"""Plant result objects.
The records :class:`~aquakin.plant.plant.Plant` *returns*, split out of
``plant.py`` so the plant module is behaviour and these are the plain data it
hands back:
- :class:`PlantCheck` -- a pre-solve wiring report (:meth:`Plant.check`).
- :class:`PlantSolution` -- the result of :meth:`Plant.solve` (carries the plant
for its accessor methods).
- :class:`SteadyStateResult` -- the result of :meth:`Plant.steady_state` /
:meth:`Plant.run_to_steady_state`.
They are re-exported from ``aquakin.plant.plant`` for backward compatibility, so
``from aquakin.plant.plant import PlantSolution`` keeps working.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
import jax.numpy as jnp
from aquakin.core.hints import did_you_mean
from aquakin.plant.errors import UnknownUnitError
if TYPE_CHECKING:
from aquakin.plant.plant import Plant
[docs]
@dataclass
class PlantCheck:
"""Result of :meth:`Plant.check` -- a pre-solve wiring report.
Attributes
----------
unfed_ports : list[str]
``"unit.port"`` input ports with no stream wired into them. These are
**errors**: the RHS sweep has no source for them, so the solve fails.
dangling_outputs : list[str]
``"unit.port"`` output ports consumed by no connection. Usually benign
-- a terminal stream that leaves the plant (final effluent, wasted
sludge, disposal cake, biogas) -- so they are reported for information,
not flagged as errors.
recycles : list[str]
``"unit.port"`` recycle (back-edge) source ports the topological sort
detected, for visibility into how the graph was cut.
"""
unfed_ports: list = field(default_factory=list)
dangling_outputs: list = field(default_factory=list)
recycles: list = field(default_factory=list)
@property
def ok(self) -> bool:
"""True when no input port is unfed (a dangling output is not an error)."""
return not self.unfed_ports
[docs]
def summary(self) -> str:
"""A short human-readable report."""
lines = [f"Plant check: {'OK' if self.ok else 'PROBLEMS'}"]
if self.unfed_ports:
lines.append(f" unfed input ports (errors): {self.unfed_ports}")
if self.dangling_outputs:
lines.append(f" unconsumed outputs (info): {self.dangling_outputs}")
if self.recycles:
lines.append(f" recycle edges: {self.recycles}")
return "\n".join(lines)
[docs]
@dataclass
class PlantSolution:
"""Result of :meth:`Plant.solve`.
Holds the integrated unit **states** (:attr:`state`), not the inter-unit
streams: the plant integrates states only, so an effluent / recycle stream is
*reconstructed on demand* from the saved states via :meth:`stream` (the whole
sweep is reconstructed once and cached on the solution). Per-unit
concentrations are read directly with :meth:`C_named` / :meth:`unit_state`.
Attributes
----------
t : jnp.ndarray
Save times, shape ``(n_t,)``.
state : jnp.ndarray
Full plant state, shape ``(n_t, total_state_size)``. This is the raw
integrated state. Where a unit's model sets ``clip_negative_states``
(on by default for ASM1, so the activated-sludge reactors in BSM1/BSM2),
entries may be **small transient negatives** -- the ``max(x, 0)`` clamp
is applied only when evaluating that unit's reaction rates, not to the
saved state. A normal numerical transient, not a solver/model error;
clip with ``jnp.maximum(state, 0.0)`` for display if needed.
plant : Plant
The plant that produced this solution; retained for accessor
methods.
events_log : list of (float, str), optional
When the solve used ``events=``, the fired events in order as
``(time, name)``. ``None`` for a plain solve.
_requested_time_unit : str, optional
The unit :attr:`t` was rescaled into when :meth:`Plant.solve` was called
with an explicit ``time_unit=``, or ``None`` for the plant's native unit.
Backs the :attr:`time_unit` property; a declared field (rather than an
attribute set after construction) so it participates in ``repr``/``eq``.
"""
t: jnp.ndarray
state: jnp.ndarray
plant: Plant
events_log: list | None = None
_requested_time_unit: str | None = None
@property
def time_unit(self) -> str | None:
"""The time unit of :attr:`t` (``"s"``, ``"d"``, ... or ``None``).
The plant's native unit (:attr:`Plant.time_unit`), or the ``time_unit=``
passed to :meth:`Plant.solve` when the times were reported in that unit.
"""
override = self._requested_time_unit
return override if override is not None else self.plant.time_unit
@property
def final_state(self) -> jnp.ndarray:
"""The full plant state at the last save time, shape ``(total_state_size,)``.
The last row of :attr:`state`. Pass it to :meth:`Plant.states_by_unit`
to read per-unit pieces without a trailing ``[-1]`` index on a 2-D
trajectory.
"""
return self.state[-1]
[docs]
def unit_state(self, unit_name: str) -> jnp.ndarray:
"""Return the trajectory of one unit's state, shape
``(n_t, unit.state_size)``. See ``plant.list_units()`` for the names."""
layout = self.plant._state_layout
if unit_name not in layout:
suffix = did_you_mean(unit_name, list(layout))
raise UnknownUnitError(
f"Unknown unit '{unit_name}'. Units (see plant.list_units()): "
f"{list(layout)}.{suffix}"
)
start, size = layout[unit_name]
return self.state[:, start : start + size]
[docs]
def available_streams(self) -> list[str]:
"""The ``"unit.port"`` output endpoints :meth:`Plant.stream` accepts
(a convenience alias for ``plant.list_ports()``)."""
return self.plant.list_ports()
[docs]
def stream(self, endpoint: str, params: jnp.ndarray | None = None):
"""Reconstruct one output stream's trajectory -- ``sol.stream("effluent")``.
A convenience for ``plant.stream(sol, endpoint, params)`` (the plant is
carried on the solution). **Inter-unit streams are not stored** -- the
plant integrates unit states, so the effluent/recycle streams are
recomputed from the saved states; the whole sweep is reconstructed once
and cached on this solution, so repeated calls for different ports are
cheap. ``endpoint`` is a semantic name (``"effluent"``, ``"ras"``, ...) or
a ``"unit.port"`` (see :meth:`available_streams` / ``plant.list_streams``).
"""
return self.plant.stream(self, endpoint, params)
[docs]
def C_named(self, unit_name: str, species: str) -> jnp.ndarray:
"""Return a single species' trajectory in a single unit.
Only valid for units whose state is a concentration vector in the unit's
model (CSTRs and other kinetic units). See ``plant.list_units()`` /
``plant.list_species(unit)`` for the valid names.
"""
unit = self.plant._unit_or_raise(unit_name)
if not self.plant._is_concentration_unit(unit):
indexable = [
n
for n in self.plant.list_units()
if self.plant._is_concentration_unit(self.plant.units[n])
]
raise KeyError(
f"Unit '{unit_name}' state is not a concentration vector, so it "
f"cannot be indexed by species (read it as a stream with "
f"plant.stream(sol, '{unit_name}.<port>') instead). "
f"Concentration units: {indexable}."
)
if species not in unit.model.species_index:
suffix = did_you_mean(species, unit.model.species)
raise KeyError(
f"Unknown species '{species}' in unit '{unit_name}'. Species "
f"(see plant.list_species('{unit_name}')): "
f"{list(unit.model.species)}.{suffix}"
)
idx = unit.model.species_index[species]
return self.unit_state(unit_name)[:, idx]
[docs]
def C_named_many(self, unit_name: str, species) -> dict[str, jnp.ndarray]:
"""Trajectories of several species in one unit, as ``{name: array}``.
The multi-species companion to :meth:`C_named` -- read a handful of a
unit's species in one call (``sol.C_named_many("tank5", ["SNH", "SNO"])``)
instead of a slice each. An unknown unit/species raises the same hinted
``KeyError`` :meth:`C_named` gives.
"""
return {sp: self.C_named(unit_name, sp) for sp in species}
[docs]
def final_named(self, unit_name: str, species=None) -> dict[str, float]:
"""Values at the **last** save time for one unit, as ``{name: float}``.
The reporting shortcut for a steady-state value: instead of
``float(sol.C_named("tank5", "SNH")[-1])`` per species,
``sol.final_named("tank5", ["SNH", "SNO"])`` returns them in one dict.
With ``species=None`` (default) every species of the unit's model is
returned (see :meth:`Plant.list_species`). Values are plain Python floats
(a post-processing read on an already-solved solution -- use
``C_named(unit, sp)[-1]`` if you need a differentiable last value).
"""
names = self.plant.list_species(unit_name) if species is None else list(species)
return {sp: float(self.C_named(unit_name, sp)[-1]) for sp in names}
[docs]
def plot(self, unit_name: str, species=None, *, ax=None, **kwargs):
"""Plot one unit's species trajectories over time.
The plant analogue of ``BatchSolution.plot``: a thin matplotlib wrapper
so ``sol.plot("tank5", "SNH")`` needs no manual ``C_named`` / unit /
axis-label boilerplate. The x-axis is labelled with the plant's time
unit; a single-species plot labels the y-axis with that species' units.
Parameters
----------
unit_name : str
A concentration-vector unit (see :meth:`Plant.list_species`).
species : str or iterable of str, optional
Species to plot; a single name, an iterable (legended), or ``None``
for every species of the unit's model.
ax : matplotlib.axes.Axes, optional
Axes to draw on; a new one is created if omitted.
**kwargs
Forwarded to ``ax.plot``.
Returns
-------
matplotlib.axes.Axes
Raises
------
ImportError
If matplotlib is not installed (``pip install aquakin[plot]``).
KeyError
For an unknown unit / species, or a non-concentration unit (hinted).
"""
import numpy as np
from aquakin.integrate._common import require_matplotlib
plt = require_matplotlib()
names = (
self.plant.list_species(unit_name)
if species is None
else [species]
if isinstance(species, str)
else list(species)
)
if ax is None:
_, ax = plt.subplots()
t = np.asarray(self.t)
net = self.plant.units[unit_name].model
for sp in names:
ax.plot(t, np.asarray(self.C_named(unit_name, sp)), label=sp, **kwargs)
unit = self.time_unit
ax.set_xlabel(f"time [{unit}]" if unit else "time")
if len(names) == 1:
ax.set_ylabel(f"{names[0]} [{net.units_of(names[0])}]")
else:
ax.set_ylabel(f"{unit_name} concentration")
ax.legend()
return ax
[docs]
def to_dataframe(self, unit: str, *, units_in_columns: bool = False):
"""Return one unit's state trajectory as a pandas ``DataFrame``.
The plant integrates a heterogeneous flat state across many units, so a
single whole-plant table is not meaningful; pick one unit. For a
kinetic unit (one with a ``model``) the columns are species names; for
any other unit they are generic ``state_0..state_{n-1}`` columns.
Parameters
----------
unit : str
Name of the unit whose trajectory to tabulate.
units_in_columns : bool, optional
If ``True``, append ``" [unit]"`` to each species column label
(kinetic units only); otherwise units are stored in
``df.attrs["units"]``.
Returns
-------
pandas.DataFrame
Time-indexed (``t``), one row per save time.
Raises
------
KeyError
If ``unit`` is not a unit of this plant.
ImportError
If pandas (an optional dependency) is not installed.
"""
from aquakin.integrate._common import build_dataframe
if unit not in self.plant.units:
raise UnknownUnitError(f"Unknown unit '{unit}'. Available: {list(self.plant.units)}")
sub = self.unit_state(unit) # (n_t, unit.state_size)
unit_obj = self.plant.units[unit]
if hasattr(unit_obj, "model"):
net = unit_obj.model
columns = [(sp, sub[:, j]) for j, sp in enumerate(net.species)]
units = {sp: net.units_of(sp) for sp in net.species}
else:
columns = [(f"state_{j}", sub[:, j]) for j in range(sub.shape[1])]
units = {}
return build_dataframe(
self.t,
columns,
index_name="t",
units=units,
units_in_columns=units_in_columns,
)
[docs]
def to_csv(self, path_or_buf=None, *, unit: str, units_in_columns: bool = True, **kwargs):
"""Write one unit's trajectory to CSV (delegates to :meth:`to_dataframe`).
``unit`` is required (the plant state is per-unit). ``units_in_columns``
defaults to ``True`` so the file is self-describing. Extra keyword
arguments are forwarded to ``pandas.DataFrame.to_csv``.
"""
return self.to_dataframe(unit, units_in_columns=units_in_columns).to_csv(
path_or_buf, **kwargs
)
@dataclass
class SteadyStateResult:
"""Result of :meth:`Plant.run_to_steady_state` or :meth:`Plant.steady_state`.
Attributes
----------
state : jnp.ndarray
The steady-state (operating-point) state vector, shape
``(total_state_size,)``. Pass it to :meth:`Plant.states_by_unit` to read
per-unit values, or as ``solve(y0=...)`` to start a dynamic run. From
:meth:`steady_state` it carries the implicit-function-theorem parameter
gradient (``jax.grad`` of a loss on it flows to the plant parameters).
converged : bool
``True`` if steady state was reached -- the dynamics died out before the
``max_time`` cap (forward) or the residual fell below ``tol`` within
``max_iter`` (algebraic PTC).
time : float or None
The time (plant units) at which a *forward* solve settled, or
``max_time`` if it did not; ``None`` for the algebraic solve.
solution : PlantSolution or None
The underlying forward solve (terminating at ``time``); ``None`` for the
algebraic solve.
method : str
How the steady state was found: ``"forward"`` (integrate-to-steady),
``"ptc"`` (pseudo-transient continuation), or ``"ptc->forward"`` (PTC did
not converge and the forward fallback was used).
iterations : int or None
PTC iteration count (``None`` for the forward method).
residual : float or None
The final scaled steady-state residual ``max_i |dy_i/dt| /
max(|y_i|, scale_floor)`` for the algebraic solve (``None`` for forward).
"""
state: jnp.ndarray
converged: bool
time: float | None = None
solution: PlantSolution | None = None
method: str = "forward"
iterations: int | None = None
residual: float | None = None
# Whether an *operating-branch* steady state exists at these parameters.
# ``True`` for a converged solve; ``False`` when pseudo-arclength continuation
# found the operating branch folds before the target (a saddle-node
# bifurcation -- the operating point is past the survival limit, e.g. digester
# washout, so only a different branch exists). ``None`` when not determined
# (the arclength layer was not invoked). A screen should *exclude* a sample
# whose operating point does not exist (``method == "past_fold"``).
operating_point_exists: bool | None = None