autocast.metrics.deterministic

Deterministic metrics.

Power-spectrum RMSE utilities in this module based on the implementation from: - Lost in Latent Space: An Empirical Study of Latent Diffusion Models for Physics Emulation (Rozet et al., 2024), https://arxiv.org/abs/2507.02608, PolymathicAI/lola - Specific code from:

• PolymathicAI/lola

• PolymathicAI/lola

class BTSCMetric(reduce_all=True, dist_sync_on_step=False)

Bases: BaseMetric[Float[Tensor, 'batch time spatial *spatial channel'] | Float[Tensor, 'batch time spatial *spatial channel ensemble'], Float[Tensor, 'batch time spatial *spatial channel']]

Base class for metrics that operate on spatial tensors.

Checks input types and shapes and converts to Tensor.

Parameters:

• reduce_all (bool) – If True, return scalar by averaging over all non-batch dims

• dist_sync_on_step (bool) – Synchronize metric state across processes at each forward()

class MSE(reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

Mean Squared Error over spatial dims.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'mse'

class MAE(reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

Mean Absolute Error over spatial dims.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'mae'

class NMAE(reduce_all=True, dist_sync_on_step=False, eps=1e-07)

Bases: BTSCMetric

Normalized Mean Absolute Error over spatial dims.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'nmae'

class NMSE(reduce_all=True, dist_sync_on_step=False, eps=1e-07)

Bases: BTSCMetric

Normalized Mean Squared Error over spatial dims.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'nmse'

class RMSE(reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

Root Mean Squared Error over spatial dims.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'rmse'

class NRMSE(eps=1e-07, reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

Normalized Root Mean Squared Error over spatial dims.

Parameters:

• eps (float)

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'nrmse'

class VMSE(eps=1e-07, reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

Variance Scaled Mean Squared Error over spatial dims.

Computes VMSE = MSE / (std(y_true)**2 + eps). Preserved for comparability with prior runs; matches The Well’s VMSE formula and default eps.

Parameters:

• eps (float)

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'vmse'

class VRMSE(eps=1e-07, reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

Variance-Scaled Root Mean Squared Error over spatial dims.

Computes VRMSE = sqrt(MSE) / (std(y_true) + eps). Preserved for comparability with prior runs. Differs from the LOLA / The Well canonical form sqrt(MSE / (var + eps)) only in the regularization regime (small std); for non-degenerate fields the two agree to <<1%.

Parameters:

• eps (float)

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'vrmse'

class VMSE2(eps=1e-06, reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

Variance-Scaled MSE matching LOLA’s evaluator.

Computes VMSE = MSE / (var(y_true) + eps) with eps = 1e-6, matching francois-rozet/lola.

Same formula as [VMSE][autocast.metrics.deterministic.VMSE], but with a larger regularizer that prevents tiny-variance fields from dominating the aggregate.

Parameters:

• eps (float)

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'vmse_v2'

class VRMSE2(eps=1e-06, reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

Variance-Scaled RMSE matching LOLA’s evaluator.

Computes VRMSE = sqrt(MSE / (var(y_true) + eps)) with eps = 1e-6, matching francois-rozet/lola.

The Well uses the same formula with eps = 1e-7. Use this for new evaluation work; [VRMSE][autocast.metrics.deterministic.VRMSE] uses a different regularizer (std + eps in the denominator) that produces inflated values for near-uniform fields.

Parameters:

• eps (float)

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'vrmse_v2'

class LInfinity(reduce_all=True, dist_sync_on_step=False)

Bases: BTSCMetric

L-Infinity Norm over spatial dims.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'l_infinity'

class PowerSpectrumRMSE(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: BTSCMetric

Average power spectrum RMSE across first three LOLA eval bands.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'psrmse'

class PowerSpectrumRMSELow(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: PowerSpectrumRMSE

Power spectrum RMSE in the low-frequency band.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'psrmse_low'

class PowerSpectrumRMSEMid(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: PowerSpectrumRMSE

Power spectrum RMSE in the mid-frequency band.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'psrmse_mid'

class PowerSpectrumRMSEHigh(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: PowerSpectrumRMSE

Power spectrum RMSE in the high-frequency band.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'psrmse_high'

class PowerSpectrumRMSETail(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: PowerSpectrumRMSE

Power spectrum RMSE in the LOLA high-frequency tail band.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'psrmse_tail'

class PowerSpectrumCCRMSE(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: BTSCMetric

Average cross-correlation RMSE across first three LOLA eval bands.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'pscc'

class PowerSpectrumCCRMSELow(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: PowerSpectrumCCRMSE

Cross-correlation RMSE in the low-frequency band.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'pscc_low'

class PowerSpectrumCCRMSEMid(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: PowerSpectrumCCRMSE

Cross-correlation RMSE in the mid-frequency band.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'pscc_mid'

class PowerSpectrumCCRMSEHigh(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: PowerSpectrumCCRMSE

Cross-correlation RMSE in the high-frequency band.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'pscc_high'

class PowerSpectrumCCRMSETail(reduce_all=True, dist_sync_on_step=False, eps=1e-06)

Bases: PowerSpectrumCCRMSE

Cross-correlation RMSE in the LOLA high-frequency tail band.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'pscc_tail'