From 8890926d89578f1f1ee98ee191014de07f459b9c Mon Sep 17 00:00:00 2001 From: araujoms Date: Fri, 7 Jun 2024 15:12:50 +0200 Subject: [PATCH] incorporate partial trace --- some_old_code/mpartiteLinAlg.jl | 109 -------------------------------- src/Ket.jl | 1 + src/basic.jl | 2 +- src/partial_tra.jl | 97 ++++++++++++++++++++++++++++ test/partial_tra.jl | 27 ++++++++ test/runtests.jl | 5 +- 6 files changed, 129 insertions(+), 112 deletions(-) delete mode 100644 some_old_code/mpartiteLinAlg.jl create mode 100644 src/partial_tra.jl create mode 100644 test/partial_tra.jl diff --git a/some_old_code/mpartiteLinAlg.jl b/some_old_code/mpartiteLinAlg.jl deleted file mode 100644 index 3494433..0000000 --- a/some_old_code/mpartiteLinAlg.jl +++ /dev/null @@ -1,109 +0,0 @@ -using LinearAlgebra -using InvertedIndices - -export ptr - -function tidx(idx::Integer, dims::Vector{<:Integer}) - #= - Converts a standard index i to a tensor index i1,i2,i3,... - Parameters: - idx - index - dims - list of dimensions of the Hilbert spaces in tensor product - Returns: - tidx - tensor index [i1,i2,i3] - Notes: - Basically converting multipartite agnostic index |i> to index taking - into account the tensor product structure |i1 i2 i3 ... > - - To comply with Julia's count from 1 we also index from 1! - (But in the actual computations we shift to 0 counting and then back again) - =# - nsys = length(dims) - tidx = Vector{Int64}(undef, nsys) - cidx = idx - 1 # Current index - dr = prod(dims) - for k = 1:nsys - # Everytime you increase a tensor index you shift by the product of remaining dimensions - dr ÷= dims[k] - tidx[k] = (cidx ÷ dr) + 1 - cidx %= dr - end - return tidx -end - -function idx(tidx::Vector{<:Integer}, dims::Vector{<:Integer}) - #= - Converts a tensor index i1,i2,i3,... to a standard index i - Parameters: - tidx - tensor index i1,i2,i3,... - dims - list of dimensions of the Hilbert spaces in tensor product - Returns: - idx - standard index i - Notes: - To comply with Julia's count from 1 we also index from 1! - (But in the actual computations we shift to 0 counting and then back again) - =# - # nsys = length(dims); - # idx = 1; - # dr = prod(dims); - # for k = 1:nsys - # dr = dr ÷ dims[k] - # idx += (tidx[k]-1) * dr - # end - # return idx - i = 1 - shift = 1 - - for k in length(tidx):-1:1 - i += (tidx[k] - 1) * shift - shift *= dims[k] - end - return i -end - -function ptr(X::AbstractMatrix, sys::Vector{<:Integer}, dims::Vector{<:Integer}) - #= - Traces out systems indicated by ksys. - Parameters: - X - Square matrix of appropriate size - ksys - list of {0,1} with 1 indicating the corresponding system is traced out - dims - list of dimensions of the systems - Returns: - Y - Matrix which is partial trace of X - =# - - sysidx = 1:length(dims) - sysKeep = sysidx[Not(sys)] # Systems kept - dimsY = dims[sysKeep] # The tensor dimensions of Y - dimsR = dims[sys] # The tensor dimensions of the traced out systems - dY = prod(dimsY) # Dimension of Y - dR = prod(dimsR) # Dimension of system traced out - - Y = Matrix{eltype(X)}(undef, (dY, dY)) # Final output Y - tXi = Vector{Int64}(undef, length(dims)) # Tensor indexing of X for column - tXj = Vector{Int64}(undef, length(dims)) # Tensor indexing of X for row - - # We loop through Y and find the corresponding element - for i = 1:dY - # Find current column tensor index for Y - ti = tidx(i, dimsY) - tXi[sysKeep] = ti - for j = 1:dY - # Find current row tensor index for Y - tj = tidx(j, dimsY) - tXj[sysKeep] = tj - # Now loop through the diagonal of the traced out systems - val = 0 - for k = 1:dR - tk = tidx(k, dimsR) - tXi[sys], tXj[sys] = tk, tk - - # Find (i,j) index of X that we are currently on and add it to total - Xi, Xj = idx(tXi, dims), idx(tXj, dims) - val += X[Xi, Xj] - end - Y[i, j] = val - end - end - return Y -end diff --git a/src/Ket.jl b/src/Ket.jl index 5c776df..d3523fe 100755 --- a/src/Ket.jl +++ b/src/Ket.jl @@ -13,6 +13,7 @@ include("games.jl") include("measurements.jl") include("entropy.jl") include("norms.jl") +include("partial_tra.jl") import Requires function __init__() diff --git a/src/basic.jl b/src/basic.jl index 8d91c6b..9f22dba 100644 --- a/src/basic.jl +++ b/src/basic.jl @@ -185,7 +185,7 @@ _realeps(::Type{<:Real}) = 0 Zeroes out real or imaginary parts of `M` that are smaller than `tol`. """ -function cleanup!(M::AbstractArray{T}; tol = _eps(T)) where {T<:Number} # SD: is it type stable? +function cleanup!(M::AbstractArray{T}; tol = _eps(T)) where {T<:Number} wrapper = Base.typename(typeof(M)).wrapper cleanup!(parent(M); tol) return wrapper(M) diff --git a/src/partial_tra.jl b/src/partial_tra.jl new file mode 100644 index 0000000..be5a19d --- /dev/null +++ b/src/partial_tra.jl @@ -0,0 +1,97 @@ +""" + _tidx(idx::Integer, dims::Vector) + +Converts a standard index `idx` to a tensor index [i₁, i₂, ...] with subsystems dimensions `dims`. +""" +function _tidx(idx::Integer, dims::Vector{<:Integer}) + result = Vector{Int64}(undef, length(dims)) + _tidx!(result, idx, dims) + return result +end + +function _tidx!(tidx::AbstractVector{<:Integer}, idx::Integer, dims::Vector{<:Integer}) + nsys = length(dims) + cidx = idx - 1 # Current index + dr = prod(dims) + for k = 1:nsys + # Everytime you increase a tensor index you shift by the product of remaining dimensions + dr ÷= dims[k] + tidx[k] = (cidx ÷ dr) + 1 + cidx %= dr + end + return tidx +end + +""" + _idx(tidx::Vector, dims::Vector) + +Converts a tensor index `tidx` = [i₁, i₂, ...] with subsystems dimensions `dims` to a standard index. +""" +function _idx(tidx::Vector{<:Integer}, dims::Vector{<:Integer}) + i = 1 + shift = 1 + + for k in length(tidx):-1:1 + i += (tidx[k] - 1) * shift + shift *= dims[k] + end + return i +end + +""" + partial_trace(X::AbstractMatrix, remove::AbstractVector, dims:: +AbstractVector) + +Takes the partial trace of matrix `X` with subsystem dimensions `dims` over the subsystems in `remove`. +""" +function partial_trace(X::AbstractMatrix{T}, remove::Vector{<:Integer}, dims::Vector{<:Integer}) where {T} + isempty(remove) && return X + length(remove) == length(dims) && return fill(T(LA.tr(X)), 1, 1) + + keep = Vector{eltype(remove)}(undef, length(dims) - length(remove)) # Systems kept + counter = 0 + for i = 1:length(dims) + if !(i in remove) + counter += 1 + keep[counter] = i + end + end + dimsY = dims[keep] # The tensor dimensions of Y + dimsR = dims[remove] # The tensor dimensions of the traced out systems + dY = prod(dimsY) # Dimension of Y + dR = prod(dimsR) # Dimension of system traced out + + Y = Matrix{T}(undef, (dY, dY)) # Final output Y + tXi = Vector{Int64}(undef, length(dims)) # Tensor indexing of X for column + tXj = Vector{Int64}(undef, length(dims)) # Tensor indexing of X for row + + @views tXikeep = tXi[keep] + @views tXiremove = tXi[remove] + @views tXjkeep = tXj[keep] + @views tXjremove = tXj[remove] + + # We loop through Y and find the corresponding element + @inbounds for i = 1:dY + # Find current column tensor index for Y + _tidx!(tXikeep, i, dimsY) + for j = 1:dY + # Find current row tensor index for Y + _tidx!(tXjkeep, j, dimsY) + + # Now loop through the diagonal of the traced out systems + val = zero(T) + for k = 1:dR + _tidx!(tXiremove, k, dimsR) + _tidx!(tXjremove, k, dimsR) + + # Find (i,j) index of X that we are currently on and add it to total + Xi, Xj = _idx(tXi, dims), _idx(tXj, dims) + val += X[Xi, Xj] + end + Y[i, j] = val + end + end + return Y +end +partial_trace(X::AbstractMatrix, sys::Integer, dims::Vector{<:Integer}) = partial_trace(X, [sys], dims) +export partial_trace diff --git a/test/partial_tra.jl b/test/partial_tra.jl new file mode 100644 index 0000000..c3d6aab --- /dev/null +++ b/test/partial_tra.jl @@ -0,0 +1,27 @@ +@testset "Partial trace" begin + d1, d2, d3 = 2, 3, 4 + for R in [Float64, Double64, Float128, BigFloat] + for T in [R, Complex{R}] + a = randn(T, d1, d1) + b = randn(T, d2, d2) + c = randn(T, d3, d3) + ab = kron(a, b) + ac = kron(a, c) + bc = kron(b, c) + abc = kron(ab, c) + @test partial_trace(ab, [1, 2], [d1, d2])[1] ≈ tr(ab) + @test partial_trace(ab, 2, [d1, d2]) ≈ a * tr(b) + @test partial_trace(ab, 1, [d1, d2]) ≈ b * tr(a) + @test partial_trace(ab, Int64[], [d1, d2]) ≈ ab + @test partial_trace(abc, [1, 2, 3], [d1, d2, d3])[1] ≈ tr(abc) + @test partial_trace(abc, [2, 3], [d1, d2, d3]) ≈ a * tr(b) * tr(c) + @test partial_trace(abc, [1, 3], [d1, d2, d3]) ≈ b * tr(a) * tr(c) + @test partial_trace(abc, [1, 2], [d1, d2, d3]) ≈ c * tr(a) * tr(b) + @test partial_trace(abc, 3, [d1, d2, d3]) ≈ ab * tr(c) + @test partial_trace(abc, 2, [d1, d2, d3]) ≈ ac * tr(b) + @test partial_trace(abc, 1, [d1, d2, d3]) ≈ bc * tr(a) + @test partial_trace(abc, Int64[], [d1, d2, d3]) ≈ abc + end + end +end +#TODO add test with JuMP variables diff --git a/test/runtests.jl b/test/runtests.jl index 3f55a7f..596f5a3 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,8 +6,9 @@ using Quadmath using Test include("basic.jl") +include("entropy.jl") include("measurements.jl") +include("nonlocal.jl") include("norms.jl") +include("partial_tra.jl") include("random.jl") -include("nonlocal.jl") -include("entropy.jl")