From 67274a0a422a738d03de78f7825586a27fec7818 Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Thu, 23 Dec 2021 19:42:29 +0000
Subject: [PATCH 1/3] Save poloidal distances along psi-contours to grid files

Also save the total poloidal distance in the core region. Since guard
cells at the branch cut are not saved to the grid file, the total
distance around the core region(s) would not be saved if we did not
create a special variable for it.
---
 doc/whats-new.md       |   2 +
 hypnotoad/core/mesh.py | 114 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 114 insertions(+), 2 deletions(-)

diff --git a/doc/whats-new.md b/doc/whats-new.md
index 903a761a..81e9487e 100644
--- a/doc/whats-new.md
+++ b/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)
diff --git a/hypnotoad/core/mesh.py b/hypnotoad/core/mesh.py
index 98cce634..fdb1c645 100644
--- a/hypnotoad/core/mesh.py
+++ b/hypnotoad/core/mesh.py
@@ -738,9 +738,9 @@ 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
@@ -748,6 +748,10 @@ def geometry1(self):
                 _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")

From dc623fd069b06539ff81d876991f9ccf6cabf6ba Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 24 Dec 2021 22:39:45 +0000
Subject: [PATCH 2/3] Update expected results with new variables

---
 .../expected_nonorthogonal.grd.nc                             | 4 ++--
 .../expected_orthogonal.grd.nc                                | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/integrated_tests/connected_doublenull_nonorthogonal/expected_nonorthogonal.grd.nc b/integrated_tests/connected_doublenull_nonorthogonal/expected_nonorthogonal.grd.nc
index afbeb105..21766ef1 100644
--- a/integrated_tests/connected_doublenull_nonorthogonal/expected_nonorthogonal.grd.nc
+++ b/integrated_tests/connected_doublenull_nonorthogonal/expected_nonorthogonal.grd.nc
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:12115fb65fbd8614a16a7a08b8d2970b56a8cb6156c95055d3b427c0edfe68c8
-size 419977
+oid sha256:b1969367565c41955666f94b49d31a111f8794fececbd37ceae47ed901372ac8
+size 411761
diff --git a/integrated_tests/connected_doublenull_orthogonal/expected_orthogonal.grd.nc b/integrated_tests/connected_doublenull_orthogonal/expected_orthogonal.grd.nc
index c1294aa7..14ac844e 100644
--- a/integrated_tests/connected_doublenull_orthogonal/expected_orthogonal.grd.nc
+++ b/integrated_tests/connected_doublenull_orthogonal/expected_orthogonal.grd.nc
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0aa8bdc34360c2f4c24d032fcaf5a87f93ba89bafb48a28f19e91baeedb2d033
-size 431439
+oid sha256:380d8c2984c6312b504dd8d4a1d3b03c723b46f404f955b8282ed95bbf4637c1
+size 420837

From 6ab4bc4a0b270cb2430eda648653dbfa6bc23ffa Mon Sep 17 00:00:00 2001
From: John Omotani <john.omotani@ukaea.uk>
Date: Fri, 24 Dec 2021 23:11:12 +0000
Subject: [PATCH 3/3] Add Python-3.9 runs to CI tests

---
 .github/workflows/pythonpackage.yml | 10 +++++-----
 .github/workflows/pythonpublish.yml |  6 +++---
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/pythonpackage.yml b/.github/workflows/pythonpackage.yml
index 49205fe5..e4fcf5aa 100644
--- a/.github/workflows/pythonpackage.yml
+++ b/.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
diff --git a/.github/workflows/pythonpublish.yml b/.github/workflows/pythonpublish.yml
index 3a107a68..5bc088c2 100644
--- a/.github/workflows/pythonpublish.yml
+++ b/.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