earth2studio.statistics.rmse

class earth2studio.statistics.rmse(reduction_dimensions, weights=None, batch_update=False)

Statistic for calculating the root mean squared error of two tensors over a set of given dimensions.

Parameters:

• reduction_dimensions (List[str]) – A list of names corresponding to dimensions to perform the statistical reduction over. Example: [‘lat’, ‘lon’]

• weights (torch.Tensor = None) – A tensor containing weights to assign to the reduction dimensions. Note that these weights must have the same number of dimensions as passed in reduction_dimensions. Example: if reduction_dimensions = [‘lat’, ‘lon’] then assert weights.ndim == 2.

• batch_update (bool = False) – Whether to applying batch updates to the rmse with each invocation of __call__. This is particularly useful when data is recieved in a stream of batches. Each invocation of __call__ will return the running rmse. In particular, it will apply the square root operation after calculating the running mean squared error.

__call__(x, x_coords, y, y_coords)

Apply metric to data x and y, checking that their coordinates are broadcastable. While reducing over reduction_dims.

If batch_update was passed True upon metric initialization then this method returns the running sample RMSE over all seen batches.

Parameters:

• x (torch.Tensor) – Input tensor, typically the forecast or prediction tensor, but RMSE is symmetric with respect to x and y.

• x_coords (CoordSystem) – Ordered dict representing coordinate system that describes the x tensor. reduction_dimensions must be in coords.

• y (torch.Tensor) – Input tensor #2 intended to be used as validation data, but ACC is symmetric with respect to x and y.

• y_coords (CoordSystem) – Ordered dict representing coordinate system that describes the y tensor. reduction_dimensions must be in coords.

Returns:

Returns root mean squared error tensor with appropriate reduced coordinates.

Return type:

tuple[torch.Tensor, CoordSystem]