Merge branch 'master' into dp/order_rho_sq_NAE

.github/workflows/benchmark.yml

@@ -1,9 +1,10 @@
 name: Benchmarks
 
 on:
-  pull_request:
+  pull_request_target:
     branches:
       - master
+    types: [opened, synchronize]
   workflow_dispatch:
     inputs:
       debug_enabled:
@@ -73,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/changelog_update.yml

@@ -0,0 +1,35 @@
+name: Check changelog updated
+
+on:
+  pull_request:
+    branches:
+      - master
+    types: [opened, synchronize, labeled, unlabeled]
+
+jobs:
+  check_changelog_updated:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Filter changes
+        id: changes
+        uses: dorny/paths-filter@v3
+        with:
+          filters: |
+            has_changes:
+              - 'desc/**'
+              - 'requirements.txt'
+              - 'requirements_conda.yml'
+              - '.github/workflows/changelog_update.yml'
+
+      - name: Check for relevant changes
+        id: check_changes
+        run: echo "has_changes=${{ steps.changes.outputs.has_changes }}" >> $GITHUB_ENV
+
+      - uses: actions/checkout@v4
+
+      - if: ${{ !contains(github.event.pull_request.labels.*.name, 'skip_changelog') && env.has_changes == 'true'}}
+        uses: danieljimeneznz/[email protected]
+        with:
+          require-changes-to: |
+            CHANGELOG.md
+          token: ${{ secrets.GITHUB_TOKEN }}

.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

@@ -8,7 +8,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/basis.py

@@ -1440,7 +1440,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/magnetic_fields/_core.py

@@ -1816,10 +1816,10 @@ def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False):
         ir = int(mgrid["ir"][()])  # number of grid points in the R coordinate
         jz = int(mgrid["jz"][()])  # number of grid points in the Z coordinate
         kp = int(mgrid["kp"][()])  # number of grid points in the phi coordinate
-        Rmin = mgrid["rmin"][()]  # Minimum R coordinate (m)
-        Rmax = mgrid["rmax"][()]  # Maximum R coordinate (m)
-        Zmin = mgrid["zmin"][()]  # Minimum Z coordinate (m)
-        Zmax = mgrid["zmax"][()]  # Maximum Z coordinate (m)
+        Rmin = mgrid["rmin"][()].filled()  # Minimum R coordinate (m)
+        Rmax = mgrid["rmax"][()].filled()  # Maximum R coordinate (m)
+        Zmin = mgrid["zmin"][()].filled()  # Minimum Z coordinate (m)
+        Zmax = mgrid["zmax"][()].filled()  # Maximum Z coordinate (m)
         nfp = int(mgrid["nfp"][()])  # Number of field periods
         Rgrid = np.linspace(Rmin, Rmax, ir)
         Zgrid = np.linspace(Zmin, Zmax, jz)
@@ -1831,9 +1831,15 @@ def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False):
         bz = np.zeros([kp, jz, ir, nextcur])
         for i in range(nextcur):
             coil_id = "%03d" % (i + 1,)
-            br[:, :, :, i] += mgrid["br_" + coil_id][()]  # B_R radial magnetic field
-            bp[:, :, :, i] += mgrid["bp_" + coil_id][()]  # B_phi toroidal field (T)
-            bz[:, :, :, i] += mgrid["bz_" + coil_id][()]  # B_Z vertical magnetic field
+            br[:, :, :, i] += mgrid["br_" + coil_id][
+                ()
+            ].filled()  # B_R radial magnetic field
+            bp[:, :, :, i] += mgrid["bp_" + coil_id][
+                ()
+            ].filled()  # B_phi toroidal field (T)
+            bz[:, :, :, i] += mgrid["bz_" + coil_id][
+                ()
+            ].filled()  # B_Z vertical magnetic field
 
         # shift axes to correct order
         br = np.moveaxis(br, (0, 1, 2), (1, 2, 0))
