Sampling

Structured sampling yields batches that preserve axis meaning (via named axes), so that operators and constraints can keep shape semantics without manual broadcasting.

For a conceptual overview (product structures, components, and coord-separable grids), see Guides → Domains and sampling.

Paired vs coord-separable sampling

Phydrax supports two complementary structured sampling modes:

• Paired sampling (PointsBatch): samples points in each block of a ProductStructure. This is the default mode used by most pointwise PDE residual constraints.
• Coord-separable sampling (CoordSeparableBatch): samples 1D coordinate axes for selected unary labels (geometry and/or scalar intervals) and evaluates on the implied Cartesian grid (with an interior mask). This is the natural mode for FFT/basis/spectral operators and neural operators (FNO, DeepONet).

Coord-separable sampling is driven by DomainComponent.sample_coord_separable(...), which takes:

• coord_separable: a mapping from unary label (e.g. "x" or "t") to either counts (int / Sequence[int]) or basis-aware axis specs (AbstractAxisSpec implementations / GridSpec);
• dense_structure + num_points: how to sample any remaining non-fixed, non-separable labels (e.g. "data" for operator-learning datasets).

Example

Coord-separable grid evaluation on an interval:

import jax.random as jr
import phydrax as phx

geom = phx.domain.Interval1d(0.0, 1.0)
component = geom.component()

batch = component.sample_coord_separable(
    {"x": phx.domain.FourierAxisSpec(64)},
    key=jr.key(0),
)

Note

A CoordSeparableBatch stores:

• coord_axes_by_label: per-label axis names (for shape/dims inference),
• coord_mask_by_label: per-label interior masks on the Cartesian grid,
• axis_discretization_by_axis: optional per-axis metadata (nodes/weights/basis), used by quadrature and basis backends.

phydrax.domain.ProductStructure

Describes how a product domain is sampled into named axes.

Phydrax uses labeled product domains, e.g. \(\Omega = \Omega_x \times \Omega_t\) with labels "x" and "t". When sampling points from a DomainComponent, we often want to control which variables are sampled jointly and which are sampled in separate blocks.

A ProductStructure defines a partition of the (non-fixed) labels into blocks (B_1, \dots, B_k). For each block \(B_j\) we sample \(n_j\) joint points in \(\prod_{\ell\in B_j}\Omega_\ell\). Each block corresponds to one named sampling axis (e.g. __phydra_blk__x__t), and values evaluated on the resulting PointsBatch are coordax.Fields carrying those axis names.

For example:

• ProductStructure((("x", "t"),)) samples paired space-time points.
• ProductStructure((("x",), ("t",))) samples space and time independently, producing a Cartesian product grid in evaluation.

__init__(blocks: tuple[tuple[str, ...], ...], *, axis_names: tuple[str, ...] | None = None)

canonicalize(domain_labels: tuple[str, ...], *, fixed_labels: frozenset[str] = frozenset()) -> ProductStructure

Return a structure with blocks ordered canonically for a domain.

Canonicalization enforces:

• Every non-fixed domain label appears in exactly one block.
• Each block is sorted in the order it appears in domain_labels.
• axis_names is set (generated if missing).

This is required before constructing a PointsBatch, since named axes are derived from the canonical blocks.

axis_for(label: str) -> str | None

Return the sampling axis name corresponding to label (or None if fixed).

---

phydrax.domain.PointsBatch

A labeled batch of sampled points with explicit axis semantics.

A PointsBatch is a mapping {label: coordax.Field} paired with a canonicalized ProductStructure. Each label's point array is stored as a coordax.Field whose named axes correspond to the sampling block(s) that include that label.

If a label is in a block with axis name a, then its sampled points carry a named dimension a of length equal to the number of samples in that block.

__init__(points: phydrax._frozendict.frozendict[str, PyTree[coordax.Field]], structure: ProductStructure)

Construct a PointsBatch from sampled points and a canonical structure.

---

phydrax.domain.QuadratureBatch

Quadrature weights associated with a PointsBatch.

A QuadratureBatch stores one 1D weight field per sampling axis. The total tensor-product weight is

\[ w(z) = \prod_{a \in \mathcal{A}} w_a, \]

where each \(w_a\) is a coordax.Field with dims (a,). These weights are used by integral estimators such as phydrax.operators.integral.

__init__(batch: PointsBatch, *, weights_by_axis: phydrax._frozendict.frozendict[str, coordax.Field] | collections.abc.Mapping[str, coordax.Field])

total_weight() -> Field

Return the product weight field.

If weights_by_axis = {a: w_a}, returns \(w=\prod_a w_a\) as a coordax.Field.

---

phydrax.domain.CoordSeparableBatch

A batch that separates coordinate axes for selected geometry labels.

For some geometries it is efficient to sample each coordinate axis independently, e.g. draw \((x_1,\dots,x_{n_1})\) and \((y_1,\dots,y_{n_2})\) separately and evaluate on the implied grid.

A CoordSeparableBatch stores: - coord_axes_by_label: which named axes correspond to each coordinate component, e.g. ("x0", "x1") for a 2D geometry label. - coord_mask_by_label: a mask coordax.Field that can be used to exclude coordinate combinations outside an irregular geometry (e.g. AABB grid masking). - dense_structure: a normal ProductStructure for any remaining (non-separable) labels sampled in paired blocks.

