Merge pull request #168 from boutproject/save-all-corners

Save R,Z locations for all four corners of each cell

doc/grid-file.rst

@@ -154,6 +154,15 @@ Spatial positions
      - Major radius and height of the lower-left corner of each grid cell. Not
        needed by BOUT++, but may be useful for post-processing.
 
+   * - ``Rxy_lower_right_corners``, ``Zxy_lower_right_corners``,
+       ``Rxy_upper_right_corners``, ``Zxy_upper_right_corners``,
+       ``Rxy_upper_left_corners``, ``Zxy_upper_left_corners``
+
+     - Major radius and height of the other three corners of each grid cell.
+       Mostly redundant information with ``Rxy_corners`` and ``Zxy_corners``,
+       but may make handling branch cuts and upper/outer boundaries more
+       convenient. Not needed by BOUT++, but may be useful for post-processing.
+
 Grid spacings
 +++++++++++++
 

doc/whats-new.md

@@ -12,6 +12,8 @@ Release history
 
 ### New features
 
+- Save R,Z locations for all four corners of each cell (#168)\
+  By [John Omotani](https://github.com/johnomotani)
 - `PsiContour.plot()` and `FineContour.plot()` can be called with an `ax`
   argument, and passing `psi` is optional (#163)\
   By [Ben Dudson](https://github.com/bendudson)

hypnotoad/cases/torpex.py

@@ -141,7 +141,6 @@ class TORPEXMagneticField(Equilibrium):
     )
 
     def __init__(self, equilibOptions, meshOptions):
-
         # Set up options read from user input
         self.user_options = self.user_options_factory.create(meshOptions)
 

hypnotoad/core/equilibrium.py

@@ -564,7 +564,6 @@ def __init__(self, parentContour, settings, *, psi):
         self.equaliseSpacing(psi=psi)
 
     def extend(self, *, psi, extend_lower=0, extend_upper=0):
-
         Nfine = self.user_options.finecontour_Nfine
 
         parentCopy = self.parentContour.newContourFromSelf()
@@ -732,7 +731,6 @@ def equaliseSpacing(self, *, psi, reallocate=False):
         # endInd unchanged - makes iteration more stable.
         count = 1
         while ds_error > self.user_options.finecontour_atol:
-
             if (
                 self.user_options.finecontour_maxits
                 and count > self.user_options.finecontour_maxits
@@ -1862,7 +1860,6 @@ def checkFineContourExtend(self, *, psi):
                 (fine_contour.positions[1, :] - fine_contour.positions[0, :]) ** 2
             )
         ):
-
             ds = fine_contour.distance[1] - fine_contour.distance[0]
             n_extend_lower = max(int(numpy.ceil(distances[0] / ds)), 1)
         else:
@@ -1882,7 +1879,6 @@ def checkFineContourExtend(self, *, psi):
                 (fine_contour.positions[-1, :] - fine_contour.positions[-2, :]) ** 2
             )
         ):
-
             ds = fine_contour.distance[-1] - fine_contour.distance[-2]
             n_extend_upper = max(int(numpy.ceil(distances[-1] / ds)), 1)
         else:
@@ -2995,7 +2991,6 @@ def combineSfuncs(
             spacings["nonorthogonal_range_lower"] is not None
             and spacings["nonorthogonal_range_upper"] is not None
         ):
-
             if sfunc_orthogonal is None:
                 # Define new_sfunc in a sensible way to create the initial distribution
                 # of points on the separatrix that is then used to create the orthogonal
@@ -3082,7 +3077,6 @@ def new_sfunc(i):
                     )
 
         elif spacings["nonorthogonal_range_lower"] is not None:
-
             if sfunc_orthogonal is None:
                 # Fix spacing so that if we call combineSfuncs again for this contour
                 # with sfunc_orthogonal from self.contourSfunc() then we get the same
@@ -3114,7 +3108,6 @@ def new_sfunc(i):
                     return (weight_lower) * sfixed_lower + (1.0 - weight_lower) * sorth
 
         elif spacings["nonorthogonal_range_upper"] is not None:
-
             if sfunc_orthogonal is None:
                 # Fix spacing so that if we call combineSfuncs again for this contour
                 # with sfunc_orthogonal from self.contourSfunc() then we get the same
@@ -4410,7 +4403,6 @@ def wallIntersection(self, p1, p2):
                     and numpy.abs(intersect.Z - second_intersect.Z)
                     < intersect_tolerance
                 ):
-
                     print("Multiple intersections with the wall")
 
                     import matplotlib.pyplot as plt

hypnotoad/core/mesh.py

@@ -146,7 +146,6 @@ def __init__(
         settings,
         parallel_map,
     ):
-
         self.user_options = self.user_options_factory.create(settings)
 
         self.name = equilibriumRegion.name + "(" + str(radialIndex) + ")"
@@ -3240,7 +3239,6 @@ class BoutMesh(Mesh):
     )
 
     def __init__(self, equilibrium, settings):
