co2-eos 0.3.0__tar.gz → 0.6.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. {co2_eos-0.3.0 → co2_eos-0.6.0}/PKG-INFO +23 -6
  2. {co2_eos-0.3.0 → co2_eos-0.6.0}/README.md +21 -4
  3. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos/__init__.py +129 -48
  4. co2_eos-0.6.0/co2_eos/backends.py +56 -0
  5. co2_eos-0.6.0/co2_eos/core.py +348 -0
  6. co2_eos-0.6.0/co2_eos/data/chi_ref_cheb_nist2026.npz +0 -0
  7. co2_eos-0.6.0/co2_eos/data/saturation_table_nist2026.npz +0 -0
  8. co2_eos-0.6.0/co2_eos/data/seed_table_nist2026.npz +0 -0
  9. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos/helmholtz.py +10 -10
  10. co2_eos-0.6.0/co2_eos/inversions.py +837 -0
  11. co2_eos-0.6.0/co2_eos/nist2026.py +278 -0
  12. co2_eos-0.6.0/co2_eos/nist2026_data.py +142 -0
  13. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos/saturation.py +65 -42
  14. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos/span_wagner.py +6 -6
  15. co2_eos-0.6.0/co2_eos/transport.py +594 -0
  16. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos.egg-info/PKG-INFO +23 -6
  17. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos.egg-info/SOURCES.txt +18 -1
  18. {co2_eos-0.3.0 → co2_eos-0.6.0}/pyproject.toml +2 -2
  19. {co2_eos-0.3.0 → co2_eos-0.6.0}/tests/test_accuracy_vs_coolprop.py +1 -1
  20. co2_eos-0.6.0/tests/test_bulk_viscosity.py +172 -0
  21. co2_eos-0.6.0/tests/test_density_from_PT_branch.py +128 -0
  22. {co2_eos-0.3.0 → co2_eos-0.6.0}/tests/test_dome_detection.py +2 -2
  23. {co2_eos-0.3.0 → co2_eos-0.6.0}/tests/test_fastpath_validation.py +1 -1
  24. {co2_eos-0.3.0 → co2_eos-0.6.0}/tests/test_map_fresh.py +2 -2
  25. co2_eos-0.6.0/tests/test_nist2026_backend.py +100 -0
  26. co2_eos-0.6.0/tests/test_nist2026_derivs.py +116 -0
  27. co2_eos-0.6.0/tests/test_nist2026_near_critical.py +110 -0
  28. co2_eos-0.6.0/tests/test_nist2026_saturation.py +92 -0
  29. co2_eos-0.6.0/tests/test_nist2026_verification.py +86 -0
  30. co2_eos-0.6.0/tests/test_nist2026_vs_span_wagner.py +114 -0
  31. co2_eos-0.6.0/tests/test_no_em_dash.py +31 -0
  32. co2_eos-0.6.0/tests/test_pressure_derivatives.py +124 -0
  33. co2_eos-0.6.0/tests/test_sw_default_bit_for_bit.py +64 -0
  34. co2_eos-0.3.0/co2_eos/core.py +0 -271
  35. co2_eos-0.3.0/co2_eos/inversions.py +0 -695
  36. co2_eos-0.3.0/co2_eos/transport.py +0 -352
  37. {co2_eos-0.3.0 → co2_eos-0.6.0}/LICENSE +0 -0
  38. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos/data/chi_ref_cheb.npz +0 -0
  39. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos/data/saturation_table.npz +0 -0
  40. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos/data/seed_table.npz +0 -0
  41. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos.egg-info/dependency_links.txt +0 -0
  42. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos.egg-info/requires.txt +0 -0
  43. {co2_eos-0.3.0 → co2_eos-0.6.0}/co2_eos.egg-info/top_level.txt +0 -0
  44. {co2_eos-0.3.0 → co2_eos-0.6.0}/setup.cfg +0 -0
  45. {co2_eos-0.3.0 → co2_eos-0.6.0}/tests/test_analytic_derivs.py +0 -0
  46. {co2_eos-0.3.0 → co2_eos-0.6.0}/tests/test_inversion_gradients.py +0 -0
  47. {co2_eos-0.3.0 → co2_eos-0.6.0}/tests/test_no_coolprop_runtime.py +0 -0
  48. {co2_eos-0.3.0 → co2_eos-0.6.0}/tests/test_rho_u_inversion.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: co2-eos
3
- Version: 0.3.0
3
+ Version: 0.6.0
4
4
  Summary: Differentiable CO₂ thermodynamic properties in JAX
5
5
  Author-email: John Patrick Therrien <john@flux-tech.de>
6
6
  License-Expression: Apache-2.0
@@ -8,7 +8,7 @@ Project-URL: Homepage, https://github.com/FluxTech-UG/co2-eos
8
8
  Project-URL: Repository, https://github.com/FluxTech-UG/co2-eos
9
9
  Project-URL: Documentation, https://github.com/FluxTech-UG/co2-eos#readme
10
10
  Project-URL: Issues, https://github.com/FluxTech-UG/co2-eos/issues
11
- Keywords: CO2,equation-of-state,thermodynamics,JAX,autodiff,Span-Wagner
11
+ Keywords: CO2,equation-of-state,thermodynamics,JAX,autodiff,Span-Wagner,NIST IR 8608
12
12
  Classifier: Development Status :: 4 - Beta
13
13
  Classifier: Intended Audience :: Science/Research
14
14
  Classifier: Operating System :: OS Independent
@@ -46,7 +46,7 @@ import co2_eos as co2
46
46
  import jax, jax.numpy as jnp
47
47
 
48
48
  # Simulation hot path: full state from the conserved variables (ρ, u).
49
- # Batched/array-native — solves T once and reuses the α-derivatives for
49
+ # Batched/array-native: solves T once and reuses the α-derivatives for
50
50
  # every property. This is the call a finite-volume code makes each RHS eval.
51
51
  rho = jnp.array([300.0, 350.0, 420.0]) # kg/m³
