From 8ec1b54da3075d4ae1ba3c9d3e4f6690726f72cf Mon Sep 17 00:00:00 2001 From: buschow1 Date: Wed, 22 Apr 2026 11:47:59 +0200 Subject: [PATCH 1/5] initial fix with no real coordinate --- .../weathergen/evaluate/export/export_core.py | 19 +++++++++++-------- .../evaluate/export/parsers/netcdf_parser.py | 7 +++++-- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index 2bd6de36a..0a156814b 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -62,20 +62,23 @@ def get_data_worker(args: tuple) -> tuple[int, int, xr.DataArray]: npoints = data_arr.shape[0] # Handle optional ensemble dimension: squeeze it out if present. - if data_arr.ndim == 3 and data_arr.shape[2] == 1: - data_arr = data_arr[:, :, 0] - - da_result = xr.DataArray( - data_arr, - dims=["ipoint", "channel"], - coords={ + data_dims = ["ipoint", "channel"] + data_coords = { "ipoint": np.arange(npoints), "channel": channels, "forecast_step": fstep, "valid_time": ("ipoint", times_arr), "lat": ("ipoint", coords_arr[:, 0]), "lon": ("ipoint", coords_arr[:, 1]), - }, + } + if data_arr.ndim == 3: + data_dims.append("mem") + #data_arr = np.mean(data_arr, axis=2) + + da_result = xr.DataArray( + data_arr, + dims=data_dims, + coords=data_coords, ) return (sample, fstep, da_result) 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 27cfc4bc3..4e116b959 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -134,17 +134,20 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: data_vars = {} 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: old_vars = [f"{new_var}_{p}" for p in pls] data_vars[new_var] = xr.DataArray( data.sel(channel=old_vars).values, - dims=["ipoint", "pressure_level"], + dims=[*data_dims, "pressure_level"], coords={"pressure_level": pls}, ) else: data_vars[new_var] = xr.DataArray( data.sel(channel=new_var).values, - dims=["ipoint"], + dims=data_dims, ) reshaped_dataset = xr.Dataset(data_vars) From bea50a0b3a0f42f850f5548cb6aa05a61975ef61 Mon Sep 17 00:00:00 2001 From: buschow1 Date: Wed, 22 Apr 2026 12:33:05 +0200 Subject: [PATCH 2/5] hacked in the coordinate --- .../evaluate/src/weathergen/evaluate/export/export_core.py | 6 ++++-- .../weathergen/evaluate/export/parsers/netcdf_parser.py | 7 ++++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index 0a156814b..ecfc420d8 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -72,8 +72,10 @@ def get_data_worker(args: tuple) -> tuple[int, int, xr.DataArray]: "lon": ("ipoint", coords_arr[:, 1]), } if data_arr.ndim == 3: - data_dims.append("mem") - #data_arr = np.mean(data_arr, axis=2) + if data_arr.shape[2] == 1: + data_arr = data_arr[:, :, 0] + else: + data_dims.append("mem") da_result = xr.DataArray( data_arr, 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 4e116b959..415d3d2e2 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -158,6 +158,10 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: if "pressure_level" in reshaped_dataset.coords: reshaped_dataset = reshaped_dataset.sortby("pressure_level") + if "mem" in reshaped_dataset.dims: + reshaped_dataset = reshaped_dataset.assign_coords( + mem=np.arange(reshaped_dataset.sizes["mem"]) + ) if grid_type == "regular": # Use original reshape logic for regular grids # This is safe for regular grids @@ -461,7 +465,8 @@ def _build_coordinate_mapping( ds.coords[coord].values, attrs[new_name], ) - + if "mem" in ds.coords: + coords["mem"] = ("mem", ds.mem.values) return coords def _add_grid_attrs(self, ds: xr.Dataset, grid_info: dict | None = None) -> xr.Dataset: From be8b09505480016b5e33d0709ae833f327eab347 Mon Sep 17 00:00:00 2001 From: buschow1 Date: Mon, 18 May 2026 13:18:46 +0200 Subject: [PATCH 3/5] coordinate ordering --- .../src/weathergen/evaluate/export/reshape.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/export/reshape.py b/packages/evaluate/src/weathergen/evaluate/export/reshape.py index d0d0bea6e..3fd4c10c9 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/reshape.py +++ b/packages/evaluate/src/weathergen/evaluate/export/reshape.py @@ -267,15 +267,9 @@ def gaussian_regular_da(self, data: xr.DataArray) -> xr.DataArray: pos = dims.index("ncells") dims[pos : pos + 1] = ["latitude", "longitude"] dims = tuple(dims) - ordered_dims = ( - ["valid_time", "pressure", "latitude", "longitude"] - if len(dims) == 4 - else ["valid_time", "latitude", "longitude"] - ) - permutation_indices = [dims.index(o_dim) for o_dim in ordered_dims] - regridded_values = np.transpose(regridded_values, axes=permutation_indices) + regrid_data = xr.DataArray( - data=regridded_values, dims=ordered_dims, coords=new_coords, attrs=attrs, name=data.name + data=regridded_values, dims=dims, coords=new_coords, attrs=attrs, name=data.name ) return regrid_data @@ -549,7 +543,10 @@ def regrid_ds( regrid_vars[var] = self.regrid_da(ds[var]) regrid_ds = xr.Dataset(regrid_vars) regrid_ds = self.add_attrs(regrid_ds) - + regrid_ds = regrid_ds.transpose( + "valid_time", "pressure", "latitude", "longitude", "mem", ..., + missing_dims="ignore" + ) return regrid_ds def regrid_da(self, da: xr.DataArray) -> xr.DataArray: From 1a5c9b35ac96e25fb1d8aab1dedd53f54670d747 Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Fri, 22 May 2026 10:36:57 +0200 Subject: [PATCH 4/5] sugggestions --- config/evaluate/config_zarr2cf.yaml | 9 ++++++++- .../src/weathergen/evaluate/export/export_core.py | 15 +++++++++------ .../evaluate/export/parsers/netcdf_parser.py | 6 ------ 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/config/evaluate/config_zarr2cf.yaml b/config/evaluate/config_zarr2cf.yaml index 75677858b..7d0e7163a 100644 --- a/config/evaluate/config_zarr2cf.yaml +++ b/config/evaluate/config_zarr2cf.yaml @@ -106,6 +106,7 @@ coordinates: forecast_step: forecast_period forecast_reference_time: forecast_reference_time ncells: ncells + mem: mem pl: pressure_level: pressure valid_time: valid_time @@ -114,6 +115,7 @@ coordinates: forecast_step: forecast_period forecast_reference_time: forecast_reference_time ncells: ncells + mem: mem dimensions: valid_time: @@ -141,4 +143,9 @@ dimensions: std_unit: hours ncells: wg: ncells - std: ncells \ No newline at end of file + std: ncells + mem: + wg: mem + std: realization + long: realization + std_unit: 1 \ No newline at end of file diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index c4dff6924..33fa7f63c 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -62,7 +62,14 @@ def get_data_worker(args: tuple) -> tuple[int, int, xr.DataArray]: npoints = data_arr.shape[0] # Handle optional ensemble dimension: squeeze it out if present. - data_dims = ["ipoint", "channel"] + if data_arr.ndim == 3: + if data_arr.shape[2] == 1: + data_arr = data_arr[:, :, 0] + data_dims = ["ipoint", "channel"] + else: + data_dims = ["ipoint", "channel", "mem"] + + data_coords = { "ipoint": np.arange(npoints), "channel": channels, @@ -71,11 +78,7 @@ def get_data_worker(args: tuple) -> tuple[int, int, xr.DataArray]: "lat": ("ipoint", coords_arr[:, 0]), "lon": ("ipoint", coords_arr[:, 1]), } - if data_arr.ndim == 3: - if data_arr.shape[2] == 1: - data_arr = data_arr[:, :, 0] - else: - data_dims.append("mem") + da_result = xr.DataArray( data_arr, 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 33e946d11..8767d6df3 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -165,10 +165,6 @@ def reshape(self, data: xr.DataArray) -> xr.Dataset: if "pressure_level" in reshaped_dataset.coords: reshaped_dataset = reshaped_dataset.sortby("pressure_level") - if "mem" in reshaped_dataset.dims: - reshaped_dataset = reshaped_dataset.assign_coords( - mem=np.arange(reshaped_dataset.sizes["mem"]) - ) if grid_type == "regular": # Use original reshape logic for regular grids # This is safe for regular grids @@ -478,8 +474,6 @@ def _build_coordinate_mapping( f"Coordinate '{coord}' will be skipped for " f"variable '{var_cfg.get('var', 'unknown')}'." ) - if "mem" in ds.coords: - coords["mem"] = ("mem", ds.mem.values) return coords def _add_grid_attrs(self, ds: xr.Dataset, grid_info: dict | None = None) -> xr.Dataset: From 4dcb7d36f9f9246fb43e4ddd7e6a6d89b74c48a5 Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Fri, 22 May 2026 10:37:47 +0200 Subject: [PATCH 5/5] linting --- .../weathergen/evaluate/export/export_core.py | 24 +++++++------------ .../src/weathergen/evaluate/export/reshape.py | 3 +-- 2 files changed, 10 insertions(+), 17 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index 33fa7f63c..fb691db93 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -69,22 +69,16 @@ def get_data_worker(args: tuple) -> tuple[int, int, xr.DataArray]: else: data_dims = ["ipoint", "channel", "mem"] - data_coords = { - "ipoint": np.arange(npoints), - "channel": channels, - "forecast_step": fstep, - "valid_time": ("ipoint", times_arr), - "lat": ("ipoint", coords_arr[:, 0]), - "lon": ("ipoint", coords_arr[:, 1]), - } - - - da_result = xr.DataArray( - data_arr, - dims=data_dims, - coords=data_coords - ) + "ipoint": np.arange(npoints), + "channel": channels, + "forecast_step": fstep, + "valid_time": ("ipoint", times_arr), + "lat": ("ipoint", coords_arr[:, 0]), + "lon": ("ipoint", coords_arr[:, 1]), + } + + da_result = xr.DataArray(data_arr, dims=data_dims, coords=data_coords) return (sample, fstep, da_result) diff --git a/packages/evaluate/src/weathergen/evaluate/export/reshape.py b/packages/evaluate/src/weathergen/evaluate/export/reshape.py index 3fd4c10c9..9ec334214 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/reshape.py +++ b/packages/evaluate/src/weathergen/evaluate/export/reshape.py @@ -544,8 +544,7 @@ def regrid_ds( regrid_ds = xr.Dataset(regrid_vars) regrid_ds = self.add_attrs(regrid_ds) regrid_ds = regrid_ds.transpose( - "valid_time", "pressure", "latitude", "longitude", "mem", ..., - missing_dims="ignore" + "valid_time", "pressure", "latitude", "longitude", "mem", ..., missing_dims="ignore" ) return regrid_ds