fix constant_offset_surface function

PlasmaControl · Oct 19, 2024 · e10b2e4 · e10b2e4
1 parent 92bf3d1
commit e10b2e4
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/desc/equilibrium/coords.py b/desc/equilibrium/coords.py
@@ -367,6 +367,7 @@ def fixup(x, *args):
                 fixup=fixup,
                 tol=tol,
                 maxiter=maxiter,
+                full_output=full_output,
                 **kwargs,
             )
         )
@@ -479,6 +480,7 @@ def fixup(x, *args):
                 fixup=fixup,
                 tol=tol,
                 maxiter=maxiter,
+                full_output=full_output,
                 **kwargs,
             )
         )

diff --git a/desc/geometry/surface.py b/desc/geometry/surface.py
@@ -741,10 +741,19 @@ def fun_jax(zeta_hat, theta, zeta):
             n, r, r_offset = n_and_r_jax(nodes)
             return jnp.arctan(r_offset[0, 1] / r_offset[0, 0]) - zeta
 
-        vecroot = jit(vmap(lambda x0, *p: root_scalar(fun_jax, x0, jac=None, args=p)))
-        zetas, (res, niter) = vecroot(
-            grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2]
+        vecroot = jit(
+            vmap(
+                lambda x0, *p: root_scalar(
+                    fun_jax, x0, jac=None, args=p, full_output=full_output
+                )
+            )
         )
+        if full_output:
+            zetas, (res, niter) = vecroot(
+                grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2]
+            )
+        else:
+            zetas = vecroot(grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2])
 
         zetas = np.asarray(zetas)
         nodes = np.vstack((np.ones_like(grid.nodes[:, 1]), grid.nodes[:, 1], zetas)).T