Skip to content

Operators

Inside an equation closure, q.ops is the operator bank bound to the field's domain. The operators overview explains the conventions.

The methods, grouped by kind:

  • Differential operators: laplacian, div_flux, gradient (alias grad), divergence (alias div), curl, advect_upwind, div_advect, face_flux.
  • Mass and integration weights: mass, lumped_mass.
  • Solid and fluid mechanics: strain, linear_elastic_stress, deformation_gradient, cauchy_stress_from_S_F, strain_rate, viscous_stress_newtonian.
  • Porous media: tortuosity_corrected.
  • Mesh boundaries: boundary_integral, boundary_trace, surface_flux.
  • Pointwise tensor algebra: trace, det, inv, transpose, sym, skew, outer, tensor_contract, double_dot, scale.
  • Finite element quadrature: grad_quad, interp_to_quad, integrate, quad_gradient, integrate_shape, integrate_shape_gradient, wrap_quad_value, wrap_quad_field, unwrap_field.

diff_soup.domains.Operators

The built-in operator bank, reached in an equation closure as q.ops.

On a grid each operator returns pointwise values, and on a mesh its weak form, integrated against each basis function. Subclass it to add an operator, listing in weak_form any that returns an integrated term, and register the class with spec.domain(name, domain, operators=...).

Parameters:

Name Type Description Default
domain

The Domain to act on.

required

discretization_type property

discretization_type: str

The discretization of the bound domain, "fvm", "fdm" or "fem".

unwrap_field

unwrap_field(value)

Return the nodal values of a field view such as q.u, or value unchanged otherwise.

The built-in operators apply it to their inputs, so they accept either form.

Parameters:

Name Type Description Default
value

A field view or an array.

required

Returns:

Type Description

The nodal values.

wrap_quad_field

wrap_quad_field(val)

Wrap nodal values in the field view that q.<name> returns on a mesh.

Arithmetic on the view samples the field at the quadrature points, and .nodal gives the nodal values. Indexing the view raises TypeError, so index .nodal instead. On a grid it returns val unchanged.

Parameters:

Name Type Description Default
val

The nodal values.

required

Returns:

Type Description

The view on a mesh, val on a grid.

laplacian

laplacian(u: _NodealOrField, *, bcs: 'Sequence | None' = None, t: 'Array | None' = None) -> Float[Array, '...']

The Laplacian ∇²u, pointwise on a grid and the weak form -∫ ∇u·∇N_i on a mesh.

On a curvilinear grid it includes the metric.

Parameters:

Name Type Description Default
u _NodealOrField

A scalar field.

required
bcs 'Sequence | None'

Flux and periodic conditions applied on a grid, which problems declare with spec.bc.

None
t 'Array | None'

The time passed to a time-dependent flux condition, zero if None.

None

Returns:

Type Description
Float[Array, '...']

One value per node.

div_flux

div_flux(D: Any, u: _NodealOrField, *, bcs: 'Sequence | None' = None, t: 'Array | None' = None) -> Float[Array, '...']

The diffusion term ∇·(D∇u), pointwise on a grid and the weak form -∫ D∇u·∇N_i on a mesh.

On a grid a tuple (Dx, Dy[, Dz]) gives each axis its own diffusivity, a diagonal tensor. A full tensor is refused.

Parameters:

Name Type Description Default
D Any

The diffusivity, a scalar or one value per node, on a mesh also one per quadrature point, on a grid also one entry per axis.

required
u _NodealOrField

A scalar field.

required
bcs 'Sequence | None'

Flux and periodic conditions applied on a grid, which problems declare with spec.bc.

None
t 'Array | None'

The time passed to a time-dependent flux condition, zero if None.

None

Returns:

Type Description
Float[Array, '...']

One value per node.

advect_upwind

advect_upwind(field: _NodealOrField, velocity: Any, *, bcs: 'Sequence | None' = None, t: 'Array | None' = None) -> Float[Array, '...']

First-order upwind advection v·∇u, on a grid.

The sign of the velocity at each cell picks the upwind side. The result is the operator itself, so a transport equation writes -q.ops.advect_upwind(q.u, v). Boundaries that are not periodic are zero-gradient outflow.

Parameters:

Name Type Description Default
field _NodealOrField

The advected field.

required
velocity Any

One constant speed for every axis, or a sequence with one entry per axis, each a constant or one value per cell.

required
bcs 'Sequence | None'

Periodic conditions only, which problems declare with spec.bc.periodic.

