def spatial_rmse(true_da: xr.DataArray, pred_da: xr.DataArray) -> xr.DataArray:
"""Calculate the RMSE collapsing the time dimension returning
a DataArray of the rmse values (spatially)
"""
true_da_shape = (true_da.lat.shape[0], true_da.lon.shape[0])
pred_da_shape = (pred_da.lat.shape[0], pred_da.lon.shape[0])
assert true_da_shape == pred_da_shape
assert tuple(true_da.dims) == tuple(pred_da.dims), (
f"Expect"
"the dimensions to be the same. Currently: "
f"True: {tuple(true_da.dims)} Preds: {tuple(pred_da.dims)}. "
'Have you tried da.transpose("time", "lat", "lon")')
# sort the lat/lons correctly just to be sure
pred_da = _sort_lat_lons(pred_da)
true_da = _sort_lat_lons(true_da)
vals = np.sqrt(
np.nansum(
(true_da.values - pred_da.values)**2, axis=0) / pred_da.shape[0])
da = xr.ones_like(pred_da).isel(time=0)
da.values = vals
# reapply the mask
da = da.where(~get_ds_mask(pred_da))
return da
def spatial_r2(true_da: xr.DataArray, pred_da: xr.DataArray) -> xr.DataArray:
true_da_shape = (true_da.lat.shape[0], true_da.lon.shape[0])
pred_da_shape = (pred_da.lat.shape[0], pred_da.lon.shape[0])
assert true_da_shape == pred_da_shape
r2_vals = 1 - (np.nansum(
(true_da.values - pred_da.values)**2, axis=0)) / (np.nansum(
(true_da.values - np.nanmean(pred_da.values))**2, axis=0))
da = xr.ones_like(pred_da).isel(time=0)
da.values = r2_vals
# reapply the mask
da = da.where(~get_ds_mask(pred_da))
return da
def spatial_rmse(true_da: xr.DataArray, pred_da: xr.DataArray) -> xr.DataArray:
"""Calculate the RMSE collapsing the time dimension returning
a DataArray of the rmse values (spatially)
"""
true_da_shape = (true_da.lat.shape[0], true_da.lon.shape[0])
pred_da_shape = (pred_da.lat.shape[0], pred_da.lon.shape[0])
assert true_da_shape == pred_da_shape
vals = np.sqrt(
np.nansum(
(true_da.values - pred_da.values)**2, axis=0) / pred_da.shape[0])
da = xr.ones_like(pred_da).isel(time=0)
da.values = vals
# reapply the mask
da = da.where(~get_ds_mask(pred_da))
return da
def spatial_r2(true_da: xr.DataArray, pred_da: xr.DataArray) -> xr.DataArray:
true_da_shape = (true_da.lat.shape[0], true_da.lon.shape[0])
pred_da_shape = (pred_da.lat.shape[0], pred_da.lon.shape[0])
assert true_da_shape == pred_da_shape
# sort the dimensions so that no inversions (e.g. lat)
true_da = true_da.sortby(["time", "lat", "lon"])
pred_da = pred_da.sortby(["time", "lat", "lon"])
r2_vals = 1 - (np.nansum(
(true_da.values - pred_da.values)**2, axis=0)) / (np.nansum(
(true_da.values - np.nanmean(pred_da.values))**2, axis=0))
da = xr.ones_like(pred_da).isel(time=0)
da.values = r2_vals
# reapply the mask
da = da.where(~get_ds_mask(pred_da))
return da
if __name__ == "__main__":
# run_administrative_region_analysis()
# run_landcover_region_analysis()
data_dir = get_data_path()
experiment = "one_month_forecast"
# ---------------------------------------- #
# Read in the data #
# ---------------------------------------- #
# read train/test data
X_train, y_train, X_test, y_test = read_all_data(data_dir,
static=False,
experiment=experiment)
mask = get_ds_mask(X_train.VCI)
test_da = y_test[list(y_test.data_vars)[0]]
# read the predicted data
predictions = read_all_available_pred_data(data_dir, experiment=experiment)
# check that the shapes of all predictions are the same
assert set([predictions[m].shape
for m in predictions.keys()]).__len__() == 1
preds = predictions[list(m for m in predictions.keys())[0]]
# get the matching shapes from predictions and test data
test_da = test_da.sel(lat=preds.lat, lon=preds.lon, time=preds.time)
# read in the models (pytorch models)
models_str = [m for m in predictions.keys()]
models = {}
