"""The public prepared-simulation object.
``gkx.prepare`` used to return the solver's internal
``PreparedExplicitNonlinearDiagnostics`` and to refuse linear cases outright,
so the reusable-execution concept existed only for half the product and only
as a private type. :class:`PreparedSimulation` is the public object the API
contract describes: it accepts linear and nonlinear cases, carries the static
topology and the compilation and cache metadata that decide whether a second
call is cheap, and exposes typed ``solve``, ``scan``, ``value_and_grad``,
``warmup``, ``estimate_memory`` and ``summary``.
It owns no numerics. Every method delegates to the runtime path the case would
have taken anyway, so preparing a case and solving it cannot disagree with
solving it directly.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Callable, Sequence, cast
import numpy as np
from gkx.core_ky_layout import FULL, KyLayout, rows_for_layout
def _case_kind(case: Any) -> str:
"""Return ``"nonlinear"`` or ``"linear"`` for a validated case."""
return "nonlinear" if bool(case.physics.nonlinear) else "linear"
def _state_shape(case: Any, *, n_laguerre: int, n_hermite: int) -> tuple[int, ...]:
"""Return the distribution-array shape this case will allocate."""
grid = case.grid
n_species = sum(1 for s in case.species if getattr(s, "kinetic", False))
# How many ky rows the run allocates, which is a question about the
# layout and not about the resolution (gkx.core_ky_layout, plan 5.3 N3).
# A two-sided state holds Ny rows and rebuilds the negative half by the
# reality condition; a half-spectrum state holds Nyc = 1 + Ny // 2 and
# does not. What this line must never do is report a layout the run did
# not build: reporting Nyc for a two-sided run halves the array it
# advertises and costs, and reporting Ny for a half-spectrum run doubles
# it. Reading the same key the grid was built from is what keeps the two
# in step.
n_ky = rows_for_layout(
int(grid.Ny), cast(KyLayout, str(getattr(grid, "ky_layout", FULL)).lower())
)
return (
max(n_species, 1),
int(n_laguerre),
int(n_hermite),
n_ky,
int(grid.Nx),
int(grid.Nz),
)
[docs]
@dataclass(frozen=True)
class PreparedSimulation:
"""A reusable, compiled simulation for one case.
Construct with :func:`gkx.prepare`. The object is frozen: deriving a
different topology means preparing a different case, which is what keeps
the compiled callables and the reported metadata describing the same run.
"""
case: Any
kind: str
state_shape: tuple[int, ...]
n_laguerre: int
n_hermite: int
options: dict[str, Any] = field(default_factory=dict)
_backend: Any = None
_warmed: list[bool] = field(default_factory=lambda: [False])
# -- execution -------------------------------------------------------
[docs]
def solve(self, parameters: Any = None, initial_state: Any = None) -> Any:
"""Run this simulation and return its typed result.
``parameters`` is reserved for the differentiable path and must be
``None`` here; a prepared object must not silently solve a different
physics problem than the one it reports.
A nonlinear ``initial_state`` on a linked (twist-shift) deck is
projected onto the linked chain cover before it is integrated, which is
the same rule the runtime applies at intake (queue rows Q19 and Q23):
the rows outside the chains are decoupled, and the ExB bracket would
alias whatever sits there back onto the chain rows. Periodic and
full-cover decks are untouched.
"""
if parameters is not None:
raise NotImplementedError(
"solve(parameters=...) requires the differentiable prepared path; "
"use value_and_grad for parameter sensitivities"
)
if self.kind == "nonlinear" and self._backend is not None:
if initial_state is None:
return self._backend.run()
return self._backend.run(initial_state)
from gkx import runtime as _runtime
if initial_state is not None:
raise NotImplementedError(
"linear prepared solves do not accept an initial state yet"
)
return _runtime.solve(self.case, **self.options)
[docs]
def scan(
self,
parameter: str,
values: Sequence[float],
*,
parallel: str = "auto",
) -> Any:
"""Solve this case once per value of a named scalar parameter."""
from gkx import runtime as _runtime
if parameter != "ky":
raise NotImplementedError(
f"prepared scan supports parameter='ky'; got {parameter!r}"
)
workers = 1 if parallel in {"auto", "serial", "off"} else int(parallel)
return _runtime.run_runtime_scan(self.case, list(values), workers=workers)
[docs]
def value_and_grad(self, objective: Callable[..., Any], parameters: Any) -> Any:
"""Return ``(value, gradient)`` of ``objective`` at ``parameters``."""
import jax
return jax.value_and_grad(objective)(parameters)
# -- introspection ---------------------------------------------------
[docs]
def warmup(self) -> "PreparedSimulation":
"""Force compilation now, so a later timed call measures execution.
Preparing a nonlinear case builds its scan closure but does not compile
it: XLA compiles when the scan first executes. Measured on the shipped
five-step KBM deck, ``prepare`` took 4.3 s, the first ``solve`` 19.6 s
and the second 0.000 s. This method used to return ``self`` without
doing anything, so code that called it and then timed ``solve`` was
timing the compile.
Forcing compilation means executing the prepared closure once, because
the scan exposes no separately lowerable handle. That costs one run --
cheap next to the compile it triggers, but not free, so it happens only
when asked for. A second call is a no-op.
A linear case has nothing to compile ahead of time: the linear runtime
chooses its solver per call, so warming one would compile a solver the
next call may not use. ``summary()["warmed"]`` reports which happened.
Returns ``self`` so a caller can chain.
"""
if self.kind == "nonlinear" and self._backend is not None:
if not self._warmed[0]:
self._backend.run()
self._warmed[0] = True
return self
[docs]
def estimate_memory(self) -> dict[str, Any]:
"""Estimate the resident bytes this simulation's arrays require.
The estimate covers the distribution array and the working copies an
explicit stage holds live at once. It is a floor, not a ceiling: it
does not model the diagnostic buffers a long sampled run accumulates,
and it says so rather than pretending to a precision it lacks.
"""
elements = int(np.prod(self.state_shape))
itemsize = 8 if self.precision == "float64" else 4
# A complex spectral state, and the stages an explicit step holds live.
state_bytes = elements * itemsize * 2
stages = 4 if str(self.case.time.method).lower() in {"rk4", "rk45"} else 3
return {
"state_shape": self.state_shape,
"elements": elements,
"state_bytes": state_bytes,
"working_set_bytes": state_bytes * stages,
"precision": self.precision,
"is_floor_not_ceiling": True,
}
@property
def precision(self) -> str:
"""Return the JAX floating precision this simulation will run in."""
import jax
return "float64" if jax.config.read("jax_enable_x64") else "float32"
@property
def devices(self) -> tuple[str, ...]:
"""Return the visible JAX devices as their platform strings."""
import jax
return tuple(f"{d.platform}:{d.id}" for d in jax.devices())
[docs]
def summary(self) -> dict[str, Any]:
"""Return the topology, precision, and cache facts that identify this."""
memory = self.estimate_memory()
return {
"kind": self.kind,
"state_shape": self.state_shape,
"n_laguerre": self.n_laguerre,
"n_hermite": self.n_hermite,
"geometry_model": self.case.geometry.model,
"precision": self.precision,
"devices": self.devices,
"working_set_bytes": memory["working_set_bytes"],
**self.compilation_metadata(),
}
[docs]
def print_summary(self, *, stream: Any = None) -> None:
"""Print the summary one field per line."""
for key, value in self.summary().items():
print(f"{key}: {value}", file=stream)
def _runtime_resolution_defaults() -> dict[str, tuple[int, int]]:
"""Mirror the runtime's own per-kind velocity fallback, from its owner.
The runtime resolves ``[run]`` Nl/Nm per call and its fallback differs by
kind: a linear call takes ``startup._RUNTIME_LINEAR_HL_FALLBACK``, a
nonlinear one (4, 8). See ``_CASE_LINEAR_SPECS`` and
``_CASE_NONLINEAR_SPECS`` in ``gkx.workflows.runtime.commands``. This reads
the linear pair from that owner rather than repeating it, because a summary
that disagrees with the runtime describes a calculation that will not run.
"""
from gkx.workflows.runtime.startup import (
_RUNTIME_LINEAR_HL_FALLBACK,
_RUNTIME_NONLINEAR_HL_FALLBACK,
)
return {
"linear": _RUNTIME_LINEAR_HL_FALLBACK,
"nonlinear": _RUNTIME_NONLINEAR_HL_FALLBACK,
}
def _resolve_velocity_resolution(
case: Any, kind: str, options: dict[str, Any]
) -> tuple[int, int]:
"""Return the ``(Nl, Nm)`` a later ``solve`` will actually build.
Precedence is explicit argument, then the deck's ``[run]`` table, then the
runtime's own default for this kind of case. Reporting anything else makes
``print_summary`` describe a calculation that will not be run: before the
``[run]`` table was carried on the case, preparing the shipped Cyclone deck
reported ``4, 8`` for a deck that asks for ``16, 48``.
"""
default_l, default_m = _runtime_resolution_defaults().get(kind, (4, 8))
run = getattr(case, "run", None)
deck_l = getattr(run, "Nl", None)
deck_m = getattr(run, "Nm", None)
n_laguerre = options.get("Nl") or deck_l or default_l
n_hermite = options.get("Nm") or deck_m or default_m
return int(n_laguerre), int(n_hermite)
def prepare_simulation(case: Any, **options: Any) -> PreparedSimulation:
"""Build a :class:`PreparedSimulation` for ``case``.
A nonlinear case builds its scan closure here, which is what makes a later
``solve`` cheap once that closure has been compiled -- XLA does that when
the scan first executes, so call ``warmup()`` if a later ``solve`` is being
timed. A linear case is validated and described but not built, because the
linear runtime chooses its solver per call. ``summary()`` reports both
facts rather than hiding them.
"""
case.validate()
kind = _case_kind(case)
n_laguerre, n_hermite = _resolve_velocity_resolution(case, kind, options)
backend = None
if kind == "nonlinear":
from gkx import runtime as _runtime
backend = _runtime.prepare(case, **options)
return PreparedSimulation(
case=case,
kind=kind,
state_shape=_state_shape(case, n_laguerre=n_laguerre, n_hermite=n_hermite),
n_laguerre=n_laguerre,
n_hermite=n_hermite,
options=dict(options),
_backend=backend,
)