Recipe API

ds.problem(name) opens a ProblemSpec, which collects a recipe's declarations until build(). The Recipe API overview shows the pattern with examples.

Recommended defaults

Most recipes need only domain, field, equation, bc and build. Set each field's scale to its typical magnitude and leave offset at its default. build runs the same static validation as spec.validate().

diff_soup.problem.problem

problem(name: str) -> ProblemSpec

Open a ProblemSpec builder for a new problem.

Parameters:

Name	Type	Description	Default
name	str	Problem name.	required

Returns:

Type	Description
ProblemSpec	The empty builder.

diff_soup.problem.ProblemSpec

Mutable builder that collects a problem's domains, unknowns, equations and conditions, then builds it.

Open one with problem. Fields, globals, parameters and functions share one pool of names. spec.bc, spec.constraint, spec.observe and spec.control hold the remaining declaration methods.

equation property

equation

Decorator that registers a closure as the equation of one field.

Use it as @spec.equation(writes=..., reads=..., name=None), where writes names the field, reads lists every field, global and coupled alias the closure reads, and name defaults to the function's name. The closure returns d(field)/dt, or an expression containing q.dt(field) that equals zero.

equations property

equations

Decorator that registers one closure as the equations of several fields.

Use it as @spec.equations(writes=(...), reads=..., name=None), with at least two fields in writes. The closure returns a dict keyed by field name.

domain

domain(name: str, domain: Any, *, operators: Any = None) -> 'ProblemSpec'

Register a domain under a name.

Parameters:

Name	Type	Description	Default
name	str	Domain name, unique in the recipe.	required
domain	Any	A grid or mesh from diff_soup.domains.	required
operators	Any	Operator-bank class or top-level named function of the domain, or None for the built-in operators.	None

Returns:

Type	Description
'ProblemSpec'	The builder, for chaining.

field

field(name: str, *, mass: str, shape: Optional[tuple[int, ...]] = None, components: Optional[int] = None, bounds: Optional[tuple[Optional[float], Optional[float]]] = None, domain: Optional[str] = None, scale: Optional[float] = None, offset: Optional[float] = None, residual_scale: Optional[float] = None, init: Optional[Callable[..., Any]] = None, quad: Optional[bool] = None) -> FieldSpec

Declare a field, an unknown with values on a domain.

The dimensional value is offset + scale * internal. Omitting scale defaults it to 1.0 with a warning at build time.

Parameters:

Name	Type	Description	Default
name	str	Field name, unique among fields, globals, parameters and functions.	required
mass	str	"differential" for a field with a time derivative, "algebraic" for one without.	required
shape	Optional[tuple[int, ...]]	Array shape of the field, or None for one value per node.	None
components	Optional[int]	Number of values per node of a vector field, exclusive with shape.	None
bounds	Optional[tuple[Optional[float], Optional[float]]]	Reserved, and any value other than None raises NotImplementedError.	None
domain	Optional[str]	Name of the field's domain, required when more than one domain is registered.	None
scale	Optional[float]	State scale, nonzero.	None
offset	Optional[float]	State offset.	None
residual_scale	Optional[float]	Divisor of the field's residual rows, nonzero.	None
init	Optional[Callable[..., Any]]	Default initial condition, a callable taking one coordinate array per spatial dimension.	None
quad	Optional[bool]	On a finite-element domain, whether q.<name> holds quadrature-point values rather than the nodal array, None meaning true there.	None

Returns:

Type	Description
FieldSpec	The declared field, which spec.couple takes as a source.

global_var

global_var(name: str, *, initial: Any = 0.0, mass: str = 'algebraic', scale: Optional[float] = None) -> GlobalVarSpec

Declare a global, a scalar or one-dimensional unknown with no domain.

An algebraic global is closed by exactly one constraint or driven by a control. A differential global is written by exactly one equation.

Parameters:

Name	Type	Description	Default
name	str	Global name, unique among fields, globals, parameters and functions.	required
initial	Any	Initial value, a scalar or one-dimensional array, which seeds the solve.	0.0
mass	str	"algebraic" or "differential".	'algebraic'
scale	Optional[float]	State scale, nonzero, or None for 1.	None

Returns:

Type	Description
GlobalVarSpec	The declared global.

param

param(name: str, *, value: Any, learnable: bool = False, bounds: Optional[tuple[Optional[float], Optional[float]]] = None, prior: Any = None) -> ParamSpec

Declare a parameter, read in closures as q.p.<name>.

Parameters:

Name	Type	Description	Default
name	str	Parameter name, unique among fields, globals, parameters and functions.	required
value	Any	Nominal value, a scalar or an array, never a callable.	required
learnable	bool	Whether a fit may select the parameter, which then needs a floating-point value.	False
bounds	Optional[tuple[Optional[float], Optional[float]]]	Bounds (lo, hi) for a fit, with None for an open side.	None
prior	Any	Prior for MAP estimation.	None

Returns:

Type	Description
ParamSpec	The declared parameter, by which a fit selects it.

geometry

geometry(name: str, *, domain: Optional[str] = None, value: float = 1.0, learnable: bool = False, bounds: Optional[tuple[Optional[float], Optional[float]]] = None) -> ParamSpec

