Fix thermal_dm and enr_thermal_dm for extreme cases
#614
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #614 +/- ##
==========================================
+ Coverage 92.28% 93.26% +0.98%
==========================================
Files 50 50
Lines 3615 3624 +9
==========================================
+ Hits 3336 3380 +44
+ Misses 279 244 -35 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Updated calculation to improve numerical stability and changed division operator for normalization.
src/qobj/states.jl
Outdated
| data = [exp(-β * ComplexF64(j)) for j in 0:(N-1)] | ||
| function thermal_dm(N::Int, n::T; sparse::Union{Bool,Val} = Val(false)) where {T<:Real} | ||
| exp_minus_β = @. 1 - (1 / (n + 1)) # use "1-1/(n+1)" instead of "n/(n+1)" for numerical stability | ||
| P_n = _complex_float_type(T)[exp_minus_β^j for j in 0:(N-1)] |
There was a problem hiding this comment.
Unfortunately this is way slower than the exp version. thermal_dm(3000, 0.2; sparse=Val(true)) becomes 50% slower than before.
There was a problem hiding this comment.
But this implementation gives correct results for all known extreme cases. I think this is more important.
There was a problem hiding this comment.
And that's why I was talking about branching on the finiteness of beta. Since precomputing the normalization factor is problematic for dealing with nbar=Inf, simply replacing the first element in the original array already works and is faster, i.e.
data = P_n = _complex_float_type(T)[j == 0 && !isfinite(β) ? one(β) : exp(-β * j) for j in 0:(N-1)]This is doing basically the same thing as the slower version since the pow function fallbacks to log and exp.
There was a problem hiding this comment.
Ok, I didn't fully understand your point in the beginning
There was a problem hiding this comment.
In the other PR @yuyichao what saying that perhaps it is better to preallocate the vector and then filling, to avoid comprehension. Is it still true here?
I also would avoid the use of the Boltzmann function we have just introduced.
There was a problem hiding this comment.
The avoiding comprehension part was mainly to not write too much code in the comprehension (for coding style only). It is true that comprehension can basically only be as performant as loop and not better but that's not too important.
The code in the loop (and therefore would've been in comprehension) was more complex in the other PR since I wanted to reuse the result for the summation computation. However, since the analytical expression to calculate the sum had issue with infinite nbar and is therefore disgarded, that's less of a problem (granted 1 + 1/n is already computed but that has always been the case) and the code in the loop is fairly simple for the current version.
src/qobj/states.jl
Outdated
| data = [exp(-β * ComplexF64(j)) for j in 0:(N-1)] | ||
| function thermal_dm(N::Int, n::T; sparse::Union{Bool,Val} = Val(false)) where {T<:Real} | ||
| exp_minus_β = @. 1 - (1 / (n + 1)) # use "1-1/(n+1)" instead of "n/(n+1)" for numerical stability | ||
| P_n = _complex_float_type(T)[exp_minus_β^j for j in 0:(N-1)] |
There was a problem hiding this comment.
And that's why I was talking about branching on the finiteness of beta. Since precomputing the normalization factor is problematic for dealing with nbar=Inf, simply replacing the first element in the original array already works and is faster, i.e.
data = P_n = _complex_float_type(T)[j == 0 && !isfinite(β) ? one(β) : exp(-β * j) for j in 0:(N-1)]This is doing basically the same thing as the slower version since the pow function fallbacks to log and exp.
src/qobj/states.jl
Outdated
| function thermal_dm(N::Int, n::T; sparse::Union{Bool,Val} = Val(false)) where {T<:Real} | ||
| exp_minus_β = 1 - (1 / (n + 1)) # use "1-1/(n+1)" instead of "n/(n+1)" for numerical stability (works for n=0 and Inf) | ||
| P_n = _complex_float_type(T)[exp_minus_β^j for j in 0:(N-1)] | ||
| P_n /= sum(P_n) |
There was a problem hiding this comment.
What's wrong with the in-place modification?
There was a problem hiding this comment.
JET.jl reports some issues, so I keep it as P_n /= sum(P_n) instead of ./=
There was a problem hiding this comment.
What’s the issue? Why is it recommending allocating a new array?
There was a problem hiding this comment.
It relates to something else in our package
I prefer to keep it as it is for now.
There was a problem hiding this comment.
So you mean that this use is perfectly fine?
I still don't understand why such a clear cut and simple optimization is problematic. I literally can't see any reason for the non-inplace version to be better. If there is one, it would be good to state it and maybe mention in a comment.
There was a problem hiding this comment.
Yeah, I also don't understand. It somehow makes the compiler feel like we are making the Qobj as a 3D-array.
I even think that the issue reported by JET.jl is not a problem from our package.
Anyway, I would prefer to merge this bug fix first, we can improve those things in a future and separate PR if necessary.
There was a problem hiding this comment.
I agree that we should avoid allocating extra vectors unnecessary. We should understand what is going on here.
|
And ref #614 (comment) which was prematurely marked as resolved. Can you make sure |
yes, I've just confirmed in my local repo. |
By which you mean it didn't work and had to be fixed.
How did you concluded that it's the line I was talking about that's triggering the issue? I don't know if you had any other local JET runs but from the CI I assume you were talking about https://github.com/qutip/QuantumToolbox.jl/actions/runs/20274818814 and https://github.com/qutip/QuantumToolbox.jl/actions/runs/20274756492. Both of which only shows issue for There was also not an attempt to apply only the |
Yes, but I tried for the latest commit and still faced that issue, and after I changed |
|
The issue of |
|
There’s simply not an issue and fwiw the previous version attempted to use the broadcast version (it’s doing ./) as well so the inplace one is closer to what the previous one was doing. It is the version you have right now that’s making unnecessary and unrelated change by introducing an array division rather than a broadcast compared to what the code used to do. And what I was saying was that you have not even tried to make the trivial change, and because of that draw the wrong conclusion that there’s an issue at all. (I.e. it seems that JET have never raised any issues about |
|
And while in general I agree with not introducing change to unrelated parts, this is not that. The previous version of code does not contain a call to vector scalar division, but an element wise broadcast. It is therefore closer to the broadcasted version that I was asking about. Alternatively, if you want to restore it to the previous version, I.e, |
|
It seems that @yuyichao is looking for performance improvement and support to Would the following case be better? function thermal_dm(N::Int, n::Tn; sparse::Union{Bool,Val} = Val(false)) where {Tn<:Real}
(n > 0) || throw(DomainError(n, "The argument n must be positive."))
T = _complex_float_type(Tn)
# Directly compute exp(-β) = n/(1+n) without logarithm
# This is more efficient and avoids numerical issues
# Derived from: β = log(1/n + 1) => exp(-β) = n/(1+n)
# For n → ∞ (infinite temperature), exp(-β) → 1 (maximally mixed state)
exp_minus_β = isinf(n) ? one(T) : T(n) / (one(T) + T(n))
# Compute partition function using geometric series formula
# Z = sum_{j=0}^{N-1} exp(-β*j) = (1 - exp(-β*N)) / (1 - exp(-β))
Z = if abs(one(real(T)) - real(exp_minus_β)) < eps(Tn)
T(N)
else
(one(T) - exp_minus_β^N) / (one(T) - exp_minus_β)
end
# Create diagonal elements with normalized probabilities
scale = one(T) / Z
data = [exp_minus_β^j * scale for j in 0:(N-1)]
if getVal(sparse)
return QuantumObject(spdiagm(0 => data), Operator(), N)
else
return QuantumObject(diagm(0 => data), Operator(), N)
end
endIt is faster than the It also handles |
|
But the current code already supports all different input cases: |
Yeah that's right. Apart from the non-existing |
Yes, what I mean is that, since we also want to optimize. Let's choose the most optimal one. Which one is faster? |
|
I have opened an issue on JET.jl about that weird behavior aviatesk/JET.jl#790 Let's see if they can give us an answer, otherwise we can proceed I would say |
| return QuantumObject(diagm(0 => P), Operator(), N) | ||
| end | ||
| end | ||
| Boltzmann_weight(β::T, E::Int) where {T<:Real} = (E != 0 || isfinite(β)) ? exp(-β * E) : one(T) |
There was a problem hiding this comment.
Could you rename is with boltzmann_weight or even _boltzmann_weight?
There was a problem hiding this comment.
ok
I will do it tomorrow.
| _complex_float_type(::Type{Complex{Float32}}) = ComplexF32 | ||
| _complex_float_type(::Type{Complex{Float64}}) = ComplexF64 | ||
| _complex_float_type(T::Type{<:Complex}) = T # Allow other untracked Complex types, like ForwardDiff.Dual | ||
| _complex_float_type(T::Type{<:Real}) = Complex{T} # Allow other untracked Complex types, like ForwardDiff.Dual |
There was a problem hiding this comment.
Have we tested this case here? I'm not sure doing just Complex{T} would work.
There was a problem hiding this comment.
I test it by myself and it works.
Do you want me to add ForwardDiff.Dual into thermal_dm test
in this case, we need to add ForwardDiff into test's Project.toml
There was a problem hiding this comment.
We already have it in the autodiff tests. We could add some tests about the derivative of the photon number of a thermal state.
3 participants
Checklist
Thank you for contributing to
QuantumToolbox.jl! Please make sure you have finished the following tasks before opening the PR.make test.juliaformatted by running:make format.docs/folder) related to code changes were updated and able to build locally by running:make docs.CHANGELOG.mdshould be updated (regarding to the code changes) and built by running:make changelog.Request for a review after you have completed all the tasks. If you have not finished them all, you can also open a Draft Pull Request to let the others know this on-going work.
Description
This is an extended fix for PR #613 .