Skip to content

Commit 3147a4f

Browse files
AngelEzquerraAngelEzquerra
committed
Add support for more mixed Complex-Real operations
This change adds support (almost all of) the same mixed Complex-Real operations that nim's complex module support. This also changes a bit some of the existing mixed ops. In particular, we used to support mixed ops of Complex64 with _any_ kind of number (including integers) but they did not support Complex32 - float32 ops. While being able to mix Complex64 with ints (for example) was nice, it was not consistent with nim's own complex library and also it would have required adding many more operator overloads, so I decided to just be consistent at the cost of a little usability in a small number of cases.
1 parent d6acd3a commit 3147a4f

File tree

3 files changed

+301
-59
lines changed

3 files changed

+301
-59
lines changed

‎src/arraymancer/tensor/operators_blas_l1.nim

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,21 +45,45 @@ proc `+`*[T: SomeNumber|Complex[float32]|Complex[float64]](a, b: Tensor[T]): Ten
4545
## Tensor addition
4646
map2_inline(a, b, x + y)
4747

48+
proc `+`*[T: SomeNumber](a: Tensor[Complex[T]], b: Tensor[T]): Tensor[Complex[T]] {.noinit.} =
49+
## Mixed Complex[T] + Real Tensor addition
50+
map2_inline(a, b, x + y)
51+
52+
proc `+`*[T: SomeNumber](a: Tensor[T], b: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit.} =
53+
## Mixed Real + Complex[T] Tensor addition
54+
map2_inline(a, b, x + y)
55+
4856
proc `-`*[T: SomeNumber|Complex[float32]|Complex[float64]](a, b: Tensor[T]): Tensor[T] {.noinit.} =
4957
## Tensor substraction
5058
map2_inline(a, b, x - y)
5159

60+
proc `-`*[T: SomeNumber](a: Tensor[Complex[T]], b: Tensor[T]): Tensor[Complex[T]] {.noinit.} =
61+
## Mixed Complex[T] - Real Tensor subtraction
62+
map2_inline(a, b, x - y)
63+
64+
proc `-`*[T: SomeNumber](a: Tensor[T], b: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit.} =
65+
## Mixed Real - Complex[T] Tensor subtraction
66+
map2_inline(a, b, x - y)
67+
5268
# #########################################################
5369
# # Tensor-Tensor in-place linear algebra
5470
5571
proc `+=`*[T: SomeNumber|Complex[float32]|Complex[float64]](a: var Tensor[T], b: Tensor[T]) =
5672
## Tensor in-place addition
5773
a.apply2_inline(b, x + y)
5874

75+
proc `+=`*[T: SomeNumber](a: var Tensor[Complex[T]], b: Tensor[T]) =
76+
## Mixed Complex + Real Tensor in-place addition
77+
a.apply2_inline(b, x + y)
78+
5979
proc `-=`*[T: SomeNumber|Complex[float32]|Complex[float64]](a: var Tensor[T], b: Tensor[T]) =
6080
## Tensor in-place substraction
6181
a.apply2_inline(b, x - y)
6282

83+
proc `-=`*[T: SomeNumber](a: var Tensor[Complex[T]], b: Tensor[T]) =
84+
## Mixed Complex - Real Tensor in-place substraction
85+
a.apply2_inline(b, x - y)
86+
6387
# #########################################################
6488
# # Tensor-scalar linear algebra
6589
@@ -68,10 +92,28 @@ proc `*`*[T: SomeNumber|Complex[float32]|Complex[float64]](a: T, t: Tensor[T]):
6892
returnEmptyIfEmpty(t)
6993
t.map_inline(x * a)
7094

95+
proc `*`*[T: SomeNumber](a: T, t: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit.} =
96+
## Element-wise mixed Real * Complex multiplication by a scalar
97+
returnEmptyIfEmpty(t)
98+
t.map_inline(x * a)
99+
100+
proc `*`*[T: SomeNumber](a: Complex[T], t: Tensor[T]): Tensor[Complex[T]] {.noinit.} =
101+
## Element-wise mixed Complex * Real multiplication by a scalar
102+
returnEmptyIfEmpty(t)
103+
t.map_inline(x * a)
104+
71105
proc `*`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: Tensor[T], a: T): Tensor[T] {.noinit.} =
72106
## Element-wise multiplication by a scalar
73107
a * t
74108