Declare a parameter that sets the size of a one-dimensional grid.

The grid is declared with outer coordinate 1, and its coordinates are multiplied by the parameter at every evaluation. Closures read the size as q.p.<name>, and a learnable size needs a finite positive lower bound.

Parameters:

Name	Type	Description	Default
name	str	Parameter name.	required
domain	Optional[str]	Name of the grid, required when more than one domain is registered.	None
value	float	Size, finite and positive.	1.0
learnable	bool	Whether a fit may select the size.	False
bounds	Optional[tuple[Optional[float], Optional[float]]]	Bounds (lo, hi) for a fit.	None

Returns:

Type	Description
ParamSpec	The declared parameter, by which a fit selects it.

function

function(name: str, fn: Any, *, learnable: bool = False, prior: Any = None, pointwise: bool = False) -> FunctionSpec

Declare a named function, called in closures as q.fn.<name>(...).

Parameters:

Name	Type	Description	Default
name	str	Function name, unique among fields, globals, parameters and functions.	required
fn	Any	The callable, such as a plain function or an eqx.Module.	required
learnable	bool	Whether a fit may select the function's array leaves.	False
prior	Any	Prior for MAP estimation.	None
pointwise	bool	Whether to apply fn at every node, for a callable that does not broadcast.	False

Returns:

Type	Description
FunctionSpec	The declared function, by which a fit selects it.

family

family(name: str, *, count: int, domain: Any, params: Optional[dict[str, Any]] = None, scale: Optional[dict[str, float]] = None, residual_scale: Optional[dict[str, float]] = None, operators: Any = None, stencil_radius: Optional[int] = None)

Declare a population of count identical members that differ only in parameter values.

The returned family declares the member template. Member equations take (q, m), with m.<name> the member's entry of a parameter in params.

Parameters:

Name	Type	Description	Default
name	str	Family name, distinct from every domain and field.	required
count	int	Number of members, positive.	required
domain	Any	Member domain, or a sequence of count domains for a ragged family.	required
params	Optional[dict[str, Any]]	Per-member parameter arrays of shape (count,), or shared scalars.	None
scale	Optional[dict[str, float]]	State scale of each member field, keyed by field name.	None
residual_scale	Optional[dict[str, float]]	Residual divisor of each member field, keyed by field name.	None
operators	Any	Operator-bank class or top-level named function of the domain, or None for the built-in operators.	None
stencil_radius	Optional[int]	Jacobian reach of a member kernel in cells per axis, or None for one cell.	None

Returns:

Type	Description
	The family, whose methods declare the member fields, equations and boundary conditions.

couple

couple(source: Any = None, *, target: Optional[FieldSpec] = None, via: Any = None, onto: Optional[str] = None, at: Any = None, reduce: Optional[str] = None, as_: Optional[str] = None, name: Optional[str] = None) -> 'ProblemSpec'

Couple a field's values into another component or family.

A field coupled onto a family gives member m the value source[at[m]]. A member projection coupled onto a domain places one value per member at the nodes at. With as_, the domain onto reads the coupled value as q.<as_>.

Parameters:

Name	Type	Description	Default
source	Any	A field returned by spec.field, or a member projection such as family.field("c").surface().	None
target	Optional[FieldSpec]	Receiving field returned by spec.field, in the explicit form without onto.	None
via	Any	A TransferOp to apply instead of the default index transfer.	None
onto	Optional[str]	Name of the receiving domain or family.	None
at	Any	Node indices, or a boolean mask, of the nodes the values come from or go to.	None
reduce	Optional[str]	How members landing on one node merge, "sum" (default), "mean", or "none" to refuse a collision.	None
as_	Optional[str]	Name under which the domain onto reads the value as q.<as_>.	None
name	Optional[str]	Coupling label, or None for "couple_<n>".	None

Returns:

Type	Description
'ProblemSpec'	The builder, for chaining.

time_scale

time_scale(value: float) -> 'ProblemSpec'

Set the time scale, so the model integrates on internal time t / value.

Equation closures read internal time, and authored callbacks such as controls and constraint targets read dimensional time. A second call with a different value raises ValueError.

Parameters:

Name	Type	Description	Default
value	float	Time scale, nonzero.	required

Returns:

Type	Description
'ProblemSpec'	The builder, for chaining.

scales

scales(*, state: Optional[dict[str, float]] = None, state_offset: Optional[dict[str, float]] = None, time: Optional[float] = None, integrand: Optional[dict[str, float]] = None) -> 'ProblemSpec'

Set scales centrally by name, overriding the inline declarations.

A central value that differs from an inline one wins, with a warning at build time.

Parameters:

Name	Type	Description	Default
state	Optional[dict[str, float]]	State scale per field, global or member field name, or per (component, field) key.	None
state_offset	Optional[dict[str, float]]	State offset, keyed like state.	None
time	Optional[float]	Time scale.	None
integrand	Optional[dict[str, float]]	Target divisor per constraint name.	None

Returns:

Type	Description
'ProblemSpec'	The builder, for chaining.

equation_scale

equation_scale(field: str, value: float) -> 'ProblemSpec'

Set the divisor of a field's residual rows, overriding field(residual_scale=...).

Parameters:

Name	Type	Description	Default
field	str	Field name.	required
value	float	Divisor, nonzero.	required

Returns:

Type	Description
'ProblemSpec'	The builder, for chaining.

build

build(*, layout=None, precision=None)

Validate the recipe and lower it into a solve-ready problem.

Recipes equal in structure and values return the same problem.

Parameters:

Name	Type	Description	Default
layout		Device placement, a Layout, or None for the default.	None
precision		Precision, or None for the default.	None

Returns:

Type	Description
	The LoweredProblem.

Raises:

Type	Description
ValueError	If the recipe fails static validation, listing every problem found.

to_equation_set

to_equation_set()

Validate the recipe and freeze it into an EquationSet.

Returns:

Type	Description
	The equation set, ready for lower.

Raises:

Type	Description
ValueError	If the recipe fails static validation, listing every problem found.

assemble_state

assemble_state(values, *, default=0.0)

Assemble the flat initial state from the declared initial conditions and overrides.

An entry in values wins over the field's init, and default fills a field with neither. Globals are not included.

Parameters:

Name	Type	Description	Default
values		Mapping from field name, or "<family>.<field>", to a scalar, an array, or a callable of the node coordinates.	required
default		Value of every field with no entry and no init.	0.0

Returns:

Type	Description
	The state vector, component fields first, then family members.

validate

validate(*, level: str = 'static')

Run the static validation pass, which build also runs.

Parameters:

Name	Type	Description	Default
level	str	Validation level, only "static".	'static'

Raises:

Type	Description
ValueError	If the recipe fails validation, listing every problem found.

---

Boundary conditions

spec.bc holds one method per boundary condition.

dirichlet

dirichlet(var: str, *, on: Optional[Any] = None, tag: Optional[str] = None, node_indices: Any = None, axis: Optional[int] = None, side: Optional[str] = None, component: Optional[int] = None, value: Any = None, value_fn: Optional[Callable] = None, fn: Optional[Callable] = None) -> None

Fix a field's value on a boundary, u = g.

Pass exactly one location (on, tag, node_indices, or axis with side), and exactly one of value, value_fn and fn.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
on	Optional[Any]	Boundary selector, a tag name, a ds.tags(...) expression, node indices, or a boolean node mask.	None
tag	Optional[str]	Boundary tag name, the single-tag spelling of on.	None
node_indices	Any	Node indices, the index spelling of on.	None
axis	Optional[int]	Grid axis of the face, with side, on a structured grid.	None
side	Optional[str]	"low" or "high", the end of axis the face sits at.	None
component	Optional[int]	Index of the one value to fix at each node of a vector field, or None for every value.	None
value	Any	Constant value, a scalar or one entry per selected location, in storage order with a face axis dropped.	None
value_fn	Optional[Callable]	Value as a function of dimensional time, lambda t: ....	None
fn	Optional[Callable]	Value as a function of the boundary view and dimensional time, lambda q, t: ....	None

flux

flux(var: str, *, on: Optional[Any] = None, tag: Optional[str] = None, node_indices: Any = None, axis: Optional[int] = None, side: Optional[str] = None, value: Any = None, value_fn: Optional[Callable] = None, fn: Optional[Callable] = None, name: Optional[str] = None, cross_facet: bool = False) -> None

Fix a field's flux density through a boundary, -n . F = g, positive into the domain.

Pass exactly one location (on, tag, node_indices, or axis with side), and exactly one of value, value_fn and fn. On a grid of two or three dimensions the cells must lie on one face, which weights them.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
on	Optional[Any]	Boundary selector, a tag name, a ds.tags(...) expression, node indices, or a boolean node mask.	None
tag	Optional[str]	Boundary tag name, the single-tag spelling of on.	None
node_indices	Any	Node indices, the index spelling of on.	None
axis	Optional[int]	Grid axis of the face, with side, on a structured grid.	None
side	Optional[str]	"low" or "high", the end of axis the face sits at.	None
value	Any	Constant value, a scalar or one entry per selected location, in storage order with a face axis dropped.	None
value_fn	Optional[Callable]	Value as a function of dimensional time, lambda t: ....	None
fn	Optional[Callable]	Value as a function of the boundary view and dimensional time, lambda q, t: ....	None
name	Optional[str]	Name that spec.constraint.flux_total(from_bc=...) refers to, on a finite-element mesh only.	None
cross_facet	bool	If True, the fn closure sees the whole boundary at once instead of one facet.	False

neumann

neumann(var: str, *, on: Optional[Any] = None, tag: Optional[str] = None, node_indices: Any = None, axis: Optional[int] = None, side: Optional[str] = None, value: Any = None, value_fn: Optional[Callable] = None, fn: Optional[Callable] = None) -> None

Fix a field's outward normal derivative on a boundary, n . grad(u) = g.