@@ -1849,13 +1855,13 @@ def from_mgrid(cls, mgrid_file, extcur=None, method="cubic", extrap=False):
                 coil_id = "%03d" % (i + 1,)
                 ar[:, :, :, i] += mgrid["ar_" + coil_id][
                     ()
-                ]  # A_R radial mag. vec. potential
+                ].filled()  # A_R radial mag. vec. potential
                 ap[:, :, :, i] += mgrid["ap_" + coil_id][
                     ()
-                ]  # A_phi toroidal mag. vec. potential
+                ].filled()  # A_phi toroidal mag. vec. potential
                 az[:, :, :, i] += mgrid["az_" + coil_id][
                     ()
-                ]  # A_Z vertical mag. vec. potential
+                ].filled()  # A_Z vertical mag. vec. potential
 
             # shift axes to correct order
             ar = np.moveaxis(ar, (0, 1, 2), (1, 2, 0))

desc/objectives/linear_objectives.py

@@ -12,7 +12,7 @@
 from termcolor import colored
 
 from desc.backend import execute_on_cpu, jnp, tree_leaves, tree_map, tree_structure
-from desc.basis import zernike_radial, zernike_radial_coeffs
+from desc.basis import zernike_radial
 from desc.geometry import FourierRZCurve
 from desc.utils import broadcast_tree, errorif, setdefault
 
@@ -774,8 +774,8 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixLambdaGauge(_Objective):
-    """Fixes gauge freedom for lambda: lambda(theta=0,zeta=0)=0.
+class FixLambdaGauge(FixParameters):
+    """Fixes gauge freedom for lambda, which sets the flux surface avg of lambda to 0.
 
     Note: this constraint is automatically applied when needed, and does not need to be
     included by the user.
