zoomy_core.model.operations module

zoomy_core.model.operations module#

General-purpose INS equation framework with composable projection operations.

Class hierarchy:

StateSpace — shared symbols (coordinates, fields, stress tensor, parameters) SymbolicBase — name + display ├── Expression — PDE terms with .project(), .ibp(), .apply(), .terms, [i] └── Relation — lhs = rhs substitution rules with .apply_to()

├── Assumption — physical conditions (kinematic BCs, hydrostatic) └── Material — constitutive models (Newtonian, inviscid)

FullINS(state) — builds INS equations from a StateSpace

Design: the user drives every step. No hidden assumptions.

class zoomy_core.model.operations.StateSpace(dimension=2)#

Bases: object

Shared symbolic state: coordinates, velocity fields, pressure, stress tensor, bathymetry, and physical parameters.

dimension is the physical space dimension:: 2 = xz plane (1D horizontal shallow water) 3 = xyz space (2D horizontal shallow water)

The vertical coordinate is always z. Horizontal coordinates are x (and y if dimension=3). The horizontal_dim property gives the number of horizontal directions (1 or 2).

property horizontal_dim#

property has_y#

class zoomy_core.model.operations.SymbolicBase(name='')#

Bases: object

Base for all symbolic objects. Provides name and notebook display.

class zoomy_core.model.operations.Expression(expr, name='', term_groups=None, solver_groups=None, term_tags=None, tag_order=None)#

Bases: SymbolicBase

Symbolic expression for PDE terms.

Supports: - Term access: .terms, [i], len(), iteration - Projection: .project(test, var, domain, …) - Integration by parts: .ibp(var, test_weight, domain) -> IBPResult - Apply conditions: .apply({old: new}, …) or .apply(relation) - SymPy: .subs(), .simplify(), .expand(), .doit() - Notebook: _repr_latex_

property terms#

apply_to_term(index, *operations)#

Apply operations to a specific term and return the full expression.

Usage:

xmom.apply_to_term(5, ProductRule())
xmom.apply_to_term(5, {old: new})

Note

Public callers should prefer xmom[5].apply(...) — Expression.__getitem__ returns a live accessor whose .apply mutates the parent in place. This method is kept as the internal split-apply-merge primitive that the accessor uses; library code may continue to call it directly.

keep_groups(*group_names)#: Return a new Expression keeping only the specified tags’ terms.

drop_groups(*group_names)#: Return a new Expression dropping the specified tags’ terms.

tag(**named_groups)#

Attach physical tags to terms for display.

Each name=value kwarg is split into additive terms and each term is recorded as belonging to name in the per-term tag map.

Tags follow the terms through .apply / .subs / .simplify: when a transformation replaces a tagged term, the new term(s) inherit the tag. Terms that cancel (via simplify) silently disappear from the view — no parallel-sum math, no cross-tag contamination.

Usage:

xmom = xmom.tag(pressure=grad_p, stress=div_tau)

property tags#

Dict of current-expr terms grouped by physical tag.

Computed on demand by walking Add.make_args(self.expr) and assigning each term to the tag recorded in _term_tags. Terms not in the main expression (cancelled by simplify, etc.) do not appear — the view stays in sync with the single source of truth.

property untagged#: Sum of current-expr terms that have no physical tag.

solver_tag(**named_groups)#

Attach/replace solver tag groups for operator routing.

Tag names are normalized via the solver-tag alias catalog, so .solver_tag(convection=...) and .solver_tag(flux=...) both write to the canonical flux slot.

Solver tags are dropped per-tag whenever apply() / simplify() changes the tagged sub-expression (structural equality check).

Usage:

xmom = xmom.solver_tag(flux=F_x, source=g * h * dbdx)

property solver_tags#: Dict of solver tags (canonical names), or empty dict if none set.

get_solver_tag(name)#

Return the sp.Expr for a solver tag, or None if not set.

