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.
- {co2_eos-0.1.0 → co2_eos-0.2.0}/PKG-INFO +67 -19
- {co2_eos-0.1.0 → co2_eos-0.2.0}/README.md +62 -14
- co2_eos-0.2.0/co2_eos/__init__.py +225 -0
- co2_eos-0.2.0/co2_eos/core.py +245 -0
- co2_eos-0.2.0/co2_eos/data/seed_table.npz +0 -0
- co2_eos-0.2.0/co2_eos/helmholtz.py +337 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/transport.py +37 -37
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/PKG-INFO +67 -19
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/SOURCES.txt +6 -1
- {co2_eos-0.1.0 → co2_eos-0.2.0}/pyproject.toml +5 -5
- co2_eos-0.2.0/tests/test_analytic_derivs.py +87 -0
- co2_eos-0.2.0/tests/test_rho_u_inversion.py +101 -0
- co2_eos-0.1.0/co2_eos/__init__.py +0 -265
- {co2_eos-0.1.0 → co2_eos-0.2.0}/LICENSE +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/data/saturation_table.npz +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/inversions.py +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/saturation.py +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos/span_wagner.py +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/dependency_links.txt +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/requires.txt +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/co2_eos.egg-info/top_level.txt +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/setup.cfg +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/tests/test_accuracy_vs_coolprop.py +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/tests/test_dome_detection.py +0 -0
- {co2_eos-0.1.0 → co2_eos-0.2.0}/tests/test_inversion_gradients.py +0 -0
- {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.
|
|
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/
|
|
8
|
-
Project-URL: Repository, https://github.com/
|
|
9
|
-
Project-URL: Documentation, https://github.com/
|
|
10
|
-
Project-URL: Issues, https://github.com/
|
|
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
|
-
[](https://colab.research.google.com/github/
|
|
37
|
+
[](https://github.com/FluxTech-UG/co2-eos/actions/workflows/test.yml)
|
|
38
|
+
[](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
|
-
#
|
|
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
|
|
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.
|
|
78
|
-
- **Fast.** JIT-compiled and vectorisable
|
|
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
|
|
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
|
-
|
|
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/
|
|
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
|
-
[](https://colab.research.google.com/github/
|
|
3
|
+
[](https://github.com/FluxTech-UG/co2-eos/actions/workflows/test.yml)
|
|
4
|
+
[](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
|
-
#
|
|
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
|
|
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.
|
|
44
|
-
- **Fast.** JIT-compiled and vectorisable
|
|
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
|
|
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
|
-
|
|
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/
|
|
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
|
+
]
|