libsalty

Fast JAX thermodynamics of water and electrolyte solutions via automatic differentiation.

libsalty encodes each thermodynamic potential once, as a pure scalar function, and derives every property from it with JAX AD. There are no hand-coded derivative tables anywhere: density is ∂g/∂p, heat capacity is −T ∂²g/∂T², chemical potentials are salinity gradients, and equilibrium curves are implicitly differentiated root solves. Everything composes with jax.jit, jax.vmap, and jax.grad end-to-end, so you can differentiate straight through seawater properties — or through an osmotic-pressure solve — into a process-model optimization.

The v0 target is desalination modeling: IAPWS industrial formulations for liquid water, steam, and Standard Seawater, vectorized and differentiable. The API is formulation-agnostic (water, seawater, gibbs.derive); the IAPWS-specific pieces live one level down (water.if97, water.region1).

Quick start

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)  # required (guarded at trace time)

from libsalty import seawater, water

# Pure water (IF97), batched: arrays broadcast like NumPy
T = jnp.linspace(300.0, 500.0, 5)
water.if97.h(T, 3.0e6)            # J/kg, auto-dispatched region 1/2
water.region1.rho(T, 3.0e6)       # kg/m^3, region-explicit
water.region4.p_sat(T)            # Pa, saturation line

# Standard Seawater (T [K], p [Pa], S [kg/kg absolute salinity])
seawater.rho(298.15, 101325.0, 0.035)
seawater.state(298.15, 101325.0, 0.035)   # all 12 properties in one pass
seawater.osmotic_coefficient(298.15, 101325.0, 0.035)

# Colligative equilibria: implicitly differentiated root solves
seawater.osmotic_pressure(298.15, 101325.0, 0.035)     # ~2.6 MPa
seawater.boiling_temperature(3.0e4, 0.035)             # boiling-point elevation
jax.grad(seawater.osmotic_pressure, argnums=2)(298.15, 101325.0, 0.035)
# (atmospheric seawater boiling ~374 K sits above AN5's 353 K validity cap,
#  so boiling_temperature(101325.0, S) is honestly NaN / converged=False)

# Everything composes: batched salinity sensitivities, compiled
sens = jax.jit(jax.vmap(jax.grad(seawater.h, argnums=2)))

What's implemented (v0)

Every property function: v, rho, s, h, u, f, cp, cv, w, alpha_v, kappa_T, gibbs, plus a state bundle that evaluates all of them from one derivative pass. seawater adds mu, mu_w, mu_s, and osmotic_coefficient.

Units and conventions

Strict SI, everywhere: T [K], p [Pa], S [kg/kg] (absolute salinity, Reference Composition), outputs in J/kg, m³/kg, kg/m³, J/(kg K), m/s, 1/K, 1/Pa. Uniform argument order (T, p[, S]). No unit switches, no kJ.

Reference states are the formulations' own (IF97: u = s = 0 for liquid at the triple point; seawater: h = s = 0 at 273.15 K, 101325 Pa, S = 0.03516504). Only differences of g, f, h, u, s are physical.

Validity and NaN semantics

Public functions enforce their formulation's nominal domain and return NaN outside — in the value and in AD derivatives of every order. A plain where(valid, x, nan) would silently report zero gradients for invalid states; libsalty's masking is built so that an invalid state is loudly NaN under grad and hessian too, while valid lanes in the same vmap batch are bit-identical to unmasked evaluation (no NaN cross-poisoning).

• water.if97.validity_mask(T, p) and region_id(T, p) make status inspectable; the saturation line belongs to region 1 (liquid).
• The seawater water part uses IF97 region 1 as a metastable liquid below 273.15 K and above the pure-water saturation pressure, exactly as AN5 prescribes — through a private unchecked evaluator, never through the public IF97 validation.
• In the high-T/high-S corner (T/K + 450·S > 362, R13-08's "region F") the extrapolated density derivatives make w, cv, alpha_v, kappa_T unreliable (e.g. w ≈ 4015 m/s at 353 K, S = 0.1 — a value AN5's own check table prints). Values there reproduce the formulation faithfully; treat them with the formulation's caveats.

Zero-salinity contracts

The saline potential behaves like A·S·ln S + B·S + C·S^{3/2} + … near S = 0, so limits are handled per property, analytically — never clamped:

• Every T/p-derived property at S = 0 equals its pure-water value bit-for-bit, with correct T/p derivatives.
• mu_w(T, p, 0) == water.region1.gibbs(T, p) exactly (the log cancellation in g − S·g_S is done analytically); osmotic_coefficient(T, p, 0) == 1.0 exactly; osmotic_pressure(T, p, 0) == 0.0 exactly.
• mu and mu_s genuinely diverge (∝ ln S) and are non-finite at S = 0 by design; salinity derivatives at exactly S = 0 read NaN rather than a misleading 0 (φ and Π have finite one-sided S→0⁺ limits — evaluate at S > 0).

Equilibrium solvers

vapor_pressure(T, S), boiling_temperature(p, S), and osmotic_pressure(T, p, S) solve chemical-potential balances (AN5 Eqs. 6 and 14) with a fixed-count safeguarded bracketed Newton wrapped in lax.custom_root: gradients come from the implicit function theorem — exact, cheap, and independent of the iteration count. Brackets and seeds come from the pure-water saturation line and van 't Hoff.

Each solver has a *_result variant returning an EquilibriumResult(root, residual, converged) pytree (jit/vmap-friendly). The plain functions return NaN wherever converged is False — e.g. boiling above the 353 K seawater limit, vapor pressure pushed below the triple-point pressure, or osmotic pressure beyond the 100 MPa cap. Check the residual, not just the value: converged requires a sign-change bracket and |residual| ≤ 1e-6 J/kg.

Bring your own potential

gibbs.derive(g) builds the full property suite for any pure scalar g(T, p, *comp) -> J/kg — nothing assumes a polynomial, so fitted Gibbs surrogates (the long-term direction of this library) get the identical API, including state bundles and transform composability.

from libsalty import gibbs
props = gibbs.derive(my_potential)   # props.h(T, p, x), props.state(...), ...

Verification

• IAPWS check tables at ~1e-8 relative: R7-97 Tables 5/15/35/36, AN5 Table A1 (water/saline/total columns at all three state points), AN5 Table A2 (all 63 boiling temperatures), R13-08 saline-part checks.
• Independent lineage: every coefficient set was cross-diffed against a second transcription and/or the iapws package before use; an optional grid cross-check against iapws (regions 1/2, 8 properties, rtol 1e-7) runs when iapws is installed.
• AD discipline: jit≡eager, vmap≡scalar consistency for every public function; first- and second-derivative behavior tested at region interiors, boundaries, domain edges, and in mixed-validity batches; named tests for every S→0 limit and every intentional singularity; implicit solver gradients checked against central differences.

204 tests. Maxwell relations hold by construction — AD guarantees them.

Performance

benchmarks/benchmark.py, Apple Silicon CPU, float64, jax 0.11. Cold JIT (compile + first call) is a one-time cost per function signature: 0.2–1.5 s here. Warm, batched:

The iapws package's scalar _Region1 call — a practical pure-Python baseline, not a claim against optimized IF97 implementations — costs ~26,000 ns per state on the same machine: roughly 250× for batched enthalpy, with derivatives included nearly for free. Single scalar calls through jit are dispatch-dominated (~10–60 µs); batch when you can.

Development

uv sync                      # includes dev deps (pytest, ruff, pyrefly, iapws)
uv run pytest                # 204 tests
uv run ruff check src tests && uv run ruff format src tests
uv run pyrefly check src tests
uv run python benchmarks/benchmark.py

Requires Python ≥ 3.12 and JAX ≥ 0.11. libsalty never mutates global JAX config; it requires x64 mode and raises an instructive error at trace time if it is off.

Roadmap

IF97 regions 3/5 and backward equations; freezing/ice Ih (IAPWS R10); transport properties; a helmholtz module (density solve behind the same property surface) as the gateway to SAFT-family and fitted surrogate potentials for arbitrary saline solutions.

Acknowledgements

The implementation was written by Claude (Fable 5, via Claude Code), working from a human-authored design plan, with every phase human-reviewed and gated on the IAPWS check tables and the full test suite before merge.

References

• IAPWS R7-97(2012), Revised Release on the IAPWS Industrial Formulation 1997 for the Thermodynamic Properties of Water and Steam.
• IAPWS R13-08, Release on the IAPWS Formulation 2008 for the Thermodynamic Properties of Seawater.
• IAPWS AN5-13, Advisory Note No. 5: Industrial Calculation of the Thermodynamic Properties of Seawater (2013).