Equations

Equation functions use JAX arrays and a generic model callable, not a concrete model class. Each public PDE accepts an ArrayMapping batch and returns named residual groups for ResidualTerm. One equation callable can compute all of its coupled components together.

Residual metadata

Attach static local names to every reusable equation. When a ResidualTerm omits explicit names, it prefixes these with the term's batch_key.

@residual_equation(names=("heat",))
def heat_equation(model_apply, model_state, batch, *, stream="residual"):
    return ((residual,),)

residual_equation [source]

residual_equation(
    *,
    names: Sequence[str],
    default_ntk_stream: ResidualStream = "residual",
) -> Callable[[Callable[P, R]], Callable[P, R]]

Attach ordered local residual names and balancing metadata to an equation callable.

The names describe the outer residual groups returned by the equation. They are local to the equation; a :class:phijax.objectives.ResidualTerm prefixes them with its configured batch_key when explicit loss names are omitted.

Parameters:

Name	Type	Description	Default
names	Sequence[str]	Non-empty unique local names ordered like the equation's outer residual groups.	required
default_ntk_stream	ResidualStream	Stream used by derivative-based balancing when a term does not override it.	'residual'

Returns:

Type	Description
Callable[[Callable[P, R]], Callable[P, R]]	Decorator that preserves the equation callable and attaches its validated residual names.

Raises:

Type	Description
TypeError	If a name is not a string.
ValueError	If names are empty, duplicated, contain /, or the stream is unsupported.

get_residual_names [source]

get_residual_names(
    residual_fn: Callable[..., Any],
) -> tuple[str, ...]

Read static residual names from an equation or configured partial.

Parameters:

Name	Type	Description	Default
residual_fn	Callable[..., Any]	Decorated equation callable or a nested :class:functools.partial wrapping one.	required

Returns:

Type	Description
tuple[str, ...]	Ordered local residual names attached by :func:residual_equation.

Raises:

Type	Description
ValueError	If the callable does not define residual-name metadata.

get_default_ntk_stream [source]

get_default_ntk_stream(
    residual_fn: Callable[..., Any],
) -> ResidualStream

Read the preferred balancing stream from an equation or configured partial.

Undecorated callables default to the residual stream so explicit :class:phijax.ResidualTerm construction remains usable when custom names are supplied.

Parameters:

Name	Type	Description	Default
residual_fn	Callable[..., Any]	Equation callable or a nested :class:functools.partial wrapping one.	required

Returns:

Type	Description
ResidualStream	Decorated default stream, or "residual" when no stream metadata is attached.

Data fidelity

base_data_fidelity supports direct component supervision, optional sample weights, vector projection, and target negation. phase_wrapped_fidelity supports the same options and represents periodic observations with cosine and sine residuals in one loss group.

Set weight_key="confidence" on either equation to read residual weights from batch["confidence"]. The default is "weight"; a missing field or weight_key=None disables weighting. Weights multiply residuals before loss reduction, including both cosine and sine residuals for phase-wrapped fidelity. Use shape (N, 1) for sample-wise weights that broadcast across residual components. The output stream remains unweighted.

Set projection_key="direction" on either equation to read projection directions from batch["direction"]. The default is "projection"; a missing field or projection_key=None disables projection. Directions must match the shape and component order selected by output_indices. Projection takes the dot product along the last axis and keeps a trailing singleton dimension, so target_indices must select one component. Phase-wrapped fidelity projects predictions before converting them to phases.

Phase-wrapped fidelity also accepts period_key="wrapping_period" to read sample-wise periods from batch["wrapping_period"]. It defaults to "period". The configured field is required for residuals; the output stream does not read it.

For example, configure phase-wrapped fidelity with custom weight, projection, and period fields:

from functools import partial

from phijax.equations import phase_wrapped_fidelity

