`decode_psc` crashes for `gauss` output

[JamesMcClung]

pscpy/src/pscpy/psc.py

Lines 94 to 122 in 1b14e73

	def decode_psc(
	ds: xr.Dataset,
	species_names: Iterable[str],
	length: ArrayLike | None = None,
	corner: ArrayLike | None = None,
	) -> xr.Dataset:
	da = ds[next(iter(ds))] # first dataset
	if da.dims[0] == "dim_0_1":
	# for compatibility, if dimensions weren't saved as attribute in the .bp file,
	# fix them up here
	ds = ds.rename_dims(
	{
	da.dims[0]: "step",
	da.dims[1]: f"comp_{da.name}",
	da.dims[2]: "z",
	da.dims[3]: "y",
	da.dims[4]: "x",
	}
	)
	ds = ds.squeeze("step")
	field_to_component = get_field_to_component(species_names)
	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({f"comp_{var_name}": component})
	ds = ds.drop_vars([var_name])
	ds = ds.assign(data_vars)

In the snippet above, ds initially looks like something like this when decoding a gauss output:

<xarray.Dataset> Size: 640kB
Dimensions:  (dim_0_1: 1, dim_1_1: 1, dim_2_40: 40, dim_3_1000: 1000, dim_4_1: 1)
Dimensions without coordinates: dim_0_1, dim_1_1, dim_2_40, dim_3_1000, dim_4_1
Data variables:
    dive     (dim_0_1, dim_1_1, dim_2_40, dim_3_1000, dim_4_1) float64 320kB ...
    rho      (dim_0_1, dim_1_1, dim_2_40, dim_3_1000, dim_4_1) float64 320kB ...
Attributes:
    length:    [  1. 500.  20.]
    dive::im:  [   1 1000   40]
    rho::ib:   [0 0 0]
    dive::ib:  [0 0 0]
    corner:    [0. 0. 0.]
    step:      [0]
    time:      [0.]
    rho::im:   [   1 1000   40]

Notably, dive and rho are separate arrays. da thus becomes the dive array on line 100, and the 2nd dimension is renamed to "comp_dive" on line 107. However, when iterating over var_name on line 117, dive and rho are both present, so line 120 crashes when there is no dimension named "comp_rho".

Ideally, decode_psc would return a dataset containing two dataarrays, dive and rho. For the gauss case, all that decode_psc really needs to do is rename the spatial dimensions and drop the first two dimensions.