name can be any alias; it is normalized to canonical form.

untagged_remainder()#

Return self.expr - sum(solver_tags.values()), simplified.

Zero means every term has been routed to a solver tag; a non-zero remainder is what drives the model’s untagged_policy (warn/error).

auto_solver_tag(*, state_vars, time_var, coords, parameters=())#

Algorithmically assign solver tags by structural matching.

Each additive term of self.expr is classified:

coeff * Derivative(q, t) with q a state variable → time_derivative.
coeff * Derivative(F, x_i) where coeff is state-variable- free → flux (conservative form).
coeff * Derivative(q_j, x_i) with q_j a state variable and coeff containing state vars → nonconservative_flux.
coeff * Derivative(p, x_i) where p is a declared parameter and coeff contains state vars → also nonconservative_flux (the term couples a state coefficient to a parameter gradient — e.g. bathymetry-gradient forcing).
No derivative and purely algebraic in state_vars ∪ parameters space → source.
Anything else (unclosed fields like u(t, x, z), unmatched Derivative shapes) stays in the untagged remainder.

Parameters:

state_vars (iterable of sympy Functions / Symbols) – The evolution variables (α_k(t, x), interface heights, etc.).
time_var (sympy Symbol) – The time coordinate (for time_derivative detection).
coords (iterable of sympy Symbols) – Spatial coordinates (x, y, …) for flux / NCP.
parameters (iterable, optional) – Known external Function/Symbol atoms (bathymetry b(t, x), prescribed fluxes, etc.). These may freely appear in any term without the term being marked “unclosed”, and they may appear as the inner of a spatial Derivative in the NCP tagging.
populated. (Returns a new Expression with solver_groups) –
preserved. (Physical tags are) –

project(test_func, var, domain=(0, 1), weight=1, scale=1, numerical=False, order=4)#: Galerkin projection: integral(expr * test * weight * scale, (var, a, b)).

ibp(var, test_weight, domain=(0, 1), scale=1)#: Integration by parts on the outermost Derivative w.r.t. var. Returns IBPResult(integrate, boundary_upper, boundary_lower).

apply(*conditions, rhs=<object object>)#

Apply substitutions / Relations / Operations — per term.

The transformation is applied to every additive term of self.expr individually. Each output term inherits the tag of its parent term (follow-the-term), so the per-term tag dict stays in sync with the evolving main expression automatically.

Shorthand: expr.apply(lhs, rhs=0) is equivalent to expr.apply({lhs: 0}). Valid only with a single positional argument which is a sympy expression (or Expression).

After Relation/dict-style applies the result is lightly simplified via _simplify_preserve_integrals (linearity-only derivatives, cancellations across the full expr). Terms that cancel vanish from both self.expr and the tag dict.

Operations are invoked via their unified __call__ interface; when the op produces multiple output pieces per input term (DepthIntegrate etc.) all pieces inherit the parent tag.

copy()#: Return an independent Expression with the same expr, tags, relation.

solve_for(variable)#

Return a new Expression whose _as_relation maps variable to the solution of self.expr == 0.

Pipeline-style counterpart to _NodeProxy.solve_for: does not mutate a system tree. Downstream model.apply(result) picks up the relation and performs the substitution.

depth_integrate(lower, upper, var, method='auto')#

Depth-integrate this expression over [lower, upper] w.r.t. var.

Parameters:

lower (sympy expressions for the integration bounds) – (typically b and b+H, both functions of t and x)
upper (sympy expressions for the integration bounds) – (typically b and b+H, both functions of t and x)
var (Symbol) – The vertical coordinate (z)
method (str) –
‘auto’ : detect derivative direction and choose method ‘leibniz’ : pull horizontal derivative outside:

int df/dx dz = d/dx[int f dz] - f(upper)*d(upper)/dx + f(lower)*d(lower)/dx

’fundamental_theorem’for vertical derivatives:
int df/dz dz = f(upper) - f(lower)

