Merge pull request #88 from ddudt/dev

Fix plotting vartheta contours

desc/configuration.py

@@ -1224,9 +1224,85 @@ def compute_dW(self, grid=None):
         dW = obj.hess_x(x, self.Rb_lmn, self.Zb_lmn, self.p_l, self.i_l, self.Psi)
         return dW
 
+    def compute_theta_coords(self, flux_coords, tol=1e-6, maxiter=20):
+        """Find the theta coordinates (rho, theta, phi) that correspond to a set of
+        straight field-line coordinates (rho, theta*, zeta).
+
+        Parameters
+        ----------
+        flux_coords : ndarray, shape(k,3)
+            2d array of flux coordinates [rho,theta*,zeta]. Each row is a different
+            coordinate.
+        tol : float
+            Stopping tolerance.
+        maxiter : int > 0
+            maximum number of Newton iterations
+
+        Returns
+        -------
+        coords : ndarray, shape(k,3)
+            coordinates [rho,theta,zeta]. If Newton method doesn't converge for
+            a given coordinate nan will be returned for those values
+
+        """
+        rho = flux_coords[:, 0]
+        theta_star = flux_coords[:, 1]
+        zeta = flux_coords[:, 2]
+        if maxiter <= 0:
+            raise ValueError(f"maxiter must be a positive integer, got{maxiter}")
+        if jnp.any(rho) <= 0:
+            raise ValueError("rho values must be positive")
+
+        # Note: theta* (also known as vartheta) is the poloidal straight field-line
+        # angle in PEST-like flux coordinates
+
+        theta_k = theta_star
+        grid = Grid(jnp.vstack([rho, theta_k, zeta]).T, sort=False)
+
+        transform = Transform(
+            grid,
+            self.L_basis,
+            derivs=np.array([[0, 0, 0], [0, 1, 0]]),
+            method="direct1",
+        )
+
+        # theta* = theta + lambda
+        theta_star_k = theta_k + transform.transform(self.L_lmn)
+        err = theta_star - theta_star_k
+
+        # Newton method for root finding
+        k = 0
+        while jnp.any(abs(err) > tol) and k < maxiter:
+            lmbda = transform.transform(self.L_lmn, 0, 0, 0)
+            lmbda_t = transform.transform(self.L_lmn, 0, 1, 0)
+            f = theta_star - theta_k - lmbda
+            df = -1 - lmbda_t
+
+            theta_k = theta_k - f / df
+
+            grid = Grid(jnp.vstack([rho, theta_k, zeta]).T, sort=False)
+            transform = Transform(
+                grid,
+                self.L_basis,
+                derivs=np.array([[0, 0, 0], [0, 1, 0]]),
+                method="direct1",
+            )
+
+            theta_star_k = theta_k + transform.transform(self.L_lmn)
+            err = theta_star - theta_star_k
+            k += 1
+
+        if k >= maxiter:  # didn't converge for all, mark those as nan
+            i = np.where(abs(err) > tol)
+            rho = put(rho, i, np.nan)
+            theta_k = put(theta_k, i, np.nan)
+            zeta = put(zeta, i, np.nan)
+
+        return jnp.vstack([rho, theta_k, zeta]).T
+
     def compute_flux_coords(self, real_coords, tol=1e-6, maxiter=20, rhomin=1e-6):
-        """Finds the flux coordinates (rho, theta, zeta) that correspond to a set of
-        real space coordinates (R,phi,Z)
+        """Find the flux coordinates (rho, theta, zeta) that correspond to a set of
+        real space coordinates (R, phi, Z).
 
         Parameters
         ----------
@@ -1245,8 +1321,8 @@ def compute_flux_coords(self, real_coords, tol=1e-6, maxiter=20, rhomin=1e-6):
             flux coordinates [rho,theta,zeta]. If Newton method doesn't converge for
             a given coordinate (often because it is outside the plasma boundary),
             nan will be returned for those values
-        """
 
+        """
         R = real_coords[:, 0]
         phi = real_coords[:, 1]
         Z = real_coords[:, 2]

desc/plotting.py

@@ -691,31 +691,27 @@ def plot_surfaces(eq, r_grid=None, t_grid=None, ax=None, **kwargs):
         if zeta.size != t_zeta.size or not np.allclose(zeta, t_zeta):
             raise ValueError(
                 colored(
-                    "r_grid and t_grid should have the same zeta planes, got r_grid={}, t_grid{}".format(
-                        zeta, t_zeta
-                    ),
+                    "r_grid and t_grid should have the same zeta planes, "
+                    + "got r_grid={}, t_grid{}".format(zeta, t_zeta),
                     "red",
                 )
             )
+
+    # Note: theta* (also known as vartheta) is the poloidal straight field-line anlge in
+    # PEST-like flux coordinates
+
+    v_grid = Grid(eq.compute_theta_coords(t_grid.nodes))
     rows = np.floor(np.sqrt(nzeta)).astype(int)
     cols = np.ceil(nzeta / rows).astype(int)
 
     r_coords = eq.compute_toroidal_coords(r_grid)
-    t_coords = eq.compute_toroidal_coords(t_grid)
-
-    # theta coordinates cooresponding to linearly spaced vartheta angles
-    v_nodes = t_grid.nodes
-    v_nodes[:, 1] = t_grid.nodes[:, 1] - t_coords["lambda"]
-    v_grid = Grid(v_nodes)
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        v_coords = eq.compute_toroidal_coords(v_grid)
+    v_coords = eq.compute_toroidal_coords(v_grid)
 
     # rho contours
     Rr = r_coords["R"].reshape((r_grid.M, r_grid.L, r_grid.N), order="F")
     Zr = r_coords["Z"].reshape((r_grid.M, r_grid.L, r_grid.N), order="F")
 
-    # theta contours
+    # vartheta contours
     Rv = v_coords["R"].reshape((t_grid.M, t_grid.L, t_grid.N), order="F")
     Zv = v_coords["Z"].reshape((t_grid.M, t_grid.L, t_grid.N), order="F")
 
@@ -732,25 +728,16 @@ def plot_surfaces(eq, r_grid=None, t_grid=None, ax=None, **kwargs):
 
     for i in range(nzeta):
         ax[i].plot(
-            Rv[:, :, i].T,
-            Zv[:, :, i].T,
-            color=colorblind_colors[2],
-            linestyle=":",
+            Rv[:, :, i].T, Zv[:, :, i].T, color=colorblind_colors[2], linestyle=":",
         )
         ax[i].plot(
-            Rr[:, :, i],
-            Zr[:, :, i],
-            color=colorblind_colors[0],
+            Rr[:, :, i], Zr[:, :, i], color=colorblind_colors[0],
         )
         ax[i].plot(
-            Rr[:, -1, i],
-            Zr[:, -1, i],
-            color=colorblind_colors[1],
+            Rr[:, -1, i], Zr[:, -1, i], color=colorblind_colors[1],
         )
         ax[i].scatter(
-            Rr[0, 0, i],
-            Zr[0, 0, i],
-            color=colorblind_colors[3],
+            Rr[0, 0, i], Zr[0, 0, i], color=colorblind_colors[3],
         )
 
         ax[i].set_xlabel(_axis_labels_RPZ[0])