-
         super().__init__(equilibrium, settings)
 
         # nx, ny both include boundary guard cells
@@ -3311,7 +3309,7 @@ def geometry(self):
         # Call geometry() method of base class
         super().geometry()
 
-        def addFromRegions(name):
+        def addFromRegions(name, *, all_corners=False):
             # Collect a 2d field from the regions
             self.fields_to_output.append(name)
             f = MultiLocationArray(self.nx, self.ny)
@@ -3334,6 +3332,17 @@ def addFromRegions(name):
                     f.corners[self.region_indices[region.myID]] = f_region.corners[
                         :-1, :-1
                     ]
+                if all_corners:
+                    if f_region._corners_array is not None:
+                        f.lower_right_corners[
+                            self.region_indices[region.myID]
+                        ] = f_region.corners[1:, :-1]
+                        f.upper_right_corners[
+                            self.region_indices[region.myID]
+                        ] = f_region.corners[1:, 1:]
+                        f.upper_left_corners[
+                            self.region_indices[region.myID]
+                        ] = f_region.corners[:-1, 1:]
 
             # Set 'bout_type' so it gets saved in the grid file
             f.attributes["bout_type"] = "Field2D"
@@ -3369,8 +3378,8 @@ def addFromRegionsXArray(name):
             # Set 'bout_type' so it gets saved in the grid file
             f.attributes["bout_type"] = "ArrayX"
 
-        addFromRegions("Rxy")
-        addFromRegions("Zxy")
+        addFromRegions("Rxy", all_corners=True)
+        addFromRegions("Zxy", all_corners=True)
         addFromRegions("psixy")
         addFromRegions("dx")
         addFromRegions("dy")
@@ -3436,6 +3445,18 @@ def writeCorners(self, name, array, f):
             name + "_corners",
             BoutArray(array.corners[:-1, :-1], attributes=array.attributes),
         )
+        f.write(
+            name + "_lower_right_corners",
+            BoutArray(array.lower_right_corners[:-1, :-1], attributes=array.attributes),
+        )
+        f.write(
+            name + "_upper_right_corners",
+            BoutArray(array.upper_right_corners[:-1, :-1], attributes=array.attributes),
+        )
+        f.write(
+            name + "_upper_left_corners",
+            BoutArray(array.upper_left_corners[:-1, :-1], attributes=array.attributes),
+        )
 
     def writeArrayXDirection(self, name, array, f):
         f.write(name, BoutArray(array.centre[:, 0], attributes=array.attributes))

hypnotoad/core/multilocationarray.py

@@ -6,12 +6,22 @@ class MultiLocationArray(numpy.lib.mixins.NDArrayOperatorsMixin):
     """
     Container for arrays representing points at different cell locations
     Not all have to be filled.
+
+    Note the ``lower_right_corners``, ``upper_right_corners``, and
+    ``upper_left_corners`` members are only intended to be used for the global arrays,
+    as within each region the ``corners`` member contains all the corners for every
+    cell. ``lower_right_corners``, ``upper_right_corners``, and ``upper_left_corners``
+    are therefore not set to zero in ``MultiLocationArray.zero()`` as they do not need
+    to be initialized.
     """
 
     _centre_array = None
     _xlow_array = None
     _ylow_array = None
     _corners_array = None
+    _lower_right_corners_array = None
+    _upper_right_corners_array = None
+    _upper_left_corners_array = None
 
     def __init__(self, nx, ny):
         self.nx = nx
@@ -67,6 +77,42 @@ def corners(self, value):
             self._corners_array = numpy.zeros([self.nx + 1, self.ny + 1])
         self._corners_array[...] = value
 
+    @property
+    def lower_right_corners(self):
+        if self._lower_right_corners_array is None:
+            self._lower_right_corners_array = numpy.zeros([self.nx + 1, self.ny + 1])
+        return self._lower_right_corners_array
+
+    @lower_right_corners.setter
+    def lower_right_corners(self, value):
+        if self._lower_right_corners_array is None:
+            self._lower_right_corners_array = numpy.zeros([self.nx + 1, self.ny + 1])
+        self._lower_right_corners_array[...] = value
+
+    @property
+    def upper_right_corners(self):
+        if self._upper_right_corners_array is None:
+            self._upper_right_corners_array = numpy.zeros([self.nx + 1, self.ny + 1])
+        return self._upper_right_corners_array
+
+    @upper_right_corners.setter
+    def upper_right_corners(self, value):
+        if self._upper_right_corners_array is None:
+            self._upper_right_corners_array = numpy.zeros([self.nx + 1, self.ny + 1])
+        self._upper_right_corners_array[...] = value
+
+    @property
+    def upper_left_corners(self):
+        if self._upper_left_corners_array is None:
+            self._upper_left_corners_array = numpy.zeros([self.nx + 1, self.ny + 1])
+        return self._upper_left_corners_array
+
+    @upper_left_corners.setter
+    def upper_left_corners(self, value):
+        if self._upper_left_corners_array is None:
+            self._upper_left_corners_array = numpy.zeros([self.nx + 1, self.ny + 1])
+        self._upper_left_corners_array[...] = value
+
     def copy(self):
         new_multilocationarray = MultiLocationArray(self.nx, self.ny)
         if self.centre is not None:
