co2-eos 0.1.0__tar.gz → 0.2.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 (26) hide show
  1. {co2_eos-0.1.0 → co2_eos-0.2.0}/PKG-INFO +67 -19
  2. {co2_eos-0.1.0 → co2_eos-0.2.0}/README.md +62 -14
  3. co2_eos-0.2.0/co2_eos/__init__.py +225 -0
  4. co2_eos-0.2.0/co2_eos/core.py +245 -0
  5. co2_eos-0.2.0/co2_eos/data/seed_table.npz +0 -0
  6. co2_eos-0.2.0/co2_eos/helmholtz.py +337 -0
  7. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/transport.py +37 -37
  8. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/PKG-INFO +67 -19
  9. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/SOURCES.txt +6 -1
  10. {co2_eos-0.1.0 → co2_eos-0.2.0}/pyproject.toml +5 -5
  11. co2_eos-0.2.0/tests/test_analytic_derivs.py +87 -0
  12. co2_eos-0.2.0/tests/test_rho_u_inversion.py +101 -0
  13. co2_eos-0.1.0/co2_eos/__init__.py +0 -265
  14. {co2_eos-0.1.0 → co2_eos-0.2.0}/LICENSE +0 -0
  15. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/data/saturation_table.npz +0 -0
  16. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/inversions.py +0 -0
  17. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/saturation.py +0 -0
  18. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/span_wagner.py +0 -0
  19. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/dependency_links.txt +0 -0
  20. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/requires.txt +0 -0
  21. {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/top_level.txt +0 -0
  22. {co2_eos-0.1.0 → co2_eos-0.2.0}/setup.cfg +0 -0
  23. {co2_eos-0.1.0 → co2_eos-0.2.0}/tests/test_accuracy_vs_coolprop.py +0 -0
  24. {co2_eos-0.1.0 → co2_eos-0.2.0}/tests/test_dome_detection.py +0 -0
  25. {co2_eos-0.1.0 → co2_eos-0.2.0}/tests/test_inversion_gradients.py +0 -0
  26. {co2_eos-0.1.0 → co2_eos-0.2.0}/tests/test_no_coolprop_runtime.py +0 -0
@@ -1,13 +1,13 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: co2-eos
3
- Version: 0.1.0
3
+ Version: 0.2.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
7
- Project-URL: Homepage, https://github.com/John-FluxTech/co2-eos
8
- Project-URL: Repository, https://github.com/John-FluxTech/co2-eos
9
- Project-URL: Documentation, https://github.com/John-FluxTech/co2-eos#readme
10
- Project-URL: Issues, https://github.com/John-FluxTech/co2-eos/issues
7
+ Project-URL: Homepage, https://github.com/FluxTech-UG/co2-eos
8
+ Project-URL: Repository, https://github.com/FluxTech-UG/co2-eos
9
+ Project-URL: Documentation, https://github.com/FluxTech-UG/co2-eos#readme
10
+ Project-URL: Issues, https://github.com/FluxTech-UG/co2-eos/issues
11
11
  Keywords: CO2,equation-of-state,thermodynamics,JAX,autodiff,Span-Wagner
12
12
  Classifier: Development Status :: 4 - Beta
13
13
  Classifier: Intended Audience :: Science/Research
@@ -34,18 +34,27 @@ Dynamic: license-file
34
34
 
35
35
  # CO2-EOS
36
36
 
37
- [![tests](https://github.com/John-FluxTech/co2-eos/actions/workflows/test.yml/badge.svg)](https://github.com/John-FluxTech/co2-eos/actions/workflows/test.yml)
38
- [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/John-FluxTech/co2-eos/blob/main/examples/launch_demo.ipynb)
37
+ [![tests](https://github.com/FluxTech-UG/co2-eos/actions/workflows/test.yml/badge.svg)](https://github.com/FluxTech-UG/co2-eos/actions/workflows/test.yml)
38
+ [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/FluxTech-UG/co2-eos/blob/main/examples/launch_demo.ipynb)
39
39
 
40
40
  Differentiable CO₂ thermodynamic properties in JAX.
41
41
 
42
- CO2-EOS is a pure-JAX implementation of the Span-Wagner equation of state for carbon dioxide. Every property evaluation is JIT-compiled, vectorisable with `vmap`, and fully differentiable with `jax.grad`, enabling gradient-based optimisation of any system that depends on CO₂ thermodynamics.
42
+ CO2-EOS is a pure-JAX implementation of the Span-Wagner equation of state for carbon dioxide, with hand-coded analytic α-derivatives. Every property evaluation is JIT-compiled, vectorisable with `vmap`, and fully differentiable with `jax.grad`, enabling gradient-based optimisation of any system that depends on CO₂ thermodynamics.
43
43
 
44
44
  ```python
45
45
  import co2_eos as co2
46
- import jax
47
-
48
- # Properties from (T, ρ)
46
+ import jax, jax.numpy as jnp
47
+
48
+ # Simulation hot path: full state from the conserved variables (ρ, u).
49
+ # Batched/array-native — solves T once and reuses the α-derivatives for
50
+ # every property. This is the call a finite-volume code makes each RHS eval.
51
+ rho = jnp.array([300.0, 350.0, 420.0]) # kg/m³
52
+ u = jnp.array([3.62e5, 3.58e5, 3.42e5]) # J/kg (specific internal energy)
53
+ state = co2.properties_from_rho_u(rho, u) # dict of arrays: T, P, cp, w, μ, k, …
54
+ print(state['temperature'], state['pressure'], state['speed_of_sound'])
55
+ # -> T ≈ [307.5, 315.4, 315.6] K, P ≈ [7.6, 8.8, 9.1] MPa (supercritical)
56
+
57
+ # Properties from (T, ρ) (scalar; vmap for batches)
49
58
  state = co2.properties(T=310.0, rho=350.0)
50
59
  print(state['pressure'], state['cp'], state['speed_of_sound'])
51
60
 
@@ -56,15 +65,14 @@ state = co2.state_from_PT(P=8e6, T=310.0)
56
65
  # (peaks near the Widom line at these conditions)
57
66
  drho_dT = jax.grad(lambda T: co2.state_from_PT(P=8e6, T=T)['density'])(310.0)
58
67
 
59
- # Vectorise across conditions
60
- import jax.numpy as jnp
68
+ # Vectorise the scalar functions across conditions
61
69
  T_array = jnp.linspace(280, 340, 100)
62
70
  states = jax.vmap(lambda T: co2.state_from_PT(P=8e6, T=T))(T_array)
63
71
  ```
64
72
 
65
73
  ## Try it without cloning
66
74
 
67
- [`examples/launch_demo.ipynb`](examples/launch_demo.ipynb) is the launch demo: validation against CoolProp, the CPU and GPU benchmarks cited below, and a worked example of gradient-based optimisation through `state_from_PT`. [Read it on GitHub](https://github.com/John-FluxTech/co2-eos/blob/main/examples/launch_demo.ipynb) for a static render, or [open it on Colab](https://colab.research.google.com/github/John-FluxTech/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.
75
+ [`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.
68
76
 
69
77
  ## Why this exists
70
78
 
@@ -74,14 +82,15 @@ But CoolProp is a C++ library with Python bindings. You can't `jax.grad` through
74
82
 
75
83
  CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagner 1996) directly in JAX:
76
84
 
77
- - **Differentiable.** `jax.grad` through any property, any inversion, any combination. No finite differences. Exact gradients via autodiff, including second and higher derivatives for free.
78
- - **Fast.** JIT-compiled and vectorisable. Forward state evaluation at fixed (T, ρ) runs in roughly 1.8 μs/point under `jax.vmap`. The (P, T) 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 (measured at a 10,000-point batch on an Apple M2 Pro; see `examples/launch_demo.ipynb`, section 2). On a Colab T4 GPU the same `state_from_PT` workflow 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. The ceiling is set by the iterative density solve: a `jax.while_loop` with data-dependent control flow forces lockstep execution across GPU warps, so the polynomial evaluation parallelises cleanly but is not the dominant cost. A fully feed-forward kernel on the same hardware would land closer to 50 to 100×.
85
+ - **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 analytic α-derivatives instead of a nested-`jax.grad` chain. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery — it solves `T(ρ, u)` once (table-seeded, fixed-iteration analytic-Cv Newton) and reuses the derivatives for every property. On an Apple M-class CPU at a 4096-point float64 batch this is **4.8× faster** for the `T(ρ, u)` inversion and **3.6× faster** for the full `T + P + μ + k` recovery than the v0.1 autodiff path. On an OVH Tesla V100S the same comparison is 1.3–2.2× at these (latency-bound) batch sizes and grows to **7.7× / 4.9×** once the GPU is compute-bound at large batches (`bench/compare.py`, `bench/PROFILING.md`).
87
+ - **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.
79
88
  - **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.
80
89
  - **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.
81
90
 
82
91
  ## What's included
83
92
 
84
- **Equation of state**: Full Span-Wagner (1996) formulation for CO₂. Helmholtz free energy as A(T, ρ), with all thermodynamic properties derived by autodiff:
93
+ **Equation of state**: Full Span-Wagner (1996) formulation for CO₂. Helmholtz free energy as A(T, ρ), with all thermodynamic properties derived from hand-coded analytic α-derivatives (validated against `jax.grad` to < 3e-12):
85
94
 
86
95
  - Pressure, internal energy, enthalpy, entropy
87
96
  - Isochoric and isobaric heat capacities (Cv, Cp)
@@ -100,9 +109,48 @@ CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagne
100
109
 
101
110
  **State inversions:**
102
111
 
112
+ - `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
+ - `state_from_Du(rho, u)`: scalar form of the above (vmap for batches).
103
114
  - `state_from_PT(P, T)`: Halley's method with pressure-aware critical-region initial guess, step damping, and bisection safety net
104
115
  - `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
105
- - `state_from_Du(rho, u)`: for simulation codes carrying (ρ, u) as conserved variables. Uses `custom_vjp` for clean gradient propagation through the implicit solve
116
+
117
+ ## Migrating to v0.2
118
+
119
+ v0.2 is a performance redesign. **Accuracy and gradients are unchanged** and the
120
+ scalar state functions keep their signatures, so most code needs no changes. The
121
+ one thing worth adopting is the fused hot path.
122
+
123
+ If your simulation carried `(ρ, u)` and recovered primitives with separate calls:
124
+
125
+ ```python
126
+ # v0.1 — separate calls, each re-evaluating α-derivatives
127
+ from co2_eos.inversions import temperature_from_Du, SUPERCRITICAL
128
+ from co2_eos import span_wagner as sw, transport
129
+ T = temperature_from_Du(rho, u, phase) # nested-autodiff Cv, while_loop
130
+ P = sw.pressure(T, rho)
131
+ mu = transport.viscosity(T, rho)
132
+ k = transport.thermal_conductivity(T, rho)
133
+ ```
134
+
135
+ replace it with the single fused call:
136
+
137
+ ```python
138
+ # v0.2 — one fused pass: solve T once, reuse the α-derivatives
139
+ import co2_eos as co2
140
+ s = co2.properties_from_rho_u(rho, u) # phase_hint defaults to SUPERCRITICAL
141
+ T, P, mu, k = s["temperature"], s["pressure"], s["viscosity"], s["thermal_conductivity"]
142
+ # also: s["cp"], s["speed_of_sound"], s["enthalpy"], s["entropy"], s["gibbs_energy"], …
143
+ ```
144
+
145
+ Notes:
146
+
147
+ - `properties_from_rho_u` / `temperature_from_rho_u` are **batched** (pass 1-D
148
+ arrays). The scalar `state_from_Du(rho, u)` returns the same dict for a single
149
+ point.
150
+ - `phase_hint` is accepted for API symmetry but unused by the `(ρ, u)` inversion
151
+ (`u` is monotone in `T` at fixed `ρ`, so no phase branch is needed).
152
+ - The old modules `co2_eos.span_wagner` and `co2_eos.inversions` still work; they
153
+ are the validation ground truth and the benchmark baseline.
106
154
 
107
155
  ## Valid range
108
156
 
@@ -125,7 +173,7 @@ Requires Python ≥ 3.10 and JAX ≥ 0.4. No other dependencies.
125
173
  To run the validation tests (which compare against CoolProp):
126
174
 
127
175
  ```bash
128
- git clone https://github.com/John-FluxTech/co2-eos.git
176
+ git clone https://github.com/FluxTech-UG/co2-eos.git
129
177
  cd co2-eos
130
178
  pip install -e ".[test]"
131
179
  pytest tests/ -v
@@ -1,17 +1,26 @@
1
1
  # CO2-EOS
2
2
 
3
- [![tests](https://github.com/John-FluxTech/co2-eos/actions/workflows/test.yml/badge.svg)](https://github.com/John-FluxTech/co2-eos/actions/workflows/test.yml)
4
- [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/John-FluxTech/co2-eos/blob/main/examples/launch_demo.ipynb)
3
+ [![tests](https://github.com/FluxTech-UG/co2-eos/actions/workflows/test.yml/badge.svg)](https://github.com/FluxTech-UG/co2-eos/actions/workflows/test.yml)
4
+ [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/FluxTech-UG/co2-eos/blob/main/examples/launch_demo.ipynb)
5
5
 
6
6
  Differentiable CO₂ thermodynamic properties in JAX.
7
7
 
8
- CO2-EOS is a pure-JAX implementation of the Span-Wagner equation of state for carbon dioxide. Every property evaluation is JIT-compiled, vectorisable with `vmap`, and fully differentiable with `jax.grad`, enabling gradient-based optimisation of any system that depends on CO₂ thermodynamics.
8
+ CO2-EOS is a pure-JAX implementation of the Span-Wagner equation of state for carbon dioxide, with hand-coded analytic α-derivatives. Every property evaluation is JIT-compiled, vectorisable with `vmap`, and fully differentiable with `jax.grad`, enabling gradient-based optimisation of any system that depends on CO₂ thermodynamics.
9
9
 
10
10
  ```python
11
11
  import co2_eos as co2
12
- import jax
13
-
14
- # Properties from (T, ρ)
12
+ import jax, jax.numpy as jnp
13
+
14
+ # Simulation hot path: full state from the conserved variables (ρ, u).
15
+ # Batched/array-native — solves T once and reuses the α-derivatives for
16
+ # every property. This is the call a finite-volume code makes each RHS eval.
17
+ rho = jnp.array([300.0, 350.0, 420.0]) # kg/m³
18
+ u = jnp.array([3.62e5, 3.58e5, 3.42e5]) # J/kg (specific internal energy)
19
+ state = co2.properties_from_rho_u(rho, u) # dict of arrays: T, P, cp, w, μ, k, …
20
+ print(state['temperature'], state['pressure'], state['speed_of_sound'])
21
+ # -> T ≈ [307.5, 315.4, 315.6] K, P ≈ [7.6, 8.8, 9.1] MPa (supercritical)
22
+
23
+ # Properties from (T, ρ) (scalar; vmap for batches)
15
24
  state = co2.properties(T=310.0, rho=350.0)
16
25
  print(state['pressure'], state['cp'], state['speed_of_sound'])
17
26
 
@@ -22,15 +31,14 @@ state = co2.state_from_PT(P=8e6, T=310.0)
22
31
  # (peaks near the Widom line at these conditions)
23
32
  drho_dT = jax.grad(lambda T: co2.state_from_PT(P=8e6, T=T)['density'])(310.0)
24
33
 
25
- # Vectorise across conditions
26
- import jax.numpy as jnp
34
+ # Vectorise the scalar functions across conditions
27
35
  T_array = jnp.linspace(280, 340, 100)
28
36
  states = jax.vmap(lambda T: co2.state_from_PT(P=8e6, T=T))(T_array)
29
37
  ```
30
38
 
31
39
  ## Try it without cloning
32
40
 
33
- [`examples/launch_demo.ipynb`](examples/launch_demo.ipynb) is the launch demo: validation against CoolProp, the CPU and GPU benchmarks cited below, and a worked example of gradient-based optimisation through `state_from_PT`. [Read it on GitHub](https://github.com/John-FluxTech/co2-eos/blob/main/examples/launch_demo.ipynb) for a static render, or [open it on Colab](https://colab.research.google.com/github/John-FluxTech/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.
41
+ [`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.
34
42
 
35
43
  ## Why this exists
36
44
 
@@ -40,14 +48,15 @@ But CoolProp is a C++ library with Python bindings. You can't `jax.grad` through
40
48
 
41
49
  CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagner 1996) directly in JAX:
42
50
 
43
- - **Differentiable.** `jax.grad` through any property, any inversion, any combination. No finite differences. Exact gradients via autodiff, including second and higher derivatives for free.
44
- - **Fast.** JIT-compiled and vectorisable. Forward state evaluation at fixed (T, ρ) runs in roughly 1.8 μs/point under `jax.vmap`. The (P, T) 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 (measured at a 10,000-point batch on an Apple M2 Pro; see `examples/launch_demo.ipynb`, section 2). On a Colab T4 GPU the same `state_from_PT` workflow 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. The ceiling is set by the iterative density solve: a `jax.while_loop` with data-dependent control flow forces lockstep execution across GPU warps, so the polynomial evaluation parallelises cleanly but is not the dominant cost. A fully feed-forward kernel on the same hardware would land closer to 50 to 100×.
51
+ - **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 analytic α-derivatives instead of a nested-`jax.grad` chain. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery — it solves `T(ρ, u)` once (table-seeded, fixed-iteration analytic-Cv Newton) and reuses the derivatives for every property. On an Apple M-class CPU at a 4096-point float64 batch this is **4.8× faster** for the `T(ρ, u)` inversion and **3.6× faster** for the full `T + P + μ + k` recovery than the v0.1 autodiff path. On an OVH Tesla V100S the same comparison is 1.3–2.2× at these (latency-bound) batch sizes and grows to **7.7× / 4.9×** once the GPU is compute-bound at large batches (`bench/compare.py`, `bench/PROFILING.md`).
53
+ - **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.
45
54
  - **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.
46
55
  - **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.
47
56
 
48
57
  ## What's included
49
58
 
50
- **Equation of state**: Full Span-Wagner (1996) formulation for CO₂. Helmholtz free energy as A(T, ρ), with all thermodynamic properties derived by autodiff:
59
+ **Equation of state**: Full Span-Wagner (1996) formulation for CO₂. Helmholtz free energy as A(T, ρ), with all thermodynamic properties derived from hand-coded analytic α-derivatives (validated against `jax.grad` to < 3e-12):
51
60
 
52
61
  - Pressure, internal energy, enthalpy, entropy
53
62
  - Isochoric and isobaric heat capacities (Cv, Cp)
@@ -66,9 +75,48 @@ CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagne
66
75
 
67
76
  **State inversions:**
68
77
 
78
+ - `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
+ - `state_from_Du(rho, u)`: scalar form of the above (vmap for batches).
69
80
  - `state_from_PT(P, T)`: Halley's method with pressure-aware critical-region initial guess, step damping, and bisection safety net
70
81
  - `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
71
- - `state_from_Du(rho, u)`: for simulation codes carrying (ρ, u) as conserved variables. Uses `custom_vjp` for clean gradient propagation through the implicit solve
82
+
83
+ ## Migrating to v0.2
84
+
85
+ v0.2 is a performance redesign. **Accuracy and gradients are unchanged** and the
86
+ scalar state functions keep their signatures, so most code needs no changes. The
87
+ one thing worth adopting is the fused hot path.
88
+
89
+ If your simulation carried `(ρ, u)` and recovered primitives with separate calls:
90
+
91
+ ```python
92
+ # v0.1 — separate calls, each re-evaluating α-derivatives
93
+ from co2_eos.inversions import temperature_from_Du, SUPERCRITICAL
94
+ from co2_eos import span_wagner as sw, transport
95
+ T = temperature_from_Du(rho, u, phase) # nested-autodiff Cv, while_loop
96
+ P = sw.pressure(T, rho)
97
+ mu = transport.viscosity(T, rho)
98
+ k = transport.thermal_conductivity(T, rho)
99
+ ```
100
+
101
+ replace it with the single fused call:
102
+
103
+ ```python
104
+ # v0.2 — one fused pass: solve T once, reuse the α-derivatives
105
+ import co2_eos as co2
106
+ s = co2.properties_from_rho_u(rho, u) # phase_hint defaults to SUPERCRITICAL
107
+ T, P, mu, k = s["temperature"], s["pressure"], s["viscosity"], s["thermal_conductivity"]
108
+ # also: s["cp"], s["speed_of_sound"], s["enthalpy"], s["entropy"], s["gibbs_energy"], …
109
+ ```
110
+
111
+ Notes:
112
+
113
+ - `properties_from_rho_u` / `temperature_from_rho_u` are **batched** (pass 1-D
114
+ arrays). The scalar `state_from_Du(rho, u)` returns the same dict for a single
115
+ point.
116
+ - `phase_hint` is accepted for API symmetry but unused by the `(ρ, u)` inversion
117
+ (`u` is monotone in `T` at fixed `ρ`, so no phase branch is needed).
118
+ - The old modules `co2_eos.span_wagner` and `co2_eos.inversions` still work; they
119
+ are the validation ground truth and the benchmark baseline.
72
120
 
73
121
  ## Valid range
74
122
 
@@ -91,7 +139,7 @@ Requires Python ≥ 3.10 and JAX ≥ 0.4. No other dependencies.
91
139
  To run the validation tests (which compare against CoolProp):
92
140
 
93
141
  ```bash
94
- git clone https://github.com/John-FluxTech/co2-eos.git
142
+ git clone https://github.com/FluxTech-UG/co2-eos.git
95
143
  cd co2-eos
96
144
  pip install -e ".[test]"
97
145
  pytest tests/ -v
@@ -0,0 +1,225 @@
1
+ """CO2-EOS: Differentiable CO₂ thermodynamic properties in JAX.
2
+
3
+ Pure-JAX Span-Wagner (1996) equation of state for carbon dioxide with hand-coded
4
+ analytic α-derivatives. JIT-compilable, ``vmap``-able, fully differentiable.
5
+
6
+ Public API
7
+ ==========
8
+
9
+ Simulation hot path — full state from the conserved variables (ρ, u). These are
10
+ **batched / array-native**: pass 1-D arrays, get back a dict of 1-D arrays. A
11
+ single fused pass solves T once and reuses the α-derivatives for every property::
12
+
13
+ properties_from_rho_u(rho, u, phase_hint=SUPERCRITICAL) -> dict
14
+ temperature_from_rho_u(rho, u, phase_hint=SUPERCRITICAL) -> T # T only
15
+
16
+ State functions — **scalar** in their inputs (0-d arrays or Python floats); wrap
17
+ with ``jax.vmap`` for batches::
18
+
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)
24
+
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``.
28
+
29
+ Saturation curve (cubic-spline interpolation of a precomputed Maxwell table)::
30
+
31
+ saturation_pressure(T) / saturation_temperature(P)
32
+ saturation_densities(T) / saturation_densities_P(P)
33
+ saturation_enthalpies(T) / saturation_enthalpies_P(P)
34
+ saturation_entropies(T) / saturation_entropies_P(P)
35
+
36
+ Transport (scalar; vmap for batches)::
37
+
38
+ viscosity(T, rho) → μ [Pa·s]
39
+ thermal_conductivity(T, rho) → λ [W/(m·K)]
40
+
41
+ Phase hints::
42
+
43
+ LIQUID = 0, VAPOR = 1, SUPERCRITICAL = 2 (alias: AUTO)
44
+
45
+ See README "Migrating to v0.2" for the move from the old per-quantity calls to
46
+ the fused ``properties_from_rho_u``.
47
+ """
48
+
49
+ import jax
50
+ import jax.numpy as jnp
51
+
52
+ # Enable float64 for the entire library (also done by submodules).
53
+ jax.config.update("jax_enable_x64", True)
54
+
55
+ __version__ = "0.2.0"
56
+
57
+ # ── Phase hint constants ────────────────────────────────────────────────────
58
+ LIQUID = 0
59
+ VAPOR = 1
60
+ SUPERCRITICAL = 2
61
+ AUTO = SUPERCRITICAL # alias: auto-pick in the dome / supercritical-style guess
62
+
63
+ # ── Internal modules ────────────────────────────────────────────────────────
64
+ from co2_eos import span_wagner as _sw
65
+ from co2_eos import transport as _tr
66
+ from co2_eos import saturation as _sat
67
+ from co2_eos import inversions as _inv
68
+ from co2_eos import core as _core
69
+
70
+ # ── Re-exports: saturation curve ────────────────────────────────────────────
71
+ saturation_pressure = _sat.saturation_pressure
72
+ saturation_temperature = _sat.saturation_temperature
73
+ saturation_densities = _sat.saturation_densities
74
+ saturation_densities_P = _sat.saturation_densities_P
75
+ saturation_enthalpies = _sat.saturation_enthalpies_T
76
+ saturation_enthalpies_P = _sat.saturation_enthalpies
77
+ saturation_entropies = _sat.saturation_entropies
78
+ saturation_entropies_P = _sat.saturation_entropies_P
79
+
80
+
81
+ # ── Transport (scalar form; vmap for batches) ───────────────────────────────
82
+
83
+ @jax.jit
84
+ def viscosity(T, rho):
85
+ """Dynamic viscosity μ [Pa·s] at scalar (T, ρ). vmap for batches."""
86
+ T = jnp.asarray(T, dtype=jnp.float64)
87
+ rho = jnp.asarray(rho, dtype=jnp.float64)
88
+ return _tr._scalar_viscosity(T, rho)
89
+
90
+
91
+ @jax.jit
92
+ def thermal_conductivity(T, rho):
93
+ """Thermal conductivity λ [W/(m·K)] at scalar (T, ρ). vmap for batches."""
94
+ T = jnp.asarray(T, dtype=jnp.float64)
95
+ rho = jnp.asarray(rho, dtype=jnp.float64)
96
+ return _tr._scalar_thermal_conductivity(T, rho)
97
+
98
+
99
+ # ═══════════════════════════════════════════════════════════════════════════
100
+ # Hot path: full state from conserved (ρ, u) — batched / array-native
101
+ # ═══════════════════════════════════════════════════════════════════════════
102
+
103
+ @jax.jit
104
+ def properties_from_rho_u(rho, u, phase_hint=SUPERCRITICAL):
105
+ """Full thermodynamic + transport state from (ρ, u). Batched.
106
+
107
+ The simulation hot path: solves T once from u(T,ρ)=u via the analytic-Cv
108
+ fixed-iteration Newton, then derives every property reusing the
109
+ α-derivatives. ``rho`` and ``u`` are 1-D arrays of equal length; the result
110
+ is a dict of 1-D arrays. ``density`` and ``internal_energy`` echo the
111
+ inputs exactly.
112
+
113
+ Args:
114
+ rho: Density [kg/m³], 1-D array.
115
+ u: Specific internal energy [J/kg], 1-D array.
116
+ phase_hint: accepted for API symmetry; the single-phase (ρ, u)
117
+ inversion is phase-agnostic (u is monotone in T at fixed ρ on a
118
+ single-phase branch).
119
+
120
+ Note: this is a single-phase EOS inversion. A (ρ, u) state inside the
121
+ two-phase dome (subcritical T, ρ between ρ_v and ρ_l) is not a single-phase
122
+ state; the EOS returns its unstable single-phase extrapolation there. The
123
+ near-critical supercritical operating regime has no dome.
124
+ """
125
+ rho = jnp.asarray(rho, dtype=jnp.float64)
126
+ u = jnp.asarray(u, dtype=jnp.float64)
127
+ phase = jnp.broadcast_to(jnp.asarray(phase_hint, dtype=jnp.int32),
128
+ jnp.shape(rho))
129
+ return jax.vmap(_core._state_from_rho_u)(rho, u, phase)
130
+
131
+
132
+ @jax.jit
133
+ def temperature_from_rho_u(rho, u, phase_hint=SUPERCRITICAL):
134
+ """Temperature T [K] such that u(T,ρ)=u at fixed ρ. Batched.
135
+
136
+ The lean inversion (no property derivation) — use when only T is needed.
137
+ Differentiable: correct jvp/vjp via the implicit function theorem.
138
+ """
139
+ rho = jnp.asarray(rho, dtype=jnp.float64)
140
+ u = jnp.asarray(u, dtype=jnp.float64)
141
+ phase = jnp.broadcast_to(jnp.asarray(phase_hint, dtype=jnp.int32),
142
+ jnp.shape(rho))
143
+ return jax.vmap(_core._temperature_from_Du)(rho, u, phase)
144
+
145
+
146
+ # ═══════════════════════════════════════════════════════════════════════════
147
+ # State functions — scalar in their inputs (vmap for batches)
148
+ # ═══════════════════════════════════════════════════════════════════════════
149
+
150
+ @jax.jit
151
+ def properties(T, rho):
152
+ """Full thermodynamic + transport state at scalar (T, ρ).
153
+
154
+ Returns a dict (see module docstring for keys). vmap for batches.
155
+ """
156
+ T = jnp.asarray(T, dtype=jnp.float64)
157
+ rho = jnp.asarray(rho, dtype=jnp.float64)
158
+ return _core._state_from_T_rho(T, rho)
159
+
160
+
161
+ @jax.jit
162
+ def state_from_PT(P, T, phase_hint=AUTO):
163
+ """Full state at (P, T). Phase-aware Halley solve for ρ, then derive all.
164
+
165
+ ``pressure`` in the returned dict echoes the input ``P`` exactly.
166
+ """
167
+ P = jnp.asarray(P, dtype=jnp.float64)
168
+ T = jnp.asarray(T, dtype=jnp.float64)
169
+ 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)
172
+ state["pressure"] = P
173
+ return state
174
+
175
+
176
+ @jax.jit
177
+ def density_from_PT(P, T, phase_hint=AUTO):
178
+ """Density [kg/m³] at scalar (P, T). Phase-aware Halley solve."""
179
+ P = jnp.asarray(P, dtype=jnp.float64)
180
+ T = jnp.asarray(T, dtype=jnp.float64)
181
+ phase = jnp.asarray(phase_hint, dtype=jnp.int32)
182
+ return _inv._density_from_PT(T, P, phase)
183
+
184
+
185
+ @jax.jit
186
+ def state_from_Ph(P, h, phase_hint=AUTO):
187
+ """Full state at (P, h). Dome-aware at subcritical P; else 2-D Newton.
188
+
189
+ ``pressure`` and ``enthalpy`` echo the inputs.
190
+ """
191
+ P = jnp.asarray(P, dtype=jnp.float64)
192
+ h = jnp.asarray(h, dtype=jnp.float64)
193
+ 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)
196
+ state["pressure"] = P
197
+ state["enthalpy"] = h
198
+ return state
199
+
200
+
201
+ @jax.jit
202
+ def state_from_Du(rho, u, phase_hint=AUTO):
203
+ """Full state at scalar (ρ, u). Scalar form of ``properties_from_rho_u``.
204
+
205
+ ``density`` and ``internal_energy`` echo the inputs.
206
+ """
207
+ rho = jnp.asarray(rho, dtype=jnp.float64)
208
+ u = jnp.asarray(u, dtype=jnp.float64)
209
+ phase = jnp.asarray(phase_hint, dtype=jnp.int32)
210
+ return _core._state_from_rho_u(rho, u, phase)
211
+
212
+
213
+ __all__ = [
214
+ "__version__",
215
+ "LIQUID", "VAPOR", "SUPERCRITICAL", "AUTO",
216
+ "properties_from_rho_u", "temperature_from_rho_u",
217
+ "properties",
218
+ "state_from_PT", "state_from_Ph", "state_from_Du",
219
+ "density_from_PT",
220
+ "saturation_pressure", "saturation_temperature",
221
+ "saturation_densities", "saturation_densities_P",
222
+ "saturation_enthalpies", "saturation_enthalpies_P",
223
+ "saturation_entropies", "saturation_entropies_P",
224
+ "viscosity", "thermal_conductivity",
225
+ ]