propax 0.1.0__py3-none-any.whl

propax/__init__.py

@@ -0,0 +1,12 @@
+from .core import Interface
+from .fluids.generic import (
+    register_conductivity_residual,
+    register_viscosity_higher_order,
+)
+
+# restrain what can be imported if user enters `from propax import *`
+__all__ = [
+    "Interface",
+    "register_viscosity_higher_order",
+    "register_conductivity_residual",
+]

propax/build_tables.py

@@ -0,0 +1,244 @@
+import argparse
+import logging
+import os
+import shutil
+import textwrap
+from pathlib import Path
+from typing import Optional
+
+from .utils.build_utils import (
+    EQS_REGISTRY,
+    build_adaptive_table,
+    build_vectorized_table,
+    fluid_table_registry,
+    get_table_path,
+    logger,
+)
+
+
+def create_all_tables(
+    fluid_name: str,
+    target: float,
+    tables_base_path: Optional[Path] = None,
+):
+    """
+    Create all interpolation tables defined in TABLE_REGISTRY for a given
+    fluid.
+
+    `target` switches the axes from uniform to graded: they are then refined
+    until the measured interpolation error meets it, which spends nodes only
+    where a measurement asks for them.
+
+    Please note a table does not have to meet high precision requirements, by
+    design it is more suited for typical float32 compatible tolerances (~1e-5
+    to ~1e-7 rtols).
+    """
+    if tables_base_path is None:
+        tables_base_path = get_table_path()
+
+    failures: list = []
+    for table_cfg in fluid_table_registry(fluid_name):
+        try:
+            if target is None:
+                build_vectorized_table(fluid_name, table_cfg, tables_base_path)
+            else:
+                build_adaptive_table(
+                    fluid_name, table_cfg, tables_base_path, target=target
+                )
+        except Exception as e:
+            logger.error(f"  ERROR building {table_cfg.name}: {e}", exc_info=True)
+            failures.append((table_cfg.name, e))
+
+    if failures:
+        names = ", ".join(name for name, _ in failures)
+        raise RuntimeError(
+            f"{len(failures)} table(s) of {fluid_name} failed to build: {names}. "
+            "The fluid is on disk but incomplete; see the logged tracebacks."
+        )
+
+
+def cached_fluids(tables_base_path: Optional[Path] = None) -> dict:
+    base = tables_base_path or get_table_path()
+    if not base.exists():
+        return {}
+    out = {}
+    for d in sorted(base.iterdir()):
+        if not d.is_dir():
+            continue
+        tables = sorted(t.name for t in d.iterdir() if t.is_dir())
+        size = sum(f.stat().st_size for f in d.rglob("*") if f.is_file())
+        out[d.name] = (tables, size)
+    return out
+
+
+def remove_fluid(fluid_name: str, tables_base_path: Optional[Path] = None) -> int:
+    base = tables_base_path or get_table_path()
+    target = base / fluid_name.lower().replace(" ", "")
+    if not target.is_dir():
+        raise FileNotFoundError(f"{target} is not in the cache")
+    size = sum(f.stat().st_size for f in target.rglob("*") if f.is_file())
+    shutil.rmtree(target)
+    return size
+
+
+def _megabytes(n: int) -> str:
+    return f"{n / 1e6:.0f} MB"
+
+
+def show_cache(base: Path) -> int:
+    cache = cached_fluids(base)
+    if not cache:
+        print(f"cache empty ({base})")
+        return 0
+    total = sum(size for _, size in cache.values())
+    print(f"{base}    {len(cache)} fluid(s), {_megabytes(total)}\n")
+    for name, (tables, size) in cache.items():
+        print(f"  {name:16s} {size / 1e6:8.1f} MB  {', '.join(tables) or '(empty)'}")
+    return 0
+
+
+def clear_cache(base: Path, assume_yes: bool) -> int:
+    cache = cached_fluids(base)
+    if not cache:
+        print(f"cache already empty ({base})")
+        return 0
+    total = sum(size for _, size in cache.values())
+    question = f"delete {len(cache)} fluid(s), {_megabytes(total)}, from {base}? [y/N] "
+    if not assume_yes and not _confirmed(question):
+        print("aborted")
+        return 1
+    shutil.rmtree(base)
+    print(f"removed {_megabytes(total)}")
+    return 0
+
+
+def _confirmed(question: str) -> bool:
+    try:
+        return input(question).strip().lower() == "y"
+    except EOFError:  # not answered is a no
+        return False
+
+
+def remove_fluids(names: list, base: Path) -> int:
+    status = 0
+    for name in names:
+        try:
+            freed = remove_fluid(name, base)
+        except FileNotFoundError as e:
+            print(f"{name}: {e}")
+            status = 1
+            continue
+        print(f"removed {name} ({_megabytes(freed)})")
+    return status
+
+
+def build_fluids(names: list, base: Path, target: float) -> int:
+    incomplete = []
+    for fluid in names:
+        try:
+            create_all_tables(fluid, tables_base_path=base, target=target)
+        except Exception as e:
+            logger.error(f"{fluid}: {e}")
+            incomplete.append(fluid)
+
+    if incomplete:
+        print(f"\nINCOMPLETE: {len(incomplete)} fluid(s) -> {', '.join(incomplete)}")
+        return 1
+    print(f"\nbuilt {len(names)} fluid(s) into {base}")
+    return 0
+
+
+def _positive_int(raw: str) -> int:
+    value = int(raw)
+    if value < 1:
+        raise argparse.ArgumentTypeError("must be at least 1")
+    return value
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m propax.build_tables",
+        description=textwrap.dedent("""
+            Build the interpolation tables a fluid needs, into the propax cache
+            (PROPAX_TABLE_DIR, else ~/.cache/propax).
+
+            A table is a few tens to a few hundred MB, so build only the fluids
+            you actually use  there is no reason to hold all covered fluids.
+            Grid corners outside the physical domain are filled from their nearest
+            valid neighbour; the warnings about them are only informative.
+        """),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("fluids", nargs="*", help="fluid names, e.g. hydrogen argon")
+
+    action = parser.add_mutually_exclusive_group()
+    action.add_argument("--all", action="store_true", help="build every known fluid")
+    action.add_argument("--list", action="store_true", help="show what the cache holds")
+    action.add_argument(
+        "--remove", metavar="FLUID", nargs="+", help="delete these from the cache"
+    )
+    action.add_argument("--clear", action="store_true", help="delete the whole cache")
+
+    parser.add_argument(
+        "--yes", action="store_true", help="answer yes to --clear's confirmation"
+    )
+    parser.add_argument(
+        "--table-dir", metavar="DIR", help="override the cache location"
+    )
+    parser.add_argument(
+        "--chunk-size",
+        type=_positive_int,
+        metavar="N",
+        help="map at most N nodes at a time instead of the whole grid; lowers "
+        "peak memory at some cost in speed (try 1024 if a build is swapping)",
+    )
+    parser.add_argument(
+        "--target",
+        type=float,
+        metavar="ERR",
+        help="grade the axes until the measured interpolation error meets ERR "
+        "(e.g. 1e-5) instead of using a uniform grid",
+    )
+    return parser
+
+
+def _apply_settings(args) -> None:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+    if args.chunk_size is not None:
+        os.environ["PROPAX_CHUNK_SIZE"] = str(args.chunk_size)
+
+
+def _targets(args, parser: argparse.ArgumentParser) -> list:
+    if args.all:
+        return sorted(EQS_REGISTRY)
+    if not args.fluids:
+        parser.error(
+            "name at least one fluid, or pass --all / --list / --remove / --clear"
+        )
+    names = [f.lower().replace(" ", "") for f in args.fluids]
+    unknown = [f for f in names if f not in EQS_REGISTRY]
+    if unknown:
+        parser.error(
+            f"unknown fluid(s): {', '.join(unknown)}. "
+            f"Known: {', '.join(sorted(EQS_REGISTRY))}"
+        )
+    return names
+
+
+def main(argv=None) -> int:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+    _apply_settings(args)
+    base = Path(args.table_dir) if args.table_dir else get_table_path()
+
+    if args.list:
+        return show_cache(base)
+    if args.clear:
+        return clear_cache(base, args.yes)
+    if args.remove:
+        return remove_fluids(args.remove, base)
+    return build_fluids(_targets(args, parser), base, args.target)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