fidelity = partial(
    phase_wrapped_fidelity,
    output_indices=(0, 1),
    target_indices=(0,),
    weight_key="confidence",
    projection_key="direction",
    period_key="wrapping_period",
)

Here, batch["direction"] has shape (N, 2), and both the selected targets and batch["wrapping_period"] have shape (N, 1). Both equations support stream="output" for output-based diagnostics; this stream returns selected model outputs before projection. The array-level residual functions accept the direction array directly through projection.

Phase-wrapped fidelity now uses a projection field when present. Set projection_key=None to retain direct supervision for batches that contain an unrelated field with that name.

base_data_fidelity_residual [source]

base_data_fidelity_residual(
    output: Array,
    target: Array,
    *,
    weight: Array | None = None,
    projection: Array | None = None,
    target_negation: bool = False,
) -> jax.Array

Return an unwrapped supervised residual.

Parameters:

Name	Type	Description	Default
output	Array	Predicted scalar or vector values.	required
target	Array	Observed values broadcast-compatible with the resolved output.	required
weight	Array | None	Optional multiplicative residual weights applied after projection.	None
projection	Array | None	Optional direction with the same shape as output. When supplied, the vector output is projected to a trailing singleton dimension before comparison with target.	None
target_negation	bool	Whether to negate targets before subtraction.	False

Returns:

Type	Description
Array	Weighted output - signed_target residual with the broadcast input shape.

Raises:

Type	Description
ValueError	If output and projection shapes do not match.

base_data_fidelity [source]

