From 3d2df98d5ab03b3d319430545c09165b7eb31a84 Mon Sep 17 00:00:00 2001 From: Sorcha Date: Thu, 12 Mar 2026 17:53:03 +0100 Subject: [PATCH 1/4] fixing pressure level --- .../evaluate/export/parsers/netcdf_parser.py | 16 +++++++++++----- .../src/weathergen/evaluate/export/reshape.py | 6 +++--- 2 files changed, 14 insertions(+), 8 deletions(-) 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 56a487c2f..6cfacbfa6 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -131,23 +131,29 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: var_dict, pl = find_pl(data.channel.values) data_vars = {} - for new_var, old_vars in var_dict.items(): - if len(old_vars) > 1: + for new_var, pls in var_dict.items(): + if pls[0] is not None: + old_vars = [f"{new_var}_{p}" for p in pls] + print(old_vars) data_vars[new_var] = xr.DataArray( data.sel(channel=old_vars).values, dims=["ipoint", "pressure_level"], + coords={"pressure_level": pls}, ) else: + print(new_var) data_vars[new_var] = xr.DataArray( - data.sel(channel=old_vars[0]).values, + data.sel(channel=new_var).values, dims=["ipoint"], ) reshaped_dataset = xr.Dataset(data_vars) reshaped_dataset = reshaped_dataset.assign_coords( ipoint=data.coords["ipoint"], - pressure_level=pl, ) + # order using pressure_level coord + if "pressure_level" in reshaped_dataset.coords: + reshaped_dataset = reshaped_dataset.sortby("pressure_level") if grid_type == "regular": # Use original reshape logic for regular grids @@ -274,7 +280,7 @@ def add_attrs(self, ds: xr.Dataset) -> xr.Dataset: else: variables = self._attrs_regular_grid(ds) - dataset = xr.merge(variables.values()) + dataset = xr.merge(variables.values(), compat="no_conflicts") dataset.attrs = ds.attrs return dataset diff --git a/packages/evaluate/src/weathergen/evaluate/export/reshape.py b/packages/evaluate/src/weathergen/evaluate/export/reshape.py index f32631ea3..80c865214 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/reshape.py +++ b/packages/evaluate/src/weathergen/evaluate/export/reshape.py @@ -75,10 +75,10 @@ def find_pl(vars: list) -> tuple[dict[str, list[str]], list[int]]: var_name = match.group(1) pressure_level = int(match.group(2)) pl.append(pressure_level) - var_dict.setdefault(var_name, []).append(var) + var_dict.setdefault(var_name, []).append(pressure_level) else: - var_dict.setdefault(var, []).append(var) - pl = sorted(set(pl)) + var_dict.setdefault(var, []).append(None) + return var_dict, pl From 923779bf8c8da03abdbb54c62b32f5b6484f03c9 Mon Sep 17 00:00:00 2001 From: Sorcha Date: Thu, 12 Mar 2026 17:59:39 +0100 Subject: [PATCH 2/4] remove prints --- .../src/weathergen/evaluate/export/parsers/netcdf_parser.py | 2 -- 1 file changed, 2 deletions(-) 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 6cfacbfa6..3adcc9c9a 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -134,14 +134,12 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: for new_var, pls in var_dict.items(): if pls[0] is not None: old_vars = [f"{new_var}_{p}" for p in pls] - print(old_vars) data_vars[new_var] = xr.DataArray( data.sel(channel=old_vars).values, dims=["ipoint", "pressure_level"], coords={"pressure_level": pls}, ) else: - print(new_var) data_vars[new_var] = xr.DataArray( data.sel(channel=new_var).values, dims=["ipoint"], From 67c36024f46f7a710f5460a3e77177f4aca97c2a Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Mon, 16 Mar 2026 11:33:48 +0100 Subject: [PATCH 3/4] fixing pressure level for 0 --- .../weathergen/evaluate/export/parsers/netcdf_parser.py | 2 +- .../evaluate/src/weathergen/evaluate/export/reshape.py | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) 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 3adcc9c9a..1d3f5caa0 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -128,7 +128,7 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: grid_type = self.grid_type # Original logic - var_dict, pl = find_pl(data.channel.values) + var_dict = find_pl(data.channel.values) data_vars = {} for new_var, pls in var_dict.items(): diff --git a/packages/evaluate/src/weathergen/evaluate/export/reshape.py b/packages/evaluate/src/weathergen/evaluate/export/reshape.py index 80c865214..aa22416cd 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/reshape.py +++ b/packages/evaluate/src/weathergen/evaluate/export/reshape.py @@ -68,18 +68,17 @@ def find_pl(vars: list) -> tuple[dict[str, list[str]], list[int]]: List of unique pressure levels found in the variable names. """ var_dict = {} - pl = [] for var in vars: match = re.search(r"^([a-zA-Z0-9_]+)_(\d+)$", var) if match: var_name = match.group(1) - pressure_level = int(match.group(2)) - pl.append(pressure_level) + pressure_level = int(match.group(2)) + if pressure_level == 0: + var_dict.setdefault(var, []).append(None) var_dict.setdefault(var_name, []).append(pressure_level) else: var_dict.setdefault(var, []).append(None) - - return var_dict, pl + return var_dict class Regridder: From 69315ab36f2c42f3fe9aef8a6d1db1e29d05c7f5 Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Mon, 16 Mar 2026 12:11:22 +0100 Subject: [PATCH 4/4] linting --- config/evaluate/config_zarr2cf.yaml | 2 +- .../src/weathergen/evaluate/export/parsers/netcdf_parser.py | 2 +- packages/evaluate/src/weathergen/evaluate/export/reshape.py | 3 ++- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/config/evaluate/config_zarr2cf.yaml b/config/evaluate/config_zarr2cf.yaml index 95bddd21b..431896336 100644 --- a/config/evaluate/config_zarr2cf.yaml +++ b/config/evaluate/config_zarr2cf.yaml @@ -89,7 +89,7 @@ variables: wg_unit: Pa std_unit: Pa level_type: sfc - tp_imerg: + tp_imerg_0: var: tp_imerg_0 long: imerg_total_precipitation std: precipitation_amount 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 1d3f5caa0..a0373b520 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -73,7 +73,7 @@ def process_sample( continue result = result.as_xarray().squeeze() - if "channel" not in result.indexes: + if "channel" not in result.indexes: result = result.expand_dims("channel") result = result.sel(channel=self.channels) result = self.reshape(result) diff --git a/packages/evaluate/src/weathergen/evaluate/export/reshape.py b/packages/evaluate/src/weathergen/evaluate/export/reshape.py index aa22416cd..74339bdad 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/reshape.py +++ b/packages/evaluate/src/weathergen/evaluate/export/reshape.py @@ -72,9 +72,10 @@ def find_pl(vars: list) -> tuple[dict[str, list[str]], list[int]]: match = re.search(r"^([a-zA-Z0-9_]+)_(\d+)$", var) if match: var_name = match.group(1) - pressure_level = int(match.group(2)) + pressure_level = int(match.group(2)) if pressure_level == 0: var_dict.setdefault(var, []).append(None) + return var_dict var_dict.setdefault(var_name, []).append(pressure_level) else: var_dict.setdefault(var, []).append(None)