propax/core/__init__.py

@@ -0,0 +1,3 @@
+from .interface import Interface
+
+__all__ = ["Interface"]

propax/core/_state.py

@@ -0,0 +1,122 @@
+from collections.abc import Mapping
+from typing import Iterator
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array
+
+from propax.utils.types import jaxBool, jaxFloat
+
+from .config import ThermoVar
+
+
+class PropertyMap(Mapping):
+    """A property map that cannot be mutated in place.
+
+    Built from `ThermoVar` keys, read back by either those or the internal
+    name they carry. `replace` and `discard` return a new map, so a value
+    reaching a traced call is never rebound behind it.
+    """
+
+    __slots__ = ("_entries",)
+
+    def __init__(self, entries=()):
+        stored = {}
+        for key, value in dict(entries).items():
+            if not isinstance(key, ThermoVar):
+                raise KeyError(
+                    f"{key!r} cannot key a PropertyMap, only a ThermoVar can. "
+                    "Use PropertyMap.of_internal to build one from internal names."
+                )
+            stored[key.internal_key] = value
+        self._entries = stored
+
+    @classmethod
+    def of_internal(cls, entries) -> "PropertyMap":
+        built = cls.__new__(cls)
+        built._entries = dict(entries)
+        return built
+
+    def replace(self, entries) -> "PropertyMap":
+        merged = dict(self._entries)
+        for key, value in dict(entries).items():
+            merged[key.internal_key if isinstance(key, ThermoVar) else key] = value
+        return PropertyMap.of_internal(merged)
+
+    def discard(self, key) -> "PropertyMap":
+        name = key.internal_key if isinstance(key, ThermoVar) else key
+        if name not in self._entries:
+            raise KeyError(f"nothing stored under {key!r}")
+        return PropertyMap.of_internal(
+            {k: v for k, v in self._entries.items() if k != name}
+        )
+
+    def __getitem__(self, key):
+        name = key.internal_key if isinstance(key, ThermoVar) else key
+        try:
+            return self._entries[name]
+        except KeyError:
+            held = ", ".join(sorted(self._entries)) or "nothing"
+            raise KeyError(f"{key!r} is not in this map, which holds {held}") from None
+
+    def __repr__(self) -> str:
+        return f"PropertyMap({', '.join(sorted(self._entries))})"
+
+    # functions to implement for mapping
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._entries)
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    def __reduce__(self):
+        names, values = _split(self)
+        return (_rebuild, (names, values))
+
+
+def _split(state: PropertyMap):
+    names = tuple(sorted(state))
+    return names, tuple(state[name] for name in names)
+
+
+def _rebuild(names, values) -> PropertyMap:
+    return PropertyMap.of_internal(zip(names, values))
+
+
+def _flatten(state: PropertyMap):
+    names, values = _split(state)
+    return values, names
+
+
+jax.tree_util.register_pytree_node(PropertyMap, _flatten, _rebuild)
+
+
+def _to_array(v):
+    return jnp.asarray(v)
+
+
+class TxResult(eqx.Module):
+    """A two-phase state: T or P known, and the quality with it.
+
+    `mix` carries D, U, H, S of the mixture, `L` and `V` the same on each
+    saturated branch.
+    """
+
+    mix: PropertyMap
+    L: PropertyMap
+    V: PropertyMap
+    x: jaxFloat
+    T: jaxFloat
+    P: jaxFloat
+
+
+class SaturationResult(eqx.Module):
+    """A point of the saturation line: T and P, and D, U, H, S on both
+    branches, with no quality to place the state between them."""
+
+    L: PropertyMap
+    V: PropertyMap
+    T: Array = eqx.field(converter=_to_array)
+    P: Array = eqx.field(converter=_to_array)
+    is_valid: jaxBool = eqx.field(default=True, converter=_to_array)