Pass exactly one location (on, tag, node_indices, or axis with side), and exactly one of value and value_fn. The condition enters laplacian and div_flux wherever they are applied to q.<var>, with their own diffusivity, so it holds whatever D is.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
on	Optional[Any]	Boundary selector, a tag name, a ds.tags(...) expression, node indices, or a boolean node mask.	None
tag	Optional[str]	Boundary tag name, the single-tag spelling of on.	None
node_indices	Any	Node indices, the index spelling of on.	None
axis	Optional[int]	Grid axis of the face, with side, on a structured grid.	None
side	Optional[str]	"low" or "high", the end of axis the face sits at.	None
value	Any	Constant derivative, a scalar or one entry per selected location, in storage order with a face axis dropped.	None
value_fn	Optional[Callable]	Derivative as a function of dimensional time, lambda t: ....	None
fn	Optional[Callable]	Refused, since a derivative that depends on the field is a Robin condition, spec.bc.robin.	None

inflow

inflow(var: str, *, axis: int, side: str, velocity: Any, value: Any = None, value_fn: Optional[Callable] = None) -> None

Feed a field through one grid face at a given speed and inlet value, v c - D dc/dn = v c_in.

The face carries the total flux velocity * value into the domain, so the transport equation must use q.ops.div_advect, which carries no flux in through a face. Pass exactly one of value and value_fn.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
axis	int	Grid axis of the inlet face.	required
side	str	"low" or "high", the end of axis the inlet sits at.	required
velocity	Any	Speed of the flow entering through the face, a non-negative scalar or one entry per face cell.	required
value	Any	Inlet value c_in, a scalar or one entry per face cell.	None
value_fn	Optional[Callable]	Inlet value as a function of dimensional time, lambda t: ....	None

periodic

periodic(var: str, *, axes: Optional[Any] = None) -> None

Wrap a field periodically on the given axes of its cartesian grid.

Every field on the same grid must be periodic on the same axes. The other faces take any other condition.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
axes	Optional[Any]	Grid axes that wrap, or None (default) for every axis.	None

robin

robin(var: str, *, on: Optional[Any] = None, tag: Optional[str] = None, node_indices: Any = None, axis: Optional[int] = None, side: Optional[str] = None, alpha: Any, beta: Any, value: Any = None, value_fn: Optional[Callable] = None, fn: Optional[Callable] = None) -> None

Fix a mixed condition on a boundary, alpha * u + beta * n . grad(u) = g.

Pass exactly one location (on, tag, node_indices, or axis with side), and exactly one of value, value_fn and fn.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
on	Optional[Any]	Boundary selector, a tag name, a ds.tags(...) expression, node indices, or a boolean node mask.	None
tag	Optional[str]	Boundary tag name, the single-tag spelling of on.	None
node_indices	Any	Node indices, the index spelling of on.	None
axis	Optional[int]	Grid axis of the face, with side, on a structured grid.	None
side	Optional[str]	"low" or "high", the end of axis the face sits at.	None
alpha	Any	Coefficient of u, a constant or a function of time.	required
beta	Any	Coefficient of the normal derivative, a constant or a function of time, never zero.	required
value	Any	Constant value, a scalar or one entry per selected location, in storage order with a face axis dropped.	None
value_fn	Optional[Callable]	Value as a function of dimensional time, lambda t: ....	None
fn	Optional[Callable]	Value as a function of the boundary view and dimensional time, lambda q, t: ....	None

traction

traction(var: str, *, on: Optional[Any] = None, tag: Optional[str] = None, node_indices: Any = None, value: Any = None, value_fn: Optional[Callable] = None, fn: Optional[Callable] = None) -> None

Fix the traction of a vector field on a boundary, sigma . n = t.

Pass exactly one of on, tag and node_indices, and exactly one of value, value_fn and fn.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
on	Optional[Any]	Boundary selector, a tag name, a ds.tags(...) expression, node indices, or a boolean node mask.	None
tag	Optional[str]	Boundary tag name, the single-tag spelling of on.	None
node_indices	Any	Node indices, the index spelling of on.	None
value	Any	Constant traction vector t.	None
value_fn	Optional[Callable]	Traction as a function of dimensional time, lambda t: ....	None
fn	Optional[Callable]	Traction as a function of the boundary view and dimensional time, lambda q, t: ....	None

symmetry

symmetry(var: str, *, on: Optional[Any] = None, tag: Optional[str] = None, node_indices: Any = None) -> None

Fix the normal component of a vector field to zero on a boundary, n . u = 0.

Pass exactly one of on, tag and node_indices. The field needs one component per spatial dimension.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
on	Optional[Any]	Boundary selector, a tag name, a ds.tags(...) expression, node indices, or a boolean node mask.	None
tag	Optional[str]	Boundary tag name, the single-tag spelling of on.	None
node_indices	Any	Node indices, the index spelling of on.	None

no_slip

no_slip(var: str, *, on: Optional[Any] = None, tag: Optional[str] = None, node_indices: Any = None, wall_velocity: Any = None) -> None

Fix a field to the wall velocity on a boundary, the same as dirichlet with value=wall_velocity.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
on	Optional[Any]	Boundary selector, a tag name, a ds.tags(...) expression, node indices, or a boolean node mask.	None
tag	Optional[str]	Boundary tag name, the single-tag spelling of on.	None
node_indices	Any	Node indices, the index spelling of on.	None
wall_velocity	Any	Velocity of the wall, or None for zero.	None

slip_wall

slip_wall(var: str, *, on: Optional[Any] = None, tag: Optional[str] = None, node_indices: Any = None) -> None

Fix the normal component of a vector field to zero on a boundary, an alias of symmetry.

