Fields API Reference

JAX-native field abstractions for scientific computing on structured grids.

Field abstractions for scientific computing on structured grids.

Provides immutable JAX pytree field types and differential operators inspired by PhiFlow, implemented in pure JAX.

Reference

Holl et al. "PhiFlow: A Differentiable PDE Solving Framework"

Box

Box(lower: tuple[float, ...], upper: tuple[float, ...])

Rectangular physical domain.

Stores the lower and upper bounds of an axis-aligned box in N dimensions. Immutable — all properties are derived from the bounds.

Attributes:

Name	Type	Description
lower	ndarray	Lower corner coordinates.
upper	ndarray	Upper corner coordinates.

Parameters:

Name	Type	Description	Default
lower	tuple[float, ...]	Lower corner coordinates.	required
upper	tuple[float, ...]	Upper corner coordinates.	required

Raises:

Type	Description
ValueError	If dimensions don't match or bounds are invalid.

lower property

lower: ndarray

Lower corner coordinates.

upper property

upper: ndarray

Upper corner coordinates.

spatial_dim property

spatial_dim: int

Number of spatial dimensions.

size property

size: ndarray

Box extent in each dimension.

center property

center: ndarray

Box center point.

CenteredGrid

CenteredGrid(values: ndarray, box: Box, extrapolation: Extrapolation = ZERO)

Scalar or vector field on a cell-centered uniform grid.

Values are stored at cell centers. The grid carries domain metadata (physical bounds and boundary conditions) and supports element-wise arithmetic. Registered as a JAX pytree so it can be used with jax.jit, jax.vmap, and jax.grad.

Attributes:

Name	Type	Description
values		Field data array, shape matches grid resolution.
box		Physical domain bounds.
extrapolation		Boundary condition type.

Parameters:

Name	Type	Description	Default
values	ndarray	Field values at cell centers.	required
box	Box	Physical domain bounds.	required
extrapolation	Extrapolation	Boundary condition type.	ZERO

shape property

shape: tuple[int, ...]

Grid shape.

spatial_dim property

spatial_dim: int

Number of spatial dimensions.

resolution property

resolution: tuple[int, ...]

Grid resolution (number of cells per dimension).

dx property

dx: ndarray

Cell size in each dimension.

cell_centers

cell_centers() -> ndarray

Compute physical coordinates of cell centers.

Returns:

Type	Description
ndarray	Meshgrid array of shape (*resolution, spatial_dim).

tree_flatten

tree_flatten()

Flatten for JAX pytree protocol.

tree_unflatten classmethod

tree_unflatten(aux_data, children)

Unflatten from JAX pytree protocol.

Extrapolation

Bases: Enum

Boundary extrapolation types for grid fields.

Determines how values outside the grid domain are handled.

maccormack

maccormack(field: CenteredGrid, velocity: CenteredGrid, dt: float, correction_strength: float = 1.0) -> CenteredGrid

Advect using MacCormack scheme (higher-order correction).

Performs a semi-Lagrangian step, then traces forward to estimate the error, and applies a correction. Clamped to prevent overshoots.

Parameters:

Name	Type	Description	Default
field	CenteredGrid	Scalar field to advect.	required
velocity	CenteredGrid	Velocity field.	required
dt	float	Time step.	required
correction_strength	float	Blending factor for error correction (0-1).	1.0

Returns:

Type	Description
CenteredGrid	Advected field with reduced numerical diffusion.

semi_lagrangian

semi_lagrangian(field: CenteredGrid, velocity: CenteredGrid, dt: float) -> CenteredGrid

Advect a field using semi-Lagrangian method.

Traces particles backward in time by -dt and samples the field at the departure points. Unconditionally stable for any dt.

Parameters:

Name	Type	Description	Default
field	CenteredGrid	Scalar field to advect, shape (*resolution).	required
velocity	CenteredGrid	Velocity field, shape (*resolution, ndim).	required
dt	float	Time step.	required

Returns:

Type	Description
CenteredGrid	Advected field at time t + dt.

curl_2d

curl_2d(field: CenteredGrid) -> CenteredGrid

Compute 2D curl (vorticity) of a vector field.

For v = (vx, vy), curl = ∂vy/∂x - ∂vx/∂y.

Parameters:

Name	Type	Description	Default
field	CenteredGrid	2D vector CenteredGrid with shape (Nx, Ny, 2).	required

Returns:

Type	Description
CenteredGrid	Scalar CenteredGrid with vorticity values.

divergence

divergence(field: CenteredGrid) -> CenteredGrid

Compute divergence of a vector field using central differences.

For a vector field v with shape (*resolution, spatial_dim), returns the scalar divergence ∇·v.

Parameters:

Name	Type	Description	Default
field	CenteredGrid	Vector CenteredGrid with trailing vector dimension.	required

Returns:

Type	Description
CenteredGrid	Scalar CenteredGrid with divergence values.

gradient

gradient(field: CenteredGrid) -> CenteredGrid

Compute spatial gradient using central finite differences.

For a scalar field u, returns ∇u as a vector field with an extra trailing dimension for the gradient components.

Parameters:

Name	Type	Description	Default
field	CenteredGrid	Scalar CenteredGrid.	required

Returns:

Type	Description
CenteredGrid	CenteredGrid with values shape (*resolution, spatial_dim).

laplacian

laplacian(field: CenteredGrid) -> CenteredGrid

Compute Laplacian using second-order central differences.

∇²u = Σ_d (u[i+1] - 2u[i] + u[i-1]) / dx_d²

Parameters:

Name	Type	Description	Default
field	CenteredGrid	Scalar CenteredGrid.	required

Returns:

Type	Description
CenteredGrid	Scalar CenteredGrid with Laplacian values.

pressure_solve_jacobi

pressure_solve_jacobi(velocity: CenteredGrid, n_iterations: int = 100, omega: float = 1.0) -> tuple[CenteredGrid, CenteredGrid]

Solve pressure Poisson equation using Jacobi iteration.

Works with any boundary condition. Slower than spectral but more general.

Parameters:

Name	Type	Description	Default
velocity	CenteredGrid	Vector velocity field.	required
n_iterations	int	Number of Jacobi iterations.	100
omega	float	Relaxation parameter (1.0 = standard, >1 = SOR).	1.0

Returns:

Type	Description
tuple[CenteredGrid, CenteredGrid]	Tuple of (divergence_free_velocity, pressure).

pressure_solve_spectral

pressure_solve_spectral(velocity: CenteredGrid) -> tuple[CenteredGrid, CenteredGrid]

Solve pressure Poisson equation using FFT (periodic boundaries).

Computes p such that ∇²p = ∇·v, then returns the divergence-free velocity v* = v - ∇p and the pressure field p.

Only works with periodic boundary conditions.

Parameters:

Name	Type	Description	Default
velocity	CenteredGrid	Vector velocity field, shape (*resolution, ndim).	required

Returns:

Type	Description
tuple[CenteredGrid, CenteredGrid]	Tuple of (divergence_free_velocity, pressure).

Raises:

Type	Description
ValueError	If boundary conditions are not periodic.