Merge branch 'master' into rc/coils_init

.github/workflows/benchmark.yml

@@ -74,6 +74,7 @@ jobs:
           source .venv-${{ matrix.python-version }}/bin/activate
           python -m pip install --upgrade pip
           pip install -r devtools/dev-requirements.txt
+          pip install matplotlib==3.9.2
 
       - name: Benchmark with pytest-benchmark (PR)
         if: env.has_changes == 'true'

.github/workflows/cache_dependencies.yml

@@ -41,7 +41,7 @@ jobs:
           source .venv-${{ matrix.python-version }}/bin/activate
           python -m pip install --upgrade pip
           pip install -r devtools/dev-requirements.txt
-          pip install matplotlib==3.7.2
+          pip install matplotlib==3.9.2
 
       - name: Cache Python environment
         id: cache-env

.github/workflows/jax_tests.yml

@@ -35,7 +35,6 @@ jobs:
           sed -i '1i\jax[cpu] == ${{ matrix.jax-version }}' ./requirements.txt
           cat ./requirements.txt
           pip install -r ./devtools/dev-requirements.txt
-          pip install matplotlib==3.7.2
       - name: Verify dependencies
         run: |
           python --version

.github/workflows/notebook_tests.yml

@@ -67,7 +67,7 @@ jobs:
           source .venv-${{ matrix.python-version }}/bin/activate
           python -m pip install --upgrade pip
           pip install -r devtools/dev-requirements.txt
-          pip install matplotlib==3.7.2
+          pip install matplotlib==3.9.2
 
       - name: Test notebooks with pytest and nbmake
         if: env.has_changes == 'true'

.github/workflows/regression_tests.yml

@@ -66,7 +66,7 @@ jobs:
           source .venv-${{ matrix.python-version }}/bin/activate
           python -m pip install --upgrade pip
           pip install -r devtools/dev-requirements.txt
-          pip install matplotlib==3.7.2
+          pip install matplotlib==3.9.2
 
       - name: Set Swap Space
         if: env.has_changes == 'true'

.github/workflows/unit_tests.yml

@@ -72,7 +72,7 @@ jobs:
           source .venv-${{ matrix.combos.python_version }}/bin/activate
           python -m pip install --upgrade pip
           pip install -r devtools/dev-requirements.txt
-          pip install matplotlib==3.7.2
+          pip install matplotlib==3.9.2
 
       - name: Set Swap Space
         if: env.has_changes == 'true'

.github/workflows/weekly_tests.yml

@@ -26,6 +26,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -r devtools/dev-requirements.txt
+          pip install matplotlib==3.9.2
       - name: Set Swap Space
         uses: pierotofy/set-swap-space@master
         with:

CHANGELOG.md

@@ -9,7 +9,7 @@ New Features
 Bug Fixes
 
 - Fixes bug that occurs when taking the gradient of ``root`` and ``root_scalar`` with newer versions of JAX (>=0.4.34) and unpins the JAX version
-
+- Changes ``FixLambdaGauge`` constraint to now enforce zero flux surface average for lambda, instead of enforcing lambda(rho,0,0)=0 as it was incorrectly doing before.
 
 v0.12.3
 -------

desc/backend.py

@@ -65,7 +65,7 @@
     )
 )
 
-if use_jax:  # noqa: C901 - FIXME: simplify this, define globally and then assign?
+if use_jax:  # noqa: C901
     from jax import custom_jvp, jit, vmap
 
     imap = jax.lax.map

desc/basis.py