Parameters:

Name	Type	Description	Default
var	str	Name of the field the condition applies to.	required
on	Optional[Any]	Boundary selector, a tag name, a ds.tags(...) expression, node indices, or a boolean node mask.	None
tag	Optional[str]	Boundary tag name, the single-tag spelling of on.	None
node_indices	Any	Node indices, the index spelling of on.	None

diff_soup.problem.tags

tags(*names: str) -> BoundarySet

Select boundary locations by tag name.

Several names select their union. On a mesh the names are those of its TaggedMesh, and a one-dimensional grid names its last node "surface".

Parameters:

Name	Type	Description	Default
*names	str	One or more non-empty tag names.	()

Returns:

Type	Description
BoundarySet	A BoundarySet, which combines with others through | and -.

Example

spec.bc.dirichlet("u", on=ds.tags("boundary") - ds.tags("rim"), value=1.0)

diff_soup.problem.BoundarySet dataclass

A boundary selector, built from tag names or node indices.

Build one with ds.tags or BoundarySet.from_nodes. Combine selectors with | (union) and - (difference), and pass the result as on= to a spec.bc declaration. It is frozen and hashable.

reduce

reduce(*, on_tags: Callable[[tuple[str, ...]], T], on_nodes: Callable[[tuple[int, ...]], T], on_union: Callable[[T, T], T], on_difference: Callable[[T, T], T]) -> T

Fold the selector expression into one value.

Each tag leaf maps through on_tags, each node leaf through on_nodes, and each union or difference combines the values of its operands. Use it to resolve a selector against a domain, for example into a set of node indices.

Parameters:

Name	Type	Description	Default
on_tags	Callable[[tuple[str, ...]], T]	Maps a tuple of tag names to a value.	required
on_nodes	Callable[[tuple[int, ...]], T]	Maps a tuple of node indices to a value.	required
on_union	Callable[[T, T], T]	Combines the values of the two operands of a union.	required
on_difference	Callable[[T, T], T]	Combines the values of the left and right operands of a difference.	required

Returns:

Type	Description
T	The value of the whole expression.

from_nodes classmethod

from_nodes(indices) -> 'BoundarySet'

Build a selector from node indices or a boolean mask.

Parameters:

Name	Type	Description	Default
indices		Node indices, or a boolean mask that selects its True positions, flattened in either case.	required

Returns:

Type	Description
'BoundarySet'	A selector of those nodes.

---

Observations

spec.observe declares the quantities a solution reports and a fit compares with data.

The observation declarations of a recipe, reached as spec.observe.

scalar

scalar(*, name: str, fn: Callable, domain: Optional[str] = None) -> None

Declare a scalar observation, computed from the state by fn(q).

Pass domain when fn reads the fields of more than one component.

Parameters:

Name	Type	Description	Default
name	str	Observation name.	required
fn	Callable	Closure lambda q: ... over the fields and globals, returning one array.	required
domain	Optional[str]	Domain the observation evaluates on, or None for the domain whose fields fn reads.	None

field

field(*, name: str, fn: Callable, domain: Optional[str] = None) -> None

Declare a spatially resolved observation, computed from the state by fn(q).

Pass domain when fn reads the fields of more than one component.

Parameters:

Name	Type	Description	Default
name	str	Observation name.	required
fn	Callable	Closure lambda q: ... over the fields and globals, returning one array.	required
domain	Optional[str]	Domain the observation evaluates on, or None for the domain whose fields fn reads.	None

boundary_mean

boundary_mean(*, name: str, fn: Callable, tag: str, domain: Optional[str] = None) -> None

Declare the mean of a nodal quantity over a tagged boundary.

fn(q) returns one value per node, and the observation averages the values on tag weighted by the boundary measure: face areas on a grid, facet measures on a mesh. Pass domain when fn reads the fields of more than one component.

Parameters:

Name	Type	Description	Default
name	str	Observation name.	required
fn	Callable	Closure lambda q: ... over the fields and globals, returning one value per node.	required
tag	str	Name of the boundary tag the mean is over.	required
domain	Optional[str]	Domain the observation evaluates on, or None for the domain whose fields fn reads.	None

spatial_mean

spatial_mean(*, name: str, fn: Callable, domain: Optional[str] = None) -> None

Declare the mean of a nodal quantity over a domain.

fn(q) returns one value per node, and the observation averages them weighted by the cell volumes on a grid or the lumped mass on a mesh. Pass domain when fn reads the fields of more than one component.

Parameters:

Name	Type	Description	Default
name	str	Observation name.	required
fn	Callable	Closure lambda q: ... over the fields and globals, returning one value per node.	required
domain	Optional[str]	Domain the observation evaluates on, or None for the domain whose fields fn reads.	None

---

Controls

spec.control turns a forcing protocol into a value the solve reads as it runs.

The control declarations of a recipe, reached as spec.control.

direct

direct(*, name: str, fn: Callable) -> None

Declare a direct control, a time-dependent input read in closures as q.control.<name>.

fn receives dimensional time and returns the dimensional value. pdesolve(..., control_state0={name: value}) replaces it for one solve.

Parameters:

Name	Type	Description	Default
name	str	Control name.	required
fn	Callable	Function lambda t: ... of dimensional time.	required