’direct’ : keep as Integral(expr, (var, lower, upper))

Returns:

DepthIntegralResult with .volume, .boundary_upper, .boundary_lower
or a plain Expression for ‘direct’ and ‘fundamental_theorem’.

subs(*args, **kwargs)#

split_integrals()#: Distribute each Integral(Add(...), lim) into a sum of single-term Integrals, recursing through any Derivative wrapper. Returns a new Expression whose .terms exposes one logical equation term per Integral.

merge_integrals()#: Group sibling Integrals with matching (limits, deriv-wrapper) signature into a single Integral(Σ c·f, lim). Inverse of split_integrals(); useful when cross-integrand cancellations need sympy’s Add canonicalisation to see them in one integrand.

simplify()#: Simplify: expand + cancel, preserving Integral and Derivative(Integral) terms.

expand()#

doit()#

has(*args)#

atoms(*types)#

xreplace(rule)#

property free_symbols#

map(fn)#

Apply fn to each term, reassemble into a single Expression.

fn receives an Expression (single term) and must return either an Expression or a DepthIntegralResult. DepthIntegralResults are assembled (volume + boundaries) before summing.

Example

integrated = expr.map(lambda t: t.depth_integrate(b, eta, z))

map_with_bcs(fn, bcs)#

Like map(), but collects boundary terms and applies BCs globally.

This is the correct way to depth-integrate a full equation: boundary terms from ALL terms are combined first, then BCs are applied once (so cross-term cancellations happen properly).

Parameters:

fn (callable) – Applied to each term. Must return DepthIntegralResult or Expression.
bcs (list of Relation) – Kinematic BCs etc. applied to the combined boundary expression.

Returns:

The fully depth-integrated equation with BCs applied.

Return type:

Expression

classify(t=None, x=None, z=None)#

Classify each term by its role in the PDE.

Returns a dict: {role: Expression} where role is one of: ‘temporal’, ‘convective’, ‘diffusive’, ‘source’.

Detection rules: - Has d/dt → temporal - Has d/dx (first-order) of a product → convective flux - Has d²/dz² or d/dz of d/dz → diffusive - Otherwise → source (algebraic)

property temporal#

only temporal (d/dt) terms.

Type:: View

property convective#

only convective flux (d/dx) terms.

Type:: View

project_onto_basis(basis, level, field_map, z_var, lower=None, upper=None, test_mode=None)#

Project a depth-integrated equation onto a polynomial basis.

Replaces every Integral(f(u,...), (z, b, eta)) by substituting the basis expansion u(z) = sum alpha_k phi_k(zeta) and evaluating the resulting integrals using the SymbolicIntegrator.

If test_mode is an integer, multiplies each integral by the test function phi_{test_mode}(zeta) before evaluating (Galerkin projection for a specific mode). If test_mode=None, returns the scalar integral (e.g. for the mass equation where no test function is needed).

Parameters:

basis (Basisfunction class (e.g. Legendre_shifted)) –
level (int) –
field_map (dict) – Maps the original Function name to a list of SymPy Symbols for the basis coefficients. Example: {‘u’: [alpha_0, alpha_1, alpha_2]}
z_var (Symbol) – The vertical coordinate (z) that appears in the integrals.
lower (sympy expressions (optional)) – The integration bounds (b, eta). If None, detected from the first Integral found.
upper (sympy expressions (optional)) – The integration bounds (b, eta). If None, detected from the first Integral found.
test_mode (int or None) – If int, project onto test function phi_{test_mode}.

Returns:

With all depth integrals replaced by basis matrix products.

Return type:

Expression

latex(strip_args=False, multiline=False)#

LaTeX representation.

Parameters:

strip_args (bool) – u(t,x,z) → u. Partial derivatives preserved.
multiline (bool) – Render as \begin{aligned} with one group per line (requires term_groups).

