Ensure separatrices, target plotted correctly if Dataset transposed

Create metadata variables called 'bout_tdim', 'bout_xdim', 'bout_ydim',
and 'bout_zdim' that record what 't', 'x', 'y', and 'z' respectively
were renamed to. This means that in plot_separatrices() and
plot_targets() we can find renamed dimensions corresponding to 'x' and
'y' even if the Dataset has been transposed so that we cannot rely on
the order of the dimensions.

xbout/geometries.py

@@ -102,6 +102,7 @@ def _set_default_toroidal_coordinates(coordinates):
         coordinates = {}
 
     # Replace any values that have not been passed in with defaults
+    coordinates['t'] = coordinates.get('t', 't')
     coordinates['x'] = coordinates.get('x', 'psi_poloidal')
     coordinates['y'] = coordinates.get('y', 'theta')
     coordinates['z'] = coordinates.get('z', 'zeta')
@@ -132,6 +133,9 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
                                  "open_boutdataset().")
             ds[v] = grid[v]
 
+    # Rename 't' if user requested it
+    ds = ds.rename(t=coordinates['t'])
+
     # Change names of dimensions to Orthogonal Toroidal ones
     ds = ds.rename(y=coordinates['y'])
 
@@ -146,6 +150,11 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
     ds = ds.set_coords(coordinates['x'])
     ds[coordinates['x']].attrs['units'] = 'Wb'
 
+    # Record which dimensions 't', 'x', and 'y' were renamed to.
+    ds.metadata['bout_tdim'] = coordinates['t']
+    ds.metadata['bout_xdim'] = coordinates['x']
+    ds.metadata['bout_ydim'] = coordinates['y']
+
     # If full data (not just grid file) then toroidal dim will be present
     if 'z' in ds.dims:
         ds = ds.rename(z=coordinates['z'])
@@ -154,6 +163,9 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
                            dims=coordinates['z'])
         ds = ds.assign_coords(**{coordinates['z']: phi})
 
+        # Record which dimension 'z' was renamed to.
+        ds.metadata['bout_zdim'] = coordinates['z']
+
     # Add 2D Cylindrical coordinates
     if ('R' not in ds) and ('Z' not in ds):
         ds = ds.rename(Rxy='R', Zxy='Z')

xbout/plotting/utils.py

@@ -160,8 +160,10 @@ def plot_separatrices(da, ax):
 
     j11, j12, j21, j22, ix1, ix2, nin, nx, ny, y_boundary_guards = _get_seps(da)
 
-    R = da.coords['R'].values
-    Z = da.coords['Z'].values
+    R = da.coords['R'].transpose(da.metadata['bout_xdim'],
+                                 da.metadata['bout_ydim']).values
+    Z = da.coords['Z'].transpose(da.metadata['bout_xdim'],
+                                 da.metadata['bout_ydim']).values
 
     if j22 + 1 < ny:
         # Lower X-point location
@@ -278,8 +280,10 @@ def plot_targets(da, ax, hatching=True):
 
     j11, j12, j21, j22, ix1, ix2, nin, nx, ny, y_boundary_guards = _get_seps(da)
 
-    R = da.coords['R'].values
-    Z = da.coords['Z'].values
+    R = da.coords['R'].transpose(da.metadata['bout_xdim'],
+                                 da.metadata['bout_ydim']).values
+    Z = da.coords['Z'].transpose(da.metadata['bout_xdim'],
+                                 da.metadata['bout_ydim']).values
 
     if j22 + 1 < ny:
         # lower PFR exists