add partial_transpose for hermitian/symmetric matrices

dev-ket · Jun 18, 2024 · dd87589 · dd87589
1 parent c9f3730
commit dd87589
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 36 deletions.
diff --git a/src/partial_tra.jl b/src/partial_tra.jl
@@ -112,55 +112,61 @@ Takes the partial trace of matrix `X` with subsystem dimensions `dims` over the
 partial_trace(X::AbstractMatrix, remove::Integer, dims::Vector{<:Integer}) = partial_trace(X, [remove], dims)
 export partial_trace
 
-"""
+@doc """
     partial_transpose(X::AbstractMatrix, transp::Vector, dims::Vector)
 
 Takes the partial transpose of matrix `X` with subsystem dimensions `dims` on the subsystems in `transp`.
-"""
-function partial_transpose(X::AbstractMatrix{T}, transp::Vector{<:Integer}, dims::Vector{<:Integer}) where {T}
-    isempty(transp) && return X
-    length(transp) == length(dims) && return LA.transpose(X)
-
-    keep = Vector{eltype(transp)}(undef, length(dims) - length(transp))  # Systems kept 
-    counter = 0
-    for i = 1:length(dims)
-        if !(i in transp)
-            counter += 1
-            keep[counter] = i
-        end
-    end
+""" partial_transpose
 
-    dY = prod(dims)                             # Dimension of the final output Y    
+for (T, symmetry, wrapper) in
+    [(:AbstractMatrix, :false, :identity), (:(LA.Hermitian), :true, :(LA.Hermitian)), (:(LA.Symmetric), :true, :(LA.Symmetric))]
+    @eval begin
+        function partial_transpose(X::$T, transp::Vector{<:Integer}, dims::Vector{<:Integer})
+            isempty(transp) && return X
+            length(transp) == length(dims) && return LA.transpose(X)
 
-    Y = Matrix{T}(undef, (dY, dY))              # Final output Y
+            keep = Vector{eltype(transp)}(undef, length(dims) - length(transp))  # Systems kept 
+            counter = 0
+            for i = 1:length(dims)
+                if !(i in transp)
+                    counter += 1
+                    keep[counter] = i
+                end
+            end
 
-    tXi = Vector{Int64}(undef, length(dims))    # Tensor indexing of X for row 
-    tXj = Vector{Int64}(undef, length(dims))    # Tensor indexing of X for column
+            dY = prod(dims)                             # Dimension of the final output Y    
 
-    tYi = Vector{Int64}(undef, length(dims))    # Tensor indexing of Y for row 
-    tYj = Vector{Int64}(undef, length(dims))    # Tensor indexing of Y for column
+            Y = similar(X, (dY, dY))                    # Final output Y
 
-    for i in 1:dY
-        _tidx!(tXi, i, dims)
-        for j in 1:i
-            _tidx!(tXj, j, dims)
+            tXi = Vector{Int64}(undef, length(dims))    # Tensor indexing of X for row 
+            tXj = Vector{Int64}(undef, length(dims))    # Tensor indexing of X for column
 
-            for k in keep
-                tYi[k] = tXi[k]
-                tYj[k] = tXj[k]
-            end
+            tYi = Vector{Int64}(undef, length(dims))    # Tensor indexing of Y for row 
+            tYj = Vector{Int64}(undef, length(dims))    # Tensor indexing of Y for column
 
-            for t in transp
-                tYi[t] = tXj[t]
-                tYj[t] = tXi[t]
-            end
+            for j in 1:dY
+                _tidx!(tYj, j, dims)
+                for i in 1:j
+                    _tidx!(tYi, i, dims)
+
+                    for k in keep
+                        tXi[k] = tYi[k]
+                        tXj[k] = tYj[k]
+                    end
 
-            Yi, Yj = _idx(tYi, dims), _idx(tYj, dims)
-            Y[Yi, Yj] = X[i, j]
-            Y[Yj, Yi] = X[j, i]
+                    for t in transp
+                        tXi[t] = tYj[t]
+                        tXj[t] = tYi[t]
+                    end
+
+                    Xi, Xj = _idx(tXi, dims), _idx(tXj, dims)
+                    Y[i, j] = X[Xi, Xj]
+                    !($symmetry) && (Y[j, i] = X[Xj, Xi])
+                end
+            end
+            return $wrapper(Y)
         end
     end
-    return Y
 end
 """
     partial_transpose(X::AbstractMatrix, transp::Vector, dims::Vector)

diff --git a/test/partial_tra.jl b/test/partial_tra.jl
@@ -56,6 +56,12 @@ end
             @test partial_transpose(abc, Int64[], [d1, d2, d3]) ≈ abc
         end
     end
+    for wrapper in [Symmetric, Hermitian]
+        M = wrapper(randn(ComplexF64, (d1 * d2 * d3, d1 * d2 * d3)))
+        x = Matrix(M)
+        @test partial_transpose(M, 2, [d1, d2, d3]) ≈ partial_transpose(x, 2, [d1, d2, d3])
+        @test partial_transpose(M, [1, 3], [d1, d2, d3]) ≈ partial_transpose(x, [1, 3], [d1, d2, d3])
+    end
 end
 
 #TODO add test with JuMP variables