Add and test 2-arg complex

Jutho · Nov 10, 2024 · 24c7636 · 24c7636
1 parent 426b1ab
commit 24c7636
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/src/tensors/tensor.jl b/src/tensors/tensor.jl
@@ -392,6 +392,9 @@ function Base.complex(t::AbstractTensorMap)
         return copy!(similar(t, complex(scalartype(t))), t)
     end
 end
+function Base.complex(r::AbstractTensorMap{<:Real}, i::AbstractTensorMap{<:Real})
+    return add(r, i, im * one(scalartype(i)))
+end
 
 # Conversion between TensorMap and Dict, for read and write purpose
 #------------------------------------------------------------------

diff --git a/test/tensors.jl b/test/tensors.jl
@@ -185,8 +185,21 @@ for V in spacelist
                 W = V1 ⊗ V2
                 for T in (Float64, ComplexF64, ComplexF32)
                     t = @constinferred randn(T, W, W)
-                    @test real(convert(Array, t)) == convert(Array, @constinferred real(t))
-                    @test imag(convert(Array, t)) == convert(Array, @constinferred imag(t))
+
+                    tr = @constinferred real(t)
+                    @test scalartype(tr) <: Real
+                    @test real(convert(Array, t)) == convert(Array, tr)
+
+                    ti = @constinferred imag(t)
+                    @test scalartype(ti) <: Real
+                    @test imag(convert(Array, t)) == convert(Array, ti)
+
+                    tc = @inferred complex(t)
+                    @test scalartype(tc) <: Complex
+                    @test complex(convert(Array, t)) == convert(Array, tc)
+
+                    tc2 = @inferred complex(tr, ti)
+                    @test tc2 ≈ tc
                 end
             end
         end