Fix and simplify partial trace, add more tests

src/atoms/affine/partialtrace.jl

@@ -1,116 +1,43 @@
-import Base.sign
 export partialtrace
-export sign, curvature, monotonicity, evaluate
 
-struct PartialTraceAtom <: AbstractExpr
-    head::Symbol
-    id_hash::UInt64
-    children::Tuple{AbstractExpr}
-    size::Tuple{Int, Int}
-    sys::Int
-    dims::Vector
-
-    function PartialTraceAtom(x::AbstractExpr, sys::Int, dims::Vector)
-        if x.size[1] ≠ x.size[2]
-            error("Only square matrices are supported")
-        end
-        if ! (1 ≤ sys ≤ length(dims))
-            error("Invalid system, should between 1 and ", length(dims), " got ", sys)
-        end
-        if x.size[1] ≠ prod(dims)
-            error("Dimension of system doesn't correspond to dimension of subsystems")
+# We compute the partial trace of x by summing over
+# (I ⊗ <j| ⊗ I) x (I ⊗ |j> ⊗ I) for all j's
+# in the system we want to trace out.
+# This function returns the jth term in the sum, namely
+# (I ⊗ <j| ⊗ I) x (I ⊗ |j> ⊗ I).
+function _term(x, j::Int, sys, dims)
+    a = sparse(1.0I, 1, 1)
+    b = sparse(1.0I, 1, 1)
+    for (i_sys, dim) in enumerate(dims)
+        if i_sys == sys
+            # create a vector that is only 1 at its jth component
+            v = spzeros(dim, 1);
+            v[j] = 1;
+            a = kron(a, v')
+            b = kron(b, v)
+        else
+            a = kron(a, sparse(1.0I, dim, dim))
+            b = kron(b, sparse(1.0I, dim, dim))
         end
-        children = (x, )
-        newsize = (round(Int, x.size[2]/dims[sys]), round(Int, x.size[1]/dims[sys]))
-        return new(:partialtrace, hash(children), children, newsize, sys, dims)
     end
+    return a * x * b
 end
 
-function sign(x::PartialTraceAtom)
-    return sign(x.children[1])
-end
-
-function curvature(x::PartialTraceAtom)
-    return ConstVexity()
-end
-
-function monotonicity(x::PartialTraceAtom)
-    return (Nondecreasing(),)
-end
-
-function evaluate(x::PartialTraceAtom)
-    ρ = evaluate(x.children[1])
-    dims = x.dims
+"""
+    partialtrace(x, sys::Int, dims::Vector)
 
-    subsystem = function(sys)
-        function term(ρ, j::Int)
-            a = sparse(1.0I, 1, 1)
-            b = sparse(1.0I, 1, 1)
-            i_sys = 1
-            for dim in dims
-                if i_sys == sys
-                # create a vector that is only 1 at its jth component
-                v = spzeros(dim, 1);
-                v[j] = 1;
-                a = kron(a, v')
-                b = kron(b, v)
-                else
-                    a = kron(a, sparse(1.0I, dim, dim))
-                    b = kron(b, sparse(1.0I, dim, dim))
-                end
-                i_sys += 1
-            end
-            return a * ρ * b
-        end
-        return sum([term(ρ, j) for j in 1:dims[sys]])
+Returns the partial trace of `x` over the `sys`th system, where `dims` is a vector of integers encoding the dimensions of each subsystem.
+"""
+function partialtrace(x, sys::Int, dims::Vector)
+    if size(x, 1) ≠ size(x, 2)
+        throw(ArgumentError("Only square matrices are supported"))
     end
-    sub_systems = [subsystem(i) for i in 1:length(dims)]
-    a = Matrix(1.0I, 1, 1)
-    for i in 1:length(dims)
-        if i == x.sys
-            continue
-        else
-            a = kron(a,sub_systems[i])
-        end
+    if ! (1 ≤ sys ≤ length(dims))
+        throw(ArgumentError("Invalid system index, should between 1 and $(length(dims)), got $sys"))
     end
-    return tr(sub_systems[x.sys])*a
-end
-
-
-function conic_form!(x::PartialTraceAtom, unique_conic_forms::UniqueConicForms=UniqueConicForms())
-    if !has_conic_form(unique_conic_forms, x)
-        sys = x.sys
-        dims = x.dims
-
-        # compute the partial trace of x by summing over
-        # (I ⊗ <j| ⊗ I) x (I ⊗ |j> ⊗ I) for all j's
-        # in the system we want to trace out
-        # This function returns every term in the sum
-        function term(ρ, j::Int)
-            a = sparse(1.0I, 1, 1)
-            b = sparse(1.0I, 1, 1)
-            i_sys = 1
-            for dim in dims
-                if i_sys == sys
-                    # create a vector that is only 1 at its jth component
-                    v = spzeros(dim, 1);
-                    v[j] = 1;
-                    a = kron(a, v')
-                    b = kron(b, v)
-                else
-                    a = kron(a, sparse(1.0I, dim, dim))
-                    b = kron(b, sparse(1.0I, dim, dim))
-                end
-                i_sys += 1
-            end
-            return a * ρ * b
-        end
-
-        # sum all terms described above for all j's
-        objective = conic_form!(sum([term(x.children[1], j) for j in 1:dims[sys]]), unique_conic_forms)
-        cache_conic_form!(unique_conic_forms, x, objective)
+    if size(x, 1) ≠ prod(dims)
+        throw(ArgumentError("Dimension of system doesn't correspond to dimension of subsystems"))
     end
-    return get_conic_form(unique_conic_forms, x)
-end
 
-partialtrace(x::AbstractExpr, sys::Int, dim::Vector) = PartialTraceAtom(x, sys, dim)
+    return sum(j -> _term(x, j, sys, dims), 1:dims[sys])
+end

test/test_sdp.jl

@@ -172,7 +172,33 @@
             p = satisfy(constraints)
             solve!(p, solver)
             @test evaluate(ρ) ≈ [0.09942819 0.29923607 0 0; 0.299237 0.900572 0 0; 0 0 0 0; 0 0 0 0] atol=TOL
-            @test evaluate(partialtrace(ρ, 1, [2; 2])) ≈ [1.0 0; 0 0] atol=TOL
+            @test evaluate(partialtrace(ρ, 1, [2; 2])) ≈ [0.09942819 0.29923607; 0.29923607 0.90057181] atol=TOL
+
+            function rand_normalized(n)
+                A = 5*randn(n, n) + im*5*randn(n, n)
+                A / tr(A)
+            end
+
+            As = [ rand_normalized(3) for _ = 1:5]
+            Bs = [ rand_normalized(2) for _ = 1:5]
+            p = rand(5)
+
+            AB = sum(i -> p[i]*kron(As[i],Bs[i]), 1:5)
+            @test partialtrace(AB, 2, [3, 2]) ≈ sum( p .* As )
+            @test partialtrace(AB, 1, [3, 2]) ≈ sum( p .* Bs )
+
+            A, B, C = rand(5,5), rand(4,4), rand(3,3)
+            ABC = kron(kron(A, B), C)
+            @test kron(A,B)*tr(C) ≈ partialtrace(ABC, 3, [5, 4, 3])
+
+            # Test 281
+            A = rand(6,6)
+            expr = partialtrace(Constant(A), 1, [2, 3])
+            @test size(expr) == size(evaluate(expr))
+
+            @test_throws ArgumentError partialtrace(rand(6, 6), 3, [2, 3])
+            @test_throws ArgumentError partialtrace(rand(6, 6), 1, [2, 4])
+            @test_throws ArgumentError partialtrace(rand(3, 4), 1, [2, 3])
         end
 
         @testset "Optimization with complex variables" begin