109+
proc `*`*[T: SomeNumber](t: Tensor[Complex[T]], a: T): Tensor[Complex[T]] {.noinit.} =
110+
## Element-wise mixed Complex * Real multiplication by a scalar
111+
a * t
112+
113+
proc `*`*[T: SomeNumber](t: Tensor[T], a: Complex[T]): Tensor[Complex[T]] {.noinit.} =
114+
## Element-wise mixed Real * Complex multiplication by a scalar
115+
a * t
116+
75117
proc `/`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: Tensor[T], a: T): Tensor[T] {.noinit.} =
76118
## Element-wise division by a float scalar
77119
returnEmptyIfEmpty(t)
@@ -80,6 +122,26 @@ proc `/`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: Tensor[T], a: T):
80122
else:
81123
t.map_inline(x / a)
82124

125+
proc `/`*[T: SomeNumber](a: T, t: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit.} =
126+
## Element-wise mixed Real scalar / Complex tensor division
127+
returnEmptyIfEmpty(t)
128+
t.map_inline(a / x )
129+
130+
proc `/`*[T: SomeNumber](a: Complex[T], t: Tensor[T]): Tensor[Complex[T]] {.noinit.} =
131+
## Element-wise mixed Complex scalar / Real tensor division
132+
returnEmptyIfEmpty(t)
133+
t.map_inline(a / x)
134+
135+
proc `/`*[T: SomeNumber](t: Tensor[Complex[T]], a: T): Tensor[Complex[T]] {.noinit.} =
136+
## Element-wise mixed Complex tensor / Real scalar division
137+
returnEmptyIfEmpty(t)
138+
t.map_inline(x / a)
139+
140+
proc `/`*[T: SomeNumber](t: Tensor[T], a: Complex[T]): Tensor[Complex[T]] {.noinit.} =
141+
## Element-wise mixed Real tensor / Complex scalar division
142+
returnEmptyIfEmpty(t)
143+
t.map_inline(x / a)
144+
83145
proc `div`*[T: SomeInteger](t: Tensor[T], a: T): Tensor[T] {.noinit.} =
84146
## Element-wise division by an integer
85147
returnEmptyIfEmpty(t)
@@ -123,12 +185,24 @@ proc `*=`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: var Tensor[T], a:
123185
return
124186
t.apply_inline(x * a)
125187

188+
proc `*=`*[T: SomeNumber](t: var Tensor[Complex[T]], a: T) =
189+
## Element-wise mixed Complex * Real multiplication by a scalar (in-place)
190+
if t.size == 0:
191+
return
192+
t.apply_inline(x * a)
193+
126194
proc `/=`*[T: SomeFloat|Complex[float32]|Complex[float64]](t: var Tensor[T], a: T) =
127195
## Element-wise division by a scalar (in-place)
128196
if t.size == 0:
129197
return
130198
t.apply_inline(x / a)
131199

200+
proc `/=`*[T: SomeNumber](t: var Tensor[Complex[T]], a: T) =
201+
## Element-wise mixed Complex / Real division by a scalar (in-place)
202+
if t.size == 0:
203+
return
204+
t.apply_inline(x / a)
205+
132206
proc `/=`*[T: SomeInteger](t: var Tensor[T], a: T) =
133207
## Element-wise division by a scalar (in-place)
134208
if t.size == 0:

‎src/arraymancer/tensor/operators_broadcasted.nim

Lines changed: 180 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,13 @@ proc `/.`*[T: SomeNumber|Complex[float32]|Complex[float64]](a, b: Tensor[T]): Te
5252
else:
5353
result = map2_inline(tmp_a, tmp_b, x / y )
5454

