Residuals

popinn.ResidualTerm

Module for evaluating a residual function over a grid of coordinate and auxiliary values, then reducing the residual grid to a scalar value with a specified metric.

Source code in src/popinn/loss.py

class ResidualTerm(eqx.Module):
    """
    Module for evaluating a residual function over a grid of coordinate and auxiliary values,
    then reducing the residual grid to a scalar value with a specified metric.
    """

    name: str = eqx.field(static=True)
    residual_fn: Callable = eqx.field(static=True)
    metric: str | Callable = eqx.field(static=True, default="mse")
    eval_fn: Callable = eqx.field(static=True, default=eval_grid)
    batch_size: int | None = eqx.field(static=True, default=None)

    def __call__(self, model: AbstractModel, data: eqx.Module) -> Float[Array, ""]:
        """
        Evaluate the residual over this term's coordinate and auxiliary grid and reduce to a scalar.

        Coordinates are read from `data.<name>_coords` (where `<name>` is the
        `name` specified at initialization) and the auxiliary
        inputs from `data.aux`.

        Args:
            model (AbstractModel): The network model.
            data (equinox.Module): data container with field carrying the 1D coordinate
            grid axes in `<name>_coords` and the auxiliary inputs in `aux`.

        Returns:
            (Float[Array, '']): the scalar loss for this term.
        """
        coords = getattr(data, self.name + "_coords")
        aux = data.aux
        kw = {}
        if self.batch_size is not None:
            key = "batch_size" if self.eval_fn is not eval_grid else "outer_batch_size"
            kw[key] = self.batch_size
        values = self.eval_fn(self.residual_fn(model), coords, aux, **kw)
        return _reduce(values, self.metric)

__init__(name, residual_fn, metric='mse', eval_fn=eval_grid, batch_size=None)

.

Parameters:

Name	Type	Description	Default
name	str	Name of the term. Should match the name of a coordinate field in the data container as data.<name>_coords.	required
residual_fn	collections.abc.Callable	Function with the signature residual_fn(model) -> Callable, where the returned function is the per-point residual function and has the signature r(*coords, aux) -> Scalar.	required
metric	str | collections.abc.Callable	Metric to compute the loss. Can be 'mse' for mean-squared-err, 'mae' for mean-absolute-error, or a custom Callable mapping a grid of values to a scalar.	'mse'
eval_fn	popinn.eval.eval_grid | popinn.eval.eval_grid_flat_aux	Function that evaluates the model over a grid of coordinates and auxiliary inputs. Can be eval_grid or eval_grid_flat_aux.	popinn.eval.eval_grid
batch_size	int | None	The batch sizes for the memory dial in eval_fn. Setting to None evaluates the whole grid at once (fastest, most memory).	None

__call__(model, data)

Evaluate the residual over this term's coordinate and auxiliary grid and reduce to a scalar.

Coordinates are read from data.<name>_coords (where <name> is the name specified at initialization) and the auxiliary inputs from data.aux.

Parameters:

Name	Type	Description	Default
model	popinn.models.AbstractModel	The network model.	required
data	equinox.Module	data container with field carrying the 1D coordinate	required

Returns:

Type	Description
jaxtyping.Float[jax.Array, '']	the scalar loss for this term.

Source code in src/popinn/loss.py

def __call__(self, model: AbstractModel, data: eqx.Module) -> Float[Array, ""]:
    """
    Evaluate the residual over this term's coordinate and auxiliary grid and reduce to a scalar.

    Coordinates are read from `data.<name>_coords` (where `<name>` is the
    `name` specified at initialization) and the auxiliary
    inputs from `data.aux`.

    Args:
        model (AbstractModel): The network model.
        data (equinox.Module): data container with field carrying the 1D coordinate
        grid axes in `<name>_coords` and the auxiliary inputs in `aux`.

    Returns:
        (Float[Array, '']): the scalar loss for this term.
    """
    coords = getattr(data, self.name + "_coords")
    aux = data.aux
    kw = {}
    if self.batch_size is not None:
        key = "batch_size" if self.eval_fn is not eval_grid else "outer_batch_size"
        kw[key] = self.batch_size
    values = self.eval_fn(self.residual_fn(model), coords, aux, **kw)
    return _reduce(values, self.metric)