Physical Metrics

Physical Metrics

neojax.metrics.physical_metrics.BoundaryConsistencyMetric

Boundary consistency metric.

Evaluates the prediction alignment with the boundary condition values at the locations specified by the boundary condition mask.

\[ \text{BoundaryConsistency} = \left( \frac{\sum_{t, c, x} |(\hat{y}_{t,c,x} - y_{\text{bc}, t,c,x}) \cdot M_{c,x}|^p}{\sum_{t, c, x} M_{c,x} + \epsilon} \right)^{1/p} \]

where \(y_{\text{bc}}\) is the boundary values, \(\hat{y}\) is the prediction, \(M\) is the boundary mask, and \(\epsilon = 10^{-7}\) is a small regularization constant.

Parameters:

Name	Type	Description	Default
p	float	Power of the error norm. Default is 2.0.	2.0
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.
• learnable_weight (bool): Flag indicating whether weight is learnable.
• p (float): Power of the norm.

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

Computes the boundary consistency metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained. Optional if pred is given.	None
target	Float[Array, 'b c ...'] | None	Unused in boundary consistency metric.	None
x	Float[Array, 'b in_c ...'] | None	Model input array shaped (batch, in_c, d1, ..., dN).	None
pred	Float[Array, 'b c ...'] | None	Model prediction array shaped (batch, t, c, d1, ..., dN). If None, the metric computes vmap(model)(x).	None
bc_masks	Int[Array, 'b c_bc ...'] | None	Boundary mask array shaped (batch, c_bc, d1, ..., dN).	None
bc_values	Float[Array, 'b t c_bc ...'] | None	Boundary values array shaped (batch, t, c_bc, d1, ..., dN).	None
**kwargs	Any	Additional arguments.	required

Returns:

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

Raises:

Type	Description
ValueError	If prediction is None and either model or x is None.

---

neojax.metrics.physical_metrics.ConservationMetric

Conservation law metric.

Evaluates the drift/deviation of a conserved physical quantity (like mass or energy) over a time trajectory relative to the initial state.

\[ \text{Drift} = \left( \frac{1}{T-1} \sum_{t=1}^{T-1} D_t^p \right)^{1/p} \]

where the deviation at time step \(t\) is:

\[ D_t = \begin{cases} |Q_t - Q_0|, & \text{if mode is absolute}, \\\\ \frac{|Q_t - Q_0|}{|Q_0| + \epsilon}, & \text{if mode is relative}. \end{cases} \]

and \(Q_t = \text{conservation_fn}(y_t, x_{\text{coords}})\) is the quantity at time step \(t\).

Parameters:

Name	Type	Description	Default
conservation_fn	Callable[[Float[Array, 'c ...'], Float[Array, 'd ...']], Float[Array, '']]	Callable taking (fields, coords) and returning a scalar. fields is shaped (c, *spatial) and coords is shaped (d, *spatial).	required
mode	Literal['absolute', 'relative']	Type of drift deviation, either "absolute" or "relative". Default is "absolute".	'absolute'
p	float	Power of the norm. Default is 2.0.	2.0
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.
• learnable_weight (bool): Flag indicating whether weight is learnable.
• conservation_fn (Callable): Conserved quantity evaluator.
• mode (Literal["absolute", "relative"]): Drift deviation mode.
• p (float): Power of the norm.

Info

May be used to ensure the model doesn't violate conservation of energy or mass.

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

Computes the conservation law metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained. Optional if pred is given.	None
target	Float[Array, 'b c ...'] | None	Unused in conservation metric.	None
x	Float[Array, 'b in_c ...'] | None	Model input array shaped (batch, in_c, d1, ..., dN).	None
pred	Float[Array, 'b c ...'] | None	Model prediction array shaped (batch, t, c, d1, ..., dN). If None, the metric computes vmap(model)(x).	None
coords	Float[Array, 'd ...'] | None	Physical coordinates shaped (d, d1, ..., dN).	None
**kwargs	Any	Additional arguments.	required

Returns:

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

Raises:

Type	Description
ValueError	If coords is None.
ValueError	If prediction is None and either model or x is None.

---

neojax.metrics.physical_metrics.ResidualMetric

PDE residual metric.

Evaluates the norm of the PDE residual using the user-provided residual function.

\[ \text{Residual} = \left( \frac{1}{N_{\text{res}}} \sum_{t, c, x} |R(\hat{y}, x_{\text{coords}})|^p \right)^{1/p} \]

where \(R\) is the residual evaluator.

Parameters:

Name	Type	Description	Default
residual_fn	Callable[[Float[Array, 't c ...'], Float[Array, 'd ...']], Float[Array, ...]]	Callable taking (fields, coords) and returning residual fields. fields is prediction trajectory shaped (t, c, *spatial) and coords is shaped (d, *spatial). Output must have shape (t_res, c_res, *spatial_res).	required
p	float	Power of the norm. Default is 2.0.	2.0
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.
• learnable_weight (bool): Flag indicating whether weight is learnable.
• residual_fn (Callable): Residual field evaluator.
• p (float): Power of the norm.

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

Computes the PDE residual metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained. Optional if pred is given.	None
target	Float[Array, 'b c ...'] | None	Unused in residual metric.	None
x	Float[Array, 'b in_c ...'] | None	Model input array shaped (batch, in_c, d1, ..., dN).	None
pred	Float[Array, 'b c ...'] | None	Model prediction array shaped (batch, t, c, d1, ..., dN). If None, the metric computes vmap(model)(x).	None
coords	Float[Array, 'd ...'] | None	Physical coordinates shaped (d, d1, ..., dN).	None
**kwargs	Any	Additional arguments.	required

Returns:

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

Raises:

Type	Description
ValueError	If coords is None.
ValueError	If prediction is None and either model or x is None.