55+
proc `^.`*[T: SomeNumber|Complex[float32]|Complex[float64]](a, b: Tensor[T]): Tensor[T] {.noinit.} =
56+
## Tensor element-wise exponentiation
57+
##
58+
## And broadcasted element-wise exponentiation.
59+
let (tmp_a, tmp_b) = broadcast2(a, b)
60+
result = map2_inline(tmp_a, tmp_b, pow(x, y))
61+
5562
proc `mod`*[T: SomeNumber](a, b: Tensor[T]): Tensor[T] {.noinit.} =
5663
## Tensor element-wise modulo operation
5764
##
@@ -195,81 +202,196 @@ proc `/.=`*[T: SomeNumber|Complex[float32]|Complex[float64]](t: var Tensor[T], v
195202

196203

197204
# #################################################
198-
# # Mixed Complex64 tensor - real scalar operations
205+
# # Mixed complex - real tensor operations
199206
#
200-
# It is always possible to operate on a Complex64 tensor with a real scalar
201-
# without loss of precission, so we allow it without explicit type conversion.
202-
# This makes complex tensor arithmetic more convenient to use.
207+
# Since nim's built-in complex module supports mixed complex-real operations
208+
# we allow them too (but in tensor form). This makes such mixed arithmetic
209+
# more efficient in addition to more convenient to use.
203210

204-
proc `+.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
205-
## Broadcasted addition for real scalar + Complex64 tensor
206-
##
207-
## The scalar is automatically converted to Complex64 before the operation.
208-
let complex_val = complex(float64(val))
209-
complex_val +. t
211+
proc `+.`*[T: SomeNumber](a: Tensor[Complex[T]], b: Tensor[T]): Tensor[Complex[T]] {.noinit,inline.} =
212+
## Broadcasted addition for tensors of incompatible but broadcastable shape.
213+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
214+
result = map2_inline(a, b, x + y)
210215
211-
proc `+.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
212-
## Broadcasted addition for real scalar + Complex64 tensor
213-
##
214-
## The scalar is automatically converted to Complex64 before the operation.
215-
let complex_val = complex(float64(val))
216-
t +. complex_val
216+
proc `+.`*[T: SomeNumber](a: Tensor[T], b: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit,inline.} =
217+
## Broadcasted addition for tensors of incompatible but broadcastable shape.
218+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
219+
result = map2_inline(a, b, x + y)
217220
218-
proc `-.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
219-
## Broadcasted subtraction for real scalar - Complex64 tensor
220-
##
221-
## The scalar is automatically converted to Complex64 before the operation.
222-
let complex_val = complex(float64(val))
223-
complex_val -. t
221+
proc `-.`*[T: SomeNumber](a: Tensor[Complex[T]], b: Tensor[T]): Tensor[Complex[T]] {.noinit,inline.} =
222+
## Broadcasted subtraction for tensors of incompatible but broadcastable shape.
223+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
224+
result = map2_inline(a, b, x - y)
224225
225-
proc `-.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
226-
## Broadcasted subtraction for real scalar - Complex64 tensor
227-
##
228-
## The scalar is automatically converted to Complex64 before the operation.
229-
let complex_val = complex(float64(val))
230-
t -. complex_val
226+
proc `-.`*[T: SomeNumber](a: Tensor[T], b: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit,inline.} =
227+
## Broadcasted subtraction for tensors of incompatible but broadcastable shape.
228+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
229+
result = map2_inline(a, b, x - y)
231230
232-
proc `*.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
233-
## Broadcasted multiplication for real scalar * Complex64 tensor
231+
proc `*.`*[T: SomeNumber](a: Tensor[Complex[T]], b: Tensor[T]): Tensor[Complex[T]] {.noinit.} =
232+
## Element-wise multiplication (Hadamard product).
234233
##
235-
## The scalar is automatically converted to Complex64 before the operation.
236-
let complex_val = complex(float64(val))
237-
complex_val *. t
234+
## And broadcasted element-wise multiplication.
235+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
236+
result = map2_inline(a, b, x * y)
238237
239-
proc `*.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
240-
## Broadcasted multiplication for real scalar * Complex64 tensor
238+
proc `*.`*[T: SomeNumber](a: Tensor[T], b: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit.} =
239+
## Element-wise multiplication (Hadamard product).
241240
##
242-
## The scalar is automatically converted to Complex64 before the operation.
243-
let complex_val = complex(float64(val))
244-
t *. complex_val
241+
## And broadcasted element-wise multiplication.
242+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
243+
result = map2_inline(a, b, x * y)
245244
246-
proc `/.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
247-
## Broadcasted division for real scalar / Complex64 tensor
245+
proc `/.`*[T: SomeNumber](a: Tensor[Complex[T]], b: Tensor[T]): Tensor[Complex[T]] {.noinit.} =
246+
## Tensor element-wise division
248247
##
249-
## The scalar is automatically converted to Complex64 before the operation.
250-
let complex_val = complex(float64(val))
251-
complex_val /. t
248+
## And broadcasted element-wise division.
249+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
250+
when T is SomeInteger:
251+
result = map2_inline(a, b, x div y )
252+
else:
253+
result = map2_inline(a, b, x / y )
252254
253-
proc `/.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
254-
## Broadcasted division for real scalar / Complex64 tensor
255+
proc `/.`*[T: SomeNumber](a: Tensor[T], b: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit.} =
256+
## Tensor element-wise division
255257
##
256-
## The scalar is automatically converted to Complex64 before the operation.
257-
let complex_val = complex(float64(val))
258-
t /. complex_val
258+
## And broadcasted element-wise division.
259+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
260+
when T is SomeInteger:
261+
result = map2_inline(a, b, x div y )
262+
else:
263+
result = map2_inline(a, b, x / y )
259264
260-
proc `^.`*[T: SomeNumber](val: T, t: Tensor[Complex64]): Tensor[Complex64] {.noinit, inline.} =
261-
## Broadcasted exponentiation for real scalar ^ Complex64 tensor
265+
proc `^.`*[T: SomeFloat](a: Tensor[Complex[T]], b: Tensor[T]): Tensor[Complex[T]] {.noinit.} =
266+
## Tensor element-wise exponentiation for real complex ^ scalar tensors
267+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
268+
result = map2_inline(a, b, pow(x, complex(y)))
269+
270+
proc `^.`*[T: SomeFloat](a: Tensor[T], b: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit.} =
271+
## Tensor element-wise exponentiation for real scalar ^ complex tensors
272+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
273+
result = map2_inline(a, b, pow(complex(x), y))
274+
275+
proc `mod`*[T: SomeNumber](a: Tensor[Complex[T]], b: Tensor[T]): Tensor[Complex[T]] {.noinit.} =
276+
## Tensor element-wise modulo operation
262277
##
263-
## The scalar is automatically converted to Complex64 before the operation.
264-
let complex_val = complex(float64(val))
265-
complex_val ^. t
278+
## And broadcasted element-wise modulo operation.
279+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
280+
result = map2_inline(a, b, x mod y)
266281
267-
proc `^.`*[T: SomeNumber](t: Tensor[Complex64], val: T): Tensor[Complex64] {.noinit, inline.} =
268-
## Broadcasted exponentiation for real scalar ^ Complex64 tensor
282+
proc `mod`*[T: SomeNumber](a: Tensor[T], b: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit.} =
283+
## Tensor element-wise modulo operation
269284
##
270-
## The scalar is automatically converted to Complex64 before the operation.
271-
let complex_val = complex(float64(val))
272-
t ^. complex_val
285+
## And broadcasted element-wise modulo operation.
286+
#check_shape(a, b, relaxed_rank1_check = RelaxedRankOne)
287+
result = map2_inline(a, b, x mod y)
288+
289+
# #################################################
290+
# # Mixed complex tensor - real scalar operations
291+
#
292+
# Since nim's built-in complex module supports mixed complex-real operations
293+
# we allow them too (but in tensor form). This makes such mixed arithmetic
294+
# more efficient in addition to more convenient to use.
295+
296+
proc `+.`*[T: SomeNumber](val: T, t: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit, inline.} =
297+
## Broadcasted addition for real scalar + complex tensor
298+
returnEmptyIfEmpty(t)
299+
result = t.map_inline(val + x)
300+
301+
proc `+.`*[T: SomeNumber](val: Complex[T], t: Tensor[T]): Tensor[Complex[T]] {.noinit, inline.} =
302+
## Broadcasted addition for real scalar + complex tensor
303+
returnEmptyIfEmpty(t)
304+
result = t.map_inline(val + x)
305+
306+
proc `+.`*[T: SomeNumber](t: Tensor[Complex[T]], val: T): Tensor[Complex[T]] {.noinit, inline.} =
307+
## Broadcasted addition for real scalar + complex tensor
308+
returnEmptyIfEmpty(t)
309+
result = t.map_inline(x + val)
310+
311+
proc `+.`*[T: SomeNumber](t: Tensor[T], val: Complex[T]): Tensor[Complex[T]] {.noinit, inline.} =
312+
## Broadcasted addition for real scalar + complex tensor
313+
returnEmptyIfEmpty(t)
314+
result = t.map_inline(x + val)
315+
316+
proc `-.`*[T: SomeNumber](val: T, t: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit, inline.} =
317+
## Broadcasted subtraction for real scalar - complex tensor
318+
returnEmptyIfEmpty(t)
319+
result = t.map_inline(val - x)
320+
321+
proc `-.`*[T: SomeNumber](val: Complex[T], t: Tensor[T]): Tensor[Complex[T]] {.noinit, inline.} =
322+
## Broadcasted subtraction for real scalar - complex tensor
323+
returnEmptyIfEmpty(t)
324+
result = t.map_inline(val + x)
325+
326+
proc `-.`*[T: SomeNumber](t: Tensor[Complex[T]], val: T): Tensor[Complex[T]] {.noinit, inline.} =
327+
## Broadcasted subtraction for real scalar - complex tensor
328+
returnEmptyIfEmpty(t)
329+
result = t.map_inline(x - val)
330+
331+
proc `-.`*[T: SomeNumber](t: Tensor[T], val: Complex[T]): Tensor[Complex[T]] {.noinit, inline.} =
332+
## Broadcasted subtraction for real scalar - complex tensor
333+
returnEmptyIfEmpty(t)
334+
result = t.map_inline(x - val)
335+
336+
proc `*.`*[T: SomeNumber](val: T, t: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit, inline.} =
337+
## Broadcasted multiplication for real scalar * complex tensor
338+
returnEmptyIfEmpty(t)
339+
result = t.map_inline(val * x)
340+
341+
proc `*.`*[T: SomeNumber](val: Complex[T], t: Tensor[T]): Tensor[Complex[T]] {.noinit, inline.} =
342+
## Broadcasted multiplication for real scalar * complex tensor
343+
returnEmptyIfEmpty(t)
344+
result = t.map_inline(val * x)
345+
346+
proc `*.`*[T: SomeNumber](t: Tensor[Complex[T]], val: T): Tensor[Complex[T]] {.noinit, inline.} =
347+
## Broadcasted multiplication for real scalar * complex tensor
348+
returnEmptyIfEmpty(t)
349+
result = t.map_inline(x * val)
350+
351+
proc `*.`*[T: SomeNumber](t: Tensor[T], val: Complex[T]): Tensor[Complex[T]] {.noinit, inline.} =
352+
## Broadcasted multiplication for real scalar * complex tensor
353+
returnEmptyIfEmpty(t)
354+
result = t.map_inline(x * val)
355+
356+
proc `/.`*[T: SomeNumber](val: T, t: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit, inline.} =
357+
## Broadcasted division for real scalar / complex tensor
358+
returnEmptyIfEmpty(t)
359+
result = t.map_inline(val / x)
360+
361+
proc `/.`*[T: SomeNumber](val: Complex[T], t: Tensor[T]): Tensor[Complex[T]] {.noinit, inline.} =
362+
## Broadcasted division for real scalar / complex tensor
363+
returnEmptyIfEmpty(t)
364+
result = t.map_inline(val / x)
365+
366+
proc `/.`*[T: SomeNumber](t: Tensor[Complex[T]], val: T): Tensor[Complex[T]] {.noinit, inline.} =
367+
## Broadcasted division for real scalar / complex tensor
368+
returnEmptyIfEmpty(t)
369+
result = t.map_inline(x / val)
370+
371+
proc `/.`*[T: SomeNumber](t: Tensor[T], val: Complex[T]): Tensor[Complex[T]] {.noinit, inline.} =
372+
## Broadcasted division for real scalar / complex tensor
373+
returnEmptyIfEmpty(t)
374+
result = t.map_inline(x / val)
375+
376+
proc `^.`*[T: SomeFloat](val: T, t: Tensor[Complex[T]]): Tensor[Complex[T]] {.noinit, inline.} =
377+
## Broadcasted exponentiation for real scalar ^ complex tensor
378+
returnEmptyIfEmpty(t)
379+
result = t.map_inline(pow(complex(val), x))
380+
381+
proc `^.`*[T: SomeFloat](val: Complex[T], t: Tensor[T]): Tensor[Complex[T]] {.noinit, inline.} =
382+
## Broadcasted exponentiation for real scalar ^ complex tensor
383+
returnEmptyIfEmpty(t)
384+
result = t.map_inline(pow(val, x))
385+
386+
proc `^.`*[T: SomeFloat](t: Tensor[Complex[T]], val: T): Tensor[Complex[T]] {.noinit, inline.} =
387+
## Broadcasted exponentiation for real scalar ^ complex tensor
388+
returnEmptyIfEmpty(t)
389+
result = t.map_inline(pow(x, val))
390+
391+
proc `^.`*[T: SomeFloat](t: Tensor[T], val: Complex[T]): Tensor[Complex[T]] {.noinit, inline.} =
392+
## Broadcasted exponentiation for real scalar ^ complex tensor
393+
returnEmptyIfEmpty(t)
394+
result = t.map_inline(pow(complex(x), val))
273395
274396
proc `+.=`*[T: SomeNumber](t: var Tensor[Complex64], val: T) {.inline.} =
275397
## Complex64 tensor in-place addition with a real scalar.

0 commit comments

Comments
 (0)