Skip to content

Commit

Permalink
skip use of cartesianindex to reduce allocations
Browse files Browse the repository at this point in the history
  • Loading branch information
Adria Labay committed Dec 23, 2024
1 parent 9910a34 commit ae029d4
Showing 1 changed file with 11 additions and 16 deletions.
27 changes: 11 additions & 16 deletions src/qobj/arithmetic_and_attributes.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ export proj, ptrace, purity, permute
export tidyup, tidyup!
export get_data, get_coherence, remove_coherence, mean_occupation

import Base: _ind2sub

# Broadcasting
Base.broadcastable(x::QuantumObject) = x.data
for op in (:(+), :(-), :(*), :(/), :(^))
Expand Down Expand Up @@ -775,32 +777,25 @@ function mean_occupation(ψ::QuantumObject{T,KetQuantumObject}) where {T}
if length.dims) == 1
return mapreduce(k -> abs2(ψ[k]) * (k - 1), +, 1:ρ.dims[1])
else
R = CartesianIndices((ψ.dims...,))
off = circshift.dims, 1)
off[end] = 1
t = Tuple.dims)

x = sum(R) do j
j_tuple = Tuple(j) .- 1
J = dot(j_tuple, off) + 1
return abs2(ψ[J]) * prod(j_tuple)
x = 0.0
for J in eachindex.data)
x += abs2(ψ[J]) * prod(Base._ind2sub(t, J) .- 1)
end

return x
return real(x)
end
end

function mean_occupation::QuantumObject{T,OperatorQuantumObject}) where {T}
if length.dims) == 1
return real(mapreduce(k -> ρ[k, k] * (k - 1), +, 1:ρ.dims[1]))
else
R = CartesianIndices((ρ.dims...,))
off = circshift.dims, 1)
off[end] = 1
t = Tuple.dims)

x = sum(R) do j
j_tuple = Tuple(j) .- 1
J = dot(j_tuple, off) + 1
return ρ[J, J] * prod(j_tuple)
x = 0.0im
for J in eachindex.data[:, 1])
x += ρ[J, J] * prod(Base._ind2sub(t, J) .- 1)
end

return real(x)
Expand Down

0 comments on commit ae029d4

Please sign in to comment.