describe(header=True, final_equation=True, parameters=False, strip_args=True)#

Composable description of this expression.

Returns a Description that renders as markdown in Jupyter.

Parameters:

header (bool) – Show expression name + term count.
final_equation (bool) – Show the symbolic equation.
parameters (bool) – List free symbols.
strip_args (bool) – Display u instead of u(t, x, z).

class zoomy_core.model.operations.DepthIntegralResult(volume, boundary_upper, boundary_lower)#

Bases: object

Result of depth-integrating a term over [lower, upper].

volume#: the volume integral (Expression)

boundary_upper#: boundary term at z=upper (Expression)

boundary_lower#: boundary term at z=lower (Expression)

The full integral = volume + boundary_upper - boundary_lower. Kinematic BCs can be applied to the boundary terms via .apply_bcs().

apply_bcs(bc_lower=None, bc_upper=None)#

assemble()#

Combine all terms: volume + boundary_upper + boundary_lower.

The sign convention is that each component already carries its correct sign. For Leibniz: upper = -f(eta)*d(eta)/dx, lower = +f(b)*db/dx. For fundamental theorem: upper = +f(eta), lower = -f(b).

class zoomy_core.model.operations.IBPResult(integrate, boundary_upper, boundary_lower)#

Bases: object

Structured result from integration by parts.

integrate#: the volume integral (Expression)

boundary_upper#: boundary term at upper limit (Expression)

boundary_lower#: boundary term at lower limit (Expression)

apply_bcs(bc_lower=None, bc_upper=None)#

assemble()#

class zoomy_core.model.operations.Relation(substitutions, name='')#

Bases: SymbolicBase

A symbolic relation: one or more substitution rules lhs_i = rhs_i.

Used as base for Assumption and Material. Can be applied to Expressions via expr.apply(relation), which calls relation.apply_to(expr).

Displays as a system of equations in notebooks.

apply_to(expr)#: Substitute all lhs -> rhs in the given SymPy expression.

solve_for(variable)#

Re-solve every rule lhs = rhs for variable and return a new Relation with substitution dict {variable: solution}.

Pipeline-style counterpart to Expression.solve_for(): lets a chain author flip the substitution direction of a relation without rewriting the relation by hand. Example: a kinematic BC w|_eta = ∂_t eta + u·∂_x eta re-solved for Derivative(state.h, state.t) returns {Derivative(h, t): w|_eta - Derivative(b, t) - u·∂_x eta}, which is exactly the substitution Leibniz boundary terms need to cancel against the ∂_z fundamental-theorem boundaries.

Each rule contributes one substitution; if multiple rules contain variable in their balance lhs - rhs, all are re-solved and merged into one dict (later wins on duplicate keys, matching dict.update semantics).

Raises ValueError if no rule’s balance contains variable (so the call is never silently a no-op).

class zoomy_core.model.operations.Assumption(substitutions, name='')#

Bases: Relation

Physical assumption (kinematic BC, hydrostatic, etc.).

class zoomy_core.model.operations.Material(substitutions, name='')#

Bases: Relation

Constitutive model (Newtonian, inviscid, etc.).

class zoomy_core.model.operations.Operation(name='', description=None)#

Bases: SymbolicBase

A structural transformation on a tree node.

A node is either an Expression (leaf) or a Zstruct of nodes (intermediate). The unified entry point is __call__(node):

Subclasses override _apply_leaf() to transform a leaf Expression. The default __call__ dispatches over the tree, recursing into intermediate Zstructs.
Rank-changing Operations (Multiply with outer=True on a leaf, or contraction on a Zstruct) override __call__ directly and may return a node of different shape than they received.

Every Operation carries name + description so System.apply can record what was done in the derivation history.

rank_changes_leaf = False#

whole_leaf_op = False#

system_level = False#

single_term_only = False#

class zoomy_core.model.operations.Multiply(factor, outer=False, name=None, description=None)#