52
52
  u = jnp.array([3.62e5, 3.58e5, 3.42e5]) # J/kg (specific internal energy)
@@ -70,6 +70,20 @@ T_array = jnp.linspace(280, 340, 100)
70
70
  states = jax.vmap(lambda T: co2.state_from_PT(P=8e6, T=T))(T_array)
71
71
  ```
72
72
 
73
+ ## Two equations of state
74
+
75
+ Every public function takes `eos=`: `"span_wagner"` (Span and Wagner 1996, the default)
76
+ or `"nist2026"` (Lemmon, Harvey and Hellmann, NIST IR 8608, 2026: 24 terms with no
77
+ nonanalytic terms, free of Span-Wagner's near-critical oscillations in c_p, κ_T and
78
+ (∂c_v/∂T)_ρ). The name is static under `jax.jit`, so each backend compiles once.
79
+
80
+ ```python
81
+ s = co2_eos.properties(314.0, 472.0, eos="nist2026") # Widom crossing at 9.1 MPa
82
+ ```
83
+
84
+ The transport correlations keep their parameters on either backend; their critical parts
85
+ read the selected backend's thermodynamics.
86
+
73
87
  ## Try it without cloning
74
88
 
75
89
  [`examples/launch_demo.ipynb`](examples/launch_demo.ipynb) is the launch demo: validation against CoolProp, the `state_from_PT` CPU/GPU benchmarks, and a worked example of gradient-based optimisation through `state_from_PT`. (The v0.2 hot-path `(ρ, u)` benchmarks are in [`bench/PROFILING.md`](bench/PROFILING.md).) [Read it on GitHub](https://github.com/FluxTech-UG/co2-eos/blob/main/examples/launch_demo.ipynb) for a static render, or [open it on Colab](https://colab.research.google.com/github/FluxTech-UG/co2-eos/blob/main/examples/launch_demo.ipynb) and switch the runtime to T4 GPU to reproduce the comparison in about 15 minutes, most of that being the CPU baseline.
@@ -83,7 +97,7 @@ But CoolProp is a C++ library with Python bindings. You can't `jax.grad` through
83
97
  CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagner 1996) directly in JAX:
84
98
 
85
99
  - **Differentiable.** `jax.grad` through any property, any inversion, any combination. No finite differences. The inversions carry hand-written `custom_jvp` rules (implicit function theorem), so forward- and reverse-mode gradients are exact and cheap; validated against central finite differences.
86
- - **Fast.** JIT-compiled and vectorisable, with hand-coded, transcendental-economized analytic α-derivatives: the Span-Wagner exponent structure (¼-integer τ-powers, integer δ-powers, shared exponential envelopes) is evaluated by sqrt/multiply ladders, so a full six-derivative bundle costs ~12 exps + 3 pows per point instead of ~150 array-exponent pows — the same arithmetic, re-associated, exact to round-off. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery: a dome-safe table seed plus three unrolled analytic-Cv Newton steps recover `T(ρ, u)` to 1e-12 K, and one shared derivative bundle feeds every property including the transport critical enhancement. On an Apple M2 Pro (float64) the full recovery runs **123 µs for a 64-point batch and 0.53 µs/point at large batches — 7.6–19.9× faster than the v0.1 autodiff path** and 2.6–6.3× faster than v0.2. On an OVH Tesla V100S the inversion reaches **63.7×** vs v0.1 at a 2²⁰-point batch (2 ns/point; full recovery 16 ns/point), and the launch-latency floor at small batches dropped ~1.7× (`bench/compare.py`, `bench/PROFILING.md`, validation in `docs/validation-fastpath-2026-07.md`).
100
+ - **Fast.** JIT-compiled and vectorisable, with hand-coded, transcendental-economized analytic α-derivatives: the Span-Wagner exponent structure (¼-integer τ-powers, integer δ-powers, shared exponential envelopes) is evaluated by sqrt/multiply ladders, so a full six-derivative bundle costs ~12 exps + 3 pows per point instead of ~150 array-exponent pows: the same arithmetic, re-associated, exact to round-off. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery: a dome-safe table seed plus three unrolled analytic-Cv Newton steps recover `T(ρ, u)` to 1e-12 K, and one shared derivative bundle feeds every property including the transport critical enhancement. On an Apple M2 Pro (float64) the full recovery runs **123 µs for a 64-point batch and 0.53 µs/point at large batches, 7.6–19.9× faster than the v0.1 autodiff path** and 2.6–6.3× faster than v0.2. On an OVH Tesla V100S the inversion reaches **63.7×** vs v0.1 at a 2²⁰-point batch (2 ns/point; full recovery 16 ns/point), and the launch-latency floor at small batches dropped ~1.7× (`bench/compare.py`, `bench/PROFILING.md`, validation in `docs/validation-fastpath-2026-07.md`).
87
101
  - **Also fast for `(P, T)`.** The `state_from_PT` workflow that mirrors `PropsSI` includes an iterative density solve and flattens to a few tens of microseconds per point at large batches, roughly 2.7× faster than `CoolProp.PropsSI` in a Python loop on the same inputs (10,000-point batch, Apple M2 Pro; see `examples/launch_demo.ipynb`, section 2). On a Colab T4 GPU it runs at 27,800 states/sec on a 10⁶-point batch (36 μs/pt) versus 1,600 states/sec on Colab CPU (615 μs/pt), a 17× speedup. Its ceiling is the density solve's data-dependent `while_loop` (lockstep across GPU warps); the feed-forward `(ρ, u)` hot path above does not have that bottleneck.
88
102
  - **Composable.** `jit`, `vmap`, `grad`, `custom_vjp`: the full JAX transformation stack works. Embed property evaluations inside your own JIT-compiled simulation and differentiate end-to-end.
89
103
  - **Phase-aware.** Robust inversions near the critical point using Halley's method with step damping and bisection fallback. Two-phase dome detection that avoids convergence to thermodynamically unstable spinodal states.
@@ -96,6 +110,7 @@ CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagne
96
110
  - Isochoric and isobaric heat capacities (Cv, Cp)
97
111
  - Speed of sound
98
112
  - Gibbs energy
113
+ - Pressure derivatives (∂P/∂T)_ρ [Pa/K] and (∂P/∂ρ)_T [Pa·m³/kg], keys `dP_dT_rho` and `dP_drho_T` in every full-state dict (what a code that evolves temperature needs for its energy equation)
99
114
 
100
115
  **Transport properties:**
101
116
 
@@ -111,6 +126,7 @@ CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagne
111
126
 
112
127
  - `properties_from_rho_u(rho, u)` / `temperature_from_rho_u(rho, u)`: the batched simulation hot path. Carries (ρ, u) as conserved variables; solves `T(ρ, u)` once with a table-seeded fixed-iteration analytic-Cv Newton, then derives every property (and viscosity / thermal conductivity) in one fused pass. `custom_jvp` (implicit function theorem) for exact forward- and reverse-mode gradients.
113
128
  - `state_from_Du(rho, u)`: scalar form of the above (vmap for batches).
129
+ - `properties_from_rho_T(rho, T)`: batched full state from (ρ, T), no inversion; the array-native counterpart of the scalar `properties(T, rho)` (ρ first, like `properties_from_rho_u`). One analytic α-derivative pass per point feeds every property.
114
130
  - `state_from_PT(P, T)`: Halley's method with pressure-aware critical-region initial guess, step damping, and bisection safety net
115
131
  - `state_from_Ph(P, h, phase_hint=)` (phase-aware): detects the two-phase dome at subcritical pressures and returns saturation properties directly, bypassing single-phase Newton iteration inside the dome
116
132
 
@@ -123,7 +139,7 @@ one thing worth adopting is the fused hot path.
123
139
  If your simulation carried `(ρ, u)` and recovered primitives with separate calls:
124
140
 
125
141
  ```python
