Latitude weighted RMSE not picking up the correct dimension #21
Description
Since the model is receiving input of [batch, seq, vars, lat (96), lon(144)]
it should be lat_size = y.shape[-2]
This is the case if channels_last is set to False.
def LLWeighted_RMSE_WheatherBench(preds: np.ndarray, y: np.ndarray):
"""
Weigthed RMSE taken from Wheather Bench.
Weighting to account for decreasing grid sizes towards the pole.
rmse = mean over forecasts and time of np.sqrt( mean over lon lat L(lat_j)*)MSE(preds, y)
weights = cos(latitude)/cos(latitude).mean()
"""
lat_size = y.shape[-1]
lats = np.linspace(-90, 90, lat_size)
weights = (np.cos(lats) / np.cos(lats)).mean()
rmse = np.sqrt(np.mean(weights * ((preds - y) ** 2), axis=(-1, -2))).mean()
return rmse