@@ -219,6 +265,13 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
 
     def zero(self):
         # Initialise all locations, set them to zero and return the result
+        #
+        # Note the ``lower_right_corners``, ``upper_right_corners``, and
+        # ``upper_left_corners`` members are only intended to be used for the global
+        # arrays, as within each region the ``corners`` member contains all the corners
+        # for every cell. ``lower_right_corners``, ``upper_right_corners``, and
+        # ``upper_left_corners`` are therefore not set to zero here as they do not need
+        # to be initialized.
         self.centre = 0.0
         self.xlow = 0.0
         self.ylow = 0.0

hypnotoad/gui/gui.py

@@ -624,7 +624,6 @@ def __init__(self, parent):
         self.saveFullYamlCheckBox.setChecked(self.parent.gui_options["save_full_yaml"])
 
     def accept(self):
-
         self.parent.gui_options["grid_file"] = self.defaultGridFileNameLineEdit.text()
         self.parent.gui_options["plot_xlow"] = self.plotXlowCheckBox.isChecked()
         self.parent.gui_options["plot_ylow"] = self.plotYlowCheckBox.isChecked()

hypnotoad/gui/matplotlib_widget.py

@@ -23,7 +23,6 @@
 
 class MatplotlibWidget:
     def __init__(self, parent):
-
         self.figure = Figure()
         self.canvas = FigureCanvas(self.figure)
         self.canvas.setParent(parent)

hypnotoad/test_suite/test_tokamak.py

@@ -79,7 +79,6 @@ def fpolprime_func(psi):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_read_geqdsk(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     # Number of mesh points
@@ -210,7 +209,6 @@ def test_bounding():
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_xpoint(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     nx = 65
@@ -521,7 +519,6 @@ def psi_func(R, Z):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_findlegs(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     eq = make_lower_single_null(settings)
@@ -540,7 +537,6 @@ def test_findlegs(psi_interpolation_method):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_findlegs_upper(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     eq = make_upper_single_null(settings)
@@ -559,7 +555,6 @@ def test_findlegs_upper(psi_interpolation_method):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_makeregions_lsn(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     eq = make_lower_single_null(settings)
@@ -570,7 +565,6 @@ def test_makeregions_lsn(psi_interpolation_method):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_makeregions_usn(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     eq = make_upper_single_null(settings)
@@ -581,7 +575,6 @@ def test_makeregions_usn(psi_interpolation_method):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_makeregions_cdn(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     eq = make_connected_double_null(settings)
@@ -592,7 +585,6 @@ def test_makeregions_cdn(psi_interpolation_method):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_makeregions_udn(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     eq = make_upper_double_null(settings)
@@ -603,7 +595,6 @@ def test_makeregions_udn(psi_interpolation_method):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_makeregions_ldn(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     eq = make_lower_double_null(settings)
@@ -614,7 +605,6 @@ def test_makeregions_ldn(psi_interpolation_method):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_makeregions_udn_largesep_1(psi_interpolation_method):
-
     settings = {"psi_interpolation_method": psi_interpolation_method}
 
     eq = make_upper_double_null_largesep(settings)
@@ -624,7 +614,6 @@ def test_makeregions_udn_largesep_1(psi_interpolation_method):
 
 @pytest.mark.parametrize("psi_interpolation_method", ["spline", "dct"])
 def test_makeregions_udn_largesep_2(psi_interpolation_method):
-
     settings = {
         "psinorm_sol": 1.2,
         "psi_interpolation_method": psi_interpolation_method,

hypnotoad/utils/critical.py

@@ -99,7 +99,6 @@ def find_critical(R, Z, psi, atol, maxits, discard_xpoints=False):
                 and (Bp2[i, j] < Bp2[i, j + 1])
                 and (Bp2[i, j] < Bp2[i, j - 1])
             ):
-
                 # Found local minimum
 
                 R0 = R[i, j]
@@ -112,7 +111,6 @@ def find_critical(R, Z, psi, atol, maxits, discard_xpoints=False):
 
                 count = 0
                 while True:
-
                     Br = -f(R1, Z1, dy=1, grid=False) / R1
                     Bz = f(R1, Z1, dx=1, grid=False) / R1
 

hypnotoad/utils/parallel_map.py

@@ -83,7 +83,7 @@ def __call__(self, function, args_list, **kwargs):
                     psi=self.psi,
                     f_R=self.f_R,
                     f_Z=self.f_Z,
-                    **kwargs
+                    **kwargs,
                 )
                 for args in args_list
             ]

hypnotoad/utils/utils.py

@@ -21,7 +21,6 @@
 
 
 def with_default(value, default):
-
     if value is not None:
         return value