__init__(points: phydrax._frozendict.frozendict[str, PyTree[coordax.Field]], *, dense_structure: ProductStructure, coord_axes_by_label: phydrax._frozendict.frozendict[str, tuple[str, ...]] | collections.abc.Mapping[str, tuple[str, ...]], coord_mask_by_label: phydrax._frozendict.frozendict[str, coordax.Field] | collections.abc.Mapping[str, coordax.Field], axis_discretization_by_axis: phydrax._frozendict.frozendict[str, phydrax.domain._grid.AxisDiscretization] | collections.abc.Mapping[str, phydrax.domain._grid.AxisDiscretization] | None = None)

Construct a coordinate-separable batch of points.

---

Coord-separable grids

phydrax.domain.AxisDiscretization

Materialized 1D discretization data for a single coordinate axis.

This bundles:

• nodes: 1D coordinates \(x_j\) along an axis.
• quad_weights: optional 1D quadrature weights \(w_j\) for approximating 1D integrals.
• basis: a hint describing how the axis was constructed ("fourier", "sine", "cosine", "legendre", "uniform"),
• periodic: whether the axis should be treated as periodic (useful for FFT/Fourier methods).

When quad_weights are present, they approximate

\[ \int_a^b f(x)\,dx \approx \sum_j w_j f(x_j), \]

__init__(*, nodes: Array, quad_weights: Array | None, basis: Literal[uniform, fourier, sine, cosine, legendre], periodic: bool)

---

phydrax.domain.AbstractAxisSpec

Abstract base class for 1D grid/basis axis specifications.

An AxisSpec is an instruction for how to discretize a 1D coordinate axis on an interval \([a,b]\). Calling materialize(a, b) produces an AxisDiscretization with nodes (and possibly quadrature weights) for that axis.

__init__(n: int)

---

phydrax.domain.GridSpec

A per-label grid spec: one axis spec per coordinate component.

For a geometry variable with var_dim=d, use GridSpec(axes=(spec0, ..., spec{d-1})) to specify a different AxisSpec per coordinate axis.

__init__(axes: Sequence[AbstractAxisSpec])

---

phydrax.domain.UniformAxisSpec

Uniform grid on \([a,b]\).

Uses jax.numpy.linspace(a, b, n, endpoint=...). Quadrature weights default to trapezoid weights when endpoint=True and uniform weights when the axis is treated as periodic (either periodic=True or endpoint=False).

__init__(n: int, *, endpoint: bool = True, periodic: bool = False)

---

phydrax.domain.FourierAxisSpec

Uniform periodic grid for Fourier/FFT methods (endpoint excluded).

Uses the nodes

\[ x_j = a + (b-a)\frac{j}{n},\quad j=0,\dots,n-1, \]

with uniform weights \(w_j=(b-a)/n\). The resulting axis is marked periodic=True.

---

phydrax.domain.SineAxisSpec

Uniform interior grid (cell-centered) suitable for sine-like expansions.

Uses the nodes

\[ x_j = a + (b-a)\frac{j+\tfrac12}{n},\quad j=0,\dots,n-1, \]

with uniform weights \(w_j=(b-a)/n\). The resulting axis is non-periodic.

---

phydrax.domain.CosineAxisSpec

Uniform endpoint-including grid suitable for cosine-like expansions.

Uses the nodes

\[ x_j = a + (b-a)\frac{j}{n-1},\quad j=0,\dots,n-1, \]

with trapezoid weights \(w_0=w_{n-1}=\tfrac12\Delta x\), \(w_j=\Delta x\) otherwise. The resulting axis is non-periodic.

---

phydrax.domain.LegendreAxisSpec

Legendre Gauss/Radau/Lobatto nodes and weights (via orthax).

orthax returns canonical nodes \(\xi_j\in[-1,1]\) and weights \(w_j\), which are mapped to \([a,b]\) via

\[ x_j=\tfrac{b-a}{2}\,\xi_j+\tfrac{a+b}{2},\qquad \tilde w_j=\tfrac{b-a}{2}\,w_j. \]

__init__(n: int, *, kind: Literal[gauss, radau, lobatto] = 'gauss')

---

Axis conventions (nodes + weights)

Many basis-aware operators (spectral derivatives, quadrature) want both:

• nodes \(x_j\) on an axis \([a,b]\),
• quadrature weights \(w_j\) to approximate \(\int_a^b f(x)\,dx \approx \sum_j w_j f(x_j)\).

When you sample with axis specs (AbstractAxisSpec implementations) / GridSpec, Phydrax materializes an AxisDiscretization and attaches it on the batch (so downstream operators can reuse nodes/weights).

Fourier (periodic, endpoint excluded)

For FourierAxisSpec(n):

\[ x_j = a + (b-a)\frac{j}{n},\quad j=0,\dots,n-1,\qquad w_j = \frac{b-a}{n}. \]

Sine (cell-centered interior grid)

For SineAxisSpec(n):

\[ x_j = a + (b-a)\frac{j+\tfrac12}{n},\quad j=0,\dots,n-1,\qquad w_j = \frac{b-a}{n}. \]

Cosine (endpoint grid + trapezoid weights)

For CosineAxisSpec(n):

\[ x_j = a + (b-a)\frac{j}{n-1},\quad j=0,\dots,n-1, \]

and trapezoid weights \(w_0=w_{n-1}=\tfrac12\Delta x\), \(w_j=\Delta x\) otherwise.

Legendre (orthax Gauss / Radau / Lobatto)

For LegendreAxisSpec(n), orthax produces nodes \(\xi_j\in[-1,1]\) and weights \(w_j\) for the canonical interval. Phydrax maps them to \([a,b]\) via

\[ x_j=\tfrac{b-a}{2}\,\xi_j+\tfrac{a+b}{2},\qquad \tilde w_j=\tfrac{b-a}{2}\,w_j. \]