From 3752ebd5f3bf21b8fc3e6ff6796cf4b2475aa730 Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Wed, 29 Apr 2026 16:13:14 +0200 Subject: [PATCH 1/5] tp fixed for zarr_dt = 0 --- .../evaluate/utils/derived_channels.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py b/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py index 5f893c590..9e072fecd 100644 --- a/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py +++ b/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py @@ -254,7 +254,7 @@ def compute_mslp(obs: xr.DataArray, time: np.datetime64) -> np.typing.NDArray: def compute_precip( - obs_data: xr.Dataset, zarr_dt: np.timedelta64, frt: np.datetime64 + obs_data: xr.Dataset, zarr_dt: np.timedelta64, valid_time: np.datetime64 ) -> np.typing.NDArray: """ Compute accumulated precipitation over the forecast time step. @@ -264,7 +264,7 @@ def compute_precip( Input data containing precipitation observations. zarr_dt : np.timedelta64 Time difference between forecast steps in hours. - frt : np.datetime64 + valid_time : np.datetime64 Forecast reference time for which to compute accumulated precipitation. Returns ------- @@ -272,15 +272,23 @@ def compute_precip( Accumulated precipitation values for the forecast time step.""" obs_dt = obs_data.time.values[1] - obs_data.time.values[0] obs_dt = obs_dt.astype("timedelta64[h]") - - if obs_dt >= zarr_dt: - return obs_data["precipitation_amount_1h"].values + if zarr_dt == 0: + # indicates nowcasting + return obs_data["precipitation_amount_1h"].sel(time=valid_time).values + elif obs_dt >= zarr_dt: + if obs_dt % zarr_dt == np.timedelta64(0): + return obs_data["precipitation_amount_1h"].sel(time=valid_time).values + else: + # return empty data + empty = np.empty(obs_data.location.shape[0]) + empty[:] = np.nan + return empty else: accumulate = np.zeros(obs_data.location.shape[0]) int_factor = int(zarr_dt / obs_dt) for i in range(int_factor): - back_time = frt - zarr_dt + (i + 1) * obs_dt + back_time = valid_time - zarr_dt + (i + 1) * obs_dt accumulate += ( obs_data.data_vars["precipitation_amount_1h"].sel(time=back_time).squeeze() ) From 8694c587e38e6b6481273989b7bf768d907188be Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Wed, 29 Apr 2026 16:16:47 +0200 Subject: [PATCH 2/5] ensuring shape compatibility --- .../src/weathergen/evaluate/utils/derived_channels.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py b/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py index 9e072fecd..8b03f408b 100644 --- a/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py +++ b/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py @@ -274,10 +274,10 @@ def compute_precip( obs_dt = obs_dt.astype("timedelta64[h]") if zarr_dt == 0: # indicates nowcasting - return obs_data["precipitation_amount_1h"].sel(time=valid_time).values + return obs_data["precipitation_amount_1h"].sel(time=valid_time).values.squeeze() elif obs_dt >= zarr_dt: if obs_dt % zarr_dt == np.timedelta64(0): - return obs_data["precipitation_amount_1h"].sel(time=valid_time).values + return obs_data["precipitation_amount_1h"].sel(time=valid_time).values.squeeze() else: # return empty data empty = np.empty(obs_data.location.shape[0]) From 6bb400abb4021885040b7c5526c3cc4571bfe71b Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Wed, 20 May 2026 18:20:07 +0200 Subject: [PATCH 3/5] adding default 6h for time step --- config/evaluate/config_zarr2cf.yaml | 7 +++++++ config/evaluate/config_zarr2verif.yaml | 2 +- .../src/weathergen/evaluate/export/parsers/verif_parser.py | 4 +++- .../src/weathergen/evaluate/utils/derived_channels.py | 6 ++---- 4 files changed, 13 insertions(+), 6 deletions(-) diff --git a/config/evaluate/config_zarr2cf.yaml b/config/evaluate/config_zarr2cf.yaml index 75677858b..0c9abe8b7 100644 --- a/config/evaluate/config_zarr2cf.yaml +++ b/config/evaluate/config_zarr2cf.yaml @@ -89,6 +89,13 @@ variables: wg_unit: Pa std_unit: Pa level_type: sfc + tp: + var: tp + long: total_precipitation + std: precipitation_amount + wg_unit: kg/m^2 + std_unit: kg/m^2 + level_type: sfc tp_imerg_0: var: tp_imerg_0 long: imerg_total_precipitation diff --git a/config/evaluate/config_zarr2verif.yaml b/config/evaluate/config_zarr2verif.yaml index 8cc573f40..8ca59f073 100644 --- a/config/evaluate/config_zarr2verif.yaml +++ b/config/evaluate/config_zarr2verif.yaml @@ -34,7 +34,7 @@ variables: var: tp long: Total precipitation amount wg_unit: {CERRA: kg/m^2, - MEPS: kg/m^2, + MEPS: m, NORA3: kg/m^2, ERA5: m, DEFAULT: kg/m^2} diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/verif_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/verif_parser.py index 3b33786a4..5148566e1 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/verif_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/verif_parser.py @@ -158,7 +158,9 @@ def get_zarr_dt( Time difference between source interval start and end in hours. """ zarr_dt = (source_interval_end - source_interval_start).astype("timedelta64[h]") - + #TODO: pull default value from config + if zarr_dt == np.timedelta64(0, "h"): + zarr_dt = np.timedelta64(6, "h") return zarr_dt def get_output_filename(self, variable: str) -> Path: diff --git a/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py b/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py index 8b03f408b..136ae4af6 100644 --- a/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py +++ b/packages/evaluate/src/weathergen/evaluate/utils/derived_channels.py @@ -272,10 +272,8 @@ def compute_precip( Accumulated precipitation values for the forecast time step.""" obs_dt = obs_data.time.values[1] - obs_data.time.values[0] obs_dt = obs_dt.astype("timedelta64[h]") - if zarr_dt == 0: - # indicates nowcasting - return obs_data["precipitation_amount_1h"].sel(time=valid_time).values.squeeze() - elif obs_dt >= zarr_dt: + + if obs_dt >= zarr_dt: if obs_dt % zarr_dt == np.timedelta64(0): return obs_data["precipitation_amount_1h"].sel(time=valid_time).values.squeeze() else: From e78b7b49b942ef7d188f2b54f76502b3df17a946 Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Thu, 21 May 2026 12:32:46 +0200 Subject: [PATCH 4/5] fix: number of samples not working --- .../src/weathergen/evaluate/export/export_core.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index eeea8576b..7fb066bb3 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -344,6 +344,8 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: initargs=(fname_zarr,), ) as pool: samples_written = 0 + + processed_samples = [] for batch_idx in range(n_batches): batch_start = batch_idx * batch_size @@ -377,8 +379,6 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: desc=f" Batch {batch_idx + 1}/{n_batches}", ) - processed_samples = [] - for sample, _fstep, data in pool.imap_unordered( get_data_worker, batch_tasks, chunksize=1 ): @@ -403,9 +403,6 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: del sample_results[sample] batch_written += 1 - # Only save here if need to merge samples, otherwise saved in process_sample - if processed_samples[0] is not None: - parser.save(processed_samples) pbar.close() samples_written += batch_written @@ -419,4 +416,8 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: # Free any remaining refs before next batch. del sample_results + # Only save here if need to merge samples, otherwise saved in process_sample + if processed_samples[0] is not None: + parser.save(processed_samples) + _logger.info(f"Export complete. Wrote {samples_written}/{len(samples)} samples.") From 1f5a7aa513fb454b8cb3b735b7a1719c14880d6f Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Thu, 21 May 2026 12:35:54 +0200 Subject: [PATCH 5/5] linting --- packages/evaluate/src/weathergen/evaluate/export/export_core.py | 2 +- .../src/weathergen/evaluate/export/parsers/verif_parser.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index 7fb066bb3..6281df0ff 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -344,7 +344,7 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: initargs=(fname_zarr,), ) as pool: samples_written = 0 - + processed_samples = [] for batch_idx in range(n_batches): diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/verif_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/verif_parser.py index 5148566e1..5df8e9918 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/verif_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/verif_parser.py @@ -158,7 +158,7 @@ def get_zarr_dt( Time difference between source interval start and end in hours. """ zarr_dt = (source_interval_end - source_interval_start).astype("timedelta64[h]") - #TODO: pull default value from config + # TODO: pull default value from config if zarr_dt == np.timedelta64(0, "h"): zarr_dt = np.timedelta64(6, "h") return zarr_dt