propax/core/config.py

@@ -0,0 +1,205 @@
+import os
+from dataclasses import dataclass
+from enum import IntEnum, StrEnum
+from pathlib import Path
+from types import MappingProxyType
+from typing import List, Literal, Mapping, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+def get_table_path() -> Path:
+    override = os.environ.get("PROPAX_TABLE_DIR")
+    if override:
+        return Path(override)
+    home = Path.home() / ".cache" / "propax"
+    home.mkdir(parents=True, exist_ok=True)
+    return home
+
+
+def get_chunk_size() -> Optional[int]:
+    """How many nodes a mapped build call may hold at once, None for no cap
+
+    Use with :
+
+    `python -m propax.build_tables --chunk-size`
+    """
+    raw = os.environ.get("PROPAX_CHUNK_SIZE")
+    if not raw:
+        return None
+    n = int(raw)
+    if n < 1:
+        raise ValueError(f"PROPAX_CHUNK_SIZE must be at least 1, got {n}")
+    return n
+
+
+@dataclass(frozen=True, slots=True)
+class VarSpec:
+    """What the flash needs to know about one property.
+
+    The CoolProp key is the enum member's own value, so it is not repeated here.
+    """
+
+    internal_key: str
+    can_be_input: bool = True
+    linear_in_x: bool = False
+    invert_for_mixing: bool = False
+    sat_key_L: Optional[str] = None
+    sat_key_V: Optional[str] = None
+    sat_key_unique: Optional[str] = None
+    scale: float = 1.0
+    unit: str = ""
+
+
+class ThermoVar(StrEnum):
+    """Canonical names for thermodynamic properties, for library internal use.
+
+    The member's value is the CoolProp key; everything else the flash needs
+    hangs off `spec`.
+    """
+
+    P = "P"
+    T = "T"
+    H = "H"
+    S = "S"
+    U = "U"
+    D = "D"
+    Q = "Q"
+    CVMASS = "CVMASS"
+    CPMASS = "CPMASS"
+    VISCOSITY = "VISCOSITY"
+    CONDUCTIVITY = "CONDUCTIVITY"
+
+    @property
+    def spec(self) -> VarSpec:
+        return _SPECS[self]
+
+    @property
+    def internal_key(self) -> str:
+        return _SPECS[self].internal_key
+
+
+_SPECS: Mapping[ThermoVar, VarSpec] = MappingProxyType(
+    {
+        ThermoVar.P: VarSpec("P", sat_key_unique="P_sat", scale=1e6, unit="Pa"),
+        ThermoVar.T: VarSpec("T", sat_key_unique="T_sat", scale=100.0, unit="K"),
+        ThermoVar.H: VarSpec(
+            "h",
+            sat_key_L="h_L",
+            sat_key_V="h_V",
+            linear_in_x=True,
+            scale=1e5,
+            unit="H",
+        ),
+        ThermoVar.S: VarSpec(
+            "s",
+            sat_key_L="s_L",
+            sat_key_V="s_V",
+            linear_in_x=True,
+            scale=1e3,
+            unit="J/K",
+        ),
+        ThermoVar.U: VarSpec(
+            "u",
+            sat_key_L="u_L",
+            sat_key_V="u_V",
+            linear_in_x=True,
+            scale=1e5,
+            unit="J",
+        ),
+        ThermoVar.D: VarSpec(
+            "rho",
+            sat_key_L="rho_L",
+            sat_key_V="rho_V",
+            invert_for_mixing=True,
+            scale=10.0,
+            unit="kg/m3",
+        ),
+        # An input, but only beside P or T `dispatch.SATURATED` says where.
+        # Any other partner would need a solve, and a quality is not a state
+        # variable of the EOS, so there is no bracket to lay one on.
+        ThermoVar.Q: VarSpec("x"),
+        ThermoVar.CVMASS: VarSpec("cv", can_be_input=False, unit="J/(K.kg)"),
+        ThermoVar.CPMASS: VarSpec("cp", can_be_input=False, unit="J/(K.kg)"),
+        ThermoVar.VISCOSITY: VarSpec("viscosity", can_be_input=False, unit="Pa/s"),
+        ThermoVar.CONDUCTIVITY: VarSpec("conductivity", can_be_input=False, unit="S/m"),
+    }
+)
+
+
+class PhaseID(IntEnum):
+    UNKNOWN = -1
+    TWO_PHASE = 0
+    LIQUID = 1
+    GAS = 2
+    SUPERCRITICAL_GAS = 3
+    SUPERCRITICAL = 4
+
+
+class Axis(BaseModel):
+    variable: ThermoVar
+    min_val: float
+    max_val: float
+    spacing: Literal["linear", "log"] = "linear"
+    n_points: int = 1000
+
+    @property
+    def key(self) -> str:
+        return str(self.variable.value)
+
+
+class TableSpec(BaseModel):
+    """`outputs` is whichever of the natural variables (D, T) the axes do not
+    already give; a dome table appends the quality, since below the saturation
+    line (D, T) does not name the state on its own. Every other property is
+    recovered from the EOS at that (rho, T) when it is asked for, never stored:
+    a stored value is one more thing to keep consistent, and interpolating it
+    is strictly worse than evaluating the EOS the table just located.
+    """
+
+    x_axis: Axis
+    y_axis: Axis
+    outputs: List[ThermoVar] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def _outputs_complete_the_flash(self):
+        axes = {self.x_axis.variable, self.y_axis.variable}
+        expected = {ThermoVar.D, ThermoVar.T} - axes
+        if set(self.outputs) != expected:
+            raise ValueError(
+                f"{self.name}: outputs must be exactly {sorted(v.value for v in expected)}, "
+                f"got {sorted(v.value for v in self.outputs)}. A table stores the flash "
+                "solution (D, T, and the quality under the dome); everything else comes "
+                "from the EOS."
+            )
+        return self
+
+    @property
+    def name(self) -> str:
+        return "_".join(("Table", self.x_axis.key, self.y_axis.key))
+
+
+def _table(x: Axis, y: Axis) -> TableSpec:
+    axes = {x.variable, y.variable}
+    return TableSpec(
+        x_axis=x, y_axis=y, outputs=sorted({ThermoVar.D, ThermoVar.T} - axes)
+    )
+
+
+_DENSITY = Axis(variable=ThermoVar.D, min_val=0.5, max_val=75.0, spacing="log")
+_TEMPERATURE = Axis(variable=ThermoVar.T, min_val=15.0, max_val=150.0)
+_PRESSURE = Axis(variable=ThermoVar.P, min_val=1e4, max_val=1.5e6, spacing="log")
+_ENERGY = Axis(variable=ThermoVar.U, min_val=-5e4, max_val=1e6)
+_ENTHALPY = Axis(variable=ThermoVar.H, min_val=-5e4, max_val=1.5e6)
+_ENTROPY = Axis(variable=ThermoVar.S, min_val=-1e4, max_val=5e4)
+
+TABLE_REGISTRY: List[TableSpec] = [
+    _table(_DENSITY, _ENERGY),
+    _table(_TEMPERATURE, _PRESSURE),
+    _table(_PRESSURE, _ENTHALPY),
+    _table(_PRESSURE, _ENTROPY),
+    _table(_PRESSURE, _ENERGY),
+    _table(_DENSITY, _ENTHALPY),
+    _table(_DENSITY, _ENTROPY),
+    _table(_TEMPERATURE, _ENTROPY),
+]

propax/core/domain.py

@@ -0,0 +1,38 @@
+from typing import Tuple
+
+import jax.numpy as jnp
+
+from propax.fluids.generic import HelmholtzEOS
+
+from .tolerances import TOL
+
+
+def rho_bounds(saturation) -> Tuple[float, float]:
+    """Every density the correlation describes."""
+    return (
+        TOL.domain.rho_lo_frac * saturation.eos.rho_crit_mass,
+        saturation.rho_max,
+    )
+
+
+def T_bounds(eos: HelmholtzEOS) -> Tuple[float, float]:
+    return eos.T_triple, eos.T_max
+
+
+def node_is_valid(saturation, rho, T):
+    T_lo, T_hi = T_bounds(saturation.eos)
+    return (rho > 0.0) & (rho <= saturation.rho_max) & (T > T_lo) & (T < T_hi)
+
+
+def is_two_phase(quality, sat_is_valid, slack: float = 0.0):
+    """
+    If the sat solve converged and the quality is withing [-slack, 1 + slack],
+    this is done to avoid misclassifying biphasic states.
+    """
+    return sat_is_valid & (quality >= -slack) & (quality <= 1.0 + slack)
+
+
+def clamp_quality(quality, slack: float = 0.0):
+    """The slack lets a quality out of [0, 1]; the lever rule may not see more
+    than the caller asked for."""
+    return jnp.clip(quality, -slack, 1.0 + slack)