derived

derived(*, name: str, fn: Callable, binds_to: str) -> None

Declare a derived control that drives a constrained global to fn(t).

Exactly one constraint must close the global binds_to, and the control replaces that constraint's row with global = fn(t). Each global takes one driving control.

Parameters:

Name	Type	Description	Default
name	str	Control name.	required
fn	Callable	Function lambda t: ... of dimensional time returning the global's dimensional value.	required
binds_to	str	Name of the global the control drives.	required

potentiostatic

potentiostatic(*, name: str, fn: Callable) -> None

Declare a set-voltage control, the same as direct.

Parameters:

Name	Type	Description	Default
name	str	Control name.	required
fn	Callable	Voltage as a function lambda t: ... of dimensional time.	required

galvanostatic

galvanostatic(*, name: str, fn: Callable, binds_to: str) -> None

Declare a set-current control, the same as derived.

Parameters:

Name	Type	Description	Default
name	str	Control name.	required
fn	Callable	Current as a function lambda t: ... of dimensional time.	required
binds_to	str	Name of the global current the control drives.	required

operational

operational(*, name: str, current: str | None = None, voltage: str, current_unknown: str, modes: Sequence[str] = _OPERATIONAL_MODES, seed_voltage_threshold: float = 1e-06, current_from_residual: str | None = None) -> None

Declare an operating table of current, voltage and power rows that close one current unknown.

The active row and its target come from an operating ControlState passed as control_state0, and switching them reuses the compiled solve. Pass exactly one of current and current_from_residual.

Parameters:

Name	Type	Description	Default
name	str	Table name.	required
current	str | None	Name of the scalar observation giving the current.	None
voltage	str	Name of the scalar observation giving the voltage.	required
current_unknown	str	Name of the global current the table closes.	required
modes	Sequence[str]	Rows of the table, distinct entries of "current", "voltage", "power", indexed by the mode.	_OPERATIONAL_MODES
seed_voltage_threshold	float	Voltage magnitude above which a power mode without a current seed starts from target / V.	1e-06
current_from_residual	str | None	Name of a field whose summed residual rows give the current.	None

---

Families

spec.family(...) declares a population of identical members and returns the object that declares their fields, equations and boundary conditions.

A declared family, returned by spec.family(...).

is_ragged property

is_ragged: bool

Whether the family's member domains form more than one size bucket.

equation property

equation: _FamilyEquationDecorator

Decorator that registers a member equation, as @members.equation(writes=..., reads=..., name=None).

The closure takes (q, m), with m.<name> the member's entry of a per-member parameter and m.member_index its index. writes names one member field or a tuple of them.

bc property

bc: _FamilyBCNamespace

Member boundary conditions, members.bc.flux(var, on=..., value= | value_fn= | fn=...).

The flux takes the same arguments as spec.bc.flux without name and cross_facet. Inside fn, q.m.<name> reads the member's entry of a per-member parameter.

field_decl

field_decl(name: str, *, mass: str, shape: Optional[tuple[int, ...]] = None) -> Any

Declare a member field, a profile on the member domain or a per-member scalar.

A profile has one value per member-domain node. shape=(1,) declares a per-member scalar, which must be mass="algebraic" and is closed by its own member equation.

Parameters:

Name	Type	Description	Default
name	str	Member field name, unique in the family.	required
mass	str	"differential" or "algebraic".	required
shape	Optional[tuple[int, ...]]	(1,) for a per-member scalar, or None for a profile, refused for a ragged family.	None

Returns:

Type	Description
Any	The member field's variable, or one per size bucket for a ragged family.

graph

graph(name: str, *, edges: Any, directed: bool, weights: Any = 1.0) -> FamilyGraph

Declare a graph of edges between members, for pairwise exchanges.

A ragged family cannot hold a graph.

Parameters:

Name	Type	Description	Default
name	str	Graph name, unique in the family.	required
edges	Any	Integer array of shape (2, E), source member indices in row 0 and targets in row 1.	required
directed	bool	Whether each edge has a direction.	required
weights	Any	A float scalar, a float array of length E, or a function (source, target, edge) returning each edge's weight.	1.0

Returns:

Type	Description
FamilyGraph	The FamilyGraph.

graph_exchange

graph_exchange(*, graph: str, writes: str, reads: tuple[str, ...], weight_reads: tuple[str, ...] = (), name: Optional[str] = None) -> Callable[[Callable], Callable]

Decorator that declares a pairwise exchange along the edges of a graph.

The decorated fn(source, target, edge) returns the flow into the target, with source.<field> and target.<field> the fields in reads and edge.weight the edge's weight.

Parameters:

Name	Type	Description	Default
graph	str	Name of a graph declared with graph.	required
writes	str	Member field whose rows receive the exchange.	required
reads	tuple[str, ...]	Member fields the exchange reads at both endpoints.	required
weight_reads	tuple[str, ...]	Member fields a weight function reads.	()
name	Optional[str]	Exchange name, or None for the function's name.	None

Returns:

Type	Description
Callable[[Callable], Callable]	The decorator, which returns the function unchanged.

describe

describe() -> FamilySummary

Return a summary of the family's declaration.

Returns:

Type	Description
FamilySummary	A record of the family's name, count, parameter names params, and per size bucket its count, first_member and fields.

members

members() -> tuple[FamilyMemberSummary, ...]

Return a summary of every member of the family.

Returns:

Type	Description
tuple[FamilyMemberSummary, ...]	One record per member with its index, size bucket, group, member fields, and params values as floats.

field

field(name: str)

Return a member field as the source of a coupling out of the family.

Call .surface() on the result for each member's value at its outer surface, the projection spec.couple takes. A ragged family raises NotImplementedError.

Parameters:

Name	Type	Description	Default
name	str	Member field declared with field_decl.	required

Returns:

Type	Description
	The member-field projection.

couple

couple(source: Any, **kwargs) -> Any

Couple a member field out of the family, the same as spec.couple.

Parameters:

Name	Type	Description	Default
source	Any	A field returned by spec.field, or a member projection such as field(name).surface().	required
**kwargs		Keywords of spec.couple, such as onto, at and reduce.	{}

Returns:

Type	Description
Any	The recipe builder, for chaining.

---

Parameters

spec.param declares a parameter, read in closures as q.p.<name>.

diff_soup.problem.ParameterSet dataclass

The learnable parameters of a fit, split into those shared by every experiment and those varying by cluster.

Pass it as parameters= of InferenceProblem. Every parameter must be learnable, come from one recipe, and appear once.

Parameters:

Name	Type	Description	Default
shared	Any	Parameters returned by spec.param or spec.function that take one value across every experiment.	()
varying	Mapping[str, Any] | None	Mapping from cluster name to the parameters that take one value per cluster.	None

diff_soup.problem.default_theta

default_theta(problem, values=None)

Return a problem's parameter values, with chosen parameters replaced.

Theta holds every floating-point array of the problem: the declared parameters and functions, the mass diagonal, the initial multipliers and the family member parameters. A number or an array is broadcast to each array of the parameter. A function also takes a module of its declared structure.

Parameters:

Name	Type	Description	Default
problem		A built problem, such as a LoweredProblem.	required
values		Mapping from a parameter returned by spec.param or spec.function to its value.	None

Returns:

Type	Description
	The parameter pytree, ready to pass as theta=.

Raises:

Type	Description
TypeError	If a key is not a parameter, or a value is not numeric or is complex for a real parameter.
ValueError	If a parameter is not declared on this problem or holds no array, or its value is None or does not fit.

The parameter values a solve takes as theta: every floating-point array of a built problem.

It holds the declared parameters and functions, the mass diagonal, the initial multipliers and the family member parameters. Read and set a declared parameter by the object spec.param or spec.function returned.

theta = ds.default_theta(problem)
theta[D]                    # the value of D
theta = theta.replace(D, 0.2)
ds.parameter_gradient(jax.grad(loss)(theta), D)   # dloss/dD

__getitem__

__getitem__(parameter)

Return a parameter's value, or a function's module holding its learnable arrays.

On a gradient of theta, read D with parameter_gradient, not grads[D].

Raises:

Type	Description
TypeError	If parameter is not one returned by spec.param or spec.function.
ValueError	If the problem does not declare parameter, or it holds no array.

replace

replace(parameter, value=None)

Return a copy with one parameter, or several, set to new values.

Pass replace(D, value) or a mapping replace({D: value, k: other}). The rules for values are those of default_theta.

Raises:

Type	Description
TypeError	If a key is not a parameter, or a value is not numeric or is complex for a real parameter.
ValueError	If the problem does not declare a parameter or it holds no array, or a value is None or does not fit.

diff_soup.problem.parameter_gradient

parameter_gradient(grads, handle)

Return the gradient with respect to a parameter, read from a gradient of theta.

ds.parameter_gradient(jax.grad(loss)(theta), D) is dloss/dD in the shape of D, or for a function a module of the gradients of its learnable arrays. Theta holds D once for each equation that reads it, and this sums those copies, so on theta itself it returns that many times D.

Parameters:

Name	Type	Description	Default
grads		A gradient of theta, such as jax.grad(loss)(theta).	required
handle		A parameter returned by spec.param or spec.function, or a tuple of them.	required

Returns:

Type	Description
	The gradient with respect to handle, or for a tuple a dictionary keyed by parameter.

Raises:

Type	Description
TypeError	If handle is not one returned by spec.param or spec.function.
ValueError	If the problem does not declare handle, or it holds no array.

diff_soup.inference.learnables

learnables(problem: Any, names: Iterable[str] | None = None) -> Optional[LearnableSelector]

Select the learnable parameters and functions declared on a built problem.

Parameters:

Name	Type	Description	Default
problem	Any	The built problem.	required
names	Iterable[str] | None	The declared names to select, or None to select every learnable.	None

Returns:

Type	Description
Optional[LearnableSelector]	A selector whose mask is a boolean pytree over problem marking the selected arrays, or None when nothing is learnable.

---

Mass operators

MassOperator declares a non-diagonal mass block inside a residual-form equation, as in op(q.dt("u")) - flux(q).

diff_soup.problem.MassOperator

Mass matrix M(t, y) applied to a time derivative, M @ ydot.

