co2-eos 0.1.0__tar.gz → 0.2.0__tar.gz

{co2_eos-0.1.0 → co2_eos-0.2.0}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: co2-eos
-Version: 0.1.0
+Version: 0.2.0
 Summary: Differentiable CO₂ thermodynamic properties in JAX
 Author-email: John Patrick Therrien <john@flux-tech.de>
 License-Expression: Apache-2.0
-Project-URL: Homepage, https://github.com/John-FluxTech/co2-eos
-Project-URL: Repository, https://github.com/John-FluxTech/co2-eos
-Project-URL: Documentation, https://github.com/John-FluxTech/co2-eos#readme
-Project-URL: Issues, https://github.com/John-FluxTech/co2-eos/issues
+Project-URL: Homepage, https://github.com/FluxTech-UG/co2-eos
+Project-URL: Repository, https://github.com/FluxTech-UG/co2-eos
+Project-URL: Documentation, https://github.com/FluxTech-UG/co2-eos#readme
+Project-URL: Issues, https://github.com/FluxTech-UG/co2-eos/issues
 Keywords: CO2,equation-of-state,thermodynamics,JAX,autodiff,Span-Wagner
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
@@ -34,18 +34,27 @@ Dynamic: license-file
 
 # CO2-EOS
 
-[![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)
-[![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)
+[![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)
+[![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)
 
 Differentiable CO₂ thermodynamic properties in JAX.
 
-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.
+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.
 
 ```python
 import co2_eos as co2
-import jax
-
-# Properties from (T, ρ)
+import jax, jax.numpy as jnp
+
+# Simulation hot path: full state from the conserved variables (ρ, u).
+# Batched/array-native — solves T once and reuses the α-derivatives for
+# every property. This is the call a finite-volume code makes each RHS eval.
+rho = jnp.array([300.0, 350.0, 420.0])      # kg/m³
+u   = jnp.array([3.62e5, 3.58e5, 3.42e5])   # J/kg  (specific internal energy)
+state = co2.properties_from_rho_u(rho, u)   # dict of arrays: T, P, cp, w, μ, k, …
+print(state['temperature'], state['pressure'], state['speed_of_sound'])
+# -> T ≈ [307.5, 315.4, 315.6] K, P ≈ [7.6, 8.8, 9.1] MPa  (supercritical)
+
+# Properties from (T, ρ)  (scalar; vmap for batches)
 state = co2.properties(T=310.0, rho=350.0)
 print(state['pressure'], state['cp'], state['speed_of_sound'])
 
@@ -56,15 +65,14 @@ state = co2.state_from_PT(P=8e6, T=310.0)
 # (peaks near the Widom line at these conditions)
 drho_dT = jax.grad(lambda T: co2.state_from_PT(P=8e6, T=T)['density'])(310.0)
 
-# Vectorise across conditions
-import jax.numpy as jnp
+# Vectorise the scalar functions across conditions
 T_array = jnp.linspace(280, 340, 100)
 states = jax.vmap(lambda T: co2.state_from_PT(P=8e6, T=T))(T_array)
 ```
 
 ## Try it without cloning
 
-[`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.
+[`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.
 
 ## Why this exists
 
@@ -74,14 +82,15 @@ But CoolProp is a C++ library with Python bindings. You can't `jax.grad` through
 
 CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagner 1996) directly in JAX:
 
-- **Differentiable.** `jax.grad` through any property, any inversion, any combination. No finite differences. Exact gradients via autodiff, including second and higher derivatives for free.
-- **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×.
+- **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.
+- **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`).
+- **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.
 - **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.
 - **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.
 
 ## What's included
 
-**Equation of state**: Full Span-Wagner (1996) formulation for CO₂. Helmholtz free energy as A(T, ρ), with all thermodynamic properties derived by autodiff:
+**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):
 
 - Pressure, internal energy, enthalpy, entropy
 - 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
 
 **State inversions:**
 
+- `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.
+- `state_from_Du(rho, u)`: scalar form of the above (vmap for batches).
 - `state_from_PT(P, T)`: Halley's method with pressure-aware critical-region initial guess, step damping, and bisection safety net
 - `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
-- `state_from_Du(rho, u)`: for simulation codes carrying (ρ, u) as conserved variables. Uses `custom_vjp` for clean gradient propagation through the implicit solve
+
+## Migrating to v0.2
+
+v0.2 is a performance redesign. **Accuracy and gradients are unchanged** and the
+scalar state functions keep their signatures, so most code needs no changes. The
+one thing worth adopting is the fused hot path.
+
+If your simulation carried `(ρ, u)` and recovered primitives with separate calls:
+
+```python
+# v0.1 — separate calls, each re-evaluating α-derivatives
+from co2_eos.inversions import temperature_from_Du, SUPERCRITICAL
+from co2_eos import span_wagner as sw, transport
+T  = temperature_from_Du(rho, u, phase)     # nested-autodiff Cv, while_loop
+P  = sw.pressure(T, rho)
+mu = transport.viscosity(T, rho)
+k  = transport.thermal_conductivity(T, rho)
+```
+
+replace it with the single fused call:
+
+```python
+# v0.2 — one fused pass: solve T once, reuse the α-derivatives
+import co2_eos as co2
+s = co2.properties_from_rho_u(rho, u)       # phase_hint defaults to SUPERCRITICAL
+T, P, mu, k = s["temperature"], s["pressure"], s["viscosity"], s["thermal_conductivity"]
+# also: s["cp"], s["speed_of_sound"], s["enthalpy"], s["entropy"], s["gibbs_energy"], …
+```
+
+Notes:
+
+- `properties_from_rho_u` / `temperature_from_rho_u` are **batched** (pass 1-D
+  arrays). The scalar `state_from_Du(rho, u)` returns the same dict for a single
+  point.
+- `phase_hint` is accepted for API symmetry but unused by the `(ρ, u)` inversion
+  (`u` is monotone in `T` at fixed `ρ`, so no phase branch is needed).
+- The old modules `co2_eos.span_wagner` and `co2_eos.inversions` still work; they
+  are the validation ground truth and the benchmark baseline.
 
 ## Valid range
 
@@ -125,7 +173,7 @@ Requires Python ≥ 3.10 and JAX ≥ 0.4. No other dependencies.
 To run the validation tests (which compare against CoolProp):
 
 ```bash
-git clone https://github.com/John-FluxTech/co2-eos.git
+git clone https://github.com/FluxTech-UG/co2-eos.git
 cd co2-eos
 pip install -e ".[test]"
 pytest tests/ -v

{co2_eos-0.1.0 → co2_eos-0.2.0}/README.md

@@ -1,17 +1,26 @@
 # CO2-EOS
 
-[![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)
-[![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)
+[![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)
+[![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)
 
 Differentiable CO₂ thermodynamic properties in JAX.
 
-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.
+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.
 
 ```python
 import co2_eos as co2
-import jax
-
-# Properties from (T, ρ)
+import jax, jax.numpy as jnp
+
+# Simulation hot path: full state from the conserved variables (ρ, u).
+# Batched/array-native — solves T once and reuses the α-derivatives for
+# every property. This is the call a finite-volume code makes each RHS eval.
+rho = jnp.array([300.0, 350.0, 420.0])      # kg/m³
+u   = jnp.array([3.62e5, 3.58e5, 3.42e5])   # J/kg  (specific internal energy)
+state = co2.properties_from_rho_u(rho, u)   # dict of arrays: T, P, cp, w, μ, k, …
+print(state['temperature'], state['pressure'], state['speed_of_sound'])
+# -> T ≈ [307.5, 315.4, 315.6] K, P ≈ [7.6, 8.8, 9.1] MPa  (supercritical)
+
+# Properties from (T, ρ)  (scalar; vmap for batches)
 state = co2.properties(T=310.0, rho=350.0)
 print(state['pressure'], state['cp'], state['speed_of_sound'])
 
@@ -22,15 +31,14 @@ state = co2.state_from_PT(P=8e6, T=310.0)
 # (peaks near the Widom line at these conditions)
 drho_dT = jax.grad(lambda T: co2.state_from_PT(P=8e6, T=T)['density'])(310.0)
 
-# Vectorise across conditions
-import jax.numpy as jnp
+# Vectorise the scalar functions across conditions
 T_array = jnp.linspace(280, 340, 100)
 states = jax.vmap(lambda T: co2.state_from_PT(P=8e6, T=T))(T_array)
 ```
 
 ## Try it without cloning
 
-[`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.
+[`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.
 
 ## Why this exists
 
@@ -40,14 +48,15 @@ But CoolProp is a C++ library with Python bindings. You can't `jax.grad` through
 
 CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagner 1996) directly in JAX:
 
-- **Differentiable.** `jax.grad` through any property, any inversion, any combination. No finite differences. Exact gradients via autodiff, including second and higher derivatives for free.
-- **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×.
+- **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.
+- **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`).
+- **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.
 - **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.
 - **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.
 
 ## What's included
 
-**Equation of state**: Full Span-Wagner (1996) formulation for CO₂. Helmholtz free energy as A(T, ρ), with all thermodynamic properties derived by autodiff:
+**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):
 
 - Pressure, internal energy, enthalpy, entropy
 - 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
 
 **State inversions:**
 
+- `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.
+- `state_from_Du(rho, u)`: scalar form of the above (vmap for batches).
 - `state_from_PT(P, T)`: Halley's method with pressure-aware critical-region initial guess, step damping, and bisection safety net
 - `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
-- `state_from_Du(rho, u)`: for simulation codes carrying (ρ, u) as conserved variables. Uses `custom_vjp` for clean gradient propagation through the implicit solve
+
+## Migrating to v0.2
+
+v0.2 is a performance redesign. **Accuracy and gradients are unchanged** and the
+scalar state functions keep their signatures, so most code needs no changes. The
+one thing worth adopting is the fused hot path.
+
+If your simulation carried `(ρ, u)` and recovered primitives with separate calls:
+
+```python
+# v0.1 — separate calls, each re-evaluating α-derivatives
+from co2_eos.inversions import temperature_from_Du, SUPERCRITICAL
+from co2_eos import span_wagner as sw, transport
+T  = temperature_from_Du(rho, u, phase)     # nested-autodiff Cv, while_loop
+P  = sw.pressure(T, rho)
+mu = transport.viscosity(T, rho)
+k  = transport.thermal_conductivity(T, rho)
+```
+
+replace it with the single fused call:
+
+```python
+# v0.2 — one fused pass: solve T once, reuse the α-derivatives
+import co2_eos as co2
+s = co2.properties_from_rho_u(rho, u)       # phase_hint defaults to SUPERCRITICAL
+T, P, mu, k = s["temperature"], s["pressure"], s["viscosity"], s["thermal_conductivity"]
+# also: s["cp"], s["speed_of_sound"], s["enthalpy"], s["entropy"], s["gibbs_energy"], …
+```
+
+Notes:
+
+- `properties_from_rho_u` / `temperature_from_rho_u` are **batched** (pass 1-D
+  arrays). The scalar `state_from_Du(rho, u)` returns the same dict for a single
+  point.
+- `phase_hint` is accepted for API symmetry but unused by the `(ρ, u)` inversion
+  (`u` is monotone in `T` at fixed `ρ`, so no phase branch is needed).
+- The old modules `co2_eos.span_wagner` and `co2_eos.inversions` still work; they
+  are the validation ground truth and the benchmark baseline.
 
 ## Valid range
 
@@ -91,7 +139,7 @@ Requires Python ≥ 3.10 and JAX ≥ 0.4. No other dependencies.
 To run the validation tests (which compare against CoolProp):
 
 ```bash
-git clone https://github.com/John-FluxTech/co2-eos.git
+git clone https://github.com/FluxTech-UG/co2-eos.git
 cd co2-eos
 pip install -e ".[test]"
 pytest tests/ -v

co2_eos-0.2.0/co2_eos/__init__.py

@@ -0,0 +1,225 @@
+"""CO2-EOS: Differentiable CO₂ thermodynamic properties in JAX.
+
+Pure-JAX Span-Wagner (1996) equation of state for carbon dioxide with hand-coded
+analytic α-derivatives.  JIT-compilable, ``vmap``-able, fully differentiable.
+
+Public API
+==========
+
+Simulation hot path — full state from the conserved variables (ρ, u).  These are
+**batched / array-native**: pass 1-D arrays, get back a dict of 1-D arrays.  A
+single fused pass solves T once and reuses the α-derivatives for every property::
+
+    properties_from_rho_u(rho, u, phase_hint=SUPERCRITICAL) -> dict
+    temperature_from_rho_u(rho, u, phase_hint=SUPERCRITICAL) -> T   # T only
+
+State functions — **scalar** in their inputs (0-d arrays or Python floats); wrap
+with ``jax.vmap`` for batches::
+
+    properties(T, rho)              — full state from (T, ρ)
+    state_from_PT(P, T)             — full state from (P, T)
+    state_from_Ph(P, h)             — full state from (P, h)
+    state_from_Du(rho, u)           — full state from (ρ, u)   (scalar form)
+    density_from_PT(P, T)           — ρ from (P, T)
+
+Each full-state dict has keys: ``temperature``, ``density``, ``pressure``,
+``cv``, ``cp``, ``speed_of_sound``, ``enthalpy``, ``internal_energy``,
+``entropy``, ``gibbs_energy``, ``viscosity``, ``thermal_conductivity``.
+
+Saturation curve (cubic-spline interpolation of a precomputed Maxwell table)::
+
+    saturation_pressure(T) / saturation_temperature(P)
+    saturation_densities(T) / saturation_densities_P(P)
+    saturation_enthalpies(T) / saturation_enthalpies_P(P)
+    saturation_entropies(T) / saturation_entropies_P(P)
+
+Transport (scalar; vmap for batches)::
+
+    viscosity(T, rho)               → μ  [Pa·s]
+    thermal_conductivity(T, rho)    → λ  [W/(m·K)]
+
+Phase hints::
+
+    LIQUID = 0, VAPOR = 1, SUPERCRITICAL = 2 (alias: AUTO)
+
+See README "Migrating to v0.2" for the move from the old per-quantity calls to
+the fused ``properties_from_rho_u``.
+"""
+
+import jax
+import jax.numpy as jnp
+
+# Enable float64 for the entire library (also done by submodules).
+jax.config.update("jax_enable_x64", True)
+
+__version__ = "0.2.0"
+
+# ── Phase hint constants ────────────────────────────────────────────────────
+LIQUID = 0
+VAPOR = 1
+SUPERCRITICAL = 2
+AUTO = SUPERCRITICAL  # alias: auto-pick in the dome / supercritical-style guess
+
+# ── Internal modules ────────────────────────────────────────────────────────
+from co2_eos import span_wagner as _sw
+from co2_eos import transport as _tr
+from co2_eos import saturation as _sat
+from co2_eos import inversions as _inv
+from co2_eos import core as _core
+
+# ── Re-exports: saturation curve ────────────────────────────────────────────
+saturation_pressure = _sat.saturation_pressure
+saturation_temperature = _sat.saturation_temperature
+saturation_densities = _sat.saturation_densities
+saturation_densities_P = _sat.saturation_densities_P
+saturation_enthalpies = _sat.saturation_enthalpies_T
+saturation_enthalpies_P = _sat.saturation_enthalpies
+saturation_entropies = _sat.saturation_entropies
+saturation_entropies_P = _sat.saturation_entropies_P
+
+
+# ── Transport (scalar form; vmap for batches) ───────────────────────────────
+
+@jax.jit
+def viscosity(T, rho):
+    """Dynamic viscosity μ [Pa·s] at scalar (T, ρ). vmap for batches."""
+    T = jnp.asarray(T, dtype=jnp.float64)
+    rho = jnp.asarray(rho, dtype=jnp.float64)
+    return _tr._scalar_viscosity(T, rho)
+
+
+@jax.jit
+def thermal_conductivity(T, rho):
+    """Thermal conductivity λ [W/(m·K)] at scalar (T, ρ). vmap for batches."""
+    T = jnp.asarray(T, dtype=jnp.float64)
+    rho = jnp.asarray(rho, dtype=jnp.float64)
+    return _tr._scalar_thermal_conductivity(T, rho)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Hot path: full state from conserved (ρ, u) — batched / array-native
+# ═══════════════════════════════════════════════════════════════════════════
+
+@jax.jit
+def properties_from_rho_u(rho, u, phase_hint=SUPERCRITICAL):
+    """Full thermodynamic + transport state from (ρ, u). Batched.
+
+    The simulation hot path: solves T once from u(T,ρ)=u via the analytic-Cv
+    fixed-iteration Newton, then derives every property reusing the
+    α-derivatives.  ``rho`` and ``u`` are 1-D arrays of equal length; the result
+    is a dict of 1-D arrays.  ``density`` and ``internal_energy`` echo the
+    inputs exactly.
+
+    Args:
+        rho: Density [kg/m³], 1-D array.
+        u: Specific internal energy [J/kg], 1-D array.
+        phase_hint: accepted for API symmetry; the single-phase (ρ, u)
+            inversion is phase-agnostic (u is monotone in T at fixed ρ on a
+            single-phase branch).
+
+    Note: this is a single-phase EOS inversion. A (ρ, u) state inside the
+    two-phase dome (subcritical T, ρ between ρ_v and ρ_l) is not a single-phase
+    state; the EOS returns its unstable single-phase extrapolation there. The
+    near-critical supercritical operating regime has no dome.
+    """
+    rho = jnp.asarray(rho, dtype=jnp.float64)
+    u = jnp.asarray(u, dtype=jnp.float64)
+    phase = jnp.broadcast_to(jnp.asarray(phase_hint, dtype=jnp.int32),
+                             jnp.shape(rho))
+    return jax.vmap(_core._state_from_rho_u)(rho, u, phase)
+
+
+@jax.jit
+def temperature_from_rho_u(rho, u, phase_hint=SUPERCRITICAL):
+    """Temperature T [K] such that u(T,ρ)=u at fixed ρ. Batched.
+
+    The lean inversion (no property derivation) — use when only T is needed.
+    Differentiable: correct jvp/vjp via the implicit function theorem.
+    """
+    rho = jnp.asarray(rho, dtype=jnp.float64)
+    u = jnp.asarray(u, dtype=jnp.float64)
+    phase = jnp.broadcast_to(jnp.asarray(phase_hint, dtype=jnp.int32),
+                             jnp.shape(rho))
+    return jax.vmap(_core._temperature_from_Du)(rho, u, phase)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# State functions — scalar in their inputs (vmap for batches)
+# ═══════════════════════════════════════════════════════════════════════════
+
+@jax.jit
+def properties(T, rho):
+    """Full thermodynamic + transport state at scalar (T, ρ).
+
+    Returns a dict (see module docstring for keys). vmap for batches.
+    """
+    T = jnp.asarray(T, dtype=jnp.float64)
+    rho = jnp.asarray(rho, dtype=jnp.float64)
+    return _core._state_from_T_rho(T, rho)
+
+
+@jax.jit
+def state_from_PT(P, T, phase_hint=AUTO):
+    """Full state at (P, T). Phase-aware Halley solve for ρ, then derive all.
+
+    ``pressure`` in the returned dict echoes the input ``P`` exactly.
+    """
+    P = jnp.asarray(P, dtype=jnp.float64)
+    T = jnp.asarray(T, dtype=jnp.float64)
+    phase = jnp.asarray(phase_hint, dtype=jnp.int32)
+    rho = _inv._density_from_PT(T, P, phase)
+    state = _core._state_from_T_rho(T, rho)
+    state["pressure"] = P
+    return state
+
+
+@jax.jit
+def density_from_PT(P, T, phase_hint=AUTO):
+    """Density [kg/m³] at scalar (P, T). Phase-aware Halley solve."""
+    P = jnp.asarray(P, dtype=jnp.float64)
+    T = jnp.asarray(T, dtype=jnp.float64)
+    phase = jnp.asarray(phase_hint, dtype=jnp.int32)
+    return _inv._density_from_PT(T, P, phase)
+
+
+@jax.jit
+def state_from_Ph(P, h, phase_hint=AUTO):
+    """Full state at (P, h). Dome-aware at subcritical P; else 2-D Newton.
+
+    ``pressure`` and ``enthalpy`` echo the inputs.
+    """
+    P = jnp.asarray(P, dtype=jnp.float64)
+    h = jnp.asarray(h, dtype=jnp.float64)
+    phase = jnp.asarray(phase_hint, dtype=jnp.int32)
+    T, rho = _inv._state_from_Ph(P, h, phase)
+    state = _core._state_from_T_rho(T, rho)
+    state["pressure"] = P
+    state["enthalpy"] = h
+    return state
+
+
+@jax.jit
+def state_from_Du(rho, u, phase_hint=AUTO):
+    """Full state at scalar (ρ, u). Scalar form of ``properties_from_rho_u``.
+
+    ``density`` and ``internal_energy`` echo the inputs.
+    """
+    rho = jnp.asarray(rho, dtype=jnp.float64)
+    u = jnp.asarray(u, dtype=jnp.float64)
+    phase = jnp.asarray(phase_hint, dtype=jnp.int32)
+    return _core._state_from_rho_u(rho, u, phase)
+
+
+__all__ = [
+    "__version__",
+    "LIQUID", "VAPOR", "SUPERCRITICAL", "AUTO",
+    "properties_from_rho_u", "temperature_from_rho_u",
+    "properties",
+    "state_from_PT", "state_from_Ph", "state_from_Du",
+    "density_from_PT",
+    "saturation_pressure", "saturation_temperature",
+    "saturation_densities", "saturation_densities_P",
+    "saturation_enthalpies", "saturation_enthalpies_P",
+    "saturation_entropies", "saturation_entropies_P",
+    "viscosity", "thermal_conductivity",
+]