@@ -1098,13 +1098,13 @@ def evaluate(
         if not len(modes):
             return np.array([]).reshape((len(nodes), 0))
 
-        # TODO: avoid duplicate calculations when mixing derivatives
+        # TODO(#1243): avoid duplicate calculations when mixing derivatives
         r, t, z = nodes.T
         l, m, n = modes.T
         lm = modes[:, :2]
 
         if unique:
-            # TODO: can avoid this here by using grid.unique_idx etc
+            # TODO(#1243): can avoid this here by using grid.unique_idx etc
             # and adding unique_modes attributes to basis
             _, ridx, routidx = np.unique(
                 r, return_index=True, return_inverse=True, axis=0
@@ -1364,7 +1364,6 @@ def polyval_vec(p, x, prec=None):
 def _polyval_exact(p, x, prec):
     p = np.atleast_2d(p)
     x = np.atleast_1d(x).flatten()
-    # TODO: possibly multithread this bit
     mpmath.mp.dps = prec
     y = np.array([np.asarray(mpmath.polyval(list(pi), x)) for pi in p])
     return y.astype(float)
@@ -1440,7 +1439,9 @@ def zernike_radial_coeffs(l, m, exact=True):
             # hence they are all integers. So, we can use exact arithmetic with integer
             # division instead of floating point division.
             # [1]https://en.wikipedia.org/wiki/Zernike_polynomials#Other_representations
-            coeffs[ii, s] = ((-1) ** ((ll - s) // 2) * factorial((ll + s) // 2)) // (
+            coeffs[ii, s] = (
+                int((-1) ** ((ll - s) // 2)) * factorial((ll + s) // 2)
+            ) // (
                 factorial((ll - s) // 2)
                 * factorial((s + mm) // 2)
                 * factorial((s - mm) // 2)

desc/batching.py

@@ -102,7 +102,6 @@ def f_(carry, x):
     return res_append
 
 
-# TODO in_axes a la vmap?
 def _scanmap(fun, scan_fun, argnums=0):
     """A helper function to wrap f with a scan_fun."""
 

desc/coils.py

@@ -1634,7 +1634,14 @@ def compute_magnetic_vector_potential(
 
     @classmethod
     def linspaced_angular(
-        cls, coil, current=None, axis=[0, 0, 1], angle=2 * np.pi, n=10, endpoint=False
+        cls,
+        coil,
+        current=None,
+        axis=[0, 0, 1],
+        angle=2 * np.pi,
+        n=10,
+        endpoint=False,
+        check_intersection=True,
     ):
         """Create a CoilSet by repeating a coil at equal spacing around the torus.
 
@@ -1652,6 +1659,8 @@ def linspaced_angular(
             Number of copies of original coil.
         endpoint : bool
             Whether to include a coil at final rotation angle. Default = False.
+        check_intersection : bool
+            whether to check the resulting coilsets for intersecting coils.
 
         """
         assert isinstance(coil, _Coil) and not isinstance(coil, CoilSet)
@@ -1665,11 +1674,17 @@ def linspaced_angular(
             coili.rotate(axis=axis, angle=phi[i])
             coili.current = currents[i]
             coils.append(coili)
-        return cls(*coils)
+        return cls(*coils, check_intersection=check_intersection)
 
     @classmethod
     def linspaced_linear(
-        cls, coil, current=None, displacement=[2, 0, 0], n=4, endpoint=False
+        cls,
+        coil,
+        current=None,
+        displacement=[2, 0, 0],
+        n=4,
+        endpoint=False,
+        check_intersection=True,
     ):
         """Create a CoilSet by repeating a coil at equal spacing in a straight line.
 
@@ -1686,6 +1701,8 @@ def linspaced_linear(
             Number of copies of original coil.
         endpoint : bool
             Whether to include a coil at final displacement location. Default = False.
+        check_intersection : bool
+            whether to check the resulting coilsets for intersecting coils.
 
         """
         assert isinstance(coil, _Coil) and not isinstance(coil, CoilSet)
@@ -1700,10 +1717,10 @@ def linspaced_linear(
             coili.translate(a[i] * displacement)
             coili.current = currents[i]
             coils.append(coili)
-        return cls(*coils)
+        return cls(*coils, check_intersection=check_intersection)
 
     @classmethod
-    def from_symmetry(cls, coils, NFP=1, sym=False):
+    def from_symmetry(cls, coils, NFP=1, sym=False, check_intersection=True):
         """Create a coil group by reflection and symmetry.
 
         Given coils over one field period, repeat coils NFP times between
@@ -1722,6 +1739,8 @@ def from_symmetry(cls, coils, NFP=1, sym=False):
         sym : bool (optional)
             Whether to enforce stellarator symmetry.
             If True, the coils will be duplicated 2*NFP times. Default = False.
+        check_intersection : bool
+            whether to check the resulting coilsets for intersecting coils.
 
         Returns
         -------
@@ -1764,7 +1783,7 @@ def from_symmetry(cls, coils, NFP=1, sym=False):
             rotated_coils.rotate(axis=[0, 0, 1], angle=2 * jnp.pi * k / NFP)
             coilset += rotated_coils
 
-        return cls(*coilset)
+        return cls(*coilset, check_intersection=check_intersection)
 
     @classmethod
     def from_makegrid_coilfile(cls, coil_file, method="cubic", check_intersection=True):
@@ -1902,8 +1921,8 @@ def save_in_makegrid_format(self, coilsFilename, NFP=None, grid=None):
             if None, will default to the coil compute functions's
             default grid
         """
-        # TODO: name each group based off of CoilSet name?
-        # TODO: have CoilGroup be automatically assigned based off of
+        # TODO(#1376): name each group based off of CoilSet name?
+        # TODO(#1376): have CoilGroup be automatically assigned based off of
         # CoilSet if current coilset is a collection of coilsets?
 
         NFP = 1 if NFP is None else NFP
@@ -2698,7 +2717,7 @@ def insert(self, i, new_item):
         self._coils.insert(i, new_item)
 
     @classmethod
-    def from_makegrid_coilfile(  # noqa: C901 - FIXME: simplify this
+    def from_makegrid_coilfile(  # noqa: C901
         cls, coil_file, method="cubic", ignore_groups=False, check_intersection=True
     ):
         """Create a MixedCoilSet of SplineXYZCoils from a MAKEGRID coil txtfile.

desc/compute/_geometry.py

@@ -194,7 +194,7 @@ def _A_of_z(params, transforms, profiles, data, **kwargs):
     data=["Z", "n_rho", "e_theta|r,p", "rho"],
     parameterization=["desc.geometry.surface.FourierRZToroidalSurface"],
     resolution_requirement="rt",  # just need max(rho) near 1
-    # FIXME: Add source grid requirement once omega is nonzero.
+    # TODO(#568): Add source grid requirement once omega is nonzero.
 )
 def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
     # Denote any vector v = [vᴿ, v^ϕ, vᶻ] with a tuple of its contravariant components.
@@ -213,7 +213,7 @@ def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwarg
         line_integrals(
             transforms["grid"],
             data["Z"] * n[:, 2] * safenorm(data["e_theta|r,p"], axis=-1),
-            # FIXME: Works currently for omega = zero, but for nonzero omega
+            # TODO(#568): Works currently for omega = zero, but for nonzero omega
             #  we need to integrate over theta at constant phi.
             #  Should be simple once we have coordinate mapping and source grid
             #  logic from GitHub pull request #1024.
@@ -449,7 +449,7 @@ def _perimeter_of_z(params, transforms, profiles, data, **kwargs):
         line_integrals(
             transforms["grid"],
             safenorm(data["e_theta|r,p"], axis=-1),
-            # FIXME: Works currently for omega = zero, but for nonzero omega
+            # TODO(#568): Works currently for omega = zero, but for nonzero omega
             #  we need to integrate over theta at constant phi.
             #  Should be simple once we have coordinate mapping and source grid
             #  logic from GitHub pull request #1024.

desc/compute/_omnigenity.py

@@ -47,7 +47,7 @@ def fitfun(x):
     return data
 
 
-# TODO: do math to change definition of nu so that we can just use B_zeta_mn here
+# TODO (#568): do math to change definition of nu so that we can just use B_zeta_mn here
 @register_compute_fun(
     name="B_phi_mn",
     label="B_{\\phi, m, n}",
@@ -63,7 +63,7 @@ def fitfun(x):
     data=["B_phi|r,t"],
     resolution_requirement="tz",
     grid_requirement={"is_meshgrid": True},
-    aliases="B_zeta_mn",  # TODO: remove when phi != zeta
+    aliases="B_zeta_mn",  # TODO(#568): remove when phi != zeta
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
 )

desc/compute/_profiles.py

@@ -1378,7 +1378,7 @@ def _iota_num_rrr(params, transforms, profiles, data, **kwargs):
                 - beta * data["sqrt(g)_rrr"],
                 data["sqrt(g)"],
             ),
-            # Todo: axis limit of beta_rrr
+            # TODO(#587): axis limit of beta_rrr
             #   Computed with four applications of l’Hôpital’s rule.
             #   Requires sqrt(g)_rrrr and fourth derivatives of basis vectors.
             jnp.nan,
@@ -1656,7 +1656,7 @@ def _iota_den_rrr(params, transforms, profiles, data, **kwargs):
             - gamma * data["sqrt(g)_rrr"],
             data["sqrt(g)"],
         ),
-        # Todo: axis limit
+        # TODO(#587): axis limit
         #   Computed with four applications of l’Hôpital’s rule.
         #   Requires sqrt(g)_rrrr and fourth derivatives of basis vectors.
         jnp.nan,
@@ -1713,7 +1713,7 @@ def _q(params, transforms, profiles, data, **kwargs):
     return data
 
 
-# TODO: add K(rho,theta,zeta)*grad(rho) term
+# TODO (#1381): add K(rho,theta,zeta)*grad(rho) term
 @register_compute_fun(
     name="I",
     label="I",

desc/compute/_stability.py

@@ -76,7 +76,7 @@ def _D_current(params, transforms, profiles, data, **kwargs):
                 / data["|grad(psi)|"] ** 3
                 * dot(Xi, data["B"])
             ),
-            # Todo: implement equivalent of equation 4.3 in desc coordinates
+            # TODO(#671): implement equivalent of equation 4.3 in desc coordinates
             jnp.nan,
         )
     )

desc/compute/_surface.py

@@ -3,7 +3,7 @@
 from .data_index import register_compute_fun
 from .geom_utils import rpz2xyz
 
-# TODO: review when zeta no longer equals phi
+# TODO(#568): review when zeta no longer equals phi
 
 
 @register_compute_fun(
@@ -27,7 +27,7 @@
 def _x_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
     R = transforms["R"].transform(params["R_lmn"])
     Z = transforms["Z"].transform(params["Z_lmn"])
-    # TODO: change when zeta no longer equals phi
+    # TODO(#568): change when zeta no longer equals phi
     phi = transforms["grid"].nodes[:, 2]
     coords = jnp.stack([R, phi, Z], axis=1)
     # default basis for "x" is rpz, the conversion will be done

desc/continuation.py

@@ -748,7 +748,6 @@ def solve_continuation(  # noqa: C901
         if len(deltas) > 0:
             if verbose > 0:
                 print("Perturbing equilibrium")
-            # TODO: pass Jx if available
             eqp = eqfam[ii - 1].copy()
             objective_i = get_equilibrium_objective(
                 eq=eqp, mode=objective, jac_chunk_size=jac_chunk_size

desc/equilibrium/coords.py

@@ -107,7 +107,7 @@ def map_coordinates(  # noqa: C901
 
     profiles = get_profiles(inbasis + basis_derivs, eq)
 
-    # TODO: make this work for permutations of in/out basis
+    # TODO (#1382): make this work for permutations of in/out basis
     if outbasis == ("rho", "theta", "zeta"):
         if inbasis == ("rho", "alpha", "zeta"):
             if "iota" in kwargs:
@@ -286,7 +286,7 @@ def _distance_body(i, idx):
     return yg[idx]
 
 
-# TODO: decide later whether to assume given phi instead of zeta.
+# TODO(#568): decide later whether to assume given phi instead of zeta.
 def _map_PEST_coordinates(
     coords,
     L_lmn,
@@ -395,7 +395,7 @@ def fixup(x, *args):
     return out
 
 
-# TODO: decide later whether to assume given phi instead of zeta.
+# TODO(#568): decide later whether to assume given phi instead of zeta.
 def _map_clebsch_coordinates(
     coords,
     iota,
@@ -766,7 +766,7 @@ def get_rtz_grid(
     return desc_grid
 
 
-# TODO: deprecated, remove eventually
+# TODO(#1383): deprecated, remove eventually
 def compute_theta_coords(
     eq, flux_coords, L_lmn=None, tol=1e-6, maxiter=20, full_output=False, **kwargs
 ):

desc/equilibrium/equilibrium.py

@@ -1443,7 +1443,6 @@ def axis(self, new):
     @property
     def spectral_indexing(self):
         """str: Type of indexing used for the spectral basis."""
-        # TODO: allow this to change?
         return self._spectral_indexing
 
     @property
@@ -1972,8 +1971,6 @@ def from_near_axis(
             raise ValueError("Input must be a pyQSC or pyQIC solution.") from e
 
         rho, _ = special.js_roots(L, 2, 2)
-        # TODO: could make this an OCS grid to improve fitting, need to figure out
-        # how concentric grids work with QSC
         grid = LinearGrid(rho=rho, theta=ntheta, zeta=na_eq.phi, NFP=na_eq.nfp)
         basis_R = FourierZernikeBasis(
             L=L,

desc/equilibrium/initial_guess.py

@@ -20,7 +20,7 @@
 from desc.utils import copy_coeffs, warnif
 
 
-def set_initial_guess(eq, *args, ensure_nested=True):  # noqa: C901 - FIXME: simplify
+def set_initial_guess(eq, *args, ensure_nested=True):  # noqa: C901
     """Set the initial guess for the flux surfaces, eg R_lmn, Z_lmn, L_lmn.
 
     Parameters