Build one with a factory. Inside a residual-form equation, op(q.dt("u")) applies it to a field's derivative, and a MassProblem takes one as its mass.

Attributes:

Name	Type	Description
apply	Callable	Function (t, y, theta, control, ydot) returning M @ ydot.
coefficient_sparsity	SparseTopology	SparseTopology of the possibly nonzero entries of M.
dependencies	Dependencies	Dependencies naming the inputs the entries vary with.
state_action_sparsity	Optional[SparseTopology]	Support of d(M @ ydot)/dy, required when the entries depend on the state.
state_action_values	Optional[Callable]	Values of d(M @ ydot)/dy on that support, or None.
coefficient_values	Optional[Callable]	Function (t, y, theta, control) returning the entries on coefficient_sparsity, or None.
transpose_apply	Optional[Callable]	Function (t, y, theta, control, w) returning M.T @ w, or None.

identity staticmethod

identity(n: int) -> 'MassOperator'

Return the constant n-by-n identity.

Parameters:

Name	Type	Description	Default
n	int	Side length.	required

diagonal staticmethod

diagonal(values_or_callback, *, n: Optional[int] = None, dependencies: Optional[Dependencies] = None, state_action_sparsity: Optional[SparseTopology] = None, state_action_values: Optional[Callable] = None) -> 'MassOperator'

Return a diagonal operator from a vector of entries or a callback.

A callback needs n and dependencies. A vector is constant unless dependencies says otherwise.

Parameters:

Name	Type	Description	Default
values_or_callback		1-D array of diagonal entries, or a callback (t, y, theta, control) returning them.	required
n	Optional[int]	Side length, required with a callback.	None
dependencies	Optional[Dependencies]	Inputs the entries vary with, required with a callback.	None
state_action_sparsity	Optional[SparseTopology]	Support of d(M @ ydot)/dy, required when the entries depend on the state.	None
state_action_values	Optional[Callable]	Values of d(M @ ydot)/dy on that support, or None.	None

sparse staticmethod

sparse(rows, cols, values_or_callback, *, n: int, dependencies: Optional[Dependencies] = None, state_action_sparsity: Optional[SparseTopology] = None, state_action_values: Optional[Callable] = None) -> 'MassOperator'

Return an operator with entries at given (row, col) positions.

Parameters:

Name	Type	Description	Default
rows		Row index of each entry, in [0, n).	required
cols		Column index of each entry, in [0, n), the same length as rows.	required
values_or_callback		One value per entry, or a callback (t, y, theta, control) returning them.	required
n	int	Side length.	required
dependencies	Optional[Dependencies]	Inputs the entries vary with, required with a callback.	None
state_action_sparsity	Optional[SparseTopology]	Support of d(M @ ydot)/dy, required when the entries depend on the state.	None
state_action_values	Optional[Callable]	Values of d(M @ ydot)/dy on that support, or None.	None

block_sparse staticmethod

block_sparse(blocks: Sequence = (), couplings: Sequence = (), *, n: int, dependencies: Optional[Dependencies] = None, state_action_sparsity: Optional[SparseTopology] = None, state_action_values: Optional[Callable] = None) -> 'MassOperator'

Return a sparse operator assembled from dense sub-blocks and single-row couplings.

Pass at least one block or coupling. Any callback entry needs dependencies.

Parameters:

Name	Type	Description	Default
blocks	Sequence	Triples (rows, cols, values), each a dense sub-block with values of shape (len(rows), len(cols)) or a callback.	()
couplings	Sequence	Triples (row, cols, weights), each one row reading cols with weights of shape (len(cols),) or a callback.	()
n	int	Side length.	required
dependencies	Optional[Dependencies]	Inputs the entries vary with, required when any entry is a callback.	None
state_action_sparsity	Optional[SparseTopology]	Support of d(M @ ydot)/dy, required when the entries depend on the state.	None
state_action_values	Optional[Callable]	Values of d(M @ ydot)/dy on that support, or None.	None

dense

dense(t, y, theta, control) -> jax.Array

Return M as a dense n-by-n array at the given inputs.

Parameters:

Name	Type	Description	Default
t		Time.	required
y		State vector.	required
theta		Parameters.	required
control		Control values.	required

Returns:

Type	Description
Array	The dense matrix.

from_apply staticmethod

from_apply(apply: Callable, *, coefficient_sparsity: SparseTopology, dependencies: Dependencies, state_action_sparsity: Optional[SparseTopology] = None, state_action_values: Optional[Callable] = None, transpose_apply: Optional[Callable] = None) -> 'MassOperator'

Return a matrix-free operator from its action apply(t, y, theta, control, ydot).

Parameters:

Name	Type	Description	Default
apply	Callable	Function (t, y, theta, control, ydot) returning M @ ydot.	required
coefficient_sparsity	SparseTopology	SparseTopology of the possibly nonzero entries of M.	required
dependencies	Dependencies	Inputs the entries vary with.	required
state_action_sparsity	Optional[SparseTopology]	Support of d(M @ ydot)/dy, required when the entries depend on the state.	None
state_action_values	Optional[Callable]	Values of d(M @ ydot)/dy on that support, or None.	None
transpose_apply	Optional[Callable]	Function (t, y, theta, control, w) returning M.T @ w, or None.	None