126
- # v0.1 — separate calls, each re-evaluating α-derivatives
142
+ # v0.1: separate calls, each re-evaluating α-derivatives
127
143
  from co2_eos.inversions import temperature_from_Du, SUPERCRITICAL
128
144
  from co2_eos import span_wagner as sw, transport
129
145
  T = temperature_from_Du(rho, u, phase) # nested-autodiff Cv, while_loop
@@ -135,7 +151,7 @@ k = transport.thermal_conductivity(T, rho)
135
151
  replace it with the single fused call:
136
152
 
137
153
  ```python
138
- # v0.2 — one fused pass: solve T once, reuse the α-derivatives
154
+ # v0.2: one fused pass that solves T once and reuses the α-derivatives
139
155
  import co2_eos as co2
140
156
  s = co2.properties_from_rho_u(rho, u) # phase_hint defaults to SUPERCRITICAL
141
157
  T, P, mu, k = s["temperature"], s["pressure"], s["viscosity"], s["thermal_conductivity"]
@@ -191,6 +207,7 @@ CO2-EOS is validated against CoolProp across the full valid range. The test suit
191
207
  - Single-phase properties on a dense (T, ρ) grid spanning gas, liquid, and supercritical regions
192
208
  - Saturation curve properties from triple point to near-critical
193
209
  - Transport properties (viscosity and thermal conductivity) across all phases
210
+ - Bulk viscosity (0.5.0) against a reference table at nine near-critical and supercritical states, the dilute limit against Cramer (2012), and finite forward and reverse gradients (`tests/test_bulk_viscosity.py`)
194
211
  - Inversions (P,T → ρ), (P,h → state), (ρ,u → state) including near-critical conditions
195
212
  - Comparison of all Helmholtz derivatives against CoolProp's analytical values
196
213
 
@@ -12,7 +12,7 @@ import co2_eos as co2
12
12
  import jax, jax.numpy as jnp
13
13
 
14
14
  # Simulation hot path: full state from the conserved variables (ρ, u).
15
- # Batched/array-native — solves T once and reuses the α-derivatives for
15
+ # Batched/array-native: solves T once and reuses the α-derivatives for
16
16
  # every property. This is the call a finite-volume code makes each RHS eval.
17
17
  rho = jnp.array([300.0, 350.0, 420.0]) # kg/m³
18
18
  u = jnp.array([3.62e5, 3.58e5, 3.42e5]) # J/kg (specific internal energy)
@@ -36,6 +36,20 @@ T_array = jnp.linspace(280, 340, 100)
36
36
  states = jax.vmap(lambda T: co2.state_from_PT(P=8e6, T=T))(T_array)
37
37
  ```
38
38
 
39
+ ## Two equations of state
40
+
41
+ Every public function takes `eos=`: `"span_wagner"` (Span and Wagner 1996, the default)
42
+ or `"nist2026"` (Lemmon, Harvey and Hellmann, NIST IR 8608, 2026: 24 terms with no
43
+ nonanalytic terms, free of Span-Wagner's near-critical oscillations in c_p, κ_T and
44
+ (∂c_v/∂T)_ρ). The name is static under `jax.jit`, so each backend compiles once.
45
+
46
+ ```python
47
+ s = co2_eos.properties(314.0, 472.0, eos="nist2026") # Widom crossing at 9.1 MPa
48
+ ```
49
+
50
+ The transport correlations keep their parameters on either backend; their critical parts
51
+ read the selected backend's thermodynamics.
52
+
39
53
  ## Try it without cloning
40
54
 
41
55
  [`examples/launch_demo.ipynb`](examples/launch_demo.ipynb) is the launch demo: validation against CoolProp, the `state_from_PT` CPU/GPU benchmarks, and a worked example of gradient-based optimisation through `state_from_PT`. (The v0.2 hot-path `(ρ, u)` benchmarks are in [`bench/PROFILING.md`](bench/PROFILING.md).) [Read it on GitHub](https://github.com/FluxTech-UG/co2-eos/blob/main/examples/launch_demo.ipynb) for a static render, or [open it on Colab](https://colab.research.google.com/github/FluxTech-UG/co2-eos/blob/main/examples/launch_demo.ipynb) and switch the runtime to T4 GPU to reproduce the comparison in about 15 minutes, most of that being the CPU baseline.
@@ -49,7 +63,7 @@ But CoolProp is a C++ library with Python bindings. You can't `jax.grad` through
49
63
  CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagner 1996) directly in JAX:
