aquakin 0.1.0__py3-none-any.whl

aquakin/__init__.py

@@ -0,0 +1,332 @@
+"""aquakin: reactive scalar transport kinetics for aqueous environmental systems."""
+
+import os as _os
+import sys as _sys
+import warnings as _warnings
+
+# aquakin enables JAX 64-bit (x64) mode at import: the stiff implicit ODE solves
+# need double precision (repeatedly forming small net differences in a stiff
+# Newton step loses the accurate digits at float32). This mutates GLOBAL,
+# process-wide JAX state -- any other JAX code in the same process will use
+# float64 afterward. The behaviour is deliberate and documented, but to keep it
+# from being a *silent* surprise we warn when we are overriding what looks like
+# an explicit float32 preference: JAX was already imported (so other code may be
+# relying on the default float32), or JAX_ENABLE_X64 is set to a false value.
+_jax_already_imported = "jax" in _sys.modules
+
+import jax
+
+if not jax.config.jax_enable_x64:
+    _env_x64 = _os.environ.get("JAX_ENABLE_X64", "").strip().lower()
+    _overrides_explicit_choice = _jax_already_imported or _env_x64 in ("0", "false", "off", "no")
+    if _overrides_explicit_choice:
+        _warnings.warn(
+            "aquakin is enabling JAX 64-bit (x64) mode, which it requires for "
+            "stiff ODE integration. This is GLOBAL, process-wide JAX state: "
+            "other JAX code in this process will now use float64. It overrides "
+            "an apparent float32 preference (JAX was already imported, or "
+            "JAX_ENABLE_X64 is set off). Enable x64 yourself before importing "
+            "aquakin, or run aquakin in a separate process, to silence this.",
+            stacklevel=2,
+        )
+    jax.config.update("jax_enable_x64", True)
+
+import logging as _logging
+
+# Library logging hygiene: attach a NullHandler to the package's root logger so
+# aquakin never configures logging on the application's behalf (it emits nothing
+# unless the caller opts in) and there is no "No handlers could be found" noise.
+# Progress output -- e.g. the ``progress=`` sampling loops in
+# ``plant.steady_state_dgsm`` / ``dynamic_dgsm`` -- is logged at INFO on child
+# loggers; a caller sees it by enabling INFO, e.g.
+# ``logging.basicConfig(level=logging.INFO)`` or
+# ``logging.getLogger("aquakin").setLevel(logging.INFO)`` with a handler.
+_logging.getLogger("aquakin").addHandler(_logging.NullHandler())
+
+from aquakin.core.conditions import OperatingConditions, SpatialConditions
+from aquakin.core.model import CompiledModel, compile_model
+from aquakin.core.parser import parse_rate_expression
+from aquakin.integrate.batch import BatchReactor, BatchSolution
+from aquakin.integrate.biofilm import BiofilmReactor, BiofilmSolution
+from aquakin.integrate.calibrate import (
+    CalibrationResult,
+    FreeICConfig,
+    LaplaceConfig,
+    OptimizerConfig,
+    PredictiveBand,
+    calibrate,
+)
+from aquakin.integrate.profile import ProfileResult, profile_likelihood
+from aquakin.integrate.cfd import CFDReactor
+from aquakin.integrate.particle import (
+    ParticleTrackReactor,
+    Track,
+    TrackSolution,
+    integrate_ensemble,
+)
+from aquakin.integrate.discrete_adjoint import (
+    esdirk_adjoint_solve,
+    implicit_euler_adjoint_solve,
+)
+from aquakin.integrate._common import (
+    DifferentiationConfig,
+    IntegratorConfig,
+    check_finite_gradient,
+    forward_adjoint,
+)
+from aquakin.integrate.forward_sensitivity import (
+    ForwardSensitivityResult,
+    forward_sensitivity,
+)
+from aquakin.integrate.pfr import PFRSolution, PlugFlowReactor
+from aquakin.integrate.sensitivity import SensitivityResult, sensitivity
+from aquakin.integrate.fit import FitResult, fit
+from aquakin.integrate.global_sensitivity import DGSMResult, dgsm
+from aquakin.integrate.events import Event, EventedResult, solve_with_events
+from aquakin.integrate.design import Constraint, OptimizeResult, optimize_design
+from aquakin.integrate.monte_carlo import MonteCarloResult, monte_carlo
+from aquakin.integrate.scenarios import (
+    KPIComparison,
+    ScenarioComparison,
+    compare_scenarios,
+    kpi_comparison,
+)
+from aquakin.plant import (
+    ActivatedSludgeSizing,
+    Aeration,
+    AerationDesignPoint,
+    AerationSystem,
+    AlgebraicTemperature,
+    BSM1Evaluation,
+    BSM2Evaluation,
+    CSTRUnit,
+    CarbonFootprint,
+    ChlorineContactUnit,
+    ComponentBalance,
+    CostFactors,
+    IdentityTranslator,
+    InfluentFractions,
+    InfluentSeries,
+    MassBalance,
+    NoDigesterError,
+    OperatingCost,
+    UnknownPortError,
+    UnknownUnitError,
+    WiringError,
+    characterize_influent,
+    fractionate,
+    mass_balance,
+    read_influent_csv,
+    MixerUnit,
+    Plant,
+    PlantCheck,
+    PlantSolution,
+    HeatBalanceTemperature,
+    RatioSplitter,
+    SetpointSplitter,
+    SludgeMetrics,
+    StateTranslator,
+    Stream,
+    TemperatureModel,
+    ThresholdSplitter,
+    UVUnit,
+    Unit,
+    aeration_energy,
+    blower_airflow_total,
+    blower_energy,
+    blower_power_kw,
+    carbon_footprint,
+    carbon_mass,
+    co2e_from_energy,
+    ct_log_removal,
+    ct_value,
+    design_summary,
+    derived_BOD,
+    derived_COD,
+    derived_TKN,
+    derived_TSS,
+    direct_n2o_emission,
+    effluent_averages,
+    effluent_quality_index,
+    evaluate_bsm1,
+    evaluate_bsm2,
+    heating_energy,
+    methane_to_co2e,
+    mixing_energy,
+    n2o_n_to_co2e,
+    operating_cost,
+    operational_cost_index,
+    operational_cost_index_bsm2,
+    pumping_energy,
+    pumping_energy_bsm2,
+    required_airflow,
+    size_activated_sludge,
+    sludge_metrics,
+    stripped_n2o,
+    t10_from_baffling,
+    t10_from_rtd,
+    time_average,
+    uv_dose,
+    uv_log_inactivation,
+)
+from aquakin.schema.loader import (
+    clear_model_cache,
+    load_model,
+    load_model_from_file,
+)
+from aquakin.schema.model_spec import SpeciationUnitsWarning
+
+# The utility helpers are aggregated by aquakin.utils (the single utility
+# surface); the top-level namespace re-exports the curated model-analysis subset
+# from there rather than reaching into utils.balance / .composition / .units.
+from aquakin.utils import (
+    UnitWarning,
+    canonical_content,
+    check_conservation,
+    check_model_units,
+    check_nitrogen,
+    composition_table,
+    parse_units,
+)
+
+__all__ = [
+    "ActivatedSludgeSizing",
+    "Aeration",
+    "AerationDesignPoint",
+    "AerationSystem",
+    "AlgebraicTemperature",
+    "BSM1Evaluation",
+    "BSM2Evaluation",
+    "BatchReactor",
+    "BatchSolution",
+    "BiofilmReactor",
+    "BiofilmSolution",
+    "CFDReactor",
+    "CSTRUnit",
+    "CalibrationResult",
+    "CarbonFootprint",
+    "ChlorineContactUnit",
+    "CompiledModel",
+    "ComponentBalance",
+    "Constraint",
+    "CostFactors",
+    "DGSMResult",
+    "DifferentiationConfig",
+    "Event",
+    "EventedResult",
+    "FitResult",
+    "ForwardSensitivityResult",
+    "FreeICConfig",
+    "HeatBalanceTemperature",
+    "IdentityTranslator",
+    "InfluentFractions",
+    "InfluentSeries",
+    "IntegratorConfig",
+    "KPIComparison",
+    "LaplaceConfig",
+    "MassBalance",
+    "MixerUnit",
+    "MonteCarloResult",
+    "NoDigesterError",
+    "OperatingConditions",
+    "OperatingCost",
+    "OptimizeResult",
+    "OptimizerConfig",
+    "PFRSolution",
+    "ParticleTrackReactor",
+    "Plant",
+    "PlantCheck",
+    "PlantSolution",
+    "PlugFlowReactor",
+    "PredictiveBand",
+    "ProfileResult",
+    "RatioSplitter",
+    "ScenarioComparison",
+    "SensitivityResult",
+    "SetpointSplitter",
+    "SludgeMetrics",
+    "SpatialConditions",
+    "SpeciationUnitsWarning",
+    "StateTranslator",
+    "Stream",
+    "TemperatureModel",
+    "ThresholdSplitter",
+    "Track",
+    "TrackSolution",
+    "UVUnit",
+    "Unit",
+    "UnitWarning",
+    "UnknownPortError",
+    "UnknownUnitError",
+    "WiringError",
+    "aeration_energy",
+    "blower_airflow_total",
+    "blower_energy",
+    "blower_power_kw",
+    "calibrate",
+    "canonical_content",
+    "carbon_footprint",
+    "carbon_mass",
+    "characterize_influent",
+    "check_conservation",
+    "check_finite_gradient",
+    "check_model_units",
+    "check_nitrogen",
+    "clear_model_cache",
+    "co2e_from_energy",
+    "compare_scenarios",
+    "compile_model",
+    "composition_table",
+    "ct_log_removal",
+    "ct_value",
+    "derived_BOD",
+    "derived_COD",
+    "derived_TKN",
+    "derived_TSS",
+    "design_summary",
+    "dgsm",
+    "direct_n2o_emission",
+    "effluent_averages",
+    "effluent_quality_index",
+    "esdirk_adjoint_solve",
+    "evaluate_bsm1",
+    "evaluate_bsm2",
+    "fit",
+    "forward_adjoint",
+    "forward_sensitivity",
+    "fractionate",
+    "heating_energy",
+    "implicit_euler_adjoint_solve",
+    "integrate_ensemble",
+    "kpi_comparison",
+    "load_model",
+    "load_model_from_file",
+    "mass_balance",
+    "methane_to_co2e",
+    "mixing_energy",
+    "monte_carlo",
+    "n2o_n_to_co2e",
+    "operating_cost",
+    "operational_cost_index",
+    "operational_cost_index_bsm2",
+    "optimize_design",
+    "parse_rate_expression",
+    "parse_units",
+    "profile_likelihood",
+    "pumping_energy",
+    "pumping_energy_bsm2",
+    "read_influent_csv",
+    "required_airflow",
+    "sensitivity",
+    "size_activated_sludge",
+    "sludge_metrics",
+    "solve_with_events",
+    "stripped_n2o",
+    "t10_from_baffling",
+    "t10_from_rtd",
+    "time_average",
+    "uv_dose",
+    "uv_log_inactivation",
+]
+
+__version__ = "0.1.0"

