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.

Parameters:

• eps (float)

• reduce_all (bool)

• dist_sync_on_step (bool)

name: str = 'vmse'

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

Bases: BTSCMetric

Variance-Scaled Root Mean Squared Error over spatial dims.

Computes VRMSE = RMSE / std(y_true), where std is computed over spatial dims.

Parameters:

• reduce_all (bool)

• dist_sync_on_step (bool)

• eps (float)

name: str = 'vrmse'

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'