None
t 'Array | None'

Ignored.

None

Returns:

Type Description
Float[Array, '...']

The advection term, shaped like field.

div_advect

div_advect(field: _NodealOrField, velocity: Any, *, scheme: str = 'tvd', limiter: str = 'van_leer', bcs: 'Sequence | None' = None, t: 'Array | None' = None) -> Float[Array, '...']

Advection in conservative form ∇·(v u), on a grid.

Each interior face carries the face velocity times a face value taken from the side the flow comes from. Under "tvd" the face value is the upstream value plus a limited share of the change to the downstream one: second order where u is smooth, and no new extrema at any cell Péclet number. Under "weno3" it is the third-order WENO reconstruction of Jiang and Shu. Under "upwind" it is the upstream value. Under "central" the face carries the mean of v u in its two cells instead. A face next to a boundary cell lacks the second upstream cell, and "tvd" and "weno3" drop to the upstream value there. A boundary face carries only an outgoing flux, so a transport equation writes -q.ops.div_advect(q.c, v) and feeds an inlet with spec.bc.inflow.

Parameters:

Name Type Description Default
field _NodealOrField

The advected field.

required
velocity Any

One constant speed for every axis, or a sequence with one entry per axis, each a constant or one value per cell.

required
scheme str

The face flux, "tvd" (default, flux-limited second order), "weno3" (third order), "central" (second order, unlimited) or "upwind" (first order).

'tvd'
limiter str

The "tvd" limiter, "van_leer" (default), "koren" or "minmod".

'van_leer'
bcs 'Sequence | None'

Periodic conditions only, which problems declare with spec.bc.periodic.

None
t 'Array | None'

Ignored.

None

Returns:

Type Description
Float[Array, '...']

The advection term, shaped like field.

gradient

gradient(u: _NodealOrField) -> Float[Array, '...']

The gradient ∇u, pointwise on a one-dimensional grid.

On a mesh, a scalar field gives the weak form ∫ ∇u N_i, shape (n, dim), and a vector field its Jacobian at the nodes, shape (n, components, dim). It has no periodic form, so it raises under spec.bc.periodic.

Parameters:

Name Type Description Default
u _NodealOrField

A scalar field, or on a mesh a vector field.

required

Returns:

Type Description
Float[Array, '...']

The gradient at every node.

divergence

divergence(vec: _NodealOrField) -> Float[Array, 'n']

The divergence ∇·v, pointwise on a grid and the weak form ∫ (∇·v) N_i on a mesh.

On a mesh it also takes a nodal tensor, row by row, or a symmetric tensor at the quadrature points, assembled as ∫ σ : ∇N_i. It has no periodic form, so it raises under spec.bc.periodic.

Parameters:

Name Type Description Default
vec _NodealOrField

A vector field of shape (n, dim), or on a mesh a tensor of shape (n, dim, dim) or (n_cells, n_q, dim, dim).

required

Returns:

Type Description
Float[Array, 'n']

One value per node for a vector field, one vector per node for a tensor.

curl

curl(u: _NodealOrField) -> Float[Array, '...']

The curl ∇×u of a vector field with one component per axis, in two or three dimensions.

In two dimensions it is the scalar ∂u_y/∂x - ∂u_x/∂y, in three the curl vector. It is pointwise on a grid and the weak form ∫ (∇×u) N_i on a mesh. It has no periodic form, so it raises under spec.bc.periodic.

Parameters:

Name Type Description Default
u _NodealOrField

A vector field of shape (n, dim).

required

Returns:

Type Description
Float[Array, '...']

One value per node in two dimensions, one vector per node in three.

mass

mass(u: _NodealOrField) -> Float[Array, 'n']

The mass-matrix action, the identity on a grid and the consistent mass ∫ u N_i on a mesh.

Use it in a custom residual that applies the mass weighting explicitly.

Parameters:

Name Type Description Default
u _NodealOrField

Nodal values, or on a mesh values at the quadrature points.

required

Returns:

Type Description
Float[Array, 'n']

One value per node.

lumped_mass

lumped_mass() -> Float[Array, 'n']

The integration weights, the cell volumes on a grid and the lumped nodal masses on a mesh.

jnp.sum(q.ops.lumped_mass() * f) integrates a field over the domain.

Returns:

Type Description
Float[Array, 'n']

One weight per node.

face_flux

face_flux(u: Float[Array, '...'], D: Float[Array, '...']) -> Float[Array, '...']

