Allow distributing the points on the inner surface more smoothly

.github/workflows/test-mms.yml

@@ -7,7 +7,7 @@ jobs:
     timeout-minutes: 60
     runs-on: ubuntu-latest
     container:
-      image: oi4ai/bout3d:latest-${{ matrix.config.mode }}
+      image: ghcr.io/boutproject/bout-dev:mpich-3d-${{ matrix.config.mode }}-fix3121
     strategy:
       fail-fast: false
       matrix:
@@ -36,13 +36,11 @@ jobs:
           cd /home/boutuser/zoidberg
           python3 -m ensurepip
           python3 -m pip install .
-          python3 -m pip install xarray
+          python3 -m pip install xarray typing_extensions
 
       - name: Run
         run: |
           export HOME=/home/boutuser
-          export PYTHONPATH=/home/boutuser/BOUT-dev/build/tools/pylib/
-          export LD_LIBRARY_PATH=/usr/lib64/mpich/lib
           cd /home/boutuser/zoidberg/test/mms/
           # Run test
-          make
+          PYTHONPATH=/home/boutuser/zoidberg:$PYTHONPATH make

pyproject.toml

@@ -52,3 +52,8 @@ zoidberg-screwpinch = "zoidberg.examples.screwpinch:main"
 [tool.setuptools_scm]
 write_to = "zoidberg/_version.py"
 git_describe_command = "git describe --dirty --tags --first-parent"
+
+
+[tool.setuptools.packages.find]
+include = ["zoidberg*"]  # alternatively: `exclude = ["additional*"]`
+namespaces = false

test/mms/delp_grid.py

@@ -33,7 +33,7 @@ def gen_grid(nx, ny, nz, R0, r0, r1, mode=0):
     pol_grid = zb.poloidal_grid.StructuredPoloidalGrid(R, Z)
 
     field = FixedCurvedSlab(Bz=0, Bzprime=0, Rmaj=R0)
-    grid = zb.grid.Grid(pol_grid, phi, 5, yperiodic=True)
+    grid = zb.grid.Grid(pol_grid, phi, 2 * np.pi / 5, yperiodic=True)
 
     fn = gen_name(*args)
 

test/mms/laplace_grid.py

@@ -46,7 +46,7 @@ def gen_grid(nx, ny, nz, R0, r0, r1, mode=0):
     pol_grid = zb.poloidal_grid.StructuredPoloidalGrid(R, Z)
 
     field = FixedCurvedSlab(Bz=0, Bzprime=0, Rmaj=R0)
-    grid = zb.grid.Grid(pol_grid, phi, 5, yperiodic=True)
+    grid = zb.grid.Grid(pol_grid, phi, 2 * np.pi / 5, yperiodic=True)
 
     fn = gen_name(*args)
 

test/mms/parallel_grid.py

@@ -31,7 +31,7 @@ def gen_grid(nx, ny, nz, R0, r0, r1, mode=0):
     pol_grid = zb.poloidal_grid.StructuredPoloidalGrid(R, Z)
 
     field = zb.field.CurvedSlab(Bz=0, Bzprime=0, Rmaj=R0)
-    grid = zb.grid.Grid(pol_grid, phi, 5, yperiodic=True)
+    grid = zb.grid.Grid(pol_grid, phi, 2 * np.pi / 5, yperiodic=True)
 
     fn = gen_name(*args)
 

zoidberg/field.py

@@ -1,8 +1,8 @@
 from math import gamma
 
 import numpy as np
-from sympy import (And, Piecewise, Sum, Symbol, atan2, cos, diff, factorial,
-                   lambdify, log, pi, sin, sqrt)
+from sympy import (Piecewise, Symbol, atan2, cos, diff, factorial, lambdify,
+                   log, pi, sin, sqrt)
 
 from . import boundary
 
@@ -711,12 +711,14 @@ def CN(self, m, k):
         )
         beta_st = lambda n, b: beta(n, b) * (2 * n - b)
 
-        delta = lambda n, b: (
-            alpha(n, b) * np.sum([1.0 / i + 1.0 / (b + i) for i in range(1, n + 1)]) / 2
-            if (n > 0)
-            else 0.0
-        )
-        delta_st = lambda n, b: delta(n, b) * (2 * n + b)
+        def delta(n, b):
+            if n <= 0:
+                return 0.0
+            return (
+                alpha(n, b)
+                * np.sum([1.0 / i + 1.0 / (b + i) for i in range(1, n + 1)])
+                / 2
+            )
 
         CN = log(1)
         for j in range(k + 1):
@@ -730,62 +732,62 @@ def CN(self, m, k):
 
         return CN
 
-    def D(self, m, n):
+    def D(self, n, v):
         """
         Parameters
         ----------
-        m: torodial mode number
-        n: summation index in  V
+        n: torodial mode number
+        v: summation index in  V
 
         Returns:
         --------
         Sympy function D_mn (R, Z) (Dirichlet boundary conditions)
         """
 
         D = log(1)
-        k_arr = np.arange(0, int(n / 2) + 1, 1)
+        k_arr = np.arange(0, int(v / 2) + 1, 1)
 
         for k in k_arr:
-            D += (self.Z ** (n - 2 * k)) / factorial(n - 2 * k) * self.CD(m, k)
+            D += (self.Z ** (v - 2 * k)) / factorial(v - 2 * k) * self.CD(n, k)
 
         return D
 
-    def N(self, m, n):
+    def N(self, n, v):
         """
         Parameters
         ----------
-        m: torodial mode number
-        n: summation index in V
+        n: torodial mode number
+        v: summation index in V
 
         Returns:
         --------
         Sympy function N_mn (R, Z) (Neumann boundary conditions)
         """
 
         N = log(1)
-        k_arr = np.arange(0, int(n / 2) + 1, 1)
+        k_arr = np.arange(0, int(v / 2) + 1, 1)
 
         for k in k_arr:
-            N += (self.Z ** (n - 2 * k)) / factorial(n - 2 * k) * self.CN(m, k)
+            N += (self.Z ** (v - 2 * k)) / factorial(v - 2 * k) * self.CN(n, k)
 
         return N
 
-    def V(self, m, l, a, b, c, d):
+    def V(self, n, m, a, b, c, d):
         """
         Parameters
         ----------
-        m: torodial mode number
-        l: polodial mode number
+        n: torodial mode number
+        m: polodial mode number
         a,b,c,d: Coefficients for m,l-th Dommaschk potential (elements of matrix A)
 
         Returns:
         --------
         Sympy function V_ml
         """
 
-        V = (a * cos(m * self.phi) + b * sin(m * self.phi)) * self.D(m, l) + (
-            c * cos(m * self.phi) + d * sin(m * self.phi)
-        ) * self.N(m, l - 1)
+        V = (a * cos(n * self.phi) + b * sin(n * self.phi)) * self.D(n, m) + (
+            c * cos(n * self.phi) + d * sin(n * self.phi)
+        ) * self.N(n, m - 1)
 
         return V
 
@@ -808,12 +810,12 @@ def U(self, A):
 
         return U
 
-    def V_hat(self, m, l, a, b, c, d):
+    def V_hat(self, n, m, a, b, c, d):
         """
         Parameters
         ----------
-        m: torodial mode number
-        l: polodial mode number
+        n: torodial mode number
+        m: polodial mode number
         a,b,c,d: Coefficients for m,l-th Dommaschk potential (elements of matrix A)
 
         Returns:
@@ -822,11 +824,11 @@ def V_hat(self, m, l, a, b, c, d):
         """
 
         V = (
-            a * cos(m * self.phi - np.pi / 2) + b * sin(m * self.phi - np.pi / 2)
-        ) * self.D(m, l) + (
-            c * cos(m * self.phi - np.pi / 2) + d * sin(m * self.phi - np.pi / 2)
+            a * cos(n * self.phi - np.pi / 2) + b * sin(n * self.phi - np.pi / 2)
+        ) * self.D(n, m) + (
+            c * cos(n * self.phi - np.pi / 2) + d * sin(n * self.phi - np.pi / 2)
         ) * self.N(
-            m, l - 1
+            n, m - 1
         )
 
         return V
@@ -1148,9 +1150,6 @@ def Byfunc(self, x, z, phi):
         """
         return self.field.Byfunc(x, z, phi)
 
-    def Bxfunc(self, x, z, phi):
-        pass
-
     def Rfunc(self, x, z, phi):
         return self.field.Rfunc(x, z, phi)
 
@@ -1631,9 +1630,6 @@ def magnetic_axis(self, phi_axis=0, configuration=0):
             res.axis.vertices.x2
         )  # (m) -- REAL SPACE from an arbitrary start point
         magnetic_axis_z = np.asarray(res.axis.vertices.x3)  # (m)
-        magnetic_axis_rmaj = np.sqrt(
-            magnetic_axis_x**2 + magnetic_axis_y**2 + magnetic_axis_z**2
-        )
 
         magnetic_axis_r = np.sqrt(
             np.asarray(magnetic_axis_x) ** 2 + np.asarray(magnetic_axis_y**2)
@@ -1680,7 +1676,7 @@ def _calc_chunk(cls, x, configuration):
         global tracer
         tracer = tracer or Client(
             "http://esb.ipp-hgw.mpg.de:8280/services/FieldLineProxy?wsdl",
-            osa_timeout=0.1,
+            osa_timeout=1,
         )
 
         assert x.shape[0] == 3
@@ -1700,7 +1696,7 @@ def _calc_chunk(cls, x, configuration):
                 res = tracer.service.magneticField(
                     pos, config, osa_timeout=(10 + len(pos.x1) / 1000)
                 )
-            except:
+            except:  # noqa
                 # Catch any error. Different errors might be
                 # reported, but we want to retry anyway.
                 # Do not except Exception, as that would also

zoidberg/fieldtracer.py

@@ -398,6 +398,55 @@ def trace_poincare(
     return result, y_slices
 
 
+class CachingFieldTracer:
+    def __init__(
+        self, tracer, cachename="_zoidberg_cache_tracing_%10.10s", debug=False
+    ):
+        self.tracer = tracer
+        self.cachename = cachename
+        self.debug = debug
+
+    def _myhash(self, *args):
+        """
+
+        Compute a digest over some numerical data
+
+        """
+        import hashlib
+
+        m = hashlib.md5()
+        for arg in args:
+            s = "-".join([str(x) for x in np.array([arg]).ravel()]).encode()
+            if self.debug:
+                print(s)
+            m.update(s)
+        return m.hexdigest()
+
+    def _guess_shape(self, x_values, y_values):
+        return len(y_values), *x_values.shape, 2
+
+    def follow_field_lines(self, x_values, z_values, y_values, **kwargs):
+        fn = self.cachename % self._myhash(x_values, z_values, y_values)
+        try:
+            dat = np.loadtxt(fn)
+            dat.shape = self._guess_shape(x_values, y_values)
+            return dat
+        except FileNotFoundError:
+            if self.debug:
+                print("no cache", fn)
+        ret = self.tracer.follow_field_lines(x_values, z_values, y_values, **kwargs)
+        shape = ret.shape
+        if self.debug:
+            print(f"saving {ret.shape}, {x_values.shape}")
+        ret.shape = -1, shape[-1]
+        np.savetxt(fn, ret)
+        assert shape == self._guess_shape(
+            x_values, y_values
+        ), f"Estimated {self._guess_shape(x_values, y_values)} but got {shape}"
+        ret.shape = shape
+        return ret
+
+
 def _get_value(*args):
     for k in args:
         if k is not None:
@@ -419,7 +468,7 @@ def __init__(
         self,
         config=None,
         configId=None,
-        timeout=0.1,
+        timeout=1,
         chunk=10000,
         stepsize=None,
         retry=None,