propax/core/flash/__init__.py

@@ -0,0 +1,52 @@
+"""The flash: given two properties, find the state.
+
+`Interface` is the facade; the work is here, split by responsibility
+rather than by size. Each module takes what it needs as arguments  an EOS, a
+saturation module, a table  so a signature says which parts of the interface
+a solver actually touches. Nothing here reaches back into the interface.
+
+    dispatch       which strategy a pair takes, and nothing else
+    two_phase      the lever rule, and the three ways T_sat is obtained
+    single_phase   the bracketed 1D solve and the generic 2D one
+    results        turning a solved state into the returned dictionary
+    tables         the optional interpolation accelerator, isolated
+"""
+
+from .dispatch import (
+    INVALID_PAIRS,
+    both_natural_pair,
+    check_supported,
+    is_degenerate_pair,
+    is_nested_pair,
+    is_saturated_pair,
+    one_dim_known_var,
+    route_of,
+    supported_pairs,
+)
+from .results import (
+    add_transport,
+    fill_result_dict,
+    get_sat_bounds,
+    properties_Tx,
+    undefine_two_phase,
+)
+from .tables import call_interp, pass_pair_args_to_table
+
+__all__ = [
+    "one_dim_known_var",
+    "is_nested_pair",
+    "is_saturated_pair",
+    "check_supported",
+    "INVALID_PAIRS",
+    "is_degenerate_pair",
+    "both_natural_pair",
+    "route_of",
+    "supported_pairs",
+    "add_transport",
+    "fill_result_dict",
+    "get_sat_bounds",
+    "properties_Tx",
+    "undefine_two_phase",
+    "call_interp",
+    "pass_pair_args_to_table",
+]