diff --git a/src/pscpy/postprocessing.py b/src/pscpy/postprocessing.py index c9af8fa..b3bc50b 100644 --- a/src/pscpy/postprocessing.py +++ b/src/pscpy/postprocessing.py @@ -30,7 +30,7 @@ def get_recentered( elif boundary == "zero": shifted[{dim: boundary_idx}] = 0 - return 0.5 * (da + shifted) + return 0.5 * (da + shifted) # type: ignore[no-any-return] def _rename_var(ds: xr.Dataset, old_name: str, new_name: str) -> None: diff --git a/src/pscpy/psc.py b/src/pscpy/psc.py index bf1e606..fd3859c 100644 --- a/src/pscpy/psc.py +++ b/src/pscpy/psc.py @@ -84,11 +84,11 @@ def decode_psc( ds = ds.squeeze("step") for var_name in ds: - for component_idx, component in enumerate( - iter_components(var_name, species_names) - ): + components = list(iter_components(var_name, species_names)) + for component_idx, component in enumerate(components): ds = ds.assign({component: ds[var_name][component_idx, :, :, :]}) - ds = ds.drop_vars([var_name]) + if var_name not in components: + ds = ds.drop_vars([var_name]) run_info = RunInfo(ds, length=length, corner=corner) coords = {