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 = {}