Skip to content

Performance

torch-flash is optimized for differentiable and batched thermodynamic work. It does not claim to beat compiled Fortran/C++ libraries on a one-off scalar state.

Execution paths

  • CPA site fractions broadcast over leading state dimensions. Ten damped fixed-point updates are followed by eight Newton updates using the analytic mass-action Jacobian and a positivity-limited step.
  • CPA's conservative volume grid is evaluated as one tensor workload. Explicit liquid/vapor roots first use a phase-specific Newton solve; the complete scan remains the fallback and stable-root path.
  • GERG-2008 and EOS-CG-2021 reducing functions, pure/departure Helmholtz terms, ideal terms, pressure, and caloric derivatives broadcast over state axes.
  • Batched multifluid density roots use a fixed vector workload. Root location is detached from the graph; if an input or fitted parameter requires a gradient, one exact Newton correction restores the implicit derivative.
  • Composition validation remains strict in eager calls. Compiled graphs assume that documented input contract so the numerical normalization can be fused.
import torch

from torch_flash import configure
from torch_flash.eos import gerg2008

runtime = configure(device="cpu", dtype=torch.float64, num_threads=8)
model = gerg2008(("hydrogen", "methane"))

temperature = torch.linspace(300.0, 400.0, 1010, **runtime.tensor_options)
pressure = torch.linspace(1.0e4, 100.0e6, 1010, **runtime.tensor_options)
hydrogen = torch.linspace(0.1, 0.9, 1010, **runtime.tensor_options)
composition = torch.stack((hydrogen, 1.0 - hydrogen), dim=-1)

volume = model.molar_volume(temperature, pressure, composition, "vapor")

compiled_pressure = torch.compile(
    lambda current_temperature, current_volume, current_composition: model.pressure(
        current_temperature,
        current_volume,
        current_composition,
    ),
    fullgraph=True,
)
compiled_pressure(temperature, volume, composition)

Compile a repeated homogeneous kernel, not a one-off call. torch.compile uses TorchInductor by default and its first call can take seconds. Validate untrusted compositions with an eager call before passing them to a compiled hot loop.

Set CPU intra-operation threads with configure(num_threads=...) before thermodynamic work. More threads help sufficiently large batches, but small kernels plateau quickly. Do not change the process-wide thread count inside a library call. The runtime guide also covers the stricter one-time inter-operation thread setting.

Executed benchmark scope

16_performance_backends.ipynb records hardware, package versions, cold compilation, warmed medians, pressure/root residuals, and plots. Its Apple arm64 results are approximately:

Workload torch-flash eager torch-flash compiled ThermoPack teqp NeqSim
GERG pressure, same coefficients 1.0 ms 38–40 µs 6 µs 1.1 µs n/a
CPA pressure workload 1.0 ms 105–120 µs 21–22 µs n/a n/a
GERG scalar TP root 12 ms n/a 11–12 µs 10–11 µs 0.3–0.45 ms*

* The NeqSim call includes setting the state and running TPflash, so it is broader than a single homogeneous density-root call.

The 1,010-state native GERG density solve reaches about 4,200 states/s with maximum normalized pressure residual below 3e-15. A 4,096-state pressure kernel gains about 2× from 4–8 CPU intra-operation threads relative to one thread on that host.

The executed CO2-H2O fitting notebook measures its exact 52-state autodifferentiable GERG residual graph separately. Batched eager evaluation takes 0.299 s, versus 5.437 s for an explicit scalar loop (18.17×); objective plus backward takes 0.322 s. The maximum batched/scalar residual difference is 1.94e-10. This is the workload used by the regression rather than a synthetic pressure-only kernel.

The complete CO2/N2 PR78 envelope in 23_synthetic_co2_pr78_derivatives.ipynb is a different workload: 469 ordered bubble, dew, and fixed-log-K continuation solves, each only three or four variables wide. The continuation path uses a secant predictor, two successive-substitution corrections, and an autodiff Newton solve. A failed prediction is retried with the conservative full initializer. The recorded Apple arm64 run takes 4.713 s. accelerated=False selects full initialization at every continuation point for numerical audits.

This envelope is intrinsically sequential: each physical solution selects the branch for the next state. Moving the individual tiny Newton systems to a GPU would add kernel-launch and host-synchronization overhead, while independently batching all temperatures can converge to the algebraic \(K_i=1\) solution instead of the physical branch. GPU execution remains appropriate for the large homogeneous state grids evaluated after the envelope is known.

These figures are measurements, not CI thresholds. Exact latency depends on hardware, PyTorch version, compiler cache, batch shape, composition, phase region, and requested derivative order.

GPU precision

The Pixi gpu environment targets CUDA and retains float64. GPU launch and transfer overhead mean that batching and device-resident downstream work are necessary to amortize costs. The benchmark checks CUDA float64 when available.

The Apple MPS runtime on the executed host rejects float64 tensors. A float32 MPS result is not an equal-accuracy thermodynamic comparison and is therefore not included in the primary timing ratios. Near critical loci, spinodals, and coalescing phase roots, float32 can be inadequate.