50
64
 
51
65
  - **Differentiable.** `jax.grad` through any property, any inversion, any combination. No finite differences. The inversions carry hand-written `custom_jvp` rules (implicit function theorem), so forward- and reverse-mode gradients are exact and cheap; validated against central finite differences.
52
- - **Fast.** JIT-compiled and vectorisable, with hand-coded, transcendental-economized analytic α-derivatives: the Span-Wagner exponent structure (¼-integer τ-powers, integer δ-powers, shared exponential envelopes) is evaluated by sqrt/multiply ladders, so a full six-derivative bundle costs ~12 exps + 3 pows per point instead of ~150 array-exponent pows — the same arithmetic, re-associated, exact to round-off. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery: a dome-safe table seed plus three unrolled analytic-Cv Newton steps recover `T(ρ, u)` to 1e-12 K, and one shared derivative bundle feeds every property including the transport critical enhancement. On an Apple M2 Pro (float64) the full recovery runs **123 µs for a 64-point batch and 0.53 µs/point at large batches — 7.6–19.9× faster than the v0.1 autodiff path** and 2.6–6.3× faster than v0.2. On an OVH Tesla V100S the inversion reaches **63.7×** vs v0.1 at a 2²⁰-point batch (2 ns/point; full recovery 16 ns/point), and the launch-latency floor at small batches dropped ~1.7× (`bench/compare.py`, `bench/PROFILING.md`, validation in `docs/validation-fastpath-2026-07.md`).
66
+ - **Fast.** JIT-compiled and vectorisable, with hand-coded, transcendental-economized analytic α-derivatives: the Span-Wagner exponent structure (¼-integer τ-powers, integer δ-powers, shared exponential envelopes) is evaluated by sqrt/multiply ladders, so a full six-derivative bundle costs ~12 exps + 3 pows per point instead of ~150 array-exponent pows: the same arithmetic, re-associated, exact to round-off. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery: a dome-safe table seed plus three unrolled analytic-Cv Newton steps recover `T(ρ, u)` to 1e-12 K, and one shared derivative bundle feeds every property including the transport critical enhancement. On an Apple M2 Pro (float64) the full recovery runs **123 µs for a 64-point batch and 0.53 µs/point at large batches, 7.6–19.9× faster than the v0.1 autodiff path** and 2.6–6.3× faster than v0.2. On an OVH Tesla V100S the inversion reaches **63.7×** vs v0.1 at a 2²⁰-point batch (2 ns/point; full recovery 16 ns/point), and the launch-latency floor at small batches dropped ~1.7× (`bench/compare.py`, `bench/PROFILING.md`, validation in `docs/validation-fastpath-2026-07.md`).
53
67
  - **Also fast for `(P, T)`.** The `state_from_PT` workflow that mirrors `PropsSI` includes an iterative density solve and flattens to a few tens of microseconds per point at large batches, roughly 2.7× faster than `CoolProp.PropsSI` in a Python loop on the same inputs (10,000-point batch, Apple M2 Pro; see `examples/launch_demo.ipynb`, section 2). On a Colab T4 GPU it runs at 27,800 states/sec on a 10⁶-point batch (36 μs/pt) versus 1,600 states/sec on Colab CPU (615 μs/pt), a 17× speedup. Its ceiling is the density solve's data-dependent `while_loop` (lockstep across GPU warps); the feed-forward `(ρ, u)` hot path above does not have that bottleneck.
54
68
  - **Composable.** `jit`, `vmap`, `grad`, `custom_vjp`: the full JAX transformation stack works. Embed property evaluations inside your own JIT-compiled simulation and differentiate end-to-end.
55
69
  - **Phase-aware.** Robust inversions near the critical point using Halley's method with step damping and bisection fallback. Two-phase dome detection that avoids convergence to thermodynamically unstable spinodal states.
@@ -62,6 +76,7 @@ CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagne
62
76
  - Isochoric and isobaric heat capacities (Cv, Cp)
63
77
  - Speed of sound
64
78
  - Gibbs energy
79
+ - Pressure derivatives (∂P/∂T)_ρ [Pa/K] and (∂P/∂ρ)_T [Pa·m³/kg], keys `dP_dT_rho` and `dP_drho_T` in every full-state dict (what a code that evolves temperature needs for its energy equation)
65
80
 
66
81
  **Transport properties:**
67
82
 
@@ -77,6 +92,7 @@ CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagne
77
92
 
78
93
  - `properties_from_rho_u(rho, u)` / `temperature_from_rho_u(rho, u)`: the batched simulation hot path. Carries (ρ, u) as conserved variables; solves `T(ρ, u)` once with a table-seeded fixed-iteration analytic-Cv Newton, then derives every property (and viscosity / thermal conductivity) in one fused pass. `custom_jvp` (implicit function theorem) for exact forward- and reverse-mode gradients.
79
94
  - `state_from_Du(rho, u)`: scalar form of the above (vmap for batches).
95
+ - `properties_from_rho_T(rho, T)`: batched full state from (ρ, T), no inversion; the array-native counterpart of the scalar `properties(T, rho)` (ρ first, like `properties_from_rho_u`). One analytic α-derivative pass per point feeds every property.
80
96
  - `state_from_PT(P, T)`: Halley's method with pressure-aware critical-region initial guess, step damping, and bisection safety net
81
97
  - `state_from_Ph(P, h, phase_hint=)` (phase-aware): detects the two-phase dome at subcritical pressures and returns saturation properties directly, bypassing single-phase Newton iteration inside the dome
82
98
 
@@ -89,7 +105,7 @@ one thing worth adopting is the fused hot path.
89
105
  If your simulation carried `(ρ, u)` and recovered primitives with separate calls:
90
106
 
