Define Regions and add method to get BoutDataArrays in a region

xbout/boutdataarray.py

@@ -11,6 +11,8 @@
 from .plotting.animate import animate_poloidal, animate_pcolormesh, animate_line
 from .plotting import plotfuncs
 from .plotting.utils import _create_norm
+from .region import (Region, _concat_inner_guards, _concat_outer_guards,
+                     _concat_lower_guards, _concat_upper_guards)
 
 
 @register_dataarray_accessor('bout')
@@ -97,7 +99,7 @@ def _shiftZ(self, zShift):
         for dim in phase.dims:
             extra_dims.remove(dim)
         phase = phase.expand_dims(extra_dims)
-        phase = phase.transpose(*fft_dims)
+        phase = phase.transpose(*fft_dims, transpose_coords=True)
 
         data_shifted_fft = data_fft * np.exp(phase.data)
 
@@ -131,6 +133,56 @@ def fromFieldAligned(self):
         result["direction_y"] = "Standard"
         return result
 
+    def from_region(self, name, with_guards=None):
+        """
+        Get a logically-rectangular section of data from a certain region.
+        Includes guard cells from neighbouring regions.
+
+        Parameters
+        ----------
+        name : str
+            Region to get data for
+        with_guards : int or dict of int, optional
+            Number of guard cells to include, by default use MXG and MYG from BOUT++.
+            Pass a dict to set different numbers for different coordinates.
+        """
+
+        region = self.data.regions[name]
+        xcoord = self.data.metadata['bout_xdim']
+        ycoord = self.data.metadata['bout_ydim']
+
+        if with_guards is None:
+            mxg = self.data.metadata['MXG']
+            myg = self.data.metadata['MYG']
+        else:
+            try:
+                mxg = with_guards.get(xcoord, self.data.metadata['MXG'])
+                myg = with_guards.get(ycoord, self.data.metadata['MYG'])
+            except AttributeError:
+                mxg = with_guards
+                myg = with_guards
+
+        da = self.data.isel(region.get_slices())
+        da.attrs['region'] = region
+
+        if mxg > 0:
+            if region.connection_inner is not None:
+                # get inner x-guard cells for da from the global array
+                da = _concat_inner_guards(da, self.data, mxg)
+            if region.connection_outer is not None:
+                # get outer x-guard cells for da from the global array
+                da = _concat_outer_guards(da, self.data, mxg)
+
+        if myg > 0:
+            if region.connection_lower is not None:
+                # get lower y-guard cells from the global array
+                da = _concat_lower_guards(da, self.data, mxg, myg)
+            if region.connection_upper is not None:
+                # get upper y-guard cells from the global array
+                da = _concat_upper_guards(da, self.data, mxg, myg)
+
+        return da
+
     def animate2D(self, animate_over='t', x=None, y=None, animate=True, fps=10,
                   save_as=None, ax=None, poloidal_plot=False, logscale=None, **kwargs):
         """

xbout/boutdataset.py

@@ -148,6 +148,14 @@ def dict_to_attrs(obj, key):
             except KeyError:
                 pass
 
+        # Do not need to save regions as these can be reconstructed from the metadata
+        del to_save.attrs['regions']
+        for var in chain(to_save.data_vars, to_save.coords):
+            try:
+                del to_save[var].attrs['regions']
+            except KeyError:
+                pass
+
         if separate_vars:
             # Save each major variable to a different netCDF file
 

xbout/geometries.py

@@ -4,7 +4,9 @@
 import xarray as xr
 import numpy as np
 
+from .region import Region, _create_regions_toroidal
 from .utils import _set_attrs_on_all_vars
+
 REGISTERED_GEOMETRIES = {}
 
 
@@ -65,6 +67,7 @@ def apply_geometry(ds, geometry_name, *, coordinates=None, grid=None):
         updated_ds = add_geometry_coords(ds, grid=grid)
     else:
         updated_ds = add_geometry_coords(ds)
+
     return updated_ds
 
 
@@ -143,15 +146,21 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
     ds = ds.rename(y=coordinates['y'])
 
     # Add 1D Orthogonal Toroidal coordinates
+    # Make index 'x' a coordinate, useful for handling global indexing
+    nx = ds.dims['x']
+    ds = ds.assign_coords(x=np.arange(nx))
     ny = ds.dims[coordinates['y']]
-    # dy should always be constant in x, so it is safe to slice to x=0
+    # dy should always be constant in x, so it is safe to slice to x=0.
     # [The y-coordinate has to be a 1d coordinate that labels x-z slices of the grid
     # (similarly x-coordinate is 1d coordinate that labels y-z slices and z-coordinate is
     # a 1d coordinate that labels x-y slices). A coordinate might have different values
     # in disconnected regions, but there are no branch-cuts allowed in the x-direction in
     # BOUT++ (at least for the momement), so the y-coordinate has to be 1d and
     # single-valued. Therefore similarly dy has to be 1d and single-valued.]
-    dy = ds['dy'].isel(x=0)
+    # Need drop=True so that the result does not have an x-coordinate value which
+    # prevents it being added as a coordinate.
+    dy = ds['dy'].isel(x=0, drop=True)
+
     # calculate theta at the centre of each cell
     theta = dy.cumsum(keep_attrs=True) - dy/2.
     ds = ds.assign_coords(**{coordinates['y']: theta})
@@ -197,6 +206,8 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
     except KeyError:
         pass
 
+    ds = _create_regions_toroidal(ds)
+
     return ds
 
 
@@ -224,6 +235,8 @@ def add_s_alpha_geometry_coords(ds, *, coordinates=None, grid=None):
                          "geometry='s-alpha'")
     ds['r'] = ds['hthe'].isel({coordinates['y']: 0}).squeeze(drop=True)
     ds['r'].attrs['units'] = 'm'
+    # remove x-index coordinate, don't need when we have 'r' as a radial coordinate
+    ds = ds.drop('x')
     ds = ds.set_coords('r')
     ds = ds.rename(x='r')
 

xbout/load.py

@@ -308,7 +308,7 @@ def _auto_open_mfboutdataset(datapath, chunks={}, info=True,
     ds = xr.open_mfdataset(paths_grid, concat_dim=concat_dims, combine='nested',
                            data_vars=_BOUT_TIME_DEPENDENT_META_VARS,
                            preprocess=_preprocess, engine=filetype,
-                           chunks=chunks, **kwargs)
+                           chunks=chunks, join='exact', **kwargs)
 
     # Remove any duplicate time values from concatenation
     _, unique_indices = unique(ds['t_array'], return_index=True)
@@ -621,5 +621,5 @@ def _open_grid(datapath, chunks, keep_xboundaries, keep_yboundaries, mxg=2):
                     y=slice(nin + 2 * yboundaries, None, None))
                 grid = xr.concat((grid_lower, grid_upper), dim='y',
                                  data_vars='minimal',
-                                 compat='identical')
+                                 compat='identical', join='exact')
     return grid

xbout/plotting/animate.py

@@ -100,16 +100,16 @@ def animate_poloidal(da, *, ax=None, cax=None, animate_over='t', separatrix=True
 
     ax.set_aspect(aspect)
 
-    regions = _decompose_regions(da)
+    da_regions = _decompose_regions(da)
 
     # Plot all regions on same axis
     blocks = []
-    for region in regions:
+    for da_region in da_regions.values():
         # Load values eagerly otherwise for some reason the plotting takes
         # 100's of times longer - for some reason animatplot does not deal
         # well with dask arrays!
-        blocks.append(amp.blocks.Pcolormesh(region.coords[x].values,
-                      region.coords[y].values, region.values, ax=ax, cmap=cmap,
+        blocks.append(amp.blocks.Pcolormesh(da_region.coords[x].values,
+                      da_region.coords[y].values, da_region.values, ax=ax, cmap=cmap,
                       **kwargs))
 
     ax.set_title(da.name)
@@ -121,10 +121,10 @@ def animate_poloidal(da, *, ax=None, cax=None, animate_over='t', separatrix=True
          targets = False
 
     if separatrix:
-        plot_separatrices(da, ax)
+        plot_separatrices(da_regions, ax)
 
     if targets:
-        plot_targets(da, ax, hatching=add_limiter_hatching)
+        plot_targets(da_regions, ax, hatching=add_limiter_hatching)
 
     if animate:
         timeline = amp.Timeline(np.arange(da.sizes[animate_over]), fps=fps)
@@ -220,7 +220,7 @@ def animate_pcolormesh(data, animate_over='t', x=None, y=None, animate=True,
             raise ValueError("Dimension {} is not present in the data" .format(x))
         y = spatial_dims[0]
 
-    data = data.transpose(animate_over, y, x)
+    data = data.transpose(animate_over, y, x, transpose_coords=True)
 
     # Load values eagerly otherwise for some reason the plotting takes
     # 100's of times longer - for some reason animatplot does not deal
@@ -317,7 +317,7 @@ def animate_line(data, animate_over='t', animate=True,
     if (t_read is animate_over):
         pass
     else:
-        data = data.transpose(animate_over, t_read)
+        data = data.transpose(animate_over, t_read, transpose_coords=True)
 
     # Load values eagerly otherwise for some reason the plotting takes
     # 100's of times longer - for some reason animatplot does not deal

xbout/plotting/plotfuncs.py

@@ -1,5 +1,3 @@
-import collections
-
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
@@ -26,10 +24,10 @@ def regions(da, ax=None, **kwargs):
     if ax is None:
         fig, ax = plt.subplots()
 
-    regions = _decompose_regions(da)
+    da_regions = _decompose_regions(da)
 
-    colored_regions = [xr.full_like(region, fill_value=num / len(regions))
-                       for num, region in enumerate(regions)]
+    colored_regions = [xr.full_like(da_region, fill_value=num / len(regions))
+                       for num, da_region in enumerate(da_regions.values())]
 
     return [region.plot.pcolormesh(x=x, y=y, vmin=0, vmax=1, cmap='tab20',
                                    infer_intervals=False, add_colorbar=False, ax=ax,
@@ -181,12 +179,13 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True,
         if 'infer_intervals' not in kwargs:
             kwargs['infer_intervals'] = False
 
-    regions = _decompose_regions(da)
+    da_regions = _decompose_regions(da)
 
     # Plot all regions on same axis
-    add_labels = [True] + [False] * (len(regions) - 1)
+    add_labels = [True] + [False] * (len(da_regions) - 1)
     artists = [method(region, x=x, y=y, ax=ax, add_colorbar=False, add_labels=add_label,
-               cmap=cmap, **kwargs) for region, add_label in zip(regions, add_labels)]
+               cmap=cmap, **kwargs)
+               for region, add_label in zip(da_regions.values(), add_labels)]
 
     if method is xr.plot.contour:
         # using extend='neither' guarantees that the ends of the colorbar will be
@@ -213,35 +212,34 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True,
                 if not isinstance(value, slice):
                     raise ValueError('Argument passed to gridlines must be bool, int or '
                                      'slice. Got a ' + type(value) + ', ' + str(value))
-        R_global = da['R']
-        R_global.attrs['metadata'] = da.metadata
-
-        Z_global = da['Z']
-        Z_global.attrs['metadata'] = da.metadata
 
-        R_regions = _decompose_regions(da['R'])
-        Z_regions = _decompose_regions(da['Z'])
+        R_regions = [da_region['R'] for da_region in da_regions.values()]
+        Z_regions = [da_region['Z'] for da_region in da_regions.values()]
 
         for R, Z in zip(R_regions, Z_regions):
             if (not da.metadata['bout_xdim'] in R.dims
                     and not da.metadata['bout_ydim'] in R.dims):
                 # Small regions around X-point do not have segments in x- or y-directions,
                 # so skip
+                # Currently this region does not exist, but there is a small white gap at
+                # the X-point, so we might add it back in future
                 continue
             if gridlines.get('x') is not None:
                 # transpose in case Dataset or DataArray has been transposed away from the usual
                 # form
                 dim_order = (da.metadata['bout_xdim'], da.metadata['bout_ydim'])
                 yarg = {da.metadata['bout_ydim']: gridlines['x']}
-                plt.plot(R.isel(**yarg).transpose(*dim_order),
-                         Z.isel(**yarg).transpose(*dim_order), color='k', lw=0.1)
+                plt.plot(R.isel(**yarg).transpose(*dim_order, transpose_coords=True),
+                         Z.isel(**yarg).transpose(*dim_order, transpose_coords=True),
+                         color='k', lw=0.1)
             if gridlines.get('y') is not None:
                 xarg = {da.metadata['bout_xdim']: gridlines['y']}
                 # Need to plot transposed arrays to make gridlines that go in the
                 # y-direction
                 dim_order = (da.metadata['bout_ydim'], da.metadata['bout_xdim'])
-                plt.plot(R.isel(**xarg).transpose(*dim_order),
-                         Z.isel(**yarg).transpose(*dim_order), color='k', lw=0.1)
+                plt.plot(R.isel(**xarg).transpose(*dim_order, transpose_coords=True),
+                         Z.isel(**yarg).transpose(*dim_order, transpose_coords=True),
+                         color='k', lw=0.1)
 
     ax.set_title(da.name)
 
@@ -250,9 +248,9 @@ def plot2d_wrapper(da, method, *, ax=None, separatrix=True, targets=True,
         targets = False
 
     if separatrix:
-        plot_separatrices(da, ax)
+        plot_separatrices(da_regions, ax)
 
     if targets:
-        plot_targets(da, ax, hatching=add_limiter_hatching)
+        plot_targets(da_regions, ax, hatching=add_limiter_hatching)
 
     return artists