minor fixes

psy_transect/plotters.py

@@ -209,9 +209,12 @@ class Transect(Formatoption):
         # update the data - this also updates the data for the plotter
         self.data = new_ds.psy[data.name]
 
+        cell_dim = self.transect_method.method_kws.get(
+            "cell_dim", "transect_cell"
+        )
         decoder = CFDecoder(
             new_ds,
-            x={self.data.dims[-1]},
+            x={cell_dim},
             y={self.raw_data.psy.get_coord("z").name},
         )
 

psy_transect/utils.py

@@ -533,9 +533,9 @@ def select_level(level, ds, coord, dim):
 
     for da in arrays:
         if da.name in ds.coords:
-            new_ds.coords[da.name] = da[selection]
+            new_ds.coords[da.name] = da.sel(**{dim: selection.values})
         else:
-            new_ds[da.name] = da[selection]
+            new_ds[da.name] = da.sel(**{dim: selection.values})
         remove_coordinates(da.attrs, [coord.name])
         remove_coordinates(da.encoding, [coord.name])
     new_ds.coords[coord.name] = ((), level, coord.attrs)