aquakin/core/__init__.py

@@ -0,0 +1,7 @@
+"""Core data structures and AST machinery (no Pydantic dependency)."""
+
+from aquakin.core.conditions import SpatialConditions
+from aquakin.core.context import CompileContext
+from aquakin.core.model import CompiledModel, compile_model
+
+__all__ = ["CompileContext", "CompiledModel", "SpatialConditions", "compile_model"]

aquakin/core/conditions.py

@@ -0,0 +1,162 @@
+"""Spatially varying condition fields (pH, temperature, etc.)."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+
+import jax.numpy as jnp
+
+
+@dataclass
+class SpatialConditions:
+    """
+    Container for spatially varying condition fields.
+
+    Each entry in ``fields`` is a JAX array of shape ``(n_locations,)`` giving
+    the value of that field at each spatial location. Reactors index into
+    these arrays with ``loc_idx`` at runtime.
+
+    Attributes
+    ----------
+    fields : dict[str, jnp.ndarray]
+        Mapping from field name (e.g. ``"pH"``) to a 1-D JAX array.
+    """
+
+    fields: dict[str, jnp.ndarray] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        normalised: dict[str, jnp.ndarray] = {}
+        n_locations: int | None = None
+        for name, value in self.fields.items():
+            arr = jnp.asarray(value)
+            if arr.ndim == 0:
+                arr = arr[None]
+            elif arr.ndim != 1:
+                raise ValueError(f"Condition field '{name}' must be 1-D, got shape {arr.shape}")
+            if n_locations is None:
+                n_locations = int(arr.shape[0])
+            elif int(arr.shape[0]) != n_locations:
+                raise ValueError(
+                    f"Condition field '{name}' has length {arr.shape[0]}, expected {n_locations}"
+                )
+            normalised[name] = arr
+        self.fields = normalised
+
+    @property
+    def n_locations(self) -> int:
+        """Number of spatial locations represented."""
+        if not self.fields:
+            return 0
+        return int(next(iter(self.fields.values())).shape[0])
+
+    @classmethod
+    def uniform(cls, n_locations: int = 1, **kwargs: float) -> SpatialConditions:
+        """
+        Build a spatially homogeneous ``SpatialConditions`` object.
+
+        Parameters
+        ----------
+        n_locations : int, optional
+            Number of spatial locations (default 1). Must be at least 1; the
+            default suits a 0-D batch reactor, so a batch user writes
+            ``SpatialConditions.uniform(pH=7.5, T=293.15)``. For the 0-D case the
+            :class:`OperatingConditions` alias reads more naturally still.
+        **kwargs : float
+            Field name -> scalar value. The scalar is broadcast to
+            ``(n_locations,)``.
+
+        Returns
+        -------
+        SpatialConditions
+        """
+        if n_locations < 1:
+            raise ValueError(f"n_locations must be >= 1, got {n_locations}")
+        # Broadcast (don't ``float()``-coerce) so a *traced* condition value -- a
+        # user differentiating a solve w.r.t. pH/T -- flows through instead of
+        # raising a ConcretizationTypeError. Identity for a concrete value.
+        fields = {
+            name: jnp.broadcast_to(jnp.asarray(value), (n_locations,))
+            for name, value in kwargs.items()
+        }
+        return cls(fields=fields)
+
+    def with_(self, **kwargs: float) -> SpatialConditions:
+        """Return a copy with some fields overridden (or added).
+
+        The common edit-from-defaults pattern: start from
+        ``model.default_conditions()`` and change only what differs, e.g.::
+
+            conditions = model.default_conditions().with_(T=283.15)   # cold
+
+        Scalar overrides are broadcast to this object's location count (so the
+        result keeps the same ``n_locations``); a length-``n_locations`` array
+        override is used as-is. Fields not named are carried over unchanged. The
+        original is not modified.
+
+        Parameters
+        ----------
+        **kwargs : float or array-like
+            Field name -> new value (scalar, broadcast to ``n_locations``, or a
+            length-``n_locations`` array).
+
+        Returns
+        -------
+        SpatialConditions
+            A new ``SpatialConditions`` (always the base type, so it stays valid
+            for every reactor) with the merged fields.
+        """
+        n = self.n_locations or 1
+        merged = dict(self.fields)
+        for name, value in kwargs.items():
+            arr = jnp.asarray(value)
+            # Broadcast a scalar rather than ``float()``-coercing it, so a traced
+            # override (a gradient w.r.t. a condition) is not concretized.
+            merged[name] = jnp.broadcast_to(arr, (n,)) if arr.ndim == 0 else arr
+        return SpatialConditions(fields=merged)
+
+    def validate_required(self, required: Iterable[str]) -> None:
+        """
+        Raise ``ValueError`` if any required field is missing.
+
+        Parameters
+        ----------
+        required : iterable of str
+            Field names that must be present.
+        """
+        missing = sorted(set(required) - set(self.fields.keys()))
+        if missing:
+            raise ValueError(
+                f"SpatialConditions is missing required fields: {missing}. "
+                f"Provided: {sorted(self.fields.keys())}"
+            )
+
+
+class OperatingConditions(SpatialConditions):
+    """Operating conditions for the 0-D (single stirred tank) case.
+
+    A single stirred tank has no spatial extent, so the spatially-varying
+    :class:`SpatialConditions` (with its ``n_locations`` array model) reads as
+    over-machinery for the most basic setup. ``OperatingConditions`` is the same
+    object specialised to one location, constructed directly from scalar field
+    values::
+
+        conditions = aquakin.OperatingConditions(pH=7.5, T=293.15)
+
+    It **is** a :class:`SpatialConditions` (one location), so it works unchanged
+    in every reactor; for a spatially varying PFR/CFD case use
+    :class:`SpatialConditions` (or :meth:`SpatialConditions.uniform`) directly.
+    To start from a model's declared defaults instead, use
+    ``model.default_conditions()`` and :meth:`SpatialConditions.with_` to edit.
+
+    Parameters
+    ----------
+    **kwargs : float
+        Condition field name (e.g. ``pH``, ``T``) -> scalar value.
+    """
+
+    def __init__(self, **kwargs: float) -> None:
+        # ``jnp.asarray(value)`` (not ``float(value)``) so a traced condition --
+        # a gradient w.r.t. an operating condition -- flows through rather than
+        # raising a ConcretizationTypeError. Identity for a concrete value.
+        super().__init__(fields={name: jnp.asarray(value) for name, value in kwargs.items()})

aquakin/core/context.py

@@ -0,0 +1,31 @@
+"""Compile-time context used by ASTNode.compile()."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class CompileContext:
+    """
+    Index maps used during AST compilation.
+
+    Attributes
+    ----------
+    species_index : dict[str, int]
+        Maps species name to its position in the concentration vector ``C``.
+    param_index : dict[str, int]
+        Maps fully-namespaced parameter name (e.g. ``"O3_Br_direct.k1"``) to
+        its position in the flat ``params`` vector.
+    condition_fields : frozenset[str]
+        Set of valid condition field names. Used to validate
+        ``ConditionNode`` references at parse/compile time.
+    reaction_name : str
+        The reaction that owns the expression currently being compiled. Used
+        for namespacing local parameter names.
+    """
+
+    species_index: dict[str, int]
+    param_index: dict[str, int]
+    condition_fields: frozenset[str]
+    reaction_name: str = ""

aquakin/core/hints.py

@@ -0,0 +1,32 @@
+"""Close-match suggestion for unknown-name errors.
+
+A single home for the "Did you mean: ...?" hint appended to ``KeyError`` /
+``ValueError`` messages when a species / parameter / unit / stream name is not
+found. Callers keep their own message prefix and exception type; this only
+formats the suffix, so the hint wording (and its similarity threshold) is tuned
+in one place.
+"""
+
+import difflib
+
+
+def did_you_mean(name, choices, n: int = 3) -> str:
+    """Return a ``" Did you mean: a, b?"`` suffix, or ``""`` if nothing is close.
+
+    Parameters
+    ----------
+    name : str
+        The unknown name the user supplied.
+    choices : iterable of str
+        The valid names to match against.
+    n : int, optional
+        Maximum number of suggestions (default 3).
+
+    Returns
+    -------
+    str
+        A leading-space suffix ready to append to an error message, empty when
+        no close match is found.
+    """
+    matches = difflib.get_close_matches(name, list(choices), n=n)
+    return f" Did you mean: {', '.join(matches)}?" if matches else ""