From 143a63c9410391b92432ad072e6b2c525ea3a63e Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <yitehuang.tw@gmail.com>
Date: Tue, 16 Dec 2025 11:14:34 +0800
Subject: [PATCH 01/11] fix `thermal_dm` and `enr_thermal_dm` for extreme cases

---
 CHANGELOG.md                           |  4 +++-
 src/qobj/energy_restricted.jl          | 12 ++++++------
 src/qobj/states.jl                     | 15 ++++++++++-----
 test/core-test/enr_state_operator.jl   |  9 +++++++++
 test/core-test/states_and_operators.jl | 22 +++++++++++++++++-----
 5 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 84bc5139e..b4c003421 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,7 +6,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
-- Add error message for bad input in state/operator generating functions ([#603])
+- Add error message for bad input in state/operator generating functions. ([#603])
+- Fix `thermal_dm` and `enr_thermal_dm` for extreme cases. ([#614])
 
 ## [v0.39.1]
 Release date: 2025-11-19
@@ -383,3 +384,4 @@ Release date: 2024-11-13
 [#591]: https://github.com/qutip/QuantumToolbox.jl/issues/591
 [#596]: https://github.com/qutip/QuantumToolbox.jl/issues/596
 [#603]: https://github.com/qutip/QuantumToolbox.jl/issues/603
+[#614]: https://github.com/qutip/QuantumToolbox.jl/issues/614
diff --git a/src/qobj/energy_restricted.jl b/src/qobj/energy_restricted.jl
index f65ea84ca..05c28a8d1 100644
--- a/src/qobj/energy_restricted.jl
+++ b/src/qobj/energy_restricted.jl
@@ -1,5 +1,5 @@
 #=
-This file defines the energy restricted space structure.
+This file defines the excitation number restricted space structure.
 =#
 
 export EnrSpace, enr_state_dictionaries
@@ -194,13 +194,13 @@ function enr_thermal_dm(
     D = s_enr.size
     idx2state = s_enr.idx2state
 
-    diags = ComplexF64[prod((nvec ./ (1 .+ nvec)) .^ idx2state[idx]) for idx in 1:D]
-    diags /= sum(diags)
-
+    exp_minus_β = @. 1 - (1 / (nvec + 1)) # use "1-1/(n+1)" instead of "n/(n+1)" to avoid n=Inf issues
+    P_n = ComplexF64[prod(exp_minus_β .^ idx2state[idx]) for idx in 1:D]
+    P_n /= sum(P_n)
     if getVal(sparse)
-        return QuantumObject(spdiagm(0 => diags), Operator(), s_enr)
+        return QuantumObject(spdiagm(0 => P_n), Operator(), s_enr)
     else
-        return QuantumObject(diagm(0 => diags), Operator(), s_enr)
+        return QuantumObject(diagm(0 => P_n), Operator(), s_enr)
     end
 end
 
diff --git a/src/qobj/states.jl b/src/qobj/states.jl
index 38cddd6c9..e4df050a6 100644
--- a/src/qobj/states.jl
+++ b/src/qobj/states.jl
@@ -119,12 +119,17 @@ Density matrix for a thermal state (generating thermal state probabilities) with
     If you want to keep type stability, it is recommended to use `thermal_dm(N, n, sparse=Val(sparse))` instead of `thermal_dm(N, n, sparse=sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 function thermal_dm(N::Int, n::Real; sparse::Union{Bool,Val} = Val(false))
-    β = log(1.0 / n + 1.0)
-    data = [exp(-β * ComplexF64(j)) for j in 0:(N-1)]
-    if getVal(sparse)
-        return QuantumObject(spdiagm(0 => data ./ sum(data)), Operator(), N)
+    if n == 0
+        return fock_dm(N, 0; sparse)
     else
-        return QuantumObject(diagm(0 => data ./ sum(data)), Operator(), N)
+        β = log(1.0 / n + 1.0)
+        P_n = [exp(-β * ComplexF64(j)) for j in 0:(N-1)]
+        P_n /= sum(P_n)
+        if getVal(sparse)
+            return QuantumObject(spdiagm(0 => P_n), Operator(), N)
+        else
+            return QuantumObject(diagm(0 => P_n), Operator(), N)
+        end
     end
 end
 
diff --git a/test/core-test/enr_state_operator.jl b/test/core-test/enr_state_operator.jl
index 50bed88c8..8f74a3d5e 100644
--- a/test/core-test/enr_state_operator.jl
+++ b/test/core-test/enr_state_operator.jl
@@ -39,6 +39,15 @@
         I_enr = enr_identity(space_enr)
         size_enr = space_enr.size
 
+        # enr_thermal_dm (extreme cases)
+        ρTd0 = enr_thermal_dm(space_enr, 0.0)
+        ρTs0 = enr_thermal_dm(space_enr, 0.0; sparse = Val(true))
+        ρTd∞ = enr_thermal_dm(space_enr, Inf)
+        ρTs∞ = enr_thermal_dm(space_enr, Inf; sparse = Val(true))
+        @test tr(ρTd0) ≈ tr(ρTs0) ≈ tr(ρTd∞) ≈ tr(ρTs∞) ≈ 1.0
+        @test ρTd0.data ≈ ρTs0.data ≈ fock_dm(size_enr, 0).data
+        @test ρTd∞.data ≈ ρTs∞.data ≈ maximally_mixed_dm(size_enr).data
+
         # tensor between normal and ENR space
         ρ_tot = tensor(ρ_s, ρ_enr)
         opstring = sprint((t, s) -> show(t, "text/plain", s), ρ_tot)
diff --git a/test/core-test/states_and_operators.jl b/test/core-test/states_and_operators.jl
index ad2814233..66645a3fa 100644
--- a/test/core-test/states_and_operators.jl
+++ b/test/core-test/states_and_operators.jl
@@ -36,11 +36,23 @@
     end
 
     @testset "thermal state" begin
-        ρTd = thermal_dm(5, 0.123)
-        ρTs = thermal_dm(5, 0.123; sparse = true)
+        N = 5
+
+        # extreme cases
+        ρTd0 = thermal_dm(N, 0.0)
+        ρTs0 = thermal_dm(N, 0.0; sparse = Val(true))
+        ρTd∞ = thermal_dm(N, Inf)
+        ρTs∞ = thermal_dm(N, Inf; sparse = Val(true))
+        @test tr(ρTd0) ≈ tr(ρTs0) ≈ tr(ρTd∞) ≈ tr(ρTs∞) ≈ 1.0
+        @test ρTd0 ≈ ρTs0 ≈ fock_dm(N, 0)
+        @test ρTd∞ ≈ ρTs∞ ≈ maximally_mixed_dm(N)
+
+        # general case
+        ρTd = thermal_dm(N, 0.123)
+        ρTs = thermal_dm(N, 0.123; sparse = Val(true))
         @test isoper(ρTd)
-        @test ρTd.dims == [5]
-        @test tr(ρTd) ≈ 1.0
+        @test ρTd.dims == [N]
+        @test tr(ρTd) ≈ tr(ρTs) ≈ 1.0
         @test ρTd.data ≈ spdiagm(
             0 => Float64[
                 0.8904859864731106,
@@ -50,7 +62,7 @@
                 0.00012815292116369966,
             ],
         )
-        @test typeof(ρTs.data) <: AbstractSparseMatrix
+        @test ρTs.data isa AbstractSparseMatrix
         @test ρTd ≈ ρTs
     end
 

From 9f75e3ee49e2780522b67d2201cd715c69e23ab6 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <yitehuang.tw@gmail.com>
Date: Tue, 16 Dec 2025 11:20:33 +0800
Subject: [PATCH 02/11] update changelog

---
 CHANGELOG.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b4c003421..76c7508c7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
 - Add error message for bad input in state/operator generating functions. ([#603])
-- Fix `thermal_dm` and `enr_thermal_dm` for extreme cases. ([#614])
+- Fix `thermal_dm` and `enr_thermal_dm` for extreme cases. ([#614], [#613])
 
 ## [v0.39.1]
 Release date: 2025-11-19
@@ -384,4 +384,5 @@ Release date: 2024-11-13
 [#591]: https://github.com/qutip/QuantumToolbox.jl/issues/591
 [#596]: https://github.com/qutip/QuantumToolbox.jl/issues/596
 [#603]: https://github.com/qutip/QuantumToolbox.jl/issues/603
+[#613]: https://github.com/qutip/QuantumToolbox.jl/issues/613
 [#614]: https://github.com/qutip/QuantumToolbox.jl/issues/614

From f9287ba44ad8f26738e8717f59f787fdb895964a Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <yitehuang.tw@gmail.com>
Date: Tue, 16 Dec 2025 12:23:13 +0800
Subject: [PATCH 03/11] fix typos

---
 src/qobj/energy_restricted.jl          | 2 +-
 src/qobj/states.jl                     | 4 ++--
 src/utilities.jl                       | 1 +
 test/core-test/states_and_operators.jl | 6 +++---
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/qobj/energy_restricted.jl b/src/qobj/energy_restricted.jl
index 05c28a8d1..31ac6ad32 100644
--- a/src/qobj/energy_restricted.jl
+++ b/src/qobj/energy_restricted.jl
@@ -195,7 +195,7 @@ function enr_thermal_dm(
     idx2state = s_enr.idx2state
 
     exp_minus_β = @. 1 - (1 / (nvec + 1)) # use "1-1/(n+1)" instead of "n/(n+1)" to avoid n=Inf issues
-    P_n = ComplexF64[prod(exp_minus_β .^ idx2state[idx]) for idx in 1:D]
+    P_n = _complex_float_type(nvec)[prod(exp_minus_β .^ idx2state[idx]) for idx in 1:D]
     P_n /= sum(P_n)
     if getVal(sparse)
         return QuantumObject(spdiagm(0 => P_n), Operator(), s_enr)
diff --git a/src/qobj/states.jl b/src/qobj/states.jl
index e4df050a6..88f303adf 100644
--- a/src/qobj/states.jl
+++ b/src/qobj/states.jl
@@ -118,12 +118,12 @@ Density matrix for a thermal state (generating thermal state probabilities) with
 !!! warning "Beware of type-stability!"
     If you want to keep type stability, it is recommended to use `thermal_dm(N, n, sparse=Val(sparse))` instead of `thermal_dm(N, n, sparse=sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
-function thermal_dm(N::Int, n::Real; sparse::Union{Bool,Val} = Val(false))
+function thermal_dm(N::Int, n::T; sparse::Union{Bool,Val} = Val(false)) where {T<:Real}
     if n == 0
         return fock_dm(N, 0; sparse)
     else
         β = log(1.0 / n + 1.0)
-        P_n = [exp(-β * ComplexF64(j)) for j in 0:(N-1)]
+        P_n = _complex_float_type(T)[exp(-β * j) for j in 0:(N-1)]
         P_n /= sum(P_n)
         if getVal(sparse)
             return QuantumObject(spdiagm(0 => P_n), Operator(), N)
diff --git a/src/utilities.jl b/src/utilities.jl
index 6d647c08b..c040a4401 100644
--- a/src/utilities.jl
+++ b/src/utilities.jl
@@ -190,6 +190,7 @@ _complex_float_type(::Type{Int32}) = ComplexF32
 _complex_float_type(::Type{Int64}) = ComplexF64
 _complex_float_type(::Type{Float32}) = ComplexF32
 _complex_float_type(::Type{Float64}) = ComplexF64
+_complex_float_type(::Type{BigFloat}) = Complex{BigFloat}
 _complex_float_type(::Type{Complex{Int32}}) = ComplexF32
 _complex_float_type(::Type{Complex{Int64}}) = ComplexF64
 _complex_float_type(::Type{Complex{Float32}}) = ComplexF32
diff --git a/test/core-test/states_and_operators.jl b/test/core-test/states_and_operators.jl
index 66645a3fa..38053730e 100644
--- a/test/core-test/states_and_operators.jl
+++ b/test/core-test/states_and_operators.jl
@@ -47,9 +47,9 @@
         @test ρTd0 ≈ ρTs0 ≈ fock_dm(N, 0)
         @test ρTd∞ ≈ ρTs∞ ≈ maximally_mixed_dm(N)
 
-        # general case
-        ρTd = thermal_dm(N, 0.123)
-        ρTs = thermal_dm(N, 0.123; sparse = Val(true))
+        # general case (also test BigFloat)
+        ρTd = thermal_dm(N, big(0.123))
+        ρTs = thermal_dm(N, big(0.123); sparse = Val(true))
         @test isoper(ρTd)
         @test ρTd.dims == [N]
         @test tr(ρTd) ≈ tr(ρTs) ≈ 1.0

From 36fb4d9e476de0313c2f1e73f8d96a5b18b626e6 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <44385685+ytdHuang@users.noreply.github.com>
Date: Wed, 17 Dec 2025 00:13:41 +0800
Subject: [PATCH 04/11] improve numerical stability for `thermal_dm`

---
 src/qobj/states.jl | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/src/qobj/states.jl b/src/qobj/states.jl
index 88f303adf..7b8f3067a 100644
--- a/src/qobj/states.jl
+++ b/src/qobj/states.jl
@@ -119,17 +119,13 @@ Density matrix for a thermal state (generating thermal state probabilities) with
     If you want to keep type stability, it is recommended to use `thermal_dm(N, n, sparse=Val(sparse))` instead of `thermal_dm(N, n, sparse=sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 function thermal_dm(N::Int, n::T; sparse::Union{Bool,Val} = Val(false)) where {T<:Real}
-    if n == 0
-        return fock_dm(N, 0; sparse)
+    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)]
+    P_n ./= sum(P_n)
+    if getVal(sparse)
+        return QuantumObject(spdiagm(0 => P_n), Operator(), N)
     else
-        β = log(1.0 / n + 1.0)
-        P_n = _complex_float_type(T)[exp(-β * j) for j in 0:(N-1)]
-        P_n /= sum(P_n)
-        if getVal(sparse)
-            return QuantumObject(spdiagm(0 => P_n), Operator(), N)
-        else
-            return QuantumObject(diagm(0 => P_n), Operator(), N)
-        end
+        return QuantumObject(diagm(0 => P_n), Operator(), N)
     end
 end
 

From 2b3b33b78b7a72082d8774ae7e64f13ec4b73627 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <44385685+ytdHuang@users.noreply.github.com>
Date: Wed, 17 Dec 2025 00:15:42 +0800
Subject: [PATCH 05/11] improve numerical stability for `enr_thermal_dm`

Updated calculation to improve numerical stability and changed division operator for normalization.
---
 src/qobj/energy_restricted.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/qobj/energy_restricted.jl b/src/qobj/energy_restricted.jl
index 31ac6ad32..c6160cc2f 100644
--- a/src/qobj/energy_restricted.jl
+++ b/src/qobj/energy_restricted.jl
@@ -194,9 +194,9 @@ function enr_thermal_dm(
     D = s_enr.size
     idx2state = s_enr.idx2state
 
-    exp_minus_β = @. 1 - (1 / (nvec + 1)) # use "1-1/(n+1)" instead of "n/(n+1)" to avoid n=Inf issues
+    exp_minus_β = @. 1 - (1 / (nvec + 1)) # use "1-1/(n+1)" instead of "n/(n+1)" for numerical stability
     P_n = _complex_float_type(nvec)[prod(exp_minus_β .^ idx2state[idx]) for idx in 1:D]
-    P_n /= sum(P_n)
+    P_n ./= sum(P_n)
     if getVal(sparse)
         return QuantumObject(spdiagm(0 => P_n), Operator(), s_enr)
     else

From c1fb5795df7fab38d504caa6798225e60278af06 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <44385685+ytdHuang@users.noreply.github.com>
Date: Wed, 17 Dec 2025 01:38:58 +0800
Subject: [PATCH 06/11] fix code quality tests

---
 src/qobj/energy_restricted.jl | 4 ++--
 src/qobj/states.jl            | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/qobj/energy_restricted.jl b/src/qobj/energy_restricted.jl
index c6160cc2f..d816cb66b 100644
--- a/src/qobj/energy_restricted.jl
+++ b/src/qobj/energy_restricted.jl
@@ -194,9 +194,9 @@ function enr_thermal_dm(
     D = s_enr.size
     idx2state = s_enr.idx2state
 
-    exp_minus_β = @. 1 - (1 / (nvec + 1)) # use "1-1/(n+1)" instead of "n/(n+1)" for numerical stability
+    exp_minus_β = @. 1 - (1 / (nvec + 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(nvec)[prod(exp_minus_β .^ idx2state[idx]) for idx in 1:D]
-    P_n ./= sum(P_n)
+    P_n /= sum(P_n)
     if getVal(sparse)
         return QuantumObject(spdiagm(0 => P_n), Operator(), s_enr)
     else
diff --git a/src/qobj/states.jl b/src/qobj/states.jl
index 7b8f3067a..9d51dfa04 100644
--- a/src/qobj/states.jl
+++ b/src/qobj/states.jl
@@ -119,9 +119,9 @@ Density matrix for a thermal state (generating thermal state probabilities) with
     If you want to keep type stability, it is recommended to use `thermal_dm(N, n, sparse=Val(sparse))` instead of `thermal_dm(N, n, sparse=sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 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
+    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)
+    P_n /= sum(P_n)
     if getVal(sparse)
         return QuantumObject(spdiagm(0 => P_n), Operator(), N)
     else

From 0d91305cda23c3c34c835ae45442d891e3372466 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <yitehuang.tw@gmail.com>
Date: Wed, 17 Dec 2025 12:15:34 +0800
Subject: [PATCH 07/11] improve type stability and efficiency

---
 src/qobj/energy_restricted.jl          | 12 ++++++-----
 src/qobj/states.jl                     | 11 +++++-----
 test/core-test/enr_state_operator.jl   | 30 +++++++++++++++++++++++---
 test/core-test/states_and_operators.jl | 16 ++++++--------
 4 files changed, 47 insertions(+), 22 deletions(-)

diff --git a/src/qobj/energy_restricted.jl b/src/qobj/energy_restricted.jl
index d816cb66b..61c913d48 100644
--- a/src/qobj/energy_restricted.jl
+++ b/src/qobj/energy_restricted.jl
@@ -194,13 +194,15 @@ function enr_thermal_dm(
     D = s_enr.size
     idx2state = s_enr.idx2state
 
-    exp_minus_β = @. 1 - (1 / (nvec + 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(nvec)[prod(exp_minus_β .^ idx2state[idx]) for idx in 1:D]
-    P_n /= sum(P_n)
+    β = @. log(1 + 1 / nvec)
+    P = _complex_float_type(T)[
+        prod(Boltzmann_weight(β[k], n_excite) for (k, n_excite) in pairs(idx2state[idx])) for idx in 1:D
+    ]
+    P /= sum(P)
     if getVal(sparse)
-        return QuantumObject(spdiagm(0 => P_n), Operator(), s_enr)
+        return QuantumObject(spdiagm(0 => P), Operator(), s_enr)
     else
-        return QuantumObject(diagm(0 => P_n), Operator(), s_enr)
+        return QuantumObject(diagm(0 => P), Operator(), s_enr)
     end
 end
 
diff --git a/src/qobj/states.jl b/src/qobj/states.jl
index 9d51dfa04..672b55d76 100644
--- a/src/qobj/states.jl
+++ b/src/qobj/states.jl
@@ -119,15 +119,16 @@ Density matrix for a thermal state (generating thermal state probabilities) with
     If you want to keep type stability, it is recommended to use `thermal_dm(N, n, sparse=Val(sparse))` instead of `thermal_dm(N, n, sparse=sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 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)
+    β = log(1 + 1 / n)
+    P = _complex_float_type(T)[Boltzmann_weight(β, j) for j in 0:(N-1)]
+    P /= sum(P)
     if getVal(sparse)
-        return QuantumObject(spdiagm(0 => P_n), Operator(), N)
+        return QuantumObject(spdiagm(0 => P), Operator(), N)
     else
-        return QuantumObject(diagm(0 => P_n), Operator(), N)
+        return QuantumObject(diagm(0 => P), Operator(), N)
     end
 end
+Boltzmann_weight(β::T, E::Int) where {T<:Real} = (E != 0 || isfinite(β)) ? exp(-β * E) : one(T)
 
 @doc raw"""
     maximally_mixed_dm(dimensions)
diff --git a/test/core-test/enr_state_operator.jl b/test/core-test/enr_state_operator.jl
index 8f74a3d5e..bf63e362e 100644
--- a/test/core-test/enr_state_operator.jl
+++ b/test/core-test/enr_state_operator.jl
@@ -1,5 +1,6 @@
 @testitem "Excitation number restricted state space" begin
     using StaticArraysCore
+    using SparseArrays
 
     @testset "EnrSpace" begin
         s_enr = EnrSpace((2, 2, 3), 3)
@@ -32,10 +33,9 @@
         space_s = (Space(D1), Space(D2))
 
         # EnrSpace
-        dims_enr = (2, 2, 3)
-        excitations = 3
+        dims_enr = (2, 3, 2)
+        excitations = 4
         space_enr = EnrSpace(dims_enr, excitations)
-        ρ_enr = enr_thermal_dm(space_enr, rand(3))
         I_enr = enr_identity(space_enr)
         size_enr = space_enr.size
 
@@ -48,7 +48,31 @@
         @test ρTd0.data ≈ ρTs0.data ≈ fock_dm(size_enr, 0).data
         @test ρTd∞.data ≈ ρTs∞.data ≈ maximally_mixed_dm(size_enr).data
 
+        # general case (also test BigFloat)
+        nvec = BigFloat[0.123, 0.456, 0.789]
+        ρTd = enr_thermal_dm(space_enr, nvec)
+        ρTs = enr_thermal_dm(space_enr, nvec; sparse = Val(true))
+        @test isoper(ρTd)
+        @test tr(ρTd) ≈ tr(ρTs) ≈ 1.0
+        @test diag(ρTd) ≈ Float64[
+            0.44317797863426783,
+            0.19545412249437527,
+            0.13879749880303996,
+            0.06121365374823841,
+            0.0434695463284246,
+            0.019171309140931812,
+            0.04854041974355738,
+            0.021407708875163092,
+            0.015202219370235007,
+            0.006704612120243388,
+            0.004761134637930744,
+            0.0020997961035927096,
+        ]
+        @test ρTs.data isa AbstractSparseMatrix
+        @test ρTd ≈ ρTs
+
         # tensor between normal and ENR space
+        ρ_enr = enr_thermal_dm(space_enr, rand(3))
         ρ_tot = tensor(ρ_s, ρ_enr)
         opstring = sprint((t, s) -> show(t, "text/plain", s), ρ_tot)
         datastring = sprint((t, s) -> show(t, "text/plain", s), ρ_tot.data)
diff --git a/test/core-test/states_and_operators.jl b/test/core-test/states_and_operators.jl
index 38053730e..916ddb94b 100644
--- a/test/core-test/states_and_operators.jl
+++ b/test/core-test/states_and_operators.jl
@@ -53,15 +53,13 @@
         @test isoper(ρTd)
         @test ρTd.dims == [N]
         @test tr(ρTd) ≈ tr(ρTs) ≈ 1.0
-        @test ρTd.data ≈ spdiagm(
-            0 => Float64[
-                0.8904859864731106,
-                0.09753319353178326,
-                0.010682620484781245,
-                0.0011700465891612583,
-                0.00012815292116369966,
-            ],
-        )
+        @test diag(ρTd) ≈ Float64[
+            0.8904859864731106,
+            0.09753319353178326,
+            0.010682620484781245,
+            0.0011700465891612583,
+            0.00012815292116369966,
+        ]
         @test ρTs.data isa AbstractSparseMatrix
         @test ρTd ≈ ρTs
     end

From 1eb3b38adf8f7f3c9332267703d71b50367cf2dc Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <yitehuang.tw@gmail.com>
Date: Wed, 17 Dec 2025 17:28:29 +0800
Subject: [PATCH 08/11] fix `_complex_float_type`

---
 src/utilities.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/utilities.jl b/src/utilities.jl
index c040a4401..ddbba0043 100644
--- a/src/utilities.jl
+++ b/src/utilities.jl
@@ -190,12 +190,12 @@ _complex_float_type(::Type{Int32}) = ComplexF32
 _complex_float_type(::Type{Int64}) = ComplexF64
 _complex_float_type(::Type{Float32}) = ComplexF32
 _complex_float_type(::Type{Float64}) = ComplexF64
-_complex_float_type(::Type{BigFloat}) = Complex{BigFloat}
 _complex_float_type(::Type{Complex{Int32}}) = ComplexF32
 _complex_float_type(::Type{Complex{Int64}}) = ComplexF64
 _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
+_complex_float_type(T::Type{<:Complex}) = T       # Allow other untracked Complex types, like ForwardDiff.Dual
 
 _convert_eltype_wordsize(::Type{T}, ::Val{64}) where {T<:Int} = Int64
 _convert_eltype_wordsize(::Type{T}, ::Val{32}) where {T<:Int} = Int32

From 9d99106d7307c74d7194545767a857e20c594318 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <y.t.d.huang@phys.ncku.edu.tw>
Date: Fri, 19 Dec 2025 11:33:16 +0800
Subject: [PATCH 09/11] move `_Boltzmann_weight` to utilities

---
 src/qobj/energy_restricted.jl | 2 +-
 src/qobj/states.jl            | 3 +--
 src/utilities.jl              | 2 ++
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/qobj/energy_restricted.jl b/src/qobj/energy_restricted.jl
index 61c913d48..31f7c19e2 100644
--- a/src/qobj/energy_restricted.jl
+++ b/src/qobj/energy_restricted.jl
@@ -196,7 +196,7 @@ function enr_thermal_dm(
 
     β = @. log(1 + 1 / nvec)
     P = _complex_float_type(T)[
-        prod(Boltzmann_weight(β[k], n_excite) for (k, n_excite) in pairs(idx2state[idx])) for idx in 1:D
+        prod(_Boltzmann_weight(β[k], n_excite) for (k, n_excite) in pairs(idx2state[idx])) for idx in 1:D
     ]
     P /= sum(P)
     if getVal(sparse)
diff --git a/src/qobj/states.jl b/src/qobj/states.jl
index 672b55d76..aad4ff7de 100644
--- a/src/qobj/states.jl
+++ b/src/qobj/states.jl
@@ -120,7 +120,7 @@ Density matrix for a thermal state (generating thermal state probabilities) with
 """
 function thermal_dm(N::Int, n::T; sparse::Union{Bool,Val} = Val(false)) where {T<:Real}
     β = log(1 + 1 / n)
-    P = _complex_float_type(T)[Boltzmann_weight(β, j) for j in 0:(N-1)]
+    P = _complex_float_type(T)[_Boltzmann_weight(β, j) for j in 0:(N-1)]
     P /= sum(P)
     if getVal(sparse)
         return QuantumObject(spdiagm(0 => P), Operator(), N)
@@ -128,7 +128,6 @@ function thermal_dm(N::Int, n::T; sparse::Union{Bool,Val} = Val(false)) where {T
         return QuantumObject(diagm(0 => P), Operator(), N)
     end
 end
-Boltzmann_weight(β::T, E::Int) where {T<:Real} = (E != 0 || isfinite(β)) ? exp(-β * E) : one(T)
 
 @doc raw"""
     maximally_mixed_dm(dimensions)
diff --git a/src/utilities.jl b/src/utilities.jl
index ddbba0043..d50496b2b 100644
--- a/src/utilities.jl
+++ b/src/utilities.jl
@@ -137,6 +137,8 @@ _sparse_similar(A::AbstractArray, args...) = sparse(args...)
 
 _Ginibre_ensemble(n::Int, rank::Int = n) = randn(ComplexF64, n, rank) / sqrt(n)
 
+_Boltzmann_weight(β::T, E::Int) where {T<:Real} = (E != 0 || isfinite(β)) ? exp(-β * E) : one(T)
+
 makeVal(x::Val{T}) where {T} = x
 makeVal(x) = Val(x)
 

From 7ac713d6996270c5329bb4a7865ab5c6e0a59c72 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <y.t.d.huang@phys.ncku.edu.tw>
Date: Fri, 19 Dec 2025 12:15:29 +0800
Subject: [PATCH 10/11] add ForwardDiff test for `thermal_dm`

---
 test/ext-test/cpu/autodiff/autodiff.jl | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/test/ext-test/cpu/autodiff/autodiff.jl b/test/ext-test/cpu/autodiff/autodiff.jl
index 5141e8854..a5ae3129b 100644
--- a/test/ext-test/cpu/autodiff/autodiff.jl
+++ b/test/ext-test/cpu/autodiff/autodiff.jl
@@ -85,6 +85,27 @@ end
 n_ss(Δ, F, γ) = abs2(F / (Δ + 1im * γ / 2))
 
 @testset "Autodiff" verbose=true begin
+
+    @testset "ForwardDiff for thermal_dm" begin
+        N = 100  # dimension of the system
+        N_op = num(N) # number operator
+
+        function N_expect(p) # p = [n]
+            ρT = thermal_dm(N, p[1]; sparse = Val(true))
+            return real(expect(N_op, ρT))
+        end
+
+        # Average photon number for thermal state
+        n = rand(Float64)
+        p = [n]
+
+        # Use ForwardDiff.gradient to compute the gradient
+        grad = ForwardDiff.gradient(N_expect, p)[1]
+
+        # Compare the result
+        @test isapprox(grad, 1.0; atol = 1e-6)
+    end
+
     @testset "sesolve" verbose=true begin
         Ω = 1.0
         params = [Ω]

From db7da3afa4645f1b6b24f4f97c60a96196044019 Mon Sep 17 00:00:00 2001
From: Yi-Te Huang <y.t.d.huang@phys.ncku.edu.tw>
Date: Fri, 19 Dec 2025 12:17:46 +0800
Subject: [PATCH 11/11] format files

---
 test/ext-test/cpu/autodiff/autodiff.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/ext-test/cpu/autodiff/autodiff.jl b/test/ext-test/cpu/autodiff/autodiff.jl
index a5ae3129b..c95205e77 100644
--- a/test/ext-test/cpu/autodiff/autodiff.jl
+++ b/test/ext-test/cpu/autodiff/autodiff.jl
@@ -85,7 +85,6 @@ end
 n_ss(Δ, F, γ) = abs2(F / (Δ + 1im * γ / 2))
 
 @testset "Autodiff" verbose=true begin
-
     @testset "ForwardDiff for thermal_dm" begin
         N = 100  # dimension of the system
         N_op = num(N) # number operator