@@ -793,9 +793,6 @@ class FixLambdaGauge(_Objective):
 
     """
 
-    _scalar = False
-    _linear = True
-    _fixed = False  # not "diagonal", since it is fixing a sum
     _units = "(rad)"
     _print_value_fmt = "lambda gauge error: "
 
@@ -806,75 +803,21 @@ def __init__(
         normalize_target=True,
         name="lambda gauge",
     ):
+        if eq.sym:
+            indices = False
+        else:
+            indices = np.where(
+                np.logical_and(eq.L_basis.modes[:, 1] == 0, eq.L_basis.modes[:, 2] == 0)
+            )[0]
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"L_lmn": indices},
             target=0,
-            bounds=None,
-            weight=1,
             normalize=normalize,
             normalize_target=normalize_target,
             name=name,
         )
 
-    def build(self, use_jit=False, verbose=1):
-        """Build constant arrays.
-
-        Parameters
-        ----------
-        use_jit : bool, optional
-            Whether to just-in-time compile the objective and derivatives.
-        verbose : int, optional
-            Level of output.
-
-        """
-        eq = self.things[0]
-        L_basis = eq.L_basis
-
-        if L_basis.sym:
-            self._A = np.zeros((0, L_basis.num_modes))
-        else:
-            # l(rho,0,0) = 0
-            # at theta=zeta=0, basis for lambda reduces to just a polynomial in rho
-            # what this constraint does is make all the coefficients of each power
-            # of rho equal to zero
-            # i.e. if lambda = (L_200 + 2*L_310) rho**2 + (L_100 + 2*L_210)*rho
-            # this constraint will make
-            # L_200 + 2*L_310 = 0
-            # L_100 + 2*L_210 = 0
-            L_modes = L_basis.modes
-            mnpos = np.where((L_modes[:, 1:] >= [0, 0]).all(axis=1))[0]
-            l_lmn = L_modes[mnpos, :]
-            if len(l_lmn) > 0:
-                c = zernike_radial_coeffs(l_lmn[:, 0], l_lmn[:, 1])
-            else:
-                c = np.zeros((0, 0))
-
-            A = np.zeros((c.shape[1], L_basis.num_modes))
-            A[:, mnpos] = c.T
-            self._A = A
-
-        self._dim_f = self._A.shape[0]
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute lambda gauge freedom errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            gauge freedom errors.
-
-        """
-        return jnp.dot(self._A, params["L_lmn"])
-
 
 class FixThetaSFL(FixParameters):
     """Fixes lambda=0 so that poloidal angle is the SFL poloidal angle.

desc/optimize/optimizer.py

@@ -209,9 +209,11 @@ def optimize(  # noqa: C901 - FIXME: simplify this
         objective, nonlinear_constraints = _maybe_wrap_nonlinear_constraints(
             eq, objective, nonlinear_constraints, self.method, options
         )
-        if not isinstance(objective, ProximalProjection):
-            for t in things:
-                linear_constraints = maybe_add_self_consistency(t, linear_constraints)
+        is_prox = isinstance(objective, ProximalProjection)
+        for t in things:
+            if isinstance(t, Equilibrium) and is_prox:
+                continue  # don't add Equilibrium self-consistency if proximal is used
+            linear_constraints = maybe_add_self_consistency(t, linear_constraints)
         linear_constraint = _combine_constraints(linear_constraints)
         nonlinear_constraint = _combine_constraints(nonlinear_constraints)
 

desc/plotting.py

@@ -1354,6 +1354,8 @@ def plot_section(
         * ``phi``: float, int or array-like. Toroidal angles to plot. If an integer,
           plot that number equally spaced in [0,2pi/NFP). Default 1 for axisymmetry and
           6 for non-axisymmetry
+        * ``fill`` : bool,  Whether the contours are filled, i.e. whether to use
+          `contourf` or `contour`. Default to ``fill=True``
 
     Returns
     -------
@@ -1442,7 +1444,7 @@ def plot_section(
     R = coords["R"].reshape((nt, nr, nz), order="F")
     Z = coords["Z"].reshape((nt, nr, nz), order="F")
     data = data.reshape((nt, nr, nz), order="F")
-
+    op = "contour" + ("f" if kwargs.pop("fill", True) else "")
     contourf_kwargs = {}
     if log:
         data = np.abs(data)  # ensure data is positive for log plot
@@ -1474,7 +1476,9 @@ def plot_section(
     for i in range(nphi):
         divider = make_axes_locatable(ax[i])
 
-        cntr = ax[i].contourf(R[:, :, i], Z[:, :, i], data[:, :, i], **contourf_kwargs)
+        cntr = getattr(ax[i], op)(
+            R[:, :, i], Z[:, :, i], data[:, :, i], **contourf_kwargs
+        )
         cax = divider.append_axes("right", **cax_kwargs)
         cbar = fig.colorbar(cntr, cax=cax)
         cbar.update_ticks()

devtools/dev-requirements_conda.yml

@@ -7,9 +7,9 @@ dependencies:
   - matplotlib >= 3.5.0, < 4.0.0
   - mpmath >= 1.0.0, < 2.0
   - netcdf4 >= 1.5.4, < 2.0
-  - numpy >= 1.20.0, < 2.0
+  - numpy >= 1.20.0
   - psutil
-  - scipy >= 1.7.0, < 2.0
+  - scipy >= 1.7.0
   - termcolor
   - pip
   - pip:

requirements.txt

@@ -6,11 +6,11 @@ interpax >= 0.3.3
 matplotlib >= 3.5.0, < 4.0.0
 mpmath >= 1.0.0, < 2.0
 netcdf4 >= 1.5.4, < 2.0
-numpy >= 1.20.0, < 2.0.0
+numpy >= 1.20.0
 nvgpu
 orthax
 plotly >= 5.16, < 6.0
 psutil
 pylatexenc >= 2.0, < 3.0
-scipy >= 1.7.0, < 2.0.0
+scipy >= 1.7.0
 termcolor

requirements_conda.yml

@@ -7,9 +7,9 @@ dependencies:
   - matplotlib >= 3.5.0, < 4.0.0
   - mpmath >= 1.0.0, < 2.0
   - netcdf4 >= 1.5.4, < 2.0
-  - numpy >= 1.20.0, < 2.0
+  - numpy >= 1.20.0
   - psutil
-  - scipy >= 1.7.0, < 2.0
+  - scipy >= 1.7.0
   - termcolor
   - pip
   - pip:

tests/conftest.py

@@ -284,13 +284,13 @@ def DummyMixedCoilSet(tmpdir_factory):
         vf_coil, displacement=[0, 0, 2], n=3, endpoint=True
     )
     xyz_coil = FourierXYZCoil(current=2)
-    phi = 2 * np.pi * np.linspace(0, 1, 20, endpoint=True) ** 2
+    phi = 2 * np.pi * np.linspace(0, 1, 20, endpoint=True)
     spline_coil = SplineXYZCoil(
         current=1,
         X=np.cos(phi),
         Y=np.sin(phi),
         Z=np.zeros_like(phi),
-        knots=np.linspace(0, 2 * np.pi, len(phi)),
+        knots=phi,
     )
     full_coilset = MixedCoilSet(
         (tf_coilset, vf_coilset, xyz_coil, spline_coil), check_intersection=False