Source code for xbout.region

from copy import deepcopy
import numpy as np
import xarray as xr
from .utils import _set_attrs_on_all_vars


[docs] class Region: """ Contains the global indices bounding a single topological region, i.e. a region with logically rectangular contiguous data. Also stores the names of any neighbouring regions. """
[docs] def __init__( self, *, name, ds=None, xinner_ind=None, xouter_ind=None, ylower_ind=None, yupper_ind=None, connection_inner_x=None, connection_outer_x=None, connection_lower_y=None, connection_upper_y=None, ): """ Parameters ---------- name : str Name of the region ds : BoutDataset, optional Dataset to get variables to calculate coordinates from xinner_ind : int, optional Global x-index of the inner points of this region xouter_ind : int, optional Global x-index of the points just beyond the outer edge of this region ylower_ind : int, optional Global y-index of the lower points of this region yupper_ind : int, optional Global y-index of the points just beyond the upper edge of this region connection_inner_x : str, optional The region inside this one in the x-direction connection_outer_x : str, optional The region outside this one in the x-direction connection_lower_y : str, optional The region below this one in the y-direction connection_upper_y : str, optional The region above this one in the y-direction """ self.name = name self.xinner_ind = xinner_ind self.xouter_ind = xouter_ind if xouter_ind is not None and xinner_ind is not None: self.nx = xouter_ind - xinner_ind self.ylower_ind = ylower_ind self.yupper_ind = yupper_ind if yupper_ind is not None and ylower_ind is not None: self.ny = yupper_ind - ylower_ind self.connection_inner_x = connection_inner_x self.connection_outer_x = connection_outer_x self.connection_lower_y = connection_lower_y self.connection_upper_y = connection_upper_y self.global_nx = ds.sizes[ds.metadata["bout_xdim"]] self.global_ny = ds.sizes[ds.metadata["bout_ydim"]] if ds is not None: # self.nx, self.ny should not include boundary points. # self.xinner, self.xouter, self.ylower, self.yupper if ds.metadata["keep_xboundaries"]: xbndry = ds.metadata["MXG"] if self.connection_inner_x is None: self.nx = self.nx - xbndry # used to calculate x-coordinate of inner side (self.xinner) xinner_ind = xinner_ind + xbndry if self.connection_outer_x is None: self.nx = self.nx - xbndry # used to calculate x-coordinate of outer side (self.xouter) xouter_ind = xouter_ind - xbndry if ds.metadata["keep_yboundaries"]: ybndry = ds.metadata["MYG"] if self.connection_lower_y is None: self.ny = self.ny - ybndry # used to calculate y-coordinate of lower side (self.ylower) ylower_ind = ylower_ind + ybndry if self.connection_upper_y is None: self.ny = self.ny - ybndry # used to calculate y-coordinate of upper side (self.yupper) yupper_ind = yupper_ind - ybndry # calculate start and end coordinates ##################################### self.xcoord = ds.metadata["bout_xdim"] self.ycoord = ds.metadata["bout_ydim"] # Note the global coordinates used here are defined so that they are zero at # the boundaries of the grid (where the grid includes all boundary cells), # not necessarily the physical boundaries because constant offsets do not # matter, as long as these bounds are consistent with the global coordinates # defined in apply_geometry (we will only use these coordinates for # interpolation) and it is simplest to calculate them with cumsum(). # dx is constant in any particular region in the y-direction, so convert to a # 1d array # Note that this is not the same coordinate as the 'x' coordinate that is # created by default from the x-index, as these values are set only for # particular regions, so do not need to be consistent between different # regions (e.g. core and PFR), so we are not forced to use just the index # value here. # Define ref_yind so that we avoid using values from the corner cells, which # might be NaN. if self.connection_lower_y is None and ds.metadata["keep_yboundaries"] == 1: ref_yind = ylower_ind + ds.metadata["MYG"] else: ref_yind = ylower_ind dx = ds["dx"].isel({self.ycoord: ref_yind}) dx_cumsum = dx.cumsum() self.xinner = (dx_cumsum[xinner_ind] - dx[xinner_ind]).values self.xouter = (dx_cumsum[xouter_ind - 1] + dx[xouter_ind - 1]).values # dy is constant in the x-direction, so convert to a 1d array # Define ref_xind so that we avoid using values from the corner cells, which # might be NaN. if self.connection_inner_x is None and ds.metadata["keep_xboundaries"] == 1: ref_xind = xinner_ind + ds.metadata["MXG"] else: ref_xind = xinner_ind dy = ds["dy"].isel(**{self.xcoord: ref_xind}) dy_cumsum = dy.cumsum() self.ylower = (dy_cumsum[ylower_ind] - dy[ylower_ind]).values self.yupper = (dy_cumsum[yupper_ind - 1]).values
def __repr__(self): result = "<xbout.region.Region>\n" for attr, val in vars(self).items(): result += f"\t{attr}\t{val}\n" return result
[docs] def get_slices(self, mxg=0, myg=0): """ Return x- and y-dimension slices that select this region from the global DataArray. Returns ------- xslice, yslice : slice, slice """ xi = self.xinner_ind if self.connection_inner_x is not None and xi is not None: xi = xi - mxg xo = self.xouter_ind if self.connection_outer_x is not None and xo is not None: xo = xo + mxg yl = self.ylower_ind if self.connection_lower_y is not None and yl is not None: yl = yl - myg yu = self.yupper_ind if self.connection_upper_y is not None and yu is not None: yu = yu + myg return {self.xcoord: slice(xi, xo), self.ycoord: slice(yl, yu)}
[docs] def get_inner_guards_slices(self, *, mxg, myg=0): """ Return x- and y-dimension slices that select mxg guard cells on the inner-x side of this region from the global DataArray. Parameters ---------- mxg : int Number of guard cells myg : int, optional Number of y-guard cells to include at the corners """ ylower = self.ylower_ind if self.connection_lower_y is not None: ylower = ylower - myg yupper = self.yupper_ind if self.connection_upper_y is not None: yupper = yupper + myg return { self.xcoord: slice(self.xinner_ind - mxg, self.xinner_ind), self.ycoord: slice(ylower, yupper), }
[docs] def get_outer_guards_slices(self, *, mxg, myg=0): """ Return x- and y-dimension slices that select mxg guard cells on the outer-x side of this region from the global DataArray. Parameters ---------- mxg : int Number of guard cells myg : int, optional Number of y-guard cells to include at the corners """ ylower = self.ylower_ind if self.connection_lower_y is not None: ylower = ylower - myg yupper = self.yupper_ind if self.connection_upper_y is not None: yupper = yupper + myg return { self.xcoord: slice(self.xouter_ind, self.xouter_ind + mxg), self.ycoord: slice(ylower, yupper), }
[docs] def get_lower_guards_slices(self, *, myg, mxg=0): """ Return x- and y-dimension slices that select myg guard cells on the lower-y side of this region from the global DataArray. Parameters ---------- myg : int Number of guard cells mxg : int, optional Number of x-guard cells to include at the corners """ xinner = self.xinner_ind if self.connection_inner_x is not None: xinner = xinner - mxg xouter = self.xouter_ind if self.connection_outer_x is not None: xouter = xouter + mxg ystart = self.ylower_ind - myg ystop = self.ylower_ind if ystart < 0: # Make sure negative indices wrap around correctly if ystop == 0: ystop = None elif ystop > 0: raise ValueError( f"'start={ystart}' of lower guards slice is negative, but " f"'stop={ystop}' is positive. Should not be possible." ) return { self.xcoord: slice(xinner, xouter), self.ycoord: slice(ystart, ystop), }
[docs] def get_upper_guards_slices(self, *, myg, mxg=0): """ Return x- and y-dimension slices that select myg guard cells on the upper-y side of this region from the global DataArray. Parameters ---------- myg : int Number of guard cells mxg : int, optional Number of x-guard cells to include at the corners """ xinner = self.xinner_ind if self.connection_inner_x is not None: xinner = xinner - mxg xouter = self.xouter_ind if self.connection_outer_x is not None: xouter = xouter + mxg # wrap around to the beginning of the grid if necessary ystart = self.yupper_ind ystop = self.yupper_ind + myg if ystart >= self.global_ny: # wrap around to beginning of global array ystart = ystart - self.global_ny ystop = ystop - self.global_ny return { self.xcoord: slice(xinner, xouter), self.ycoord: slice(ystart, ystop), }
def __eq__(self, other): if not isinstance(other, Region): return NotImplemented return vars(self) == vars(other)
def _in_range(val, lower, upper): if val < lower: return lower elif val > upper: return upper else: return val def _order_vars(lower, upper): if upper < lower: return upper, lower else: return lower, upper def _get_topology(ds): jys11 = ds.metadata["jyseps1_1"] jys21 = ds.metadata["jyseps2_1"] ny_inner = ds.metadata["ny_inner"] jys12 = ds.metadata["jyseps1_2"] jys22 = ds.metadata["jyseps2_2"] ny = ds.metadata["ny"] ixs1 = ds.metadata["ixseps1"] ixs2 = ds.metadata["ixseps2"] nx = ds.metadata["nx"] if jys21 == jys12: # No upper X-point if jys11 <= 0 and jys22 >= ny - 1: ix = min(ixs1, ixs2) if ix >= nx - 1: return "core" elif ix <= 0: return "sol" else: return "limiter" return "single-null" if jys11 == jys21 and jys12 == jys22: if jys11 < ny_inner - 1 and jys22 > ny_inner: return "xpoint" else: raise ValueError("Currently unsupported topology") if ixs1 == ixs2: if jys21 < ny_inner - 1 and jys12 > ny_inner: return "connected-double-null" else: raise ValueError("Currently unsupported topology") elif ixs1 < ixs2: return "lower-disconnected-double-null" else: # ixs1 > ixs2 return "upper-disconnected-double-null" def _check_connections(regions): for region in regions.values(): if region.connection_inner_x is not None: if regions[region.connection_inner_x].connection_outer_x != region.name: raise ValueError( f"Inner-x connection of {region.name} is " f"{region.connection_inner_x}, but outer-x connection of " f"{region.connection_inner_x} is " f"{regions[region.connection_inner_x].connection_outer_x}" ) if region.connection_outer_x is not None: if regions[region.connection_outer_x].connection_inner_x != region.name: raise ValueError( f"Inner-x connection of {region.name} is " f"{region.connection_outer_x}, but inner-x connection of " f"{region.connection_outer_x} is " f"{regions[region.connection_outer_x].connection_inner_x}" ) if region.connection_lower_y is not None: if regions[region.connection_lower_y].connection_upper_y != region.name: raise ValueError( f"Lower-y connection of {region.name} is " f"{region.connection_lower_y}, but upper-y connection of " f"{region.connection_lower_y} is " f"{regions[region.connection_lower_y].connection_upper_y}" ) if region.connection_upper_y is not None: if regions[region.connection_upper_y].connection_lower_y != region.name: raise ValueError( f"Upper-y connection of {region.name} is " f"{region.connection_upper_y}, but lower-y connection of " f"{region.connection_upper_y} is " f"{regions[region.connection_upper_y].connection_lower_y}" ) topologies = {}
[docs] def topology_lower_disconnected_double_null( *, ds, ixs1, ixs2, nx, jys11, jys21, ny_inner, jys12, jys22, ny, ybndry ): regions = {} regions["lower_inner_PFR"] = Region( name="lower_inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=0, yupper_ind=jys11 + 1, connection_outer_x="lower_inner_intersep", connection_upper_y="lower_outer_PFR", ) regions["lower_inner_intersep"] = Region( name="lower_inner_intersep", ds=ds, xinner_ind=ixs1, xouter_ind=ixs2, ylower_ind=0, yupper_ind=jys11 + 1, connection_inner_x="lower_inner_PFR", connection_outer_x="lower_inner_SOL", connection_upper_y="inner_intersep", ) regions["lower_inner_SOL"] = Region( name="lower_inner_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=0, yupper_ind=jys11 + 1, connection_inner_x="lower_inner_intersep", connection_upper_y="inner_SOL", ) regions["inner_core"] = Region( name="inner_core", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys11 + 1, yupper_ind=jys21 + 1, connection_outer_x="inner_intersep", connection_lower_y="outer_core", connection_upper_y="outer_core", ) regions["inner_intersep"] = Region( name="inner_intersep", ds=ds, xinner_ind=ixs1, xouter_ind=ixs2, ylower_ind=jys11 + 1, yupper_ind=jys21 + 1, connection_inner_x="inner_core", connection_outer_x="inner_SOL", connection_lower_y="lower_inner_intersep", connection_upper_y="outer_intersep", ) regions["inner_SOL"] = Region( name="inner_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=jys11 + 1, yupper_ind=jys21 + 1, connection_inner_x="inner_intersep", connection_lower_y="lower_inner_SOL", connection_upper_y="upper_inner_SOL", ) regions["upper_inner_PFR"] = Region( name="upper_inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys21 + 1, yupper_ind=ny_inner, connection_outer_x="upper_inner_intersep", connection_lower_y="upper_outer_PFR", ) regions["upper_inner_intersep"] = Region( name="upper_inner_intersep", ds=ds, xinner_ind=ixs1, xouter_ind=ixs2, ylower_ind=jys21 + 1, yupper_ind=ny_inner, connection_inner_x="upper_inner_PFR", connection_outer_x="upper_inner_SOL", connection_lower_y="upper_outer_intersep", ) regions["upper_inner_SOL"] = Region( name="upper_inner_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=jys21 + 1, yupper_ind=ny_inner, connection_inner_x="upper_inner_intersep", connection_lower_y="inner_SOL", ) regions["upper_outer_PFR"] = Region( name="upper_outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=ny_inner, yupper_ind=jys12 + 1, connection_outer_x="upper_outer_intersep", connection_upper_y="upper_inner_PFR", ) regions["upper_outer_intersep"] = Region( name="upper_outer_intersep", ds=ds, xinner_ind=ixs1, xouter_ind=ixs2, ylower_ind=ny_inner, yupper_ind=jys12 + 1, connection_inner_x="upper_outer_PFR", connection_outer_x="upper_outer_SOL", connection_upper_y="upper_inner_intersep", ) regions["upper_outer_SOL"] = Region( name="upper_outer_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=ny_inner, yupper_ind=jys12 + 1, connection_inner_x="upper_outer_intersep", connection_upper_y="outer_SOL", ) regions["outer_core"] = Region( name="outer_core", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys12 + 1, yupper_ind=jys22 + 1, connection_outer_x="outer_intersep", connection_lower_y="inner_core", connection_upper_y="inner_core", ) regions["outer_intersep"] = Region( name="outer_intersep", ds=ds, xinner_ind=ixs1, xouter_ind=ixs2, ylower_ind=jys12 + 1, yupper_ind=jys22 + 1, connection_inner_x="outer_core", connection_outer_x="outer_SOL", connection_lower_y="inner_intersep", connection_upper_y="lower_outer_intersep", ) regions["outer_SOL"] = Region( name="outer_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=jys12 + 1, yupper_ind=jys22 + 1, connection_inner_x="outer_intersep", connection_lower_y="upper_outer_SOL", connection_upper_y="lower_outer_SOL", ) regions["lower_outer_PFR"] = Region( name="lower_outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys22 + 1, yupper_ind=ny, connection_outer_x="lower_outer_intersep", connection_lower_y="lower_inner_PFR", ) regions["lower_outer_intersep"] = Region( name="lower_outer_intersep", ds=ds, xinner_ind=ixs1, xouter_ind=ixs2, ylower_ind=jys22 + 1, yupper_ind=ny, connection_inner_x="lower_outer_PFR", connection_outer_x="lower_outer_SOL", connection_lower_y="outer_intersep", ) regions["lower_outer_SOL"] = Region( name="lower_outer_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=jys22 + 1, yupper_ind=ny, connection_inner_x="lower_outer_intersep", connection_lower_y="outer_SOL", ) return regions
topologies["lower-disconnected-double-null"] = topology_lower_disconnected_double_null
[docs] def topology_upper_disconnected_double_null( *, ds, ixs1, ixs2, nx, jys11, jys21, ny_inner, jys12, jys22, ny, ybndry ): regions = {} regions["lower_inner_PFR"] = Region( name="lower_inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs2, ylower_ind=0, yupper_ind=jys11 + 1, connection_outer_x="lower_inner_intersep", connection_upper_y="lower_outer_PFR", ) regions["lower_inner_intersep"] = Region( name="lower_inner_intersep", ds=ds, xinner_ind=ixs2, xouter_ind=ixs1, ylower_ind=0, yupper_ind=jys11 + 1, connection_inner_x="lower_inner_PFR", connection_outer_x="lower_inner_SOL", connection_upper_y="lower_outer_intersep", ) regions["lower_inner_SOL"] = Region( name="lower_inner_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=0, yupper_ind=jys11 + 1, connection_inner_x="lower_inner_intersep", connection_upper_y="inner_SOL", ) regions["inner_core"] = Region( name="inner_core", ds=ds, xinner_ind=0, xouter_ind=ixs2, ylower_ind=jys11 + 1, yupper_ind=jys21 + 1, connection_outer_x="inner_intersep", connection_lower_y="outer_core", connection_upper_y="outer_core", ) regions["inner_intersep"] = Region( name="inner_intersep", ds=ds, xinner_ind=ixs2, xouter_ind=ixs1, ylower_ind=jys11 + 1, yupper_ind=jys21 + 1, connection_inner_x="inner_core", connection_outer_x="inner_SOL", connection_lower_y="outer_intersep", connection_upper_y="upper_inner_intersep", ) regions["inner_SOL"] = Region( name="inner_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=jys11 + 1, yupper_ind=jys21 + 1, connection_inner_x="inner_intersep", connection_lower_y="lower_inner_SOL", connection_upper_y="upper_inner_SOL", ) regions["upper_inner_PFR"] = Region( name="upper_inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs2, ylower_ind=jys21 + 1, yupper_ind=ny_inner, connection_outer_x="upper_inner_intersep", connection_lower_y="upper_outer_PFR", ) regions["upper_inner_intersep"] = Region( name="upper_inner_intersep", ds=ds, xinner_ind=ixs2, xouter_ind=ixs1, ylower_ind=jys21 + 1, yupper_ind=ny_inner, connection_inner_x="upper_inner_PFR", connection_outer_x="upper_inner_SOL", connection_lower_y="inner_intersep", ) regions["upper_inner_SOL"] = Region( name="upper_inner_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=jys21 + 1, yupper_ind=ny_inner, connection_inner_x="upper_inner_intersep", connection_lower_y="inner_SOL", ) regions["upper_outer_PFR"] = Region( name="upper_outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs2, ylower_ind=ny_inner, yupper_ind=jys12 + 1, connection_outer_x="upper_outer_intersep", connection_upper_y="upper_inner_PFR", ) regions["upper_outer_intersep"] = Region( name="upper_outer_intersep", ds=ds, xinner_ind=ixs2, xouter_ind=ixs1, ylower_ind=ny_inner, yupper_ind=jys12 + 1, connection_inner_x="upper_outer_PFR", connection_outer_x="upper_outer_SOL", connection_upper_y="outer_intersep", ) regions["upper_outer_SOL"] = Region( name="upper_outer_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=ny_inner, yupper_ind=jys12 + 1, connection_inner_x="upper_outer_intersep", connection_upper_y="outer_SOL", ) regions["outer_core"] = Region( name="outer_core", ds=ds, xinner_ind=0, xouter_ind=ixs2, ylower_ind=jys12 + 1, yupper_ind=jys22 + 1, connection_outer_x="outer_intersep", connection_lower_y="inner_core", connection_upper_y="inner_core", ) regions["outer_intersep"] = Region( name="outer_intersep", ds=ds, xinner_ind=ixs2, xouter_ind=ixs1, ylower_ind=jys12 + 1, yupper_ind=jys22 + 1, connection_inner_x="outer_core", connection_outer_x="outer_SOL", connection_lower_y="upper_outer_intersep", connection_upper_y="inner_intersep", ) regions["outer_SOL"] = Region( name="outer_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=jys12 + 1, yupper_ind=jys22 + 1, connection_inner_x="outer_intersep", connection_lower_y="upper_outer_SOL", connection_upper_y="lower_outer_SOL", ) regions["lower_outer_PFR"] = Region( name="lower_outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs2, ylower_ind=jys22 + 1, yupper_ind=ny, connection_outer_x="lower_outer_intersep", connection_lower_y="lower_inner_PFR", ) regions["lower_outer_intersep"] = Region( name="lower_outer_intersep", ds=ds, xinner_ind=ixs2, xouter_ind=ixs1, ylower_ind=jys22 + 1, yupper_ind=ny, connection_inner_x="lower_outer_PFR", connection_outer_x="lower_outer_SOL", connection_lower_y="lower_inner_intersep", ) regions["lower_outer_SOL"] = Region( name="lower_outer_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=jys22 + 1, yupper_ind=ny, connection_inner_x="lower_outer_intersep", connection_lower_y="outer_SOL", ) return regions
topologies["upper-disconnected-double-null"] = topology_upper_disconnected_double_null
[docs] def topology_connected_double_null( *, ds, ixs1, ixs2, nx, jys11, jys21, ny_inner, jys12, jys22, ny, ybndry ): regions = {} regions["lower_inner_PFR"] = Region( name="lower_inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=0, yupper_ind=jys11 + 1, connection_outer_x="lower_inner_SOL", connection_upper_y="lower_outer_PFR", ) regions["lower_inner_SOL"] = Region( name="lower_inner_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=0, yupper_ind=jys11 + 1, connection_inner_x="lower_inner_PFR", connection_upper_y="inner_SOL", ) regions["inner_core"] = Region( name="inner_core", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys11 + 1, yupper_ind=jys21 + 1, connection_outer_x="inner_SOL", connection_lower_y="outer_core", connection_upper_y="outer_core", ) regions["inner_SOL"] = Region( name="inner_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=jys11 + 1, yupper_ind=jys21 + 1, connection_inner_x="inner_core", connection_lower_y="lower_inner_SOL", connection_upper_y="upper_inner_SOL", ) regions["upper_inner_PFR"] = Region( name="upper_inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys21 + 1, yupper_ind=ny_inner, connection_outer_x="upper_inner_SOL", connection_lower_y="upper_outer_PFR", ) regions["upper_inner_SOL"] = Region( name="upper_inner_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=jys21 + 1, yupper_ind=ny_inner, connection_inner_x="upper_inner_PFR", connection_lower_y="inner_SOL", ) regions["upper_outer_PFR"] = Region( name="upper_outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=ny_inner, yupper_ind=jys12 + 1, connection_outer_x="upper_outer_SOL", connection_upper_y="upper_inner_PFR", ) regions["upper_outer_SOL"] = Region( name="upper_outer_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=ny_inner, yupper_ind=jys12 + 1, connection_inner_x="upper_outer_PFR", connection_upper_y="outer_SOL", ) regions["outer_core"] = Region( name="outer_core", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys12 + 1, yupper_ind=jys22 + 1, connection_outer_x="outer_SOL", connection_lower_y="inner_core", connection_upper_y="inner_core", ) regions["outer_SOL"] = Region( name="outer_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=jys12 + 1, yupper_ind=jys22 + 1, connection_inner_x="outer_core", connection_lower_y="upper_outer_SOL", connection_upper_y="lower_outer_SOL", ) regions["lower_outer_PFR"] = Region( name="lower_outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys22 + 1, yupper_ind=ny, connection_outer_x="lower_outer_SOL", connection_lower_y="lower_inner_PFR", ) regions["lower_outer_SOL"] = Region( name="lower_outer_SOL", ds=ds, xinner_ind=ixs2, xouter_ind=nx, ylower_ind=jys22 + 1, yupper_ind=ny, connection_inner_x="lower_outer_PFR", connection_lower_y="outer_SOL", ) return regions
topologies["connected-double-null"] = topology_connected_double_null
[docs] def topology_single_null( *, ds, ixs1, ixs2, nx, jys11, jys21, ny_inner, jys12, jys22, ny, ybndry ): regions = {} regions["inner_PFR"] = Region( name="inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=0, yupper_ind=jys11 + 1, connection_outer_x="inner_SOL", connection_upper_y="outer_PFR", ) regions["inner_SOL"] = Region( name="inner_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=0, yupper_ind=jys11 + 1, connection_inner_x="inner_PFR", connection_upper_y="SOL", ) regions["core"] = Region( name="core", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys11 + 1, yupper_ind=jys22 + 1, connection_outer_x="SOL", connection_lower_y="core", connection_upper_y="core", ) regions["SOL"] = Region( name="SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=jys11 + 1, yupper_ind=jys22 + 1, connection_inner_x="core", connection_lower_y="inner_SOL", connection_upper_y="outer_SOL", ) regions["outer_PFR"] = Region( name="outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys22 + 1, yupper_ind=ny, connection_outer_x="outer_SOL", connection_lower_y="inner_PFR", ) regions["outer_SOL"] = Region( name="outer_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=jys22 + 1, yupper_ind=ny, connection_inner_x="outer_PFR", connection_lower_y="SOL", ) return regions
topologies["single-null"] = topology_single_null
[docs] def topology_limiter( *, ds, ixs1, ixs2, nx, jys11, jys21, ny_inner, jys12, jys22, ny, ybndry ): regions = {} regions["core"] = Region( name="core", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=ybndry, yupper_ind=ny - ybndry, connection_outer_x="SOL", connection_lower_y="core", connection_upper_y="core", ) regions["SOL"] = Region( name="SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=0, yupper_ind=ny, connection_inner_x="core", ) return regions
topologies["limiter"] = topology_limiter
[docs] def topology_core( *, ds, ixs1, ixs2, nx, jys11, jys21, ny_inner, jys12, jys22, ny, ybndry ): regions = {} regions["core"] = Region( name="core", ds=ds, xinner_ind=0, xouter_ind=nx, ylower_ind=ybndry, yupper_ind=ny - ybndry, connection_lower_y="core", connection_upper_y="core", ) return regions
topologies["core"] = topology_core
[docs] def topology_sol( *, ds, ixs1, ixs2, nx, jys11, jys21, ny_inner, jys12, jys22, ny, ybndry ): regions = {} regions["SOL"] = Region( name="SOL", ds=ds, xinner_ind=0, xouter_ind=nx, ylower_ind=0, yupper_ind=ny ) return regions
topologies["sol"] = topology_sol
[docs] def topology_xpoint( *, ds, ixs1, ixs2, nx, jys11, jys21, ny_inner, jys12, jys22, ny, ybndry ): regions = {} regions["lower_inner_PFR"] = Region( name="lower_inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=0, yupper_ind=jys11 + 1, connection_outer_x="lower_inner_SOL", connection_upper_y="lower_outer_PFR", ) regions["lower_inner_SOL"] = Region( name="lower_inner_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=0, yupper_ind=jys11 + 1, connection_inner_x="lower_inner_PFR", connection_upper_y="upper_inner_SOL", ) regions["upper_inner_PFR"] = Region( name="upper_inner_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys11 + 1, yupper_ind=ny_inner, connection_outer_x="upper_inner_SOL", connection_lower_y="upper_outer_PFR", ) regions["upper_inner_SOL"] = Region( name="upper_inner_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=jys11 + 1, yupper_ind=ny_inner, connection_inner_x="upper_inner_PFR", connection_lower_y="lower_inner_SOL", ) regions["upper_outer_PFR"] = Region( name="upper_outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=ny_inner, yupper_ind=jys22 + 1, connection_outer_x="upper_outer_SOL", connection_upper_y="upper_inner_PFR", ) regions["upper_outer_SOL"] = Region( name="upper_outer_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=ny_inner, yupper_ind=jys22 + 1, connection_inner_x="upper_outer_PFR", connection_upper_y="lower_outer_SOL", ) regions["lower_outer_PFR"] = Region( name="lower_outer_PFR", ds=ds, xinner_ind=0, xouter_ind=ixs1, ylower_ind=jys22 + 1, yupper_ind=ny, connection_outer_x="lower_outer_SOL", connection_lower_y="lower_inner_PFR", ) regions["lower_outer_SOL"] = Region( name="lower_outer_SOL", ds=ds, xinner_ind=ixs1, xouter_ind=nx, ylower_ind=jys22 + 1, yupper_ind=ny, connection_inner_x="lower_outer_PFR", connection_lower_y="upper_outer_SOL", ) return regions
topologies["xpoint"] = topology_xpoint def _create_regions_toroidal(ds): topology = _get_topology(ds) ixs1 = ds.metadata["ixseps1"] ixs2 = ds.metadata["ixseps2"] nx = ds.metadata["nx"] jys11 = ds.metadata["jyseps1_1"] jys21 = ds.metadata["jyseps2_1"] ny_inner = ds.metadata["ny_inner"] jys12 = ds.metadata["jyseps1_2"] jys22 = ds.metadata["jyseps2_2"] ny = ds.metadata["ny"] mxg = ds.metadata["MXG"] myg = ds.metadata["MYG"] # keep_yboundaries is 1 if there are y-boundaries and 0 if there are not ybndry = ds.metadata["keep_yboundaries"] * myg if jys21 == jys12: # No upper targets ybndry_upper = 0 else: ybndry_upper = ybndry # Make sure all sizes are sensible ixs1 = _in_range(ixs1, 0, nx) ixs2 = _in_range(ixs2, 0, nx) jys11 = _in_range(jys11, 0, ny - 1) jys21 = _in_range(jys21, 0, ny - 1) jys12 = _in_range(jys12, 0, ny - 1) jys21, jys12 = _order_vars(jys21, jys12) ny_inner = _in_range(ny_inner, jys21 + 1, jys12 + 1) jys22 = _in_range(jys22, 0, ny - 1) # Adjust for boundary cells # keep_xboundaries is 1 if there are x-boundaries and 0 if there are not if not ds.metadata["keep_xboundaries"]: ixs1 = ixs1 - mxg ixs2 = ixs2 - mxg nx = nx - 2 * mxg jys11 = jys11 + ybndry jys21 = jys21 + ybndry ny_inner = ny_inner + ybndry + ybndry_upper jys12 = jys12 + ybndry + 2 * ybndry_upper jys22 = jys22 + ybndry + 2 * ybndry_upper ny = ny + 2 * ybndry + 2 * ybndry_upper # Note, include guard cells in the created regions, fill them later if topology in topologies: regions = topologies[topology]( ds=ds, ixs1=ixs1, ixs2=ixs2, nx=nx, jys11=jys11, jys21=jys21, ny_inner=ny_inner, jys12=jys12, jys22=jys22, ny=ny, ybndry=ybndry, ) else: raise NotImplementedError(f"Topology '{topology}' is not implemented") _check_connections(regions) ds = _set_attrs_on_all_vars(ds, "regions", regions) ds.metadata["topology"] = topology return ds def _create_single_region(ds, periodic_y=True): nx = ds.metadata["nx"] ny = ds.metadata["ny"] connection = "all" if periodic_y else None regions = { "all": Region( name="all", ds=ds, xouter_ind=nx, xinner_ind=0, ylower_ind=0, yupper_ind=ny, connection_lower_y=connection, connection_upper_y=connection, ) } _check_connections(regions) ds = _set_attrs_on_all_vars(ds, "regions", regions) return ds def _concat_inner_guards(da, da_global, mxg): """ Concatenate inner x-guard cells to da, which is in a single region, getting the guard cells from the global array """ if mxg <= 0: return da if len(da.bout._regions) > 1: raise ValueError("da passed should have only one region") region = list(da.bout._regions.values())[0] if region.connection_inner_x not in da_global.bout._regions: # No connection, or plotting restricted set of regions not including this # connection return da myg_da = da_global.metadata["MYG"] keep_yboundaries = da_global.metadata["keep_yboundaries"] xcoord = da_global.metadata["bout_xdim"] ycoord = da_global.metadata["bout_ydim"] da_inner = da_global.bout.from_region(region.connection_inner_x, with_guards=0) if len(da_inner.bout._regions) > 1: raise ValueError("da_inner should have only one region") region_inner = list(da_inner.bout._regions.values())[0] if ( myg_da > 0 and keep_yboundaries and region.connection_lower_y is not None and region_inner.connection_lower_y is None ): # da_inner may have more points in the y-direction, because it has an actual # boundary, not a connection. Need to remove any extra points da_inner = da_inner.isel(**{ycoord: slice(myg_da, None)}) if ( myg_da > 0 and keep_yboundaries and region.connection_upper_y is not None and region_inner.connection_upper_y is None ): # da_inner may have more points in the y-direction, because it has an actual # boundary, not a connection. Need to remove any extra points da_inner = da_inner.isel(**{ycoord: slice(None, -myg_da)}) # select just the points we need to fill the guard cells of da da_inner = da_inner.isel(**{xcoord: slice(-mxg, None)}) if ( myg_da > 0 and keep_yboundaries and region.connection_lower_y is None and region_inner.connection_lower_y is not None ): # da_inner may have fewer points in the y-direction, because it has a connection, # not an actual boundary. Need to get the extra points from its connection da_inner_lower = da_global.bout.from_region( region_inner.connection_lower_y, with_guards=0 ) da_inner_lower = da_inner_lower.isel( **{xcoord: slice(-mxg, None), ycoord: slice(-myg_da, None)} ) save_regions = da_inner.bout._regions da_inner = xr.concat((da_inner_lower, da_inner), ycoord, join="exact") # xr.concat takes attributes from the first variable, but we need da_inner's # regions da_inner.attrs["regions"] = save_regions if ( myg_da > 0 and keep_yboundaries and region.connection_upper_y is None and region_inner.connection_upper_y is not None ): # da_inner may have fewer points in the y-direction, because it has a connection, # not an actual boundary. Need to get the extra points from its connection da_inner_upper = da_global.bout.from_region( region_inner.connection_upper_y, with_guards=0 ) da_inner_upper = da_inner_upper.isel( **{xcoord: slice(-mxg, None), ycoord: slice(myg_da)} ) da_inner = xr.concat((da_inner, da_inner_upper), ycoord, join="exact") # xr.concat takes attributes from the first variable, so regions are OK if xcoord in da.coords: # Some coordinates may not be single-valued, so use local coordinates for # neighbouring region, not communicated ones. Ensures the coordinates are # continuous so that interpolation works correctly near the boundaries. slices = region.get_inner_guards_slices(mxg=mxg) new_xcoord = da_global[xcoord].isel(**{xcoord: slices[xcoord]}) new_ycoord = da_global[ycoord].isel(**{ycoord: slices[ycoord]}) # can't use commented out version, uncommented one works around xarray bug # removing attrs # https://github.com/pydata/xarray/issues/4415 # https://github.com/pydata/xarray/issues/4393 # da_inner = da_inner.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord}) da_inner[xcoord].data[...] = new_xcoord.data da_inner = da_inner.reset_index(xcoord).set_xindex(xcoord) da_inner[ycoord].data[...] = new_ycoord.data da_inner = da_inner.reset_index(ycoord).set_xindex(ycoord) save_regions = da.bout._regions da = xr.concat((da_inner, da), xcoord, join="exact") # xr.concat takes attributes from the first variable (for xarray>=0.15.0, keeps attrs # that are the same in all objects for xarray<0.15.0) da.attrs["regions"] = save_regions return da def _concat_outer_guards(da, da_global, mxg): """ Concatenate outer x-guard cells to da, which is in a single region, getting the guard cells from the global array """ if mxg <= 0: return da if len(da.bout._regions) > 1: raise ValueError("da passed should have only one region") region = list(da.bout._regions.values())[0] if region.connection_outer_x not in da_global.bout._regions: # No connection, or plotting restricted set of regions not including this # connection return da myg_da = da_global.metadata["MYG"] keep_yboundaries = da_global.metadata["keep_yboundaries"] xcoord = da_global.metadata["bout_xdim"] ycoord = da_global.metadata["bout_ydim"] da_outer = da_global.bout.from_region(region.connection_outer_x, with_guards=0) if len(da_outer.bout._regions) > 1: raise ValueError("da_outer should have only one region") region_outer = list(da_outer.bout._regions.values())[0] if ( myg_da > 0 and keep_yboundaries and region.connection_lower_y is not None and region_outer.connection_lower_y is None ): # da_outer may have more points in the y-direction, because it has an actual # boundary, not a connection. Need to remove any extra points da_outer = da_outer.isel(**{ycoord: slice(myg_da, None)}) if ( myg_da > 0 and keep_yboundaries and region.connection_upper_y is not None and region_outer.connection_upper_y is None ): # da_outer may have more points in the y-direction, because it has an actual # boundary, not a connection. Need to remove any extra points da_outer = da_outer.isel(**{ycoord: slice(None, -myg_da)}) # select just the points we need to fill the guard cells of da da_outer = da_outer.isel(**{xcoord: slice(mxg)}) if ( myg_da > 0 and keep_yboundaries and region.connection_lower_y is None and region_outer.connection_lower_y is not None ): # da_outer may have fewer points in the y-direction, because it has a connection, # not an actual boundary. Need to get the extra points from its connection da_outer_lower = da_global.bout.from_region( region_outer.connection_lower_y, with_guards=0 ) da_outer_lower = da_outer_lower.isel( **{xcoord: slice(-mxg, None), ycoord: slice(-myg_da, None)} ) save_regions = da_outer.bout._regions da_outer = xr.concat((da_outer_lower, da_outer), ycoord, join="exact") # xr.concat takes attributes from the first variable, but we need da_outer's # regions da_outer.attrs["regions"] = save_regions if ( myg_da > 0 and keep_yboundaries and region.connection_upper_y is None and region_outer.connection_upper_y is not None ): # da_outer may have fewer points in the y-direction, because it has a connection, # not an actual boundary. Need to get the extra points from its connection da_outer_upper = da_global.bout.from_region( region_outer.connection_upper_y, with_guards=0 ) da_outer_upper = da_outer_upper.isel( **{xcoord: slice(-mxg, None), ycoord: slice(myg_da)} ) da_outer = xr.concat((da_outer, da_outer_upper), ycoord, join="exact") # xr.concat takes attributes from the first variable, so regions are OK if xcoord in da.coords: # Some coordinates may not be single-valued, so use local coordinates for # neighbouring region, not communicated ones. Ensures the coordinates are # continuous so that interpolation works correctly near the boundaries. slices = region.get_outer_guards_slices(mxg=mxg) new_xcoord = da_global[xcoord].isel(**{xcoord: slices[xcoord]}) new_ycoord = da_global[ycoord].isel(**{ycoord: slices[ycoord]}) # can't use commented out version, uncommented one works around xarray bug # removing attrs # https://github.com/pydata/xarray/issues/4415 # https://github.com/pydata/xarray/issues/4393 # da_outer = da_outer.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord}) da_outer[xcoord].data[...] = new_xcoord.data da_outer = da_outer.reset_index(xcoord).set_xindex(xcoord) da_outer[ycoord].data[...] = new_ycoord.data da_outer = da_outer.reset_index(ycoord).set_xindex(ycoord) save_regions = da.bout._regions da = xr.concat((da, da_outer), xcoord, join="exact") # xarray<0.15.0 only keeps attrs that are the same on all variables passed to concat da.attrs["regions"] = save_regions return da def _concat_lower_guards(da, da_global, mxg, myg): """ Concatenate lower y-guard cells to da, which is in a single region, getting the guard cells from the global array """ if myg <= 0: return da if len(da.bout._regions) > 1: raise ValueError("da passed should have only one region") region = list(da.bout._regions.values())[0] if region.connection_lower_y not in da_global.bout._regions: # No connection, or plotting restricted set of regions not including this # connection return da xcoord = da_global.metadata["bout_xdim"] ycoord = da_global.metadata["bout_ydim"] da_lower = da_global.bout.from_region( region.connection_lower_y, with_guards={xcoord: mxg, ycoord: 0} ) # select just the points we need to fill the guard cells of da da_lower = da_lower.isel(**{ycoord: slice(-myg, None)}) if ycoord in da.coords: # Some coordinates may not be single-valued, so use local coordinates for # neighbouring region, not communicated ones. Ensures the coordinates are # continuous so that interpolation works correctly near the boundaries. slices = region.get_lower_guards_slices(mxg=mxg, myg=myg) new_xcoord = da_global[xcoord].isel(**{xcoord: slices[xcoord]}) new_ycoord = da_global[ycoord].isel(**{ycoord: slices[ycoord]}) # new_ycoord might not join smoothly onto da[ycoord] for limiter or core-only # geometries where the guard cells have to come from the opposite side of a # branch cut. Need to work around this, or it breaks parallel interpolation. ycoord_increasing = np.any( da[ycoord].isel(**{ycoord: 1}) > da[ycoord].isel(**{ycoord: 0}) ) if ycoord_increasing and np.any( new_ycoord.isel(**{ycoord: -1}) > da[ycoord].isel(**{ycoord: 0}) ): if not np.all( new_ycoord.isel(**{ycoord: -1}) > da[ycoord].isel(**{ycoord: 0}) ): raise ValueError( f"Inconsistent {ycoord} - should be either increasing everywhere " f"or decreasing everywhere" ) # Estimate what coordinate values should be by extrapolating linearly from # interior of region. This is exact if the grid spacing is constant, which # it usually is in 'y' offset = ( 2.0 * da[ycoord].isel(**{ycoord: 0}) - da[ycoord].isel(**{ycoord: 1}) - new_ycoord.isel(**{ycoord: -1}) ) new_ycoord = new_ycoord + offset elif (not ycoord_increasing) and np.any( new_ycoord.isel(**{ycoord: -1}) < da[ycoord].isel(**{ycoord: 0}) ): if not np.all( new_ycoord.isel(**{ycoord: -1}) < da[ycoord].isel(**{ycoord: 0}) ): raise ValueError( f"Inconsistent {ycoord} - should be either increasing everywhere " f"or decreasing everywhere" ) # Estimate what coordinate values should be by extrapolating linearly from # interior of region. This is exact if the grid spacing is constant, which # it usually is in 'y' offset = ( 2.0 * da[ycoord].isel(**{ycoord: 0}) - da[ycoord].isel(**{ycoord: 1}) - new_ycoord.isel(**{ycoord: -1}) ) new_ycoord = new_ycoord + offset # can't use commented out version, uncommented one works around xarray bug # removing attrs # https://github.com/pydata/xarray/issues/4415 # https://github.com/pydata/xarray/issues/4393 # da_lower = da_lower.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord}) da_lower[xcoord].data[...] = new_xcoord.data da_lower = da_lower.reset_index(xcoord).set_xindex(xcoord) da_lower[ycoord].data[...] = new_ycoord.data da_lower = da_lower.reset_index(ycoord).set_xindex(ycoord) if "poloidal_distance" in da.coords and myg > 0: # Special handling for core regions to deal with branch cut if "core" in region.name: # import pdb; pdb.set_trace() # Try to detect whether there is branch cut at lower boundary: if there is # poloidal_distance_ylow should be zero at the boundary of this region poloidal_distance_bottom = da["poloidal_distance_ylow"].isel({ycoord: 0}) if all(abs(poloidal_distance_bottom) < 1.0e-16): # Offset so that the poloidal_distance in da_lower is continuous from # the poloidal_distance in this region. # Expect there to be y-boundary cells in the Dataset, this will probably # fail if there are not. total_poloidal_distance = da["total_poloidal_distance"] da_lower["poloidal_distance"] -= total_poloidal_distance da_lower["poloidal_distance_ylow"] -= total_poloidal_distance save_regions = da.bout._regions da = xr.concat((da_lower, da), ycoord, join="exact") # xr.concat takes attributes from the first variable (for xarray>=0.15.0, keeps attrs # that are the same in all objects for xarray<0.15.0) da.attrs["regions"] = save_regions return da def _concat_upper_guards(da, da_global, mxg, myg): """ Concatenate upper y-guard cells to da, which is in a single region, getting the guard cells from the global array """ if myg <= 0: return da if len(da.bout._regions) > 1: raise ValueError("da passed should have only one region") region = list(da.bout._regions.values())[0] if region.connection_upper_y not in da_global.bout._regions: # No connection, or plotting restricted set of regions not including this # connection return da xcoord = da_global.metadata["bout_xdim"] ycoord = da_global.metadata["bout_ydim"] da_upper = da_global.bout.from_region( region.connection_upper_y, with_guards={xcoord: mxg, ycoord: 0} ) # select just the points we need to fill the guard cells of da da_upper = da_upper.isel(**{ycoord: slice(myg)}) if ycoord in da.coords: # Some coordinates may not be single-valued, so use local coordinates for # neighbouring region, not communicated ones. Ensures the coordinates are # continuous so that interpolation works correctly near the boundaries. slices = region.get_upper_guards_slices(mxg=mxg, myg=myg) new_xcoord = da_global[xcoord].isel(**{xcoord: slices[xcoord]}) new_ycoord = da_global[ycoord].isel(**{ycoord: slices[ycoord]}) # new_ycoord might not join smoothly onto da[ycoord] for limiter or core-only # geometries where the guard cells have to come from the opposite side of a # branch cut. Need to work around this, or it breaks parallel interpolation. ycoord_increasing = np.any( da[ycoord].isel(**{ycoord: -1}) > da[ycoord].isel(**{ycoord: -2}) ) if ycoord_increasing and np.any( new_ycoord.isel(**{ycoord: 0}) < da[ycoord].isel(**{ycoord: -1}) ): if not np.all( new_ycoord.isel(**{ycoord: 0}) < da[ycoord].isel(**{ycoord: -1}) ): raise ValueError( f"Inconsistent {ycoord} - should be either increasing everywhere " f"or decreasing everywhere" ) # Estimate what coordinate values should be by extrapolating linearly from # interior of region. This is exact if the grid spacing is constant, which # it usually is in 'y' offset = ( 2.0 * da[ycoord].isel(**{ycoord: -1}) - da[ycoord].isel(**{ycoord: -2}) - new_ycoord.isel(**{ycoord: 0}) ) new_ycoord = new_ycoord + offset elif (not ycoord_increasing) and np.any( new_ycoord.isel(**{ycoord: 0}) > da[ycoord].isel(**{ycoord: -1}) ): if not np.all( new_ycoord.isel(**{ycoord: 0}) > da[ycoord].isel(**{ycoord: -1}) ): raise ValueError( f"Inconsistent {ycoord} - should be either increasing everywhere " f"or decreasing everywhere" ) # Estimate what coordinate values should be by extrapolating linearly from # interior of region. This is exact if the grid spacing is constant, which # it usually is in 'y' offset = ( 2.0 * da[ycoord].isel(**{ycoord: -1}) - da[ycoord].isel(**{ycoord: -2}) - new_ycoord.isel(**{ycoord: 0}) ) new_ycoord = new_ycoord + offset # can't use commented out version, uncommented one works around xarray bug # removing attrs # https://github.com/pydata/xarray/issues/4415 # https://github.com/pydata/xarray/issues/4393 # da_upper = da_upper.assign_coords(**{xcoord: new_xcoord, ycoord: new_ycoord}) da_upper[xcoord].data[...] = new_xcoord.data da_upper = da_upper.reset_index(xcoord).set_xindex(xcoord) da_upper[ycoord].data[...] = new_ycoord.data da_upper = da_upper.reset_index(ycoord).set_xindex(ycoord) if "poloidal_distance" in da.coords and myg > 0: # Special handling for core regions to deal with branch cut if "core" in region.name: # import pdb; pdb.set_trace() # Try to detect whether there is branch cut at upper boundary: if there is # poloidal_distance_ylow should be zero at the boundary of da_upper poloidal_distance_bottom = da_upper["poloidal_distance_ylow"].isel( {ycoord: 0} ) if all(abs(poloidal_distance_bottom) < 1.0e-16): # Offset so that the poloidal_distance in da_upper is continuous from # the poloidal_distance in this region. # Expect there to be y-boundary cells in the Dataset, this will probably # fail if there are not. total_poloidal_distance = da["total_poloidal_distance"] da_upper["poloidal_distance"] += total_poloidal_distance da_upper["poloidal_distance_ylow"] += total_poloidal_distance save_regions = da.bout._regions da = xr.concat((da, da_upper), ycoord, join="exact") # xarray<0.15.0 only keeps attrs that are the same on all variables passed to concat da.attrs["regions"] = save_regions return da def _from_region(ds_or_da, name, with_guards): # ensure we do not modify the input ds_or_da = ds_or_da.copy(deep=True) region = ds_or_da.bout._regions[name] xcoord = ds_or_da.metadata["bout_xdim"] ycoord = ds_or_da.metadata["bout_ydim"] if with_guards is None: mxg = ds_or_da.metadata["MXG"] myg = ds_or_da.metadata["MYG"] else: try: mxg = with_guards.get(xcoord, ds_or_da.metadata["MXG"]) myg = with_guards.get(ycoord, ds_or_da.metadata["MYG"]) except AttributeError: mxg = with_guards myg = with_guards result = ds_or_da.isel(region.get_slices()).copy() # The returned result has only one region single_region = deepcopy(region) result.attrs["regions"] = {name: single_region} # get inner x-guard cells for result from the global array result = _concat_inner_guards(result, ds_or_da, mxg) # get outer x-guard cells for result from the global array result = _concat_outer_guards(result, ds_or_da, mxg) # get lower y-guard cells from the global array result = _concat_lower_guards(result, ds_or_da, mxg, myg) # get upper y-guard cells from the global array result = _concat_upper_guards(result, ds_or_da, mxg, myg) # If the result (which only has a single region) is passed to from_region a # second time, don't want to slice anything. single_region = list(result.bout._regions.values())[0] single_region.xinner_ind = None single_region.xouter_ind = None single_region.ylower_ind = None single_region.yupper_ind = None return result