Try to make splines more periodic

zoidberg/field.py

@@ -183,31 +183,28 @@ def field_direction(self, pos, ycoord, flatten=False):
         else:
             x, z = pos
 
-        By = self.Byfunc(x, z, ycoord)
-        Rmaj = self.Rfunc(x, z, ycoord)  # Major radius. None if Cartesian
-
-        if Rmaj is not None:
-            # In cylindrical coordinates
+        if hasattr(self, "Bxyzfunc"):
+            Bx, By, Bz = self.Bxyzfunc(x, z, ycoord)
+        else:
+            Bx = self.Bxfunc(x, z, ycoord)
+            By = self.Byfunc(x, z, ycoord)
+            Bz = self.Bzfunc(x, z, ycoord)
 
-            if np.amin(np.abs(By)) < 1e-8:
-                # Very small By
-                print(x, z, ycoord, By)
-                raise ValueError(
-                    "Small By ({}) at (x={}, y={}, z={})".format(By, x, ycoord, z)
-                )
+        Rmaj = self.Rfunc(x, z, ycoord)  # Major radius. None if Cartesian
+        if Rmaj is None:
+            Rmaj = 1
 
-            R_By = Rmaj / By
-            # Rate of change of x location [m] with y angle [radians]
-            dxdphi = R_By * self.Bxfunc(x, z, ycoord)
-            # Rate of change of z location [m] with y angle [radians]
-            dzdphi = R_By * self.Bzfunc(x, z, ycoord)
-        else:
-            # In Cartesian coordinates
+        if np.amin(np.abs(By)) < 1e-8:
+            # Very small By
+            raise ValueError(
+                "Small By ({}) at (x={}, y={}, z={})".format(By, x, ycoord, z)
+            )
 
-            # Rate of change of x location [m] with y angle [radians]
-            dxdphi = self.Bxfunc(x, z, ycoord) / By
-            # Rate of change of z location [m] with y angle [radians]
-            dzdphi = self.Bzfunc(x, z, ycoord) / By
+        R_By = Rmaj / By
+        # Rate of change of x location [m] with y angle [radians]
+        dxdphi = R_By * Bx
+        # Rate of change of z location [m] with y angle [radians]
+        dzdphi = R_By * Bz
 
         if flatten:
             result = np.column_stack((dxdphi, dzdphi)).flatten()
@@ -1843,6 +1840,12 @@ def Byfunc(self, x, z, phi):
         B = self.getB(x, z, phi)
         return -B[0] * np.sin(phi) + B[1] * np.cos(phi)
 
+    def Bxyzfunc(self, x, z, phi):
+        B = self.getB(x, z, phi)
+        Bx = B[0] * np.cos(phi) + B[1] * np.sin(phi)
+        By = -B[0] * np.sin(phi) + B[1] * np.cos(phi)
+        return Bx, By, B[2]
+
     def Bzfunc(self, *pos):
         return self.getB(*pos)[2]
 

zoidberg/poloidal_grid.py

@@ -293,10 +293,10 @@ def __init__(self, R, Z):
         self.nz = nz
 
         xinds = np.arange(nx)
-        zinds = np.arange(nz + 1)
-        # Repeat the final point in y since periodic in y
-        R_ext = np.concatenate((R, np.reshape(R[:, 0], (nx, 1))), axis=1)
-        Z_ext = np.concatenate((Z, np.reshape(Z[:, 0], (nx, 1))), axis=1)
+        zinds = np.arange(nz * 3)
+        # Repeat the data in z, to approximate periodicity
+        R_ext = np.concatenate((R, R, R), axis=1)
+        Z_ext = np.concatenate((Z, Z, Z), axis=1)
 
         self._spl_r = RectBivariateSpline(xinds, zinds, R_ext)
         self._spl_z = RectBivariateSpline(xinds, zinds, Z_ext)
@@ -330,8 +330,8 @@ def getCoordinate(self, xind, zind, dx=0, dz=0):
         # Periodic in y
         zind = np.remainder(zind, nz)
 
-        R = self._spl_r(xind, zind, dx=dx, dy=dz, grid=False)
-        Z = self._spl_z(xind, zind, dx=dx, dy=dz, grid=False)
+        R = self._spl_r(xind, zind + self.nz, dx=dx, dy=dz, grid=False)
+        Z = self._spl_z(xind, zind + self.nz, dx=dx, dy=dz, grid=False)
 
         return R, Z