Merge branch 'master' into rc/transforms

.github/dependabot.yml

@@ -0,0 +1,44 @@
+# Please see the documentation for all configuration options:
+# https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file
+
+version: 2
+updates:
+  - package-ecosystem: "pip"
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "weekly"
+    groups:
+      # group similar dependencies together into fewer PRs
+      docs:
+        patterns:
+          - "sphinx"
+      linting:
+        patterns:
+          - "flake8"
+          - "black"
+          - "pre-commit"
+      testing:
+        patterns:
+          - "pytest"
+          - "qic"
+          - "qsc"
+          - "shapely"
+      minor_packages:
+        patterns:
+          - "colorama"
+          - "nvgpu"
+          - "psutil"
+          - "pylatexenc"
+          - "termcolor"
+      jax:
+        patterns:
+          - "jax"
+          - "diffrax"
+          - "interpax"
+          - "orthax"
+          - "quadax"
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      # Check for updates to GitHub Actions every week
+      interval: "weekly"

.github/workflows/benchmark.yml

@@ -82,6 +82,7 @@ jobs:
           source .venv-${{ matrix.python-version }}/bin/activate
           pwd
           lscpu
+          pip list
           cd tests/benchmarks
           python -m pytest benchmark_cpu_small.py -vv \
             --benchmark-save='Latest_Commit' \
@@ -108,6 +109,7 @@ jobs:
           source .venv-${{ matrix.python-version }}/bin/activate
           pwd
           lscpu
+          pip list
           cd tests/benchmarks
           python -m pytest benchmark_cpu_small.py -vv \
             --benchmark-save='master' \
@@ -149,12 +151,12 @@ jobs:
 
       - name: Comment PR with the results
         if: env.has_changes == 'true'
-        uses: thollander/actions-comment-pull-request@v2
+        uses: thollander/actions-comment-pull-request@v3
         env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          github-token: ${{ secrets.GITHUB_TOKEN }}
         with:
-          filePath: tests/benchmarks/commit_msg.txt
-          comment_tag: benchmark
+          file-path: tests/benchmarks/commit_msg.txt
+          comment-tag: benchmark
 
       - name: Upload benchmark data
         if: always() && env.has_changes == 'true'

.github/workflows/jax_tests.yml

@@ -44,6 +44,7 @@ jobs:
         run: |
           pwd
           lscpu
+          pip list
           python -m pytest -m unit \
             --durations=0 \
             --mpl \

.github/workflows/mpl_tests.yml

@@ -72,4 +72,5 @@ jobs:
         run: |
           pwd
           lscpu
+          pip list
           python -m pytest tests/test_plotting.py --durations=0 --mpl --maxfail=1

.github/workflows/notebook_tests.yml

@@ -75,6 +75,7 @@ jobs:
           source .venv-${{ matrix.python-version }}/bin/activate
           pwd
           lscpu
+          pip list
           export PYTHONPATH=$(pwd)
           pytest -v --nbmake "./docs/notebooks" \
             --nbmake-timeout=2000 \

.github/workflows/regression_tests.yml

@@ -80,6 +80,7 @@ jobs:
           source .venv-${{ matrix.python-version }}/bin/activate
           pwd
           lscpu
+          pip list
           python -m pytest -v -m regression\
             --durations=0 \
             --cov-report xml:cov.xml \
@@ -106,14 +107,9 @@ jobs:
       - name: Upload coverage
         if: env.has_changes == 'true'
         id : codecov
-        uses: Wandalen/wretry.action@v3.5.0
+        uses: codecov/codecov-action@v5
         with:
-          action: codecov/codecov-action@v4
-          with: |
-            token: ${{ secrets.CODECOV_TOKEN }}
-            name: codecov-umbrella
-            files: ./cov.xml
-            fail_ci_if_error: true
-            verbose: true
-          attempt_limit: 10
-          attempt_delay: 60000 # ms, 1 min
+          name: codecov-umbrella
+          files: ./cov.xml
+          fail_ci_if_error: true
+          verbose: true

.github/workflows/release.yml

@@ -20,7 +20,7 @@ jobs:
       - name: Build package
         run: python -m build --sdist --wheel
       - name: Publish package
-        uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+        uses: pypa/gh-action-pypi-publish@67339c736fd9354cd4f8cb0b744f2b82a74b5c70
         with:
           user: __token__
           password: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/unit_tests.yml

@@ -86,6 +86,7 @@ jobs:
           source .venv-${{ matrix.combos.python_version }}/bin/activate
           pwd
           lscpu
+          pip list
           python -m pytest -v -m unit \
             --durations=0 \
             --cov-report xml:cov.xml \
@@ -112,14 +113,9 @@ jobs:
       - name: Upload coverage
         if: env.has_changes == 'true'
         id : codecov
-        uses: Wandalen/wretry.action@v3.5.0
+        uses: codecov/codecov-action@v5
         with:
-          action: codecov/codecov-action@v4
-          with: |
-            token: ${{ secrets.CODECOV_TOKEN }}
-            name: codecov-umbrella
-            files: ./cov.xml
-            fail_ci_if_error: true
-            verbose: true
-          attempt_limit: 10
-          attempt_delay: 60000 # ms, 1 min
+          name: codecov-umbrella
+          files: ./cov.xml
+          fail_ci_if_error: true
+          verbose: true

.github/workflows/update_notebook_examples.yml

@@ -0,0 +1,75 @@
+name: Update Notebooks and Examples
+on:
+  schedule:
+    # runs on the first of every month at noon
+    - cron: '00 12 1 * *'
+
+jobs:
+  create_issue:
+    name: Create issues to update notebooks and examples
+    runs-on: ubuntu-latest
+    permissions:
+      issues: write
+    steps:
+      - name: Create issue to update notebooks
+        run: |
+          if [[ $CLOSE_PREVIOUS == true ]]; then
+            previous_issue_number=$(gh issue list \
+              --label "$LABELS" \
+              --json number \
+              --jq '.[0].number')
+            if [[ -n $previous_issue_number ]]; then
+              gh issue close "$previous_issue_number"
+              gh issue unpin "$previous_issue_number"
+            fi
+          fi
+          new_issue_url=$(gh issue create \
+            --title "$TITLE" \
+            --assignee "$ASSIGNEES" \
+            --label "$LABELS" \
+            --body "$BODY")
+          if [[ $PINNED == true ]]; then
+            gh issue pin "$new_issue_url"
+          fi
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          GH_REPO: ${{ github.repository }}
+          TITLE: Update Notebooks
+          ASSIGNEES: f0uriest, dpanici, ddudt, rahulgaur104, kianorr, YigitElma
+          LABELS: monthly_update_notebooks
+          BODY: |
+            Reminder to check that notebooks evaluate correctly, and to update
+            if necessary.
+          PINNED: false
+          CLOSE_PREVIOUS: true
+      - name: Create issue to update notebooks
+        run: |
+            if [[ $CLOSE_PREVIOUS == true ]]; then
+                previous_issue_number=$(gh issue list \
+                --label "$LABELS" \
+                --json number \
+                --jq '.[0].number')
+                if [[ -n $previous_issue_number ]]; then
+                gh issue close "$previous_issue_number"
+                gh issue unpin "$previous_issue_number"
+                fi
+            fi
+            new_issue_url=$(gh issue create \
+                --title "$TITLE" \
+                --assignee "$ASSIGNEES" \
+                --label "$LABELS" \
+                --body "$BODY")
+            if [[ $PINNED == true ]]; then
+                gh issue pin "$new_issue_url"
+            fi
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          GH_REPO: ${{ github.repository }}
+          TITLE: Update Examples
+          ASSIGNEES: f0uriest, dpanici, ddudt, rahulgaur104, kianorr, YigitElma
+          LABELS: monthly_update_examples
+          BODY: |
+              Reminder to check that examples run correctly, and to update
+              if necessary.
+          PINNED: false
+          CLOSE_PREVIOUS: true

.github/workflows/weekly_tests.yml

@@ -35,6 +35,7 @@ jobs:
         run: |
           pwd
           lscpu
+          pip list
           python -m pytest -v -m unit \
             --durations=0 \
             --splits 4 \

.pre-commit-config.yaml

@@ -1,11 +1,11 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 24.10.0
     hooks:
       - id: black-jupyter
         log_file: devtools/pre-commit.log
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v5.0.0
     hooks:
       - id: check-added-large-files
         log_file: devtools/pre-commit.log
@@ -22,20 +22,20 @@ repos:
       - id: no-commit-to-branch
         log_file: devtools/pre-commit.log
   - repo: https://github.com/pycqa/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - id: isort
         name: isort (python)
         args: ["--profile", "black"]
         log_file: devtools/pre-commit.log
   - repo: https://github.com/PyCQA/flake8
-    rev: "5.0.4"
+    rev: "7.1.1"
     hooks:
       - id: flake8
         additional_dependencies: ["flake8-docstrings", "flake8-eradicate", "flake8-isort"]
         log_file: devtools/pre-commit.log
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.2.0
+    rev: v3.19.0
     hooks:
       - id: pyupgrade
         log_file: devtools/pre-commit.log

CHANGELOG.md

@@ -1,14 +1,46 @@
 Changelog
 =========
 
+New Feature
+
+- Adds a new profile class ``PowerProfile`` for raising profiles to a power.
+- Adds an option ``scaled_termination`` (defaults to True) to all of the desc optimizers to measure the norms for ``xtol`` and ``gtol`` in the scaled norm provided by ``x_scale`` (which defaults to using an adaptive scaling based on the Jacobian or Hessian). This should make things more robust when optimizing parameters with widely different magnitudes. The old behavior can be recovered by passing ``options={"scaled_termination": False}``.
+- ``desc.objectives.Omnigenity`` is now vectorized and able to optimize multiple surfaces at the same time. Previously it was required to use a different objective for each surface.
+
+Bug Fixes
+
+- Small bug fix to use the correct normalization length ``a`` in the BallooningStability objective
+
+v0.13.0
+-------
+
 New Features
-- Add ``from_input_file`` method to ``Equilibrium`` class to generate an ``Equilibrium`` object with boundary, profiles, resolution and flux specified in a given DESC or VMEC input file
+
+- Adds ``from_input_file`` method to ``Equilibrium`` class to generate an ``Equilibrium`` object with boundary, profiles, resolution and flux specified in a given DESC or VMEC input file
+- Adds function ``solve_regularized_surface_current`` to ``desc.magnetic_fields`` module that implements the REGCOIL algorithm (Landreman, (2017)) for surface current normal field optimization
+    * Can specify the tuple ``current_helicity=(M_coil, N_coil)`` to determine if resulting contours correspond to helical topology (both ``(M_coil, N_coil)`` not equal to 0), modular (``N_coil`` equal to 0 and ``M_coil`` nonzero) or windowpane/saddle (``M_coil`` and ``N_coil`` both zero)
+    * ``M_coil`` is the number of poloidal transits a coil makes before returning to itself, while ``N_coil`` is the number of toroidal transits a coil makes before returning to itself (this is sort of like the QS ``helicity``)
+    * if multiple values of the regularization parameter are input, will return a family of surface current fields (as a list) corresponding to the solution at each regularization value
+- Adds method ``to_CoilSet`` to ``FourierCurrentPotentialField`` which implements a coil cutting algorithm to discretize the surface current into coils
+    * works for both modular and helical coils
+- Adds a new objective ``SurfaceCurrentRegularization`` (which minimizes ``w*|K|``, the regularization term from surface current in the REGCOIL algorithm, with `w` being the objective weight which act as the regularization parameter)
+    * use of both this and the ``QuadraticFlux`` objective allows for REGCOIL solutions to be obtained through the optimization framework, and combined with other objectives as well.
+- Changes local area weighting of Bn in QuadraticFlux objective to be the square root of the local area element (Note that any existing optimizations using this objective may need different weights to achieve the same result now.)
+- Adds a new tutorial showing how to use``REGCOIL`` features.
+- Adds an ``NFP`` attribute to ``ScalarPotentialField``, ``VectorPotentialField`` and ``DommaschkPotentialField``, to allow ``SplineMagneticField.from_field`` and ``MagneticField.save_mgrid`` to efficiently take advantage of the discrete toroidal symmetry of these fields, if present.
+- Adds ``SurfaceQuadraticFlux`` objective which minimizes the quadratic magnetic flux through a ``FourierRZToroidalSurface`` object, allowing for optimizing for Quadratic flux minimizing (QFM) surfaces.
+- Allows ``ToroidalFlux`` objective to accept ``FourierRZToroidalSurface`` so it can be used to specify the toroidal flux through a QFM surface.
+- Adds ``eq_fixed`` flag to ``ToroidalFlux`` to allow for the equilibrium/QFM surface to vary during optimization, useful for single-stage optimizations.
+- Adds tutorial notebook showcasing QFM surface capability.
+- Adds ``rotate_zeta`` function to ``desc.compat`` to rotate an ``Equilibrium`` around Z axis.
 
 
 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
+- 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.
+- Fixes bug in ``softmin/softmax`` implementation.
+- Fixes bug that occured when using ``ProximalProjection`` with a scalar optimization algorithm.
 
 v0.12.3
 -------

desc/__main__.py

@@ -37,7 +37,7 @@ def main(cl_args=sys.argv[1:]):
         and (inputs[-1]["pres_ratio"] is None)
         and (inputs[-1]["bdry_ratio"] is None)
     ):
-        eq = Equilibrium(**inputs[-1], check_kwargs=False)
+        eq = Equilibrium(**inputs[-1], check_kwargs=False, ensure_nested=False)
         equil_fam = EquilibriaFamily.solve_continuation_automatic(
             eq,
             objective=inputs[-1]["objective"],

desc/backend.py

@@ -73,10 +73,10 @@
     from jax.lax import cond, fori_loop, scan, switch, while_loop
     from jax.nn import softmax as softargmax
     from jax.numpy import bincount, flatnonzero, repeat, take
-    from jax.numpy.fft import irfft, rfft, rfft2
+    from jax.numpy.fft import ifft, irfft, irfft2, rfft, rfft2
     from jax.scipy.fft import dct, idct
     from jax.scipy.linalg import block_diag, cho_factor, cho_solve, qr, solve_triangular
-    from jax.scipy.special import gammaln, logsumexp
+    from jax.scipy.special import gammaln
     from jax.tree_util import (
         register_pytree_node,
         tree_flatten,
@@ -434,7 +434,7 @@ def tangent_solve(g, y):
     jit = lambda func, *args, **kwargs: func
     execute_on_cpu = lambda func: func
     import scipy.optimize
-    from numpy.fft import irfft, rfft, rfft2  # noqa: F401
+    from numpy.fft import ifft, irfft, irfft2, rfft, rfft2  # noqa: F401
     from scipy.fft import dct, idct  # noqa: F401
     from scipy.integrate import odeint  # noqa: F401
     from scipy.linalg import (  # noqa: F401
@@ -445,7 +445,7 @@ def tangent_solve(g, y):
         qr,
         solve_triangular,
     )
-    from scipy.special import gammaln, logsumexp  # noqa: F401
+    from scipy.special import gammaln  # noqa: F401
     from scipy.special import softmax as softargmax  # noqa: F401
 
     trapezoid = np.trapezoid if hasattr(np, "trapezoid") else np.trapz