From f2f0ed3a480014886c68a4621ac372e915482b66 Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Tue, 2 Jun 2026 14:47:49 +0200 Subject: [PATCH 1/2] more changes --- .../src/weathergen/evaluate/export/parsers/netcdf_parser.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py index 8767d6df3..91dc7e637 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -145,6 +145,9 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: if "mem" in data.dims: data_dims.append("mem") if pls[0] is not None: + data_dims.append("pressure_level") + if "mem" in data.dims: + data_dims.append("mem") old_vars = [f"{new_var}_{p}" for p in pls] data_vars[new_var] = xr.DataArray( data.sel(channel=old_vars).values, @@ -152,6 +155,8 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: coords={"pressure_level": pls}, ) else: + if "mem" in data.dims: + data_dims.append("mem") data_vars[new_var] = xr.DataArray( data.sel(channel=new_var).values, dims=data_dims, From c0651432e22a897feddb79e2534539f8ce0b3352 Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Fri, 12 Jun 2026 15:36:52 +0200 Subject: [PATCH 2/2] testing complete --- .../src/weathergen/evaluate/export/export_core.py | 7 +++++-- .../src/weathergen/evaluate/export/export_inference.py | 2 +- .../weathergen/evaluate/export/parsers/netcdf_parser.py | 8 +++----- .../weathergen/evaluate/export/parsers/quaver_parser.py | 1 + 4 files changed, 10 insertions(+), 8 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index fb691db93..35dc9135b 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -314,13 +314,16 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: fsteps = get_fsteps(fsteps, fname_zarr) samples = get_samples(samples, fname_zarr) streams = get_streams(stream, fname_zarr) + _logger.info(f"Streams to process: {streams}") for stream in streams: grid_type = get_grid_type(data_type, stream, fname_zarr) - channels = get_channels(channels, stream, fname_zarr) + stream_channels = get_channels(channels, stream, fname_zarr) source_starts, source_ends = get_source_info(fname_zarr, stream, samples) + # prevent overwritting of channels between streams kwargs["grid_type"] = grid_type - kwargs["channels"] = channels + kwargs["channels"] = stream_channels kwargs["data_type"] = data_type + kwargs["stream"] = stream parser = CfParserFactory.get_parser(config=config, **kwargs) diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_inference.py b/packages/evaluate/src/weathergen/evaluate/export/export_inference.py index 22a674a86..0409a461d 100755 --- a/packages/evaluate/src/weathergen/evaluate/export/export_inference.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_inference.py @@ -104,7 +104,7 @@ def parse_args(args: list) -> argparse.Namespace: parser.add_argument( "--stream", type=str, - choices=["ERA5", "CERRA", "MEPS", "NORA3", "IMERG_ANEMOI"], + choices=["ERA5", "ERA5pl", "ERA5ml", "CERRA", "MEPS", "NORA3", "IMERG_ANEMOI"], help="Stream name to retrieve data for", ) diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py index 91dc7e637..e0b6bbe90 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -70,7 +70,7 @@ def process_sample( da_fs = [] for result in fstep_iterator_results: - if result is None: + if result is None or (isinstance(result, xr.DataArray) and result.size == 0): continue # result is already a materialized xarray DataArray (built in the worker). @@ -139,11 +139,9 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: # Original logic var_dict = find_pl(data.channel.values) data_vars = {} - + #order of appending upoints should be ipoint, pressure_level, mem (if mem exists) for new_var, pls in var_dict.items(): data_dims = ["ipoint"] - if "mem" in data.dims: - data_dims.append("mem") if pls[0] is not None: data_dims.append("pressure_level") if "mem" in data.dims: @@ -151,7 +149,7 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: old_vars = [f"{new_var}_{p}" for p in pls] data_vars[new_var] = xr.DataArray( data.sel(channel=old_vars).values, - dims=[*data_dims, "pressure_level"], + dims=data_dims, coords={"pressure_level": pls}, ) else: diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py index a19077ccb..f68fa7c78 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py @@ -74,6 +74,7 @@ def process_sample( ref_time: np.datetime64, source_interval_start: np.datetime64 = None, source_interval_end: np.datetime64 = None, + **kwargs ): """ Process results from get_data_worker: reshape, concatenate, add metadata, and save.