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 8767d6df3..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,19 +139,22 @@ 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: + 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, - dims=[*data_dims, "pressure_level"], + dims=data_dims, 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, 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.