The diffusive flux -D ∂u/∂x at the faces of a one-dimensional grid.

The two end faces carry zero flux. It has no periodic form, so it raises under spec.bc.periodic.

Parameters:

Name Type Description Default
u Float[Array, '...']

One value per cell.

required
D Float[Array, '...']

The diffusivity, a scalar, one value per cell, or one per face.

required

Returns:

Type Description
Float[Array, '...']

One flux per face, n + 1 values.

strain

strain(u: _NodealOrField) -> Float[Array, 'n dim dim']

The small strain sym(∇u) of a displacement field, at the nodes of a mesh.

Parameters:

Name Type Description Default
u _NodealOrField

A displacement field, shape (n, dim).

required

Returns:

Type Description
Float[Array, 'n dim dim']

The strain at each node, shape (n, dim, dim).

linear_elastic_stress

linear_elastic_stress(u: _NodealOrField, *, E, nu, eigenstrain=None) -> Float[Array, 'n dim dim']

The isotropic linear-elastic stress λ tr(ε) I + 2μ ε, at the nodes of a mesh.

The Lamé constants λ and μ follow from E and nu, and the strain ε is sym(∇u) less any eigenstrain.

Parameters:

Name Type Description Default
u _NodealOrField

A displacement field, shape (n, dim).

required
E

Young's modulus.

required
nu

Poisson's ratio.

required
eigenstrain

A nodal strain subtracted before the stress law, shape (n, dim, dim).

None

Returns:

Type Description
Float[Array, 'n dim dim']

The stress at each node, shape (n, dim, dim).

deformation_gradient

deformation_gradient(gradu)

The deformation gradient F = I + ∇u.

Parameters:

Name Type Description Default
gradu

The displacement gradient, shape (..., dim, dim).

required

Returns:

Type Description

F, of the same shape.

cauchy_stress_from_S_F

cauchy_stress_from_S_F(S, F)

The Cauchy stress F S Fᵀ / det(F) from the second Piola-Kirchhoff stress S.

Parameters:

Name Type Description Default
S

The second Piola-Kirchhoff stress, shape (..., dim, dim).

required
F

The deformation gradient, of the same shape.

required

Returns:

Type Description

The Cauchy stress, of the same shape.

strain_rate

strain_rate(u)

The strain rate sym(∇v) of a velocity field, at the nodes of a mesh.

Parameters:

Name Type Description Default
u

A velocity field, shape (n, dim).

required

Returns:

Type Description

The strain rate at each node, shape (n, dim, dim).

viscous_stress_newtonian

viscous_stress_newtonian(u, *, mu)

The Newtonian viscous stress 2μ sym(∇v) of a velocity field, at the nodes of a mesh.

Parameters:

Name Type Description Default
u

A velocity field, shape (n, dim).

required
mu

The dynamic viscosity.

required

Returns:

Type Description

The stress at each node, shape (n, dim, dim).

tortuosity_corrected

tortuosity_corrected(D, eps, beta: float = 1.5)

The Bruggeman-corrected diffusivity D ε^β of a porous medium.

Parameters:

Name Type Description Default
D

The bulk diffusivity.

required
eps

The porosity ε.

required
beta float

The Bruggeman exponent β.

1.5

Returns:

Type Description

The effective diffusivity, broadcast over D and eps.

cell_indicator

cell_indicator(name)

One per cell of a named region of a tagged mesh and zero elsewhere, shape (n_cells,).

Parameters:

Name Type Description Default
name

A region name in the mesh's cell_tag_names.

required

boundary_integral

boundary_integral(field, *, tag)

The integral of a nodal field over a tagged boundary of a mesh.

Each facet contributes the consistent share of its node values times its measure.

Parameters:

Name Type Description Default
field

A nodal field.

required
tag

The boundary, by name or by integer tag.

required

Returns:

Type Description

The integral, a scalar.

boundary_trace

boundary_trace(field, *, tag)

The values of a nodal field at the nodes of a tagged boundary of a mesh.

Parameters:

Name Type Description Default
field

A nodal field.

required
tag

The boundary, by name or by integer tag.

required

Returns:

Type Description

The field at those nodes, in ascending node order.

surface_flux

surface_flux(field, *, tag, kinetics, ctx=None, t=0.0)

A boundary flux from a kinetics law, evaluated on a field's values at a tagged boundary of a mesh.

Parameters:

Name Type Description Default
field

A nodal field.

required
tag

The boundary, by name or by integer tag.

required
kinetics

A function kinetics(trace, ctx, t), called with the boundary values from boundary_trace.

required
ctx

Passed to kinetics.

None
t

Passed to kinetics.

0.0

Returns:

Type Description

The result of kinetics.

trace

trace(T)

The trace of a tensor over its last two axes.

Parameters:

Name Type Description Default
T

A tensor, shape (..., dim, dim).

required

Returns:

Type Description

The trace, shape (...).

det

det(T)

The determinant of a tensor over its last two axes.

Parameters:

Name Type Description Default
T

A tensor, shape (..., dim, dim).

required

Returns:

Type Description

The determinant, shape (...).

inv

inv(T)

The inverse of a tensor over its last two axes.

Parameters:

Name Type Description Default
T

A tensor, shape (..., dim, dim).

required

Returns:

Type Description

The inverse, of the same shape.

transpose

transpose(T)

The transpose of a tensor over its last two axes.

Parameters:

Name Type Description Default
T

A tensor, shape (..., a, b).

required

Returns:

Type Description

The transpose, shape (..., b, a).

sym

sym(T)

The symmetric part (T + Tᵀ) / 2 of a tensor.

Parameters:

Name Type Description Default
T

A tensor, shape (..., dim, dim).

required

Returns:

Type Description

The symmetric part, of the same shape.

skew

skew(T)

The antisymmetric part (T - Tᵀ) / 2 of a tensor.

Parameters:

Name Type Description Default
T

A tensor, shape (..., dim, dim).

required

Returns:

Type Description

The antisymmetric part, of the same shape.

outer

outer(u, v)

The outer product of two vectors, batched over leading axes.

Parameters:

Name Type Description Default
u

A vector, shape (..., dim).

required
v

A vector, shape (..., dim).

required

Returns:

Type Description

The tensor u_i v_j, shape (..., dim, dim).

tensor_contract

tensor_contract(A, B, axes: int = 1)

The contraction of the last axes axes of A with the first axes axes of B, as jnp.tensordot.

Parameters:

Name Type Description Default
A

The first tensor.

required
B

The second tensor.

required
axes int

The number of axes to contract.

1

Returns:

Type Description

The contracted tensor.

double_dot

double_dot(A, B)

The double contraction A : B over the last two axes.

A rank-four A, with trailing shape (dim, dim, dim, dim), gives the tensor C : ε, as in a stress law. Otherwise the result is one scalar per point.

Parameters:

Name Type Description Default
A

A tensor, shape (..., dim, dim), or a stiffness, shape (..., dim, dim, dim, dim).

required
B

A tensor, shape (..., dim, dim).

required

Returns:

Type Description

A scalar per point, or a tensor per point for a rank-four A.

scale

scale(scalar_field, tensor)

The product of a scalar field and a tensor, broadcast over the tensor's components.

Parameters:

Name Type Description Default
scalar_field

One scalar per point.

required
tensor

A constant tensor, or one whose leading axes match scalar_field.

required

Returns:

Type Description

The scaled tensor at every point.

grad_quad

grad_quad(u_vec: _NodealOrField) -> Float[Array, 'n_cells n_q dim dim']

The gradient of a vector field at the quadrature points of a mesh.

Parameters:

Name Type Description Default
u_vec _NodealOrField

A nodal vector field, shape (n, components).

required

Returns:

Type Description
Float[Array, 'n_cells n_q dim dim']

The gradient ∂u_i/∂x_j, shape (n_cells, n_q, components, dim).

grad

grad(u)

An alias of gradient.

div

div(vec_or_tensor)

An alias of divergence.

interp_to_quad

interp_to_quad(u: _NodealOrField) -> Float[Array, 'n_cells n_q']

Sample a nodal field at the quadrature points of a mesh.

Parameters:

Name Type Description Default
u _NodealOrField

A nodal field, scalar or with trailing component axes.

required

Returns:

Type Description
Float[Array, 'n_cells n_q']

The values at the quadrature points, shape (n_cells, n_q, ...).

integrate

integrate(expr: _NodealOrField) -> Float[Array, 'n']

The weak integral ∫ f N_i of an expression on a mesh, the identity on a grid.

Parameters:

Name Type Description Default
expr _NodealOrField

Nodal values, or values at the quadrature points of shape (n_cells, n_q, ...).

required

Returns:

Type Description
Float[Array, 'n']

One value per node.

quad_gradient

quad_gradient(value)

The gradient of a nodal field at the quadrature points of a mesh, the value q.grad returns.

Parameters:

Name Type Description Default
value

A nodal scalar or vector field.

required

Returns:

Type Description

The gradient, shape (n_cells, n_q, dim) for a scalar field and (n_cells, n_q, components, dim) for a vector field.

integrate_shape

integrate_shape(value_q)

Assemble ∫ f N_i from values at the quadrature points of a mesh, the term q.source returns.

Parameters:

Name Type Description Default
value_q

The values at the quadrature points, shape (n_cells, n_q, ...).

required

Returns:

Type Description

One value per node, marked as an integrated weak-form term.

integrate_shape_gradient

integrate_shape_gradient(flux_q)

Assemble ∫ F·∇N_i from a flux at the quadrature points of a mesh, the term q.flux returns.

A flux vector gives one value per node. A symmetric stress in Voigt order gives one vector per node.

Parameters:

Name Type Description Default
flux_q

A flux of shape (n_cells, n_q, dim), or a Voigt stress of shape (n_cells, n_q, 3) in 2-D and (n_cells, n_q, 6) in 3-D.

required

Returns:

Type Description

The integrated term per node, marked as a weak-form term.

wrap_quad_value

wrap_quad_value(value)

Wrap values at the quadrature points of a mesh so arithmetic on them stays there.

Multiplying by a constant tensor broadcasts over its components. Apply a nonlinearity with .apply(fn), which keeps the wrapper. A residual integrates the result as ∫ f N_i.

Parameters:

Name Type Description Default
value

The values at the quadrature points, shape (n_cells, n_q, ...).

required

Returns:

Type Description

The wrapped values.


Abstract operator bank

The details matter only when writing a bank for a new discretization.

diff_soup.domains.AbstractOperator

The base class of an operator bank, which an equation closure reaches as q.ops.

Subclass it for a new discretization, implementing discretization_type and the operators it offers, and register it with spec.domain(name, domain, operators=...). Operators it leaves out raise NotImplementedError. A bank used under spec.bc.periodic must accept bcs= on laplacian, div_flux, advect_upwind and div_advect, and one used under spec.bc.neumann on laplacian and div_flux.

Attributes:

Name Type Description
domain Any

The domain the bank is bound to.

weak_form frozenset

The names of operators whose result on a mesh is already integrated against the test functions, merged across subclasses.

provides property

provides: frozenset

The names of the standard operators this bank implements, such as laplacian or mass.

laplacian

laplacian(u: Any, *, bcs: 'Any | None' = None, t: 'Any | None' = None) -> Any

The Laplacian ∇²u.

The base class raises NotImplementedError.

Parameters:

Name Type Description Default
u Any

The field values.

required
bcs 'Any | None'

Grid boundary conditions to apply to the result, or None.

None
t 'Any | None'

The time for time-dependent boundary conditions, or None.

None

Returns:

Type Description
Any

The Laplacian term.

div_flux

div_flux(*args, **kwargs)

The diffusion term ∇·(D∇u), called as div_flux(D, u, *, bcs=None, t=None).

The base class raises NotImplementedError.

gradient

gradient(*args, **kwargs)

The gradient ∇u, called as gradient(u).

The base class raises NotImplementedError.

divergence

divergence(*args, **kwargs)

The divergence ∇·v, called as divergence(v).

The base class raises NotImplementedError.

mass

mass(*args, **kwargs)

The mass-matrix action on a field, called as mass(u).

The base class raises NotImplementedError.

lumped_mass

lumped_mass(*args, **kwargs)

The integration weights of the domain, called as lumped_mass().

The base class raises NotImplementedError.

advect_upwind

advect_upwind(*args, **kwargs)

Upwind advection of a field, called as advect_upwind(u, velocity, *, bcs=None, t=None).

The base class raises NotImplementedError.

div_advect

div_advect(*args, **kwargs)

Conservative advection of a field, called as div_advect(u, velocity, *, scheme="tvd", limiter="van_leer", bcs=None, t=None).

The base class raises NotImplementedError.

as_weak

as_weak(value)

Mark value as an already-integrated weak-form term.

On a mesh, pass a hand-assembled residual through it so the residual is not integrated a second time. On a grid it returns value unchanged.

Parameters:

Name Type Description Default
value

One value per node, already integrated against the test functions.

required

Returns:

Type Description

The marked term on a mesh, value on a grid.

for_residual

for_residual()

Return the bank as equation closures see it, as q.ops.

On a mesh the view marks the result of every operator in weak_form as already integrated. On a grid it is the bank itself.