Save poloidal distances along psi-contours to grid files

.github/workflows/pythonpackage.yml

@@ -17,7 +17,7 @@ jobs:
     if: always()
     strategy:
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: [3.7, 3.8, 3.9]
       fail-fast: false
 
     steps:
@@ -42,7 +42,7 @@ jobs:
     if: always()
     strategy:
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: [3.7, 3.8, 3.9]
       fail-fast: false
 
     steps:
@@ -78,7 +78,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v1
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
@@ -104,7 +104,7 @@ jobs:
     if: always()
     strategy:
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: 3.9
 
     steps:
     - uses: actions/checkout@v2
@@ -128,7 +128,7 @@ jobs:
     if: always()
     strategy:
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: 3.9
 
     steps:
     - uses: actions/checkout@v2

.github/workflows/pythonpublish.yml

@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: [3.7, 3.8, 3.9]
 
     steps:
     - uses: actions/checkout@v2
@@ -37,7 +37,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: 3.9
 
     steps:
     - uses: actions/checkout@v2
@@ -60,7 +60,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: 3.9
 
     steps:
     - uses: actions/checkout@v2

doc/whats-new.md

@@ -21,6 +21,8 @@ What's new
 
 ### New features
 
+- Save poloidal distances along psi contours to grid file (#116)\
+  By [John Omotani](https://github.com/johnomotani)
 - Enable DCT interpolation scheme in TokamakEquilibrum, can be selected with new option
   psi_interpolation_method (#113)\
   By [John Omotani](https://github.com/johnomotani)

hypnotoad/core/mesh.py

@@ -738,16 +738,20 @@ def getRZBoundary(self):
             self.Rxy.corners[:, -1] = up.Rxy.corners[:, 0]
             self.Zxy.corners[:, -1] = up.Zxy.corners[:, 0]
 
-    def geometry1(self):
+    def calcDistances(self):
         """
-        Calculate geometrical quantities for this region
+        Calculate the distances for all PsiContours in the region
         """
         if self.user_options.orthogonal:
             # Distances already calculated in non-orthogonal case
             self.contours = self.parallel_map(
                 _calc_contour_distance, enumerate(self.contours)
             )
 
+    def geometry1(self):
+        """
+        Calculate geometrical quantities for this region
+        """
         self.psixy = self.meshParent.equilibrium.psi(self.Rxy, self.Zxy)
 
         self.dx = MultiLocationArray(self.nx, self.ny)
@@ -772,6 +776,8 @@ def geometry1(self):
         self.Bzxy = self.meshParent.equilibrium.Bp_Z(self.Rxy, self.Zxy)
         self.Bpxy = numpy.sqrt(self.Brxy ** 2 + self.Bzxy ** 2)
 
+        self.calcPoloidalDistance()
+
         if hasattr(
             self.meshParent.equilibrium.regions[self.equilibriumRegion.name], "pressure"
         ):
@@ -1542,6 +1548,105 @@ def integrand_func(R, Z):
                 region.zShift.corners[:, -1] - self.zShift.corners[:, 0]
             ).reshape((-1, 1))
 
+    def calcPoloidalDistance(self):
+        """
+        Calculate poloidal distance by following contours between regions.
+        """
+        # Cannot just test 'connections['lower'] is not None' because periodic regions
+        # always have a lower connection - requires us to give a yGroupIndex to each
+        # region when creating the groups.
+        if self.yGroupIndex != 0:
+            return None
+
+        region = self
+        region.poloidal_distance = MultiLocationArray(region.nx, region.ny)
+        region.poloidal_distance.centre = 0.0
+        region.poloidal_distance.ylow = 0.0
+        region.poloidal_distance.xlow = 0.0
+        region.poloidal_distance.corners = 0.0
+
+        # Initialise so that distance counts from the lower wall (for SOL/PFR) or wall
+        # (for core)
+        for i in range(self.nx):
+            c = region.contours[2 * i + 1]
+            # Cell-centre points
+            region.poloidal_distance.centre[i, :] -= c.get_distance(
+                psi=self.meshParent.equilibrium.psi
+            )[c.startInd]
+            # ylow points
+            region.poloidal_distance.ylow[i, :] -= c.get_distance(
+                psi=self.meshParent.equilibrium.psi
+            )[c.startInd]
+        for i in range(self.nx + 1):
+            c = region.contours[2 * i]
+            # Cell-centre points
+            region.poloidal_distance.xlow[i, :] -= c.get_distance(
+                psi=self.meshParent.equilibrium.psi
+            )[c.startInd]
+            # ylow points
+            region.poloidal_distance.corners[i, :] -= c.get_distance(
+                psi=self.meshParent.equilibrium.psi
+            )[c.startInd]
+
+        # Get distances from contours
+        while True:
+            for i in range(self.nx):
+                c = region.contours[2 * i + 1]
+                # Cell-centre points
+                region.poloidal_distance.centre[i, :] += c.get_distance(
+                    psi=self.meshParent.equilibrium.psi
+                )[1::2]
+                # ylow points
+                region.poloidal_distance.ylow[i, :] += c.get_distance(
+                    psi=self.meshParent.equilibrium.psi
+                )[::2]
+            for i in range(self.nx + 1):
+                c = region.contours[2 * i]
+                # Cell-centre points
+                region.poloidal_distance.xlow[i, :] += c.get_distance(
+                    psi=self.meshParent.equilibrium.psi
+                )[1::2]
+                # ylow points
+                region.poloidal_distance.corners[i, :] += c.get_distance(
+                    psi=self.meshParent.equilibrium.psi
+                )[::2]
+
+            next_region = region.getNeighbour("upper")
+            if (next_region is None) or (next_region is self):
+                # Note: If periodic, next_region is self (back to start)
+                break
+            else:
+                # Initialise with values at the lower y-boundary of next_region
+                next_region.poloidal_distance = MultiLocationArray(
+                    next_region.nx, next_region.ny
+                )
+                next_region.poloidal_distance.centre[
+                    :, :
+                ] = region.poloidal_distance.ylow[:, -1, numpy.newaxis]
+                next_region.poloidal_distance.ylow[
+                    :, :
+                ] = region.poloidal_distance.ylow[:, -1, numpy.newaxis]
+                next_region.poloidal_distance.xlow[
+                    :, :
+                ] = region.poloidal_distance.corners[:, -1, numpy.newaxis]
+                next_region.poloidal_distance.corners[
+                    :, :
+                ] = region.poloidal_distance.corners[:, -1, numpy.newaxis]
+                region = next_region
+
+        # Save total poloidal distance in core
+        self.total_poloidal_distance = MultiLocationArray(region.nx, 1)
+        if self.connections["lower"] is not None:
+            # This is a periodic region (we already checked that the self.yGroupIndex is
+            # 0).
+            # 'region' is the last region in the y-group
+            self.total_poloidal_distance.centre[:, 0] = region.poloidal_distance.ylow[
+                :, -1
+            ]
+            self.total_poloidal_distance.xlow[:, 0] = region.poloidal_distance.corners[
+                :, -1
+            ]
+
     def getNeighbour(self, face):
         if self.connections[face] is None:
             return None
@@ -2602,6 +2707,9 @@ def geometry(self):
                 self.calculateRZ()
                 break
         print("Calculate geometry", flush=True)
+        for region in self.regions.values():
+            print("Distances", region.name, flush=True)
+            region.calcDistances()
         for region in self.regions.values():
             print("1", region.name, flush=True)
             region.geometry1()
@@ -3262,6 +3370,8 @@ def addFromRegionsXArray(name):
         addFromRegions("psixy")
         addFromRegions("dx")
         addFromRegions("dy")
+        addFromRegions("poloidal_distance")
+        addFromRegionsXArray("total_poloidal_distance")
         addFromRegions("Brxy")
         addFromRegions("Bzxy")
         addFromRegions("Bpxy")