From 38f577cec845f27a03d26e3d27d4714c0b31d76c Mon Sep 17 00:00:00 2001 From: James McClung Date: Wed, 12 Mar 2025 14:04:33 -0400 Subject: [PATCH 1/2] psc: change 2nd dim name in decode_psc it gets dropped anyways now, so the name shouldn't matter --- src/pscpy/psc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pscpy/psc.py b/src/pscpy/psc.py index d2654cc..c0fcd9d 100644 --- a/src/pscpy/psc.py +++ b/src/pscpy/psc.py @@ -111,7 +111,7 @@ def decode_psc( ds = ds.rename_dims( { da.dims[0]: "step", - da.dims[1]: f"comp_{da.name}", + da.dims[1]: "component", da.dims[2]: "z", da.dims[3]: "y", da.dims[4]: "x", @@ -124,7 +124,7 @@ def decode_psc( for var_name in ds: if var_name in field_to_component: for field, component in field_to_component[var_name].items(): # type: ignore[index] - data_vars[field] = ds[var_name].isel({f"comp_{var_name}": component}) + data_vars[field] = ds[var_name].isel({"component": component}) ds = ds.drop_vars([var_name]) ds = ds.assign(data_vars) From bb6773b8e33a1cb511432c01aa9ead38e5408826 Mon Sep 17 00:00:00 2001 From: James McClung Date: Tue, 25 Mar 2025 11:39:21 -0400 Subject: [PATCH 2/2] psc: use positional indexing for component dim --- src/pscpy/psc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pscpy/psc.py b/src/pscpy/psc.py index c0fcd9d..2f67e46 100644 --- a/src/pscpy/psc.py +++ b/src/pscpy/psc.py @@ -111,7 +111,7 @@ def decode_psc( ds = ds.rename_dims( { da.dims[0]: "step", - da.dims[1]: "component", + # dims[1] is the "component" dimension, which gets removed later da.dims[2]: "z", da.dims[3]: "y", da.dims[4]: "x", @@ -123,8 +123,8 @@ def decode_psc( data_vars = {} for var_name in ds: if var_name in field_to_component: - for field, component in field_to_component[var_name].items(): # type: ignore[index] - data_vars[field] = ds[var_name].isel({"component": component}) + for field, component_idx in field_to_component[var_name].items(): # type: ignore[index] + data_vars[field] = ds[var_name][component_idx, :, :, :] ds = ds.drop_vars([var_name]) ds = ds.assign(data_vars)