base_data_fidelity(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    output_indices: Sequence[int] = (0,),
    target_indices: Sequence[int] = (0,),
    weight_key: str | None = "weight",
    projection_key: str | None = "projection",
    target_negation: bool = False,
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate direct supervised residuals or selected model outputs.

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure explicit-state model application callable.	required
model_state	Any	Differentiable model parameter PyTree.	required
batch	ArrayMapping	Arrays containing inputs, targets, and optional weight and projection fields.	required
output_indices	Sequence[int]	Model-output components compared with observations.	(0,)
target_indices	Sequence[int]	Target components aligned with output_indices, or one component for projected output.	(0,)
weight_key	str | None	Optional batch field containing multiplicative residual weights. Defaults to "weight". A missing field or None disables weighting. The output stream is unweighted.	'weight'
projection_key	str | None	Optional batch field containing directions matching the selected output shape and order. Defaults to "projection". A missing field or None disables projection.	'projection'
target_negation	bool	Whether to negate targets before subtraction.	False
stream	ResidualStream	"residual" for direct supervised errors or "output" for selected model outputs.	'residual'

Returns:

Type	Description
ResidualGroups	One residual or output group for scalar loss reduction or derivative-based balancing.

Raises:

Type	Description
ValueError	If the stream, component selection, or optional projection is invalid.

phase_wrapped_residuals [source]

phase_wrapped_residuals(
    output: Array,
    target: Array,
    period: Array,
    *,
    weight: Array | None = None,
    projection: Array | None = None,
    target_negation: bool = False,
) -> tuple[jax.Array, jax.Array]

Return cosine and sine residuals for periodic scalar observations.

Parameters:

Name	Type	Description	Default
output	Array	Predicted scalar or vector values.	required
target	Array	Observed scalar values.	required
period	Array	Sample-wise wrapping periods.	required
weight	Array | None	Optional multiplicative residual weights applied after phase conversion.	None
projection	Array | None	Optional direction with the same shape as output. Projects vector predictions to a trailing singleton dimension before phase conversion.	None
target_negation	bool	Whether to negate targets before phase conversion.	False

Returns:

Type	Description
tuple[Array, Array]	Cosine and sine residual arrays with the broadcast input shape after projection.

Raises:

Type	Description
ValueError	If output and projection shapes do not match.

phase_wrapped_fidelity [source]

phase_wrapped_fidelity(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    output_indices: Sequence[int] = (0,),
    target_indices: Sequence[int] = (0,),
    weight_key: str | None = "weight",
    projection_key: str | None = "projection",
    period_key: str = "period",
    target_negation: bool = False,
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate grouped phase-wrapped fidelity residuals or supervised outputs.

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure explicit-state model application callable.	required
model_state	Any	Differentiable model parameter PyTree.	required
batch	ArrayMapping	Arrays containing inputs, targets, the configured period field, and optional weight and projection fields.	required
output_indices	Sequence[int]	Model-output components compared with observations.	(0,)
target_indices	Sequence[int]	Target components aligned with output_indices, or one component for projected output.	(0,)
weight_key	str | None	Optional batch field containing multiplicative residual weights. Defaults to "weight". A missing field or None disables weighting. The output stream is unweighted.	'weight'
projection_key	str | None	Optional batch field containing directions matching the selected output shape and order. Defaults to "projection". A missing field or None disables projection.	'projection'
period_key	str	Batch field containing sample-wise wrapping periods. Required for the residual stream only.	'period'
target_negation	bool	Whether to negate targets before phase conversion.	False
stream	ResidualStream	"residual" for cosine and sine errors or "output" for selected model outputs.	'residual'

Returns:

Type	Description
ResidualGroups	One residual group containing cosine and sine arrays, or one selected-output group for output-based NTK
ResidualGroups	diagnostics.

Raises:

Type	Description
ValueError	If the stream, component selection, or optional projection is invalid.
KeyError	If the configured period field is missing when evaluating residuals.

Boundary conditions

no_slip_boundary constrains every selected velocity component to the wall velocity in batch["targets"]. Use zero targets for a stationary wall or prescribed velocities for a moving wall. output_indices and target_indices default to (0, 1) and must select the same number of components in matching order. For three velocity components, set both to (0, 1, 2). The equation returns one no_slip group containing the component-wise residual array. Its default NTK stream is "output", which returns selected predictions without reading targets. No wall normals are required.

free_slip_boundary accepts normals_key="wall_normals" to read wall-normal directions from batch["wall_normals"]. The default remains "normals". Normal components must follow the order selected by output_indices and have the same final width. The configured field is required for residuals; stream="output" returns the selected model outputs without reading normals.

base_boundary_residual [source]

base_boundary_residual(
    output: Array, target: Array
) -> jax.Array

Return a direct supervised boundary residual.

Parameters:

Name	Type	Description	Default
output	Array	Predicted boundary values with shape [..., components].	required
target	Array	Reference boundary values broadcast-compatible with output and the same final width.	required

Returns:

Type	Description
Array	Direct output - target residual with the broadcast input shape.

Raises:

Type	Description
ValueError	If output and target final widths do not match.

no_slip_residual [source]

no_slip_residual(
    velocity: Array, target: Array
) -> jax.Array

Constrain every velocity component to a reference wall velocity.

Parameters:

Name	Type	Description	Default
velocity	Array	Predicted velocity components with shape [..., dimensions].	required
target	Array	Reference wall velocity broadcast-compatible with velocity.	required

Returns:

Type	Description
Array	Component-wise no-slip residual with the broadcast input shape.

Raises:

Type	Description
ValueError	If the final velocity and target widths do not match.

no_slip_boundary [source]

no_slip_boundary(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    output_indices: Sequence[int] = (0, 1),
    target_indices: Sequence[int] = (0, 1),
    stream: ResidualStream = "residual",
) -> ResidualGroups

Constrain selected velocity components to a reference wall velocity.

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure explicit-state model application callable.	required
model_state	Any	Differentiable model parameter PyTree.	required
batch	ArrayMapping	Arrays containing inputs and wall-velocity targets. Use zero targets for a stationary wall. Targets are required for the residual stream only.	required
output_indices	Sequence[int]	Model-output velocity components constrained at the boundary.	(0, 1)
target_indices	Sequence[int]	Target velocity components aligned with output_indices.	(0, 1)
stream	ResidualStream	"residual" for component-wise velocity errors or "output" for selected model outputs.	'residual'

Returns:

Type	Description
ResidualGroups	One group containing a vector residual with shape (N, len(output_indices)), or one selected-output
ResidualGroups	group for output-based NTK diagnostics. Target batch dimensions may broadcast against the outputs.

Raises:

Type	Description
ValueError	If the stream or component selection is invalid.
KeyError	If inputs is missing, or targets is missing when evaluating residuals.

free_slip_residual [source]

free_slip_residual(
    velocity: Array, target: Array, normals: Array
) -> jax.Array

Project velocity mismatch onto the wall-normal direction.

Parameters:

Name	Type	Description	Default
velocity	Array	Predicted velocity components with shape [..., dimensions].	required
target	Array	Reference wall velocity broadcast-compatible with velocity.	required
normals	Array	Wall-normal components with the same final width as velocity.	required

Returns:

Type	Description
Array	Scalar normal residual per sample with shape [..., 1].

Raises:

Type	Description
ValueError	If the final velocity, target, and normal dimensions do not match.

free_slip_boundary [source]

free_slip_boundary(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    output_indices: Sequence[int] = (0, 1),
    target_indices: Sequence[int] = (0, 1),
    normals_key: str = "normals",
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate grouped free-slip boundary residuals or constrained outputs.

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure explicit-state model application callable.	required
model_state	Any	Differentiable model parameter PyTree.	required
batch	ArrayMapping	Arrays containing inputs, targets, and the configured normals field.	required
output_indices	Sequence[int]	Model-output velocity components constrained at the boundary.	(0, 1)
target_indices	Sequence[int]	Target velocity components aligned with output_indices.	(0, 1)
normals_key	str	Batch field containing wall-normal components aligned with output_indices. Defaults to "normals". Required for the residual stream only.	'normals'
stream	ResidualStream	"residual" for the normal-velocity error or "output" for selected model outputs.	'residual'

Returns:

Type	Description
ResidualGroups	One residual group for scalar loss reduction, or one selected-output group for output-based NTK diagnostics.

Raises:

Type	Description
ValueError	If the stream or component selection is invalid.
KeyError	If the configured normals field is missing when evaluating residuals.

One-dimensional Burgers equation

Inputs use [t, x]; output_index selects the scalar solution. The equation is

\[ \frac{\partial u}{\partial t} + u\frac{\partial u}{\partial x} - \nu\frac{\partial^2 u}{\partial x^2} = 0, \]

where \(\nu\) is viscosity_coefficient. The function returns the left-hand side as the residual optimized during training.

burgers_1d [source]

burgers_1d(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    viscosity_coefficient: float = 0.01 / math.pi,
    output_index: int = 0,
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate a grouped one-dimensional Burgers equation residual.

Coordinates are ordered as [t, x]. For the selected scalar model output u, the returned residual is du/dt + u * du/dx - viscosity_coefficient * d2u/dx2. A zero viscosity selects the inviscid path and avoids tracing a second spatial derivative. The viscous path computes only d2u/dx2, without tracing the unused mixed derivative d2u/(dx dt).

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure callable mapping (model_state, point) to model outputs.	required
model_state	Any	Explicit differentiable model state.	required
batch	ArrayMapping	Arrays containing rank-two [t, x] coordinates under inputs.	required
viscosity_coefficient	float	Nonnegative coefficient multiplying d2u/dx2.	0.01 / pi
output_index	int	Nonnegative scalar model-output component representing u.	0
stream	ResidualStream	Equation representation requested by the objective. Burgers supports only "residual".	'residual'

Returns:

Type	Description
ResidualGroups	One single-array residual group.

Raises:

Type	Description
ValueError	If stream requests raw model outputs or the equation settings are invalid.

Cartesian Navier--Stokes

Dimension	Input order	Output order	Residual order
2D	[x, y, t]	[u_x, u_y, p]	continuity, x momentum, y momentum
3D	[x, y, z, t]	[u_x, u_y, u_z, p]	continuity, x/y/z momentum

A zero viscosity_coefficient selects an inviscid path that does not trace second derivatives.

cartesian_2d_navier_stokes [source]

cartesian_2d_navier_stokes(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    pressure_coefficient: float = 1.0,
    viscosity_coefficient: float = 0.0,
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate grouped incompressible Navier--Stokes residuals in two Cartesian dimensions.

The coordinate order is [x, y, t], model outputs are [u_x, u_y, p], and the returned groups represent continuity, x momentum, and y momentum. With c_p=pressure_coefficient and c_nu=viscosity_coefficient, the residuals are

\[ \begin{aligned} R_c &= \frac{\partial u_x}{\partial x} + \frac{\partial u_y}{\partial y}, \\ R_x &= \frac{\partial u_x}{\partial t} + u_x\frac{\partial u_x}{\partial x} + u_y\frac{\partial u_x}{\partial y} + c_p\frac{\partial p}{\partial x} - c_\nu\left( \frac{\partial^2 u_x}{\partial x^2} + \frac{\partial^2 u_x}{\partial y^2} \right), \\ R_y &= \frac{\partial u_y}{\partial t} + u_x\frac{\partial u_y}{\partial x} + u_y\frac{\partial u_y}{\partial y} + c_p\frac{\partial p}{\partial y} - c_\nu\left( \frac{\partial^2 u_y}{\partial x^2} + \frac{\partial^2 u_y}{\partial y^2} \right). \end{aligned} \]

A zero viscosity_coefficient selects the inviscid path and avoids tracing second derivatives.

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure callable mapping (model_state, point) to [u_x, u_y, p].	required
model_state	Any	Explicit differentiable model state.	required
batch	ArrayMapping	Arrays containing rank-two physical coordinates under inputs, with columns [x, y, t].	required
pressure_coefficient	float	Coefficient multiplying Cartesian pressure gradients.	1.0
viscosity_coefficient	float	Nonnegative coefficient multiplying Cartesian velocity Laplacians. Use 0.0 for inviscid flow or 1 / Re for a nondimensional viscous equation.	0.0
stream	ResidualStream	Equation representation requested by the objective. Navier--Stokes supports only "residual".	'residual'

Returns:

Type	Description
ResidualGroups	Three single-array groups ordered as continuity, x momentum, and y momentum.

Raises:

Type	Description
ValueError	If stream requests model outputs or the equation settings are invalid.

cartesian_3d_navier_stokes [source]

cartesian_3d_navier_stokes(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    pressure_coefficient: float = 1.0,
    viscosity_coefficient: float = 0.0,
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate grouped incompressible Navier--Stokes residuals in three Cartesian dimensions.

The coordinate order is [x, y, z, t], model outputs are [u_x, u_y, u_z, p], and the returned groups represent continuity and the x, y, and z momentum equations. With c_p=pressure_coefficient and c_nu=viscosity_coefficient, the residuals are

\[ \begin{aligned} R_c &= \frac{\partial u_x}{\partial x} + \frac{\partial u_y}{\partial y} + \frac{\partial u_z}{\partial z}, \\ R_x &= \frac{\partial u_x}{\partial t} + u_x\frac{\partial u_x}{\partial x} + u_y\frac{\partial u_x}{\partial y} + u_z\frac{\partial u_x}{\partial z} + c_p\frac{\partial p}{\partial x} - c_\nu\nabla^2u_x, \\ R_y &= \frac{\partial u_y}{\partial t} + u_x\frac{\partial u_y}{\partial x} + u_y\frac{\partial u_y}{\partial y} + u_z\frac{\partial u_y}{\partial z} + c_p\frac{\partial p}{\partial y} - c_\nu\nabla^2u_y, \\ R_z &= \frac{\partial u_z}{\partial t} + u_x\frac{\partial u_z}{\partial x} + u_y\frac{\partial u_z}{\partial y} + u_z\frac{\partial u_z}{\partial z} + c_p\frac{\partial p}{\partial z} - c_\nu\nabla^2u_z, \end{aligned} \]

where \(\nabla^2u_i = \partial^2u_i/\partial x^2 + \partial^2u_i/\partial y^2 + \partial^2u_i/\partial z^2\).

A zero viscosity_coefficient selects the inviscid path and avoids tracing second derivatives.

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure callable mapping (model_state, point) to [u_x, u_y, u_z, p].	required
model_state	Any	Explicit differentiable model state.	required
batch	ArrayMapping	Arrays containing rank-two physical coordinates under inputs, with columns [x, y, z, t].	required
pressure_coefficient	float	Coefficient multiplying Cartesian pressure gradients.	1.0
viscosity_coefficient	float	Nonnegative coefficient multiplying Cartesian velocity Laplacians. Use 0.0 for inviscid flow or 1 / Re for a nondimensional viscous equation.	0.0
stream	ResidualStream	Equation representation requested by the objective. Navier--Stokes supports only "residual".	'residual'

Returns:

Type	Description
ResidualGroups	Four single-array groups ordered as continuity, x momentum, y momentum, and z momentum.

Raises:

Type	Description
ValueError	If stream requests model outputs or the equation settings are invalid.

Polar Navier--Stokes

Polar inputs use [r, th, t] and outputs use [u_r, u_th, p]. radius_epsilon protects terms that divide by the radius. Equation weighting keeps the same zero-residual solutions while reducing singular behavior near the origin.

polar_navier_stokes [source]

polar_navier_stokes(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    pressure_coefficient: float = 1.0,
    viscosity_coefficient: float = 0.0,
    radius_epsilon: float = 1e-12,
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate grouped polar Navier--Stokes equation residuals.

The coordinate order is [r, th, t], and model outputs are [u_r, u_th, p]. The returned groups represent continuity, radial momentum, and angular momentum. All reciprocal-radius terms use max(r, radius_epsilon). A zero viscosity_coefficient selects the inviscid path and avoids tracing second derivatives. A nonzero coefficient adds the polar vector-Laplacian terms used by the viscous incompressible Navier--Stokes equations. The viscous path computes only pure radial and angular second derivatives of velocity, excluding unused mixed and pressure Hessian entries.

With \(c_p\) equal to pressure_coefficient, define the continuity and nonviscous momentum terms as

\[ \begin{aligned} C &= \frac{\partial u_r}{\partial r} + \frac{u_r}{r} + \frac{1}{r}\frac{\partial u_\theta}{\partial \theta}, \\ F_r &= \frac{\partial u_r}{\partial t} + u_r\frac{\partial u_r}{\partial r} + \frac{u_\theta}{r}\frac{\partial u_r}{\partial \theta} - \frac{u_\theta^2}{r} + c_p\frac{\partial p}{\partial r}, \\ F_\theta &= \frac{\partial u_\theta}{\partial t} + u_r\frac{\partial u_\theta}{\partial r} + \frac{u_\theta}{r}\frac{\partial u_\theta}{\partial \theta} + \frac{u_ru_\theta}{r} + \frac{c_p}{r}\frac{\partial p}{\partial \theta}. \end{aligned} \]

Define the scalar polar Laplacian as

\[ \mathcal{L}(f) = \frac{\partial^2 f}{\partial r^2} + \frac{1}{r}\frac{\partial f}{\partial r} + \frac{1}{r^2}\frac{\partial^2 f}{\partial \theta^2}. \]

The viscous velocity terms are

\[ \begin{aligned} (\nabla^2\mathbf{u})_r &= \mathcal{L}(u_r) - \frac{u_r}{r^2} - \frac{2}{r^2}\frac{\partial u_\theta}{\partial \theta}, \\ (\nabla^2\mathbf{u})_\theta &= \mathcal{L}(u_\theta) - \frac{u_\theta}{r^2} + \frac{2}{r^2}\frac{\partial u_r}{\partial \theta}. \end{aligned} \]

With \(c_\nu\) equal to viscosity_coefficient, the returned residuals are

\[ R_c = C, \qquad R_r = F_r - c_\nu(\nabla^2\mathbf{u})_r, \qquad R_\theta = F_\theta - c_\nu(\nabla^2\mathbf{u})_\theta. \]

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure callable mapping (model_state, point) to [u_r, u_th, p].	required
model_state	Any	Explicit differentiable model state.	required
batch	ArrayMapping	Arrays containing rank-two [r, th, t] coordinates under inputs.	required
pressure_coefficient	float	Coefficient multiplying radial and angular pressure gradients.	1.0
viscosity_coefficient	float	Coefficient multiplying the polar vector-Laplacian terms. Use 0.0 for inviscid flow or 1 / Re for a nondimensional viscous equation. This value must be static while tracing.	0.0
radius_epsilon	float	Minimum radius used in reciprocal geometric terms.	1e-12
stream	ResidualStream	Equation representation requested by the objective. Polar Navier--Stokes supports only "residual".	'residual'

Returns:

Type	Description
ResidualGroups	Three single-array groups ordered as continuity, radial momentum, and angular momentum.

Raises:

Type	Description
ValueError	If stream requests raw model outputs or the equation settings are invalid.

Spherical Navier--Stokes

Spherical inputs use [r, th, phi, t] and outputs use [u_r, u_th, u_phi, p]. radius_epsilon and sine_epsilon guard geometric singularities. Residuals are weighted by r * sin(th) as described in the function documentation.

spherical_navier_stokes [source]

spherical_navier_stokes(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    pressure_coefficient: float = 1.0,
    viscosity_coefficient: float = 0.0,
    radius_epsilon: float = 1e-12,
    sine_epsilon: float = 1e-12,
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate grouped weighted spherical Navier--Stokes residuals.

The coordinate order is [r, th, phi, t], model outputs are [u_r, u_th, u_phi, p], and the returned groups represent continuity and radial, polar, and azimuthal momentum. Every physical residual is multiplied by r sin(th). Away from the spherical axis this preserves its zero set, while removing reciprocal-sine factors from the inviscid equations and reducing their influence near the coordinate singularity.

With \(c_p\) equal to pressure_coefficient, define the unweighted continuity and nonviscous momentum terms as

\[ \begin{aligned} C &= \frac{\partial u_r}{\partial r} + \frac{2u_r}{r} + \frac{1}{r}\frac{\partial u_\theta}{\partial \theta} + \frac{u_\theta\cot\theta}{r} + \frac{1}{r\sin\theta}\frac{\partial u_\phi}{\partial \phi}, \\ F_r &= \frac{\partial u_r}{\partial t} + u_r\frac{\partial u_r}{\partial r} + \frac{u_\theta}{r}\frac{\partial u_r}{\partial \theta} + \frac{u_\phi}{r\sin\theta}\frac{\partial u_r}{\partial \phi} - \frac{u_\theta^2 + u_\phi^2}{r} + c_p\frac{\partial p}{\partial r}, \\ F_\theta &= \frac{\partial u_\theta}{\partial t} + u_r\frac{\partial u_\theta}{\partial r} + \frac{u_\theta}{r}\frac{\partial u_\theta}{\partial \theta} + \frac{u_\phi}{r\sin\theta}\frac{\partial u_\theta}{\partial \phi} + \frac{u_ru_\theta}{r} - \frac{u_\phi^2\cot\theta}{r} + \frac{c_p}{r}\frac{\partial p}{\partial \theta}, \\ F_\phi &= \frac{\partial u_\phi}{\partial t} + u_r\frac{\partial u_\phi}{\partial r} + \frac{u_\theta}{r}\frac{\partial u_\phi}{\partial \theta} + \frac{u_\phi}{r\sin\theta}\frac{\partial u_\phi}{\partial \phi} + \frac{u_ru_\phi}{r} + \frac{u_\theta u_\phi\cot\theta}{r} + \frac{c_p}{r\sin\theta}\frac{\partial p}{\partial \phi}. \end{aligned} \]

Define the scalar spherical Laplacian as

\[ \mathcal{L}(f) = \frac{\partial^2 f}{\partial r^2} + \frac{2}{r}\frac{\partial f}{\partial r} + \frac{1}{r^2}\left( \frac{\partial^2 f}{\partial \theta^2} + \cot\theta\frac{\partial f}{\partial \theta} + \frac{1}{\sin^2\theta}\frac{\partial^2 f}{\partial \phi^2} \right). \]

The unweighted viscous velocity terms are

\[ \begin{aligned} (\nabla^2\mathbf{u})_r &= \mathcal{L}(u_r) - \frac{2u_r}{r^2} - \frac{2}{r^2}\left( \frac{\partial u_\theta}{\partial \theta} + u_\theta\cot\theta \right) - \frac{2}{r^2\sin\theta}\frac{\partial u_\phi}{\partial \phi}, \\ (\nabla^2\mathbf{u})_\theta &= \mathcal{L}(u_\theta) + \frac{2}{r^2}\frac{\partial u_r}{\partial \theta} - \frac{u_\theta}{r^2\sin^2\theta} - \frac{2\cos\theta}{r^2\sin^2\theta} \frac{\partial u_\phi}{\partial \phi}, \\ (\nabla^2\mathbf{u})_\phi &= \mathcal{L}(u_\phi) + \frac{2}{r^2\sin\theta}\frac{\partial u_r}{\partial \phi} + \frac{2\cos\theta}{r^2\sin^2\theta} \frac{\partial u_\theta}{\partial \phi} - \frac{u_\phi}{r^2\sin^2\theta}. \end{aligned} \]

With \(c_\nu\) equal to viscosity_coefficient, the returned weighted residuals are

\[ R_c = r\sin\theta\,C, \qquad R_r = r\sin\theta\left[F_r - c_\nu(\nabla^2\mathbf{u})_r\right], \qquad R_\theta = r\sin\theta\left[F_\theta - c_\nu(\nabla^2\mathbf{u})_\theta\right], \qquad R_\phi = r\sin\theta\left[F_\phi - c_\nu(\nabla^2\mathbf{u})_\phi\right]. \]

A zero viscosity_coefficient selects the inviscid path and avoids tracing all second derivatives.

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure callable mapping (model_state, point) to [u_r, u_th, u_phi, p].	required
model_state	Any	Explicit differentiable model state.	required
batch	ArrayMapping	Arrays containing rank-two [r, th, phi, t] coordinates under inputs.	required
pressure_coefficient	float	Coefficient multiplying spherical pressure gradients.	1.0
viscosity_coefficient	float	Nonnegative coefficient multiplying spherical vector-Laplacian terms. Use 0.0 for inviscid flow or 1 / Re for a nondimensional viscous equation.	0.0
radius_epsilon	float	Minimum radius used in reciprocal geometric terms.	1e-12
sine_epsilon	float	Minimum value of sin(th) used in reciprocal geometric terms.	1e-12
stream	ResidualStream	Equation representation requested by the objective. Navier--Stokes supports only "residual".	'residual'

Returns:

Type	Description
ResidualGroups	Four single-array groups ordered as continuity, radial momentum, polar momentum, and azimuthal momentum.

Raises:

Type	Description
ValueError	If stream requests raw model outputs or the equation settings are invalid.

See Building equations and objectives for residual-group design and analytic tests.