Lp-Metrics

Lp-Metrics

The standard approach for training Neural Operators is using continuous \(L^p\) norms.

• LpMetric: Computes the standard \(L^p\) norm of the error over the spatial domain.
• RelativeLpMetric: Computes the \(L^p\) error normalized by the \(L^p\) norm of the true target function. This is standard practice in operator learning because physical data often ranges across different magnitudes, and a relative metric ensures stable gradients across samples.

neojax.metrics.lp_metrics.LpMetric

General \(L^{p}\)-metric.

Computes

\[ \Vert y - \hat{y}\Vert_{L^p} = \Bigg(\int_{\Omega} \Big| y - \hat{y} \Big|^p dx \Bigg)^{\frac{1}{p}} \]

where \(y\) is the ground truth and \(\hat{y}\) is the model prediction.

Parameters:

Name	Type	Description	Default
p	Union[float, int, Literal['inf']]	Power of the norm. Can be a number or "inf" (representing L-infinity norm).	required
weight	float	(Learnable) scalar weight coefficient of the metric. The metric is computed as 'metric * weight'. Default is 1.0. If learnable_weight is True, weight has to be strictly positive and will be optimized via gradient descent.	1.0
learnable_weight	bool	Optional flag indicating whether weight is learnable. If True, this flag is used in the is_learnable_metric_weight filter function to indicate to eqx.filter_... or eqx.partition that weight should be adapted by the optimizer. Default is False.	False

Internal Attributes

These fields store the internal state of the metric.

• weight (Float[Array, ""]): Learnable metric weight. Filter during training to prevent updates.
• p (float | int | Literal["inf"]): Power of the norm.
• learnable_weight (bool): Flag indicating whether weight is learnable.
• is_p_inf (bool): Flag indicating if p is infinity.

__call__(model: equinox._module._module.Module | None = None, *, target: Float[Array, 'b c ...'], x: PyTree[Float[Array, 'b ...']] | None = None, pred: Float[Array, 'b c ...'] | None = None, **kwargs: typing.Any) -> Float[Array, '']

Computes the \(L^{p}\)-metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained. Optional if pred is given.	None
target	Float[Array, 'b c ...']	Ground truth array shaped (batch, c, d1, ..., dN).	required
x	PyTree[Float[Array, 'b ...']] | None	Model input array shaped (batch, in_c, d1, ..., dN).	None
pred	Float[Array, 'b c ...'] | None	Model prediction array. If None, the metric computes vmap(model)(x).	None

Returns:

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

Raises:

Type	Description
ValueError	If pred is None or model and x are None.

---

neojax.metrics.lp_metrics.RelativeLpMetric

General relative \(L^{p}\)-metric.

Computes

\[ \frac{\Vert y - \hat{y}\Vert_{L^p}}{\Vert y \Vert_{L^p}} \]

where \(y\) is the ground truth and \(\hat{y}\) is the model prediction.

Parameters:

Name	Type	Description	Default
p	Union[float, int, Literal['inf']]	Power of the norm. Can be a number or "inf" (representing L-infinity norm).	required
weight	float	(Learnable) scalar weight coefficient of the metric. The metric is computed as 'metric * weight'. Default is 1.0. If learnable_weight is True, weight has to be strictly positive and will be optimized via gradient descent.	1.0
learnable_weight	bool	Optional flag indicating whether weight is learnable. If True, this flag is used in the is_learnable_metric_weight filter function to indicate to eqx.filter_... or eqx.partition that weight should be adapted by the optimizer. Default is False.	False

Internal Attributes

These fields store the internal state of the metric.

• weight (Float[Array, ""]): Learnable metric weight. Filter during training to prevent updates.
• p (float | int | Literal["inf"]): Power of the norm.
• learnable_weight (bool): Flag indicating whether weight is learnable.
• is_p_inf (bool): Flag indicating if p is infinity.

__call__(model: equinox._module._module.Module | None = None, *, target: Float[Array, 'b c ...'], x: PyTree[Float[Array, 'b ...']] | None = None, pred: Float[Array, 'b c ...'] | None = None, **kwargs: typing.Any) -> Float[Array, '']

Computes relative \(L^{p}\)-metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained. Optional if pred is given.	None
target	Float[Array, 'b c ...']	Ground truth array shaped (batch, c, d1, ..., dN).	required
x	PyTree[Float[Array, 'b ...']] | None	Model input array shaped (batch, in_c, d1, ..., dN).	None
pred	Float[Array, 'b c ...'] | None	Model prediction array. If None, the metric computes vmap(model)(x).	None

Returns:

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

Raises:

Type	Description
ValueError	If pred is None or model and x are None.