diff --git a/src/arraymancer/tensor/complex.nim b/src/arraymancer/tensor/complex.nim index 036630ef..67bee011 100644 --- a/src/arraymancer/tensor/complex.nim +++ b/src/arraymancer/tensor/complex.nim @@ -43,6 +43,21 @@ proc complex*[T: SomeNumber](re: Tensor[T]): auto {.inline, noinit.} = else: map_inline(re, complex(x)) +proc complex_imag*[T: SomeNumber](im: Tensor[T]): auto {.inline, noinit.} = + ## Create a new, imaginary Tensor from a single real Tensor + ## + ## The input Tensor is copied into the imaginary part of the output Tensor, + ## while the real part is set to all zeros. + ## + ## If the input is an integer Tensor, the output will be a Tensor of + ## Complex64 to avoid any loss of precision. If you want to convert it + ## into a Tensor of Complex32, you must convert the input to float32 first + ## by using `.asType(float32)`. + when T is SomeInteger: + map_inline(im, complex(float64(0.0), float64(x))) + else: + map_inline(im, complex(T(0.0), x)) + proc real*[T: SomeFloat](t: Tensor[Complex[T]]): Tensor[T] {.inline, noinit.} = ## Get the real part of a complex Tensor (as a float Tensor) t.map_inline(x.re) diff --git a/tests/tensor/test_complex.nim b/tests/tensor/test_complex.nim index b07db158..01716c4b 100644 --- a/tests/tensor/test_complex.nim +++ b/tests/tensor/test_complex.nim @@ -56,6 +56,24 @@ proc main() = check: c64_from_int == expected_c64 check: c32 == expected_c32 + block: # Single tensor to imaginary complex conversion + var im_int = [1, -10, 20].toTensor + var im_float64 = im_int.asType(float64) + var im_float32 = im_int.asType(float32) + var c64_imag_from_int = complex_imag(im_int) + var c64_imag = complex_imag(im_float64) + var c32_imag = complex_imag(im_float32) + var expected_imag64 = [complex64(0.0, 1.0), + complex64(0.0, -10.0), + complex64(0.0, 20.0)].toTensor + var expected_imag32 = [complex32(0.0, 1.0), + complex32(0.0, -10.0), + complex32(0.0, 20.0)].toTensor + + check: c64_imag == expected_imag64 + check: c64_imag_from_int == expected_imag64 + check: c32_imag == expected_imag32 + test "Get the Real and Imaginary Components of a Complex Tensor": var c = [ complex(1.0, -300.0),