91
107
  ```python
92
- # v0.1 — separate calls, each re-evaluating α-derivatives
108
+ # v0.1: separate calls, each re-evaluating α-derivatives
93
109
  from co2_eos.inversions import temperature_from_Du, SUPERCRITICAL
94
110
  from co2_eos import span_wagner as sw, transport
95
111
  T = temperature_from_Du(rho, u, phase) # nested-autodiff Cv, while_loop
@@ -101,7 +117,7 @@ k = transport.thermal_conductivity(T, rho)
101
117
  replace it with the single fused call:
102
118
 
103
119
  ```python
104
- # v0.2 — one fused pass: solve T once, reuse the α-derivatives
120
+ # v0.2: one fused pass that solves T once and reuses the α-derivatives
105
121
  import co2_eos as co2
106
122
  s = co2.properties_from_rho_u(rho, u) # phase_hint defaults to SUPERCRITICAL
107
123
  T, P, mu, k = s["temperature"], s["pressure"], s["viscosity"], s["thermal_conductivity"]
@@ -157,6 +173,7 @@ CO2-EOS is validated against CoolProp across the full valid range. The test suit
157
173
  - Single-phase properties on a dense (T, ρ) grid spanning gas, liquid, and supercritical regions
158
174
  - Saturation curve properties from triple point to near-critical
159
175
  - Transport properties (viscosity and thermal conductivity) across all phases
176
+ - Bulk viscosity (0.5.0) against a reference table at nine near-critical and supercritical states, the dilute limit against Cramer (2012), and finite forward and reverse gradients (`tests/test_bulk_viscosity.py`)
160
177
  - Inversions (P,T → ρ), (P,h → state), (ρ,u → state) including near-critical conditions
161
178
  - Comparison of all Helmholtz derivatives against CoolProp's analytical values
162
179
 
@@ -1,30 +1,51 @@
1
1
  """CO2-EOS: Differentiable CO₂ thermodynamic properties in JAX.
2
2
 
3
- Pure-JAX Span-Wagner (1996) equation of state for carbon dioxide with hand-coded
3
+ Pure-JAX Helmholtz equations of state for carbon dioxide with hand-coded
4
4
  analytic α-derivatives. JIT-compilable, ``vmap``-able, fully differentiable.
5
5
 
6
+ Two backends, selected by the keyword ``eos`` every public function takes (a
7
+ Python string, static under ``jax.jit``: one compiled kernel per backend)::
8
+
9
+ eos="span_wagner" Span & Wagner (1996), the default, bit for bit 0.5.0
10
+ eos="nist2026" Lemmon, Harvey & Hellmann, NIST IR 8608 (2026):
11
+ 24 terms, no nonanalytic terms, free of Span-Wagner's
12
+ near-critical oscillations in c_p, κ_T, (∂c_v/∂T)_ρ
13
+
14
+ ``EOS_NAMES`` lists them. The transport correlations keep their own
15
+ parameters on either backend; their critical parts read the selected
16
+ backend's thermodynamics.
17
+
6
18
  Public API
7
19
  ==========
8
20
 
9
- Simulation hot path — full state from the conserved variables (ρ, u). These are
21
+ Simulation hot path: full state from the conserved variables (ρ, u). These are
10
22
  **batched / array-native**: pass 1-D arrays, get back a dict of 1-D arrays. A
11
23
  single fused pass solves T once and reuses the α-derivatives for every property::
12
24
 
13
25
  properties_from_rho_u(rho, u, phase_hint=SUPERCRITICAL) -> dict
14
26
  temperature_from_rho_u(rho, u, phase_hint=SUPERCRITICAL) -> T # T only
15
27
 
16
- State functions — **scalar** in their inputs (0-d arrays or Python floats); wrap
28
+ Full state from (ρ, T), for a code that evolves temperature: also batched, no
29
+ inversion (a direct evaluation of the Helmholtz surface)::
30
+
31
+ properties_from_rho_T(rho, T) -> dict
32
+
33
+ State functions: **scalar** in their inputs (0-d arrays or Python floats); wrap
17
34
  with ``jax.vmap`` for batches::
18
35
 
19
- properties(T, rho) — full state from (T, ρ)
20
- state_from_PT(P, T) — full state from (P, T)
21
- state_from_Ph(P, h) — full state from (P, h)
22
- state_from_Du(rho, u) — full state from (ρ, u) (scalar form)
23
- density_from_PT(P, T) — ρ from (P, T)
36
+ properties(T, rho) full state from (T, ρ)
37
+ state_from_PT(P, T) full state from (P, T)
38
+ state_from_Ph(P, h) full state from (P, h)
39
+ state_from_Du(rho, u) full state from (ρ, u) (scalar form)
40
+ density_from_PT(P, T) ρ from (P, T)
24
41
 
25
- Each full-state dict has keys: ``temperature``, ``density``, ``pressure``,
26
- ``cv``, ``cp``, ``speed_of_sound``, ``enthalpy``, ``internal_energy``,
27
- ``entropy``, ``gibbs_energy``, ``viscosity``, ``thermal_conductivity``.
42
+ Each full-state dict has keys: ``temperature`` [K], ``density`` [kg/m³],
43
+ ``pressure`` [Pa], ``cv`` and ``cp`` [J/(kg·K)], ``speed_of_sound`` [m/s],
44
+ ``enthalpy``, ``internal_energy`` and ``gibbs_energy`` [J/kg], ``entropy``
45
+ [J/(kg·K)], ``dP_dT_rho`` ((∂P/∂T)_ρ [Pa/K]), ``dP_drho_T`` ((∂P/∂ρ)_T
46
+ [Pa·m³/kg]), ``viscosity`` [Pa·s], ``thermal_conductivity`` [W/(m·K)],
47
+ ``bulk_viscosity`` [Pa·s] (ζ = ζ_vib + ζ_crit) and
48
+ ``bulk_viscosity_vibrational`` [Pa·s] (its vibrational part ζ_vib alone).
28
49
 
29
50
  Saturation curve (cubic-spline interpolation of a precomputed Maxwell table)::
30
51
 
@@ -35,8 +56,11 @@ Saturation curve (cubic-spline interpolation of a precomputed Maxwell table)::
35
56
 
36
57
  Transport (scalar; vmap for batches)::
37
58
 
38
- viscosity(T, rho) → μ [Pa·s]
59
+ viscosity(T, rho) → μ [Pa·s] (no EOS input)
39
60
  thermal_conductivity(T, rho) → λ [W/(m·K)]
61
+ bulk_viscosity(T, rho) → ζ [Pa·s] (vibrational + critical)
62
+
63
+ Every function above also takes ``eos=`` (default ``"span_wagner"``).
40
64
 
41
65
  Phase hints::
42
66
 
@@ -46,13 +70,15 @@ See README "Migrating to v0.2" for the move from the old per-quantity calls to
46
70
  the fused ``properties_from_rho_u``.
47
71
  """
48
72
 
73
+ from functools import partial
74
+
49
75
  import jax
50
76
  import jax.numpy as jnp
51
77
 
52
78
  # Enable float64 for the entire library (also done by submodules).
53
79
  jax.config.update("jax_enable_x64", True)
54
80
 
55
- __version__ = "0.3.0"
81
+ __version__ = "0.6.0"
56
82
 
57
83
  # ── Phase hint constants ────────────────────────────────────────────────────
58
84
  LIQUID = 0
@@ -61,12 +87,16 @@ SUPERCRITICAL = 2
61
87
  AUTO = SUPERCRITICAL # alias: auto-pick in the dome / supercritical-style guess
62
88
 
63
89
  # ── Internal modules ────────────────────────────────────────────────────────
90
+ from co2_eos.backends import EOS_NAMES, DEFAULT_EOS, SPAN_WAGNER, NIST2026
64
91
  from co2_eos import span_wagner as _sw
65
92
  from co2_eos import transport as _tr
66
93
  from co2_eos import saturation as _sat
67
94
  from co2_eos import inversions as _inv
68
95
  from co2_eos import core as _core
69
96
 
97
+ # jit with the backend name static: one compiled kernel per backend.
98
+ _jit = partial(jax.jit, static_argnames=("eos",))
99
+
70
100
  # ── Re-exports: saturation curve ────────────────────────────────────────────
71
101
  saturation_pressure = _sat.saturation_pressure
72
102
  saturation_temperature = _sat.saturation_temperature
@@ -80,28 +110,47 @@ saturation_entropies_P = _sat.saturation_entropies_P
80
110
 
81
111
  # ── Transport (scalar form; vmap for batches) ───────────────────────────────
82
112
 
83
- @jax.jit
84
- def viscosity(T, rho):
85
- """Dynamic viscosity μ [Pa·s] at scalar (T, ρ). vmap for batches."""
113
+ @_jit
114
+ def viscosity(T, rho, eos=DEFAULT_EOS):
115
+ """Dynamic viscosity μ [Pa·s] at scalar (T, ρ). vmap for batches.
116
+
117
+ Laesecke-Muzny reads no EOS property; ``eos`` is accepted (and checked)
118
+ so every property function takes the same keyword.
119
+ """
86
120
  T = jnp.asarray(T, dtype=jnp.float64)
87
121
  rho = jnp.asarray(rho, dtype=jnp.float64)
88
- return _tr._scalar_viscosity(T, rho)
122
+ return _tr.for_eos(eos)._scalar_viscosity(T, rho)
89
123
 
90
124
 
91
- @jax.jit
92
- def thermal_conductivity(T, rho):
125
+ @_jit
126
+ def thermal_conductivity(T, rho, eos=DEFAULT_EOS):
93
127
  """Thermal conductivity λ [W/(m·K)] at scalar (T, ρ). vmap for batches."""
94
128
  T = jnp.asarray(T, dtype=jnp.float64)
95
129
  rho = jnp.asarray(rho, dtype=jnp.float64)
96
- return _tr._scalar_thermal_conductivity(T, rho)
130
+ return _tr.for_eos(eos)._scalar_thermal_conductivity(T, rho)
131
+
132
+
133
+ @_jit
134
+ def bulk_viscosity(T, rho, eos=DEFAULT_EOS):
135
+ """Bulk viscosity ζ [Pa·s] at scalar (T, ρ). vmap for batches.
136
+
137
+ ζ = ζ_vib + ζ_crit, the zero-frequency coefficient: vibrational relaxation
138
+ (Lin & Trusler 2014, pure CO₂) plus the critical part (Onuki 2007,
139
+ R_B = 0.03) on the Olchowy-Sengers correlation length the conductivity
140
+ enhancement uses. Valid while ω·τ₁₁ ≪ 1, single phase, 270 to 500 K.
141
+ """
142
+ T = jnp.asarray(T, dtype=jnp.float64)
143
+ rho = jnp.asarray(rho, dtype=jnp.float64)
144
+ return _tr.for_eos(eos)._scalar_bulk_viscosity(T, rho)
97
145
 
98
146
 
99
147
  # ═══════════════════════════════════════════════════════════════════════════
100
- # Hot path: full state from conserved (ρ, u) — batched / array-native
148
+ # Hot path: full state from conserved (ρ, u), batched / array-native
101
149
  # ═══════════════════════════════════════════════════════════════════════════
102
150
 
103
- @jax.jit
104
- def properties_from_rho_u(rho, u, phase_hint=SUPERCRITICAL):
151
+ @_jit
152
+ def properties_from_rho_u(rho, u, phase_hint=SUPERCRITICAL,
153
+ eos=DEFAULT_EOS):
105
154
  """Full thermodynamic + transport state from (ρ, u). Batched.
106
155
 
107
156
  The simulation hot path: solves T once from u(T,ρ)=u via the analytic-Cv
@@ -126,40 +175,71 @@ def properties_from_rho_u(rho, u, phase_hint=SUPERCRITICAL):
126
175
  u = jnp.asarray(u, dtype=jnp.float64)
127
176
  phase = jnp.broadcast_to(jnp.asarray(phase_hint, dtype=jnp.int32),
128
177
  jnp.shape(rho))
129
- return jax.vmap(_core._state_from_rho_u)(rho, u, phase)
178
+ return jax.vmap(_core.for_eos(eos)._state_from_rho_u)(rho, u, phase)
179
+
180
+
181
+ @_jit
182
+ def properties_from_rho_T(rho, T, eos=DEFAULT_EOS):
183
+ """Full thermodynamic + transport state from (ρ, T). Batched.
184
+
185
+ The array-native counterpart of the scalar ``properties(T, rho)`` (note the
186
+ argument order: ρ first, matching ``properties_from_rho_u``). A direct
187
+ evaluation of the Helmholtz surface: no inversion, so no phase hint. One
188
+ analytic α-derivative pass per point feeds every property, including
189
+ ``dP_dT_rho`` and ``dP_drho_T``. ``rho`` and ``T`` are 1-D arrays of
190
+ equal length; the result is a dict of 1-D float64 arrays.
191
+
192
+ This is ``jax.vmap`` of the same scalar kernel ``properties`` uses, so the
193
+ two agree bit for bit on Span-Wagner (on nist2026, whose kernel sums its
194
+ 24 terms as one vector reduction, to round-off); it exists as the
195
+ documented batched entry point.
196
+
197
+ Note: inside the two-phase dome (subcritical T, ρ between ρ_v and ρ_l) the
198
+ EOS returns its unstable single-phase extrapolation (there
199
+ ``dP_drho_T`` can be negative); the caller keeps its states out of the dome.
200
+ """
201
+ rho = jnp.asarray(rho, dtype=jnp.float64)
202
+ T = jnp.asarray(T, dtype=jnp.float64)
203
+ if rho.ndim != 1 or T.shape != rho.shape:
204
+ raise ValueError(
205
+ f"properties_from_rho_T expects 1-D rho and T of equal shape; "
206
+ f"got rho {rho.shape}, T {T.shape}")
207
+ return jax.vmap(_core.for_eos(eos)._state_from_T_rho)(T, rho)
130
208
 
131
209
 
132
- @jax.jit
133
- def temperature_from_rho_u(rho, u, phase_hint=SUPERCRITICAL):
210
+ @_jit
211
+ def temperature_from_rho_u(rho, u, phase_hint=SUPERCRITICAL,
212
+ eos=DEFAULT_EOS):
134
213
  """Temperature T [K] such that u(T,ρ)=u at fixed ρ. Batched.
135
214
 
136
- The lean inversion (no property derivation) — use when only T is needed.
215
+ The lean inversion (no property derivation); use when only T is needed.
137
216
  Differentiable: correct jvp/vjp via the implicit function theorem.
138
217
  """
139
218
  rho = jnp.asarray(rho, dtype=jnp.float64)
140
219
  u = jnp.asarray(u, dtype=jnp.float64)
141
220
  phase = jnp.broadcast_to(jnp.asarray(phase_hint, dtype=jnp.int32),
142
221
  jnp.shape(rho))
143
- return jax.vmap(_core._temperature_from_Du)(rho, u, phase)
222
+ return jax.vmap(_core.for_eos(eos)._temperature_from_Du)(rho, u,
223
+ phase)
144
224
 
145
225
 
146
226
  # ═══════════════════════════════════════════════════════════════════════════
147
- # State functions — scalar in their inputs (vmap for batches)
227
+ # State functions: scalar in their inputs (vmap for batches)
148
228
  # ═══════════════════════════════════════════════════════════════════════════
149
229
 
150
- @jax.jit
151
- def properties(T, rho):
230
+ @_jit
231
+ def properties(T, rho, eos=DEFAULT_EOS):
152
232
  """Full thermodynamic + transport state at scalar (T, ρ).
153
233
 
154
234
  Returns a dict (see module docstring for keys). vmap for batches.
155
235
  """
156
236
  T = jnp.asarray(T, dtype=jnp.float64)
157
237
  rho = jnp.asarray(rho, dtype=jnp.float64)
158
- return _core._state_from_T_rho(T, rho)
238
+ return _core.for_eos(eos)._state_from_T_rho(T, rho)
159
239
 
160
240
 
161
- @jax.jit
162
- def state_from_PT(P, T, phase_hint=AUTO):
241
+ @_jit
242
+ def state_from_PT(P, T, phase_hint=AUTO, eos=DEFAULT_EOS):
163
243
  """Full state at (P, T). Phase-aware Halley solve for ρ, then derive all.
164
244
 
165
245
  ``pressure`` in the returned dict echoes the input ``P`` exactly.
@@ -167,23 +247,23 @@ def state_from_PT(P, T, phase_hint=AUTO):
167
247
  P = jnp.asarray(P, dtype=jnp.float64)
168
248
  T = jnp.asarray(T, dtype=jnp.float64)
169
249
  phase = jnp.asarray(phase_hint, dtype=jnp.int32)
170
- rho = _inv._density_from_PT(T, P, phase)
171
- state = _core._state_from_T_rho(T, rho)
250
+ rho = _inv.for_eos(eos)._density_from_PT(T, P, phase)
251
+ state = _core.for_eos(eos)._state_from_T_rho(T, rho)
172
252
  state["pressure"] = P
173
253
  return state
174
254
 
175
255
 
176
- @jax.jit
177
- def density_from_PT(P, T, phase_hint=AUTO):
256
+ @_jit
257
+ def density_from_PT(P, T, phase_hint=AUTO, eos=DEFAULT_EOS):
178
258
  """Density [kg/m³] at scalar (P, T). Phase-aware Halley solve."""
179
259
  P = jnp.asarray(P, dtype=jnp.float64)
180
260
  T = jnp.asarray(T, dtype=jnp.float64)
181
261
  phase = jnp.asarray(phase_hint, dtype=jnp.int32)
182
- return _inv._density_from_PT(T, P, phase)
262
+ return _inv.for_eos(eos)._density_from_PT(T, P, phase)
183
263
 
184
264
 
185
- @jax.jit
186
- def state_from_Ph(P, h, phase_hint=AUTO):
265
+ @_jit
266
+ def state_from_Ph(P, h, phase_hint=AUTO, eos=DEFAULT_EOS):
187
267
  """Full state at (P, h). Dome-aware at subcritical P; else 2-D Newton.
188
268
 
189
269
  ``pressure`` and ``enthalpy`` echo the inputs.
@@ -191,15 +271,15 @@ def state_from_Ph(P, h, phase_hint=AUTO):
191
271
  P = jnp.asarray(P, dtype=jnp.float64)
192
272
  h = jnp.asarray(h, dtype=jnp.float64)
193
273
  phase = jnp.asarray(phase_hint, dtype=jnp.int32)
194
- T, rho = _inv._state_from_Ph(P, h, phase)
195
- state = _core._state_from_T_rho(T, rho)
274
+ T, rho = _inv.for_eos(eos)._state_from_Ph(P, h, phase)
275
+ state = _core.for_eos(eos)._state_from_T_rho(T, rho)
196
276
  state["pressure"] = P
197
277
  state["enthalpy"] = h
198
278
  return state
199
279
 
200
280
 
201
- @jax.jit
202
- def state_from_Du(rho, u, phase_hint=AUTO):
281
+ @_jit
282
+ def state_from_Du(rho, u, phase_hint=AUTO, eos=DEFAULT_EOS):
203
283
  """Full state at scalar (ρ, u). Scalar form of ``properties_from_rho_u``.
204
284
 
205
285
  ``density`` and ``internal_energy`` echo the inputs.
@@ -207,13 +287,14 @@ def state_from_Du(rho, u, phase_hint=AUTO):
207
287
  rho = jnp.asarray(rho, dtype=jnp.float64)
208
288
  u = jnp.asarray(u, dtype=jnp.float64)
209
289
  phase = jnp.asarray(phase_hint, dtype=jnp.int32)
210
- return _core._state_from_rho_u(rho, u, phase)
290
+ return _core.for_eos(eos)._state_from_rho_u(rho, u, phase)
211
291
 
212
292
 
213
293
  __all__ = [
214
294
  "__version__",
295
+ "EOS_NAMES", "DEFAULT_EOS", "SPAN_WAGNER", "NIST2026",
215
296
  "LIQUID", "VAPOR", "SUPERCRITICAL", "AUTO",
216
- "properties_from_rho_u", "temperature_from_rho_u",
297
+ "properties_from_rho_u", "properties_from_rho_T", "temperature_from_rho_u",
217
298
  "properties",
218
299
  "state_from_PT", "state_from_Ph", "state_from_Du",
219
300
  "density_from_PT",
@@ -221,5 +302,5 @@ __all__ = [
221
302
  "saturation_densities", "saturation_densities_P",
222
303
  "saturation_enthalpies", "saturation_enthalpies_P",
223
304
  "saturation_entropies", "saturation_entropies_P",
224
- "viscosity", "thermal_conductivity",
305
+ "viscosity", "thermal_conductivity", "bulk_viscosity",
225
306
  ]
@@ -0,0 +1,56 @@
1
+ """The equation-of-state backends and the static switch that selects one.
2
+
3
+ Every public function of co2-eos takes a keyword ``eos``, a Python string that
4
+ is static under ``jax.jit`` (one compiled kernel per backend; no traced
5
+ branching), with the default ``DEFAULT_EOS``:
6
+
7
+ * ``"span_wagner"``: Span & Wagner (1996), the default. Its surface is
8
+ ``span_wagner`` (constants, α, the autodiff property functions) and its
9
+ analytic kernel ``helmholtz``; its tables carry no suffix.
10
+ * ``"nist2026"``: Lemmon, Harvey & Hellmann, NIST IR 8608 (2026), 24 terms
11
+ with no nonanalytic terms. Surface and kernel are both ``nist2026``; its
12
+ tables carry the suffix ``_nist2026``.
13
+
14
+ A backend's surface module provides ``TC, RHOC, RHOC_MOLAR, R, R_MOLAR, M,
15
+ PC, T_TRIPLE`` and ``_scalar_pressure``, ``_scalar_enthalpy``,
16
+ ``_scalar_internal_energy``, ``_scalar_entropy``; its kernel module provides
17
+ ``ideal_derivs``, ``ideal_tau_only``, ``residual_derivs``,
18
+ ``residual_tau_prep``, ``residual_tau_fast``, ``residual_tau_derivs``. Each
19
+ backend carries its own gas constant and molar mass; nothing shared between
20
+ backends assumes Span-Wagner's.
21
+ """
22
+
23
+ from pathlib import Path
24
+ from types import SimpleNamespace
25
+
26
+ from co2_eos import helmholtz as _helmholtz
27
+ from co2_eos import nist2026 as _nist2026
28
+ from co2_eos import span_wagner as _span_wagner
29
+
30
+ SPAN_WAGNER = "span_wagner"
31
+ NIST2026 = "nist2026"
32
+ EOS_NAMES = (SPAN_WAGNER, NIST2026)
33
+ DEFAULT_EOS = SPAN_WAGNER
34
+
35
+ DATA_DIR = Path(__file__).parent / "data"
36
+
37
+ _BACKENDS = {
38
+ SPAN_WAGNER: SimpleNamespace(name=SPAN_WAGNER, surface=_span_wagner,
39
+ kernel=_helmholtz, suffix=""),
40
+ NIST2026: SimpleNamespace(name=NIST2026, surface=_nist2026,
41
+ kernel=_nist2026, suffix="_nist2026"),
42
+ }
43
+
44
+
45
+ def get(eos):
46
+ """The backend record for ``eos``; raises on an unknown name."""
47
+ try:
48
+ return _BACKENDS[eos]
49
+ except (KeyError, TypeError):
50
+ raise ValueError(
51
+ f"unknown eos {eos!r}; choose one of {EOS_NAMES}") from None
52
+
53
+
54
+ def table_path(eos, stem):
55
+ """Path of the backend's table ``<stem><suffix>.npz`` in ``co2_eos/data``."""
56
+ return DATA_DIR / f"{stem}{get(eos).suffix}.npz"