Cleanup code

Author: dschwoerer

zoidberg/field.py

@@ -1888,6 +1888,8 @@ def Rfunc(self, x, z, phi):
 
 
 class EMC3(MagneticField):
+    """Field based on a EMC3 grid file"""
+
     def __init__(self, ds):
         self.ds = ds
         assert hasattr(ds, "emc3"), "Expected an xemc3 dataset. Is xemc3 imported?"
@@ -1899,7 +1901,6 @@ def Bxfunc(self, x, z, phi):
         raise NotImplementedError("Use maybe EMC3 tracer?")
 
     def Byfunc(self, x, z, phi):
-        # raise NotImplementedError("Use maybe EMC3 tracer?")
         return self.Bmag(x, z, phi)
 
     def Bzfunc(self, x, z, phi):
@@ -1914,45 +1915,20 @@ def Bmag(self, x, z, phi):
         ].values
         nans = np.isnan(vals)
         if np.any(nans):
-            if 1:
-                from scipy.interpolate import CubicSpline as CS
-
-                for i in range(x.shape[1]):
-                    ni = nans[:, i]
-                    if not np.any(ni):
-                        continue
-                    xi = x[:, i]
-                    zi = z[:, i]
-                    vi = vals[:, i]
-                    si = np.zeros_like(xi)
-                    si[1:] = np.cumsum(
-                        np.sqrt((xi[1:] - xi[:-1]) ** 2 + (zi[1:] - zi[:-1]) ** 2)
-                    )
-                    interp = CS(si[~ni], vi[~ni])
-                    vals[ni, i] = interp(si[ni])
-            else:
-                from scipy.interpolate import LinearNDInterpolator as LinInter
-
-                if phi.shape == ():
-                    pos = np.array((x, z))
-                else:
-                    pos = np.array((x, z, phi))
-                inter = LinInter(pos[:, ~nans].T, vals[~nans])
-                print(vals[~nans])
-
-                print(inter(pos[:, nans].T))
-                vals[nans] = inter(pos[:, nans].T)
-
-                import matplotlib.pyplot as plt
-
-                plt.figure()
-                # plt.pcolormesh(X, Y, Z, shading='auto')
-                plt.plot(*pos[:, ~nans], "ok", label="input point")
-                plt.plot(*pos[:, nans], "or", label="evaled")
-                print(*pos[:, nans].T)
-                plt.legend()
-                # plt.colorbar()
-                plt.axis("equal")
-                plt.show()
+            from scipy.interpolate import CubicSpline as CS
+
+            for i in range(x.shape[1]):
+                ni = nans[:, i]
+                if not np.any(ni):
+                    continue
+                xi = x[:, i]
+                zi = z[:, i]
+                vi = vals[:, i]
+                si = np.zeros_like(xi)
+                si[1:] = np.cumsum(
+                    np.sqrt((xi[1:] - xi[:-1]) ** 2 + (zi[1:] - zi[:-1]) ** 2)
+                )
+                interp = CS(si[~ni], vi[~ni])
+                vals[ni, i] = interp(si[ni])
 
         return vals

zoidberg/fieldtracer.py

@@ -730,7 +730,7 @@ def __init__(self, x, y):
 
 
 class EMC3FieldTracer(FieldTracer):
-    """A class for following magnetic field lines
+    """A class for following magnetic field lines provided by an EMC3 grid.
 
     Parameters
     ----------
@@ -760,9 +760,6 @@ def __init__(self, field):
         self.firstlast = (first, last)
 
     def follow_field_lines(self, x_values, z_values, y_values, rtol=None):
-        # assert np.all(y_values >= self.ds.phi.min().values) and np.all(
-        #     y_values <= self.ds.phi.max().values
-        # ), f"The condition is not fulfilled: {self.ds.phi.min()} <= {y_values} <= {self.ds.phi.max()}"
         meshes = [self.makeMeshes(phi) for phi in y_values]
         assert x_values.shape == z_values.shape
         out = np.empty((len(y_values), *x_values.shape, 2))
@@ -796,10 +793,6 @@ def _rz_to_ab(self, rz, grid, ij):
         nz -= 1
         i, j = ij // nz, ij % nz
         ABCD = grid[i : i + 2, j : j + 2]
-        # if j + 1 == nz:
-        #    ABCD[:, 1] = grid[i : i + 2, 0]
-        # print(ABCD)
-        # print(i, j, nz, ij)
         if ABCD.shape != (2, 2, 2):
             print(grid.shape)
             print(rz)
@@ -836,8 +829,6 @@ def J(albe):
         for i in range(100):
             assert np.all(np.isfinite(albe))
             albe = albe - np.linalg.inv(J(albe).T) @ fun(albe)
-            # if i > 20:
-            #    print(f"Failing to converge! {albe} {fun(albe)}")
             res = np.sum(fun(albe) ** 2)
             if res < tol:
                 return albe
@@ -870,7 +861,6 @@ def rz_to_ab(self, rz, meshes, ij, zid):
         for mind in 1e-3, 1e-2, 1:
             mind = mind**2
             for mesh in itertools.chain(meshes[zid:], meshes[:zid]):
-                # print(f"checking mesh {mesh.zid}")
                 l2 = (rz[0] - mesh.r) ** 2 + (rz[1] - mesh.z) ** 2
                 ij = np.argmin(l2)
                 if l2.flat[ij] < mind: