Merge branch 'master' into dependabot/pip/jax-1cca3be8ee

PlasmaControl · Dec 11, 2024 · f31d8cf · f31d8cf
2 parents 0b61ab5 + 8804180
commit f31d8cf
Show file tree

Hide file tree

Showing 43 changed files with 155 additions and 65 deletions.
diff --git a/.github/workflows/update_notebook_examples.yml b/.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
diff --git a/desc/__main__.py b/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"],

diff --git a/desc/equilibrium/utils.py b/desc/equilibrium/utils.py
@@ -141,6 +141,7 @@ def parse_axis(axis, NFP=1, sym=True, surface=None):
             axis[:, 1],
             axis[:, 2],
             axis[:, 0].astype(int),
+            axis[:, 0].astype(int),
             NFP=NFP,
             sym=sym,
             name="axis",

diff --git a/desc/examples/ARIES-CS_output.h5 b/desc/examples/ARIES-CS_output.h5
diff --git a/desc/examples/ATF b/desc/examples/ATF
@@ -34,10 +34,6 @@ l:   0	p =   5.00000000E+05	i =   -3.50000000E-01
 l:   2	p =  -1.00000000E+06	i =   -6.50000000E-01
 l:   4	p =   5.00000000E+05	i =   0.00000000E+00
 
-# magnetic axis initial guess
-n:   0	R0 =   2.10000000E+00	Z0 =   0.00000000E+00
-n:   1	R0 =   0.10000000E+00	Z0 =   0.10000000E+00
-
 # fixed-boundary surface shape
 m:  -2	n:  -1	R1 =  -2.50000000E-02	Z1 =   0.00000000E+00
 m:  -2	n:   0	R1 =   0.00000000E+00	Z1 =  -6.75000000E-03

diff --git a/desc/examples/ATF_output.h5 b/desc/examples/ATF_output.h5
diff --git a/desc/examples/DSHAPE_CURRENT_output.h5 b/desc/examples/DSHAPE_CURRENT_output.h5
diff --git a/desc/examples/DSHAPE_output.h5 b/desc/examples/DSHAPE_output.h5
diff --git a/desc/examples/ESTELL_output.h5 b/desc/examples/ESTELL_output.h5
diff --git a/desc/examples/HELIOTRON b/desc/examples/HELIOTRON
@@ -14,15 +14,15 @@ M_grid = 12, 18
 N_grid =  0,  0,  6
 
 # continuation parameters
-bdry_ratio = 0, 0, 1
+bdry_ratio = 0, 0, 0.5, 1
 pres_ratio = 0, 1, 1
 pert_order = 2
 
 # solver tolerances
 ftol = 1e-2
 xtol = 1e-6
 gtol = 1e-6
-maxiter = 100
+maxiter = 50
 
 # solver methods
 optimizer         = lsq-exact

diff --git a/desc/examples/HELIOTRON_output.h5 b/desc/examples/HELIOTRON_output.h5
diff --git a/desc/examples/HSX_output.h5 b/desc/examples/HSX_output.h5
diff --git a/desc/examples/NCSX b/desc/examples/NCSX
@@ -20,7 +20,7 @@ pert_order = 2.0
 # solver tolerances
 ftol = 1e-2
 xtol = 1e-6
-gtol = 1e-2
+gtol = 1e-8
 maxiter = 50
 
 # solver methods

diff --git a/desc/examples/NCSX_output.h5 b/desc/examples/NCSX_output.h5
diff --git a/desc/examples/SOLOVEV_output.h5 b/desc/examples/SOLOVEV_output.h5
diff --git a/desc/examples/W7-X b/desc/examples/W7-X
@@ -6,25 +6,19 @@ NFP =  5
 Psi = -2.13300000E+00
 
 # spectral resolution
-L_rad = 12,16,24
+L_rad  = 12
 M_pol  = 12
-N_tor  =   12
-M_grid =  16
-N_grid =   16
+N_tor  = 12
+M_grid = 16
+N_grid = 16
 
 # continuation parameters
-# because we specify these and a sequence for L_rad, it will not use the automatic continuation
-bdry_ratio=1.0
 pert_order = 2
 
 
 # solver methods
 optimizer         = lsq-exact
 objective         = force
-spectral_indexing = fringe
-ftol=1e-4
-xtol= 1e-6
-maxiter=150
 
 # pressure and rotational transform profiles
 l:   0	p =   1.85596929e+05	i =  -8.56047021e-01

diff --git a/desc/examples/W7-X_output.h5 b/desc/examples/W7-X_output.h5
diff --git a/desc/examples/WISTELL-A_output.h5 b/desc/examples/WISTELL-A_output.h5
diff --git a/desc/examples/precise_QA_output.h5 b/desc/examples/precise_QA_output.h5
diff --git a/desc/examples/precise_QH_output.h5 b/desc/examples/precise_QH_output.h5
diff --git a/desc/examples/regenerate_all_equilibria.py b/desc/examples/regenerate_all_equilibria.py
@@ -4,43 +4,67 @@
 
 python3 regenerate_all_equilibria.py
 
-Copy the two lines below to the main code to run it on a GPU
-from desc import set_device
-set_device("gpu")
 """
+import argparse
 import glob
 import os
 import subprocess as spr
+import sys
+from pathlib import Path
 
-from desc.__main__ import main
-from desc.io import load
-from desc.vmec import VMECIO
+parser = argparse.ArgumentParser(
+    prog="regenerate_all_equilibria",
+)
+parser.add_argument(
+    "-o",
+    "--output",
+    metavar="output_dir",
+    help="Path to output files. If not specified, defaults to input dir",
+)
+parser.add_argument(
+    "--gpu",
+    "-g",
+    action="store_true",
+    help="Use GPU if available. If more than one are available, selects the "
+    + "GPU with most available memory. ",
+)
+parser.add_argument(
+    "--noprecise",
+    action="store_true",
+    help="Don't regenerate precise QS equilibria.",
+)
+
+
+args = parser.parse_args(sys.argv[1:])
 
 pwd = os.getcwd()
+if args.output:
+    output_path = os.path.abspath(args.output)
+else:
+    output_path = pwd
+
+Path(output_path).mkdir(parents=True, exist_ok=True)
+
+files = sorted(glob.glob(pwd + "/*.h5"))
+print("Regenerating files:")
+for f in files:
+    print(f)
+print("saving to: ", output_path)
+names = [f.split("/")[-1].split(".")[0].replace("_output", "") for f in files]
+names = [f for f in names if "precise" not in f]
+
+for fname in names:
+    print(f"Running the input file {fname} \n")
+    cargs = [
+        "desc",
+        fname,
+        "-vv",
+        f"-o {os.path.join(output_path, fname + '_output.h5')}",
+    ]
+    if args.gpu:
+        cargs += ["-g"]
+    spr.run(cargs)
 
-for fname in glob.glob(pwd + "/*.h5"):
-    if (
-        fname.split(".")[-1] == "h5"
-        and os.path.isfile(fname.split(".")[-1] + ".nc") is True
-        and fname.split(".")[0].split("/")[-1].split("_")[0] != "precise"
-    ):
-        finputname = fname.split(".")[0].split("/")[-1].split("_")[0]
-        print(f"Running the input file {finputname} \n")
-        main(cl_args=[str(f"{finputname}"), "-vv"])
-        # save wout file
-        eq = load(f"{pwd}/{finputname}")[-1]
-        VMECIO.save(eq, "wout_" + fname + ".nc", surfs=256)
-    elif (
-        fname.split(".")[-1] == "h5"
-        and os.path.isfile(fname + ".nc") is False
-        and fname.split(".")[0].split("/")[-1].split("_")[0] != "precise"
-    ):
-        finputname = fname.split(".")[0].split("/")[-1].split("_")[0]
-        print(f"Running the input file {finputname} \n")
-        main(cl_args=[str(f"{finputname}"), "-vv"])
-    else:
-        continue
-
-main(cl_args=["DSHAPE_CURRENT", "-vv"])
-spr.call(["python3 -u precise_QA.py"], shell=True)
-spr.call(["python3 -u precise_QH.py"], shell=True)
+if not args.noprecise:
+    spr.run(["python3", "-u", "precise_QA.py"])
+    spr.run(["python3", "-u", "precise_QH.py"])
diff --git a/tests/baseline/test_1d_elongation.png b/tests/baseline/test_1d_elongation.png
diff --git a/tests/baseline/test_2d_g_tz.png b/tests/baseline/test_2d_g_tz.png
diff --git a/tests/baseline/test_2d_logF.png b/tests/baseline/test_2d_logF.png
diff --git a/tests/baseline/test_effective_ripple.png b/tests/baseline/test_effective_ripple.png
diff --git a/tests/baseline/test_fsa_F_normalized.png b/tests/baseline/test_fsa_F_normalized.png
diff --git a/tests/baseline/test_plot_b_mag.png b/tests/baseline/test_plot_b_mag.png
diff --git a/tests/baseline/test_plot_boozer_modes_breaking_only.png b/tests/baseline/test_plot_boozer_modes_breaking_only.png
diff --git a/tests/baseline/test_plot_boozer_modes_no_norm.png b/tests/baseline/test_plot_boozer_modes_no_norm.png
diff --git a/tests/baseline/test_plot_boozer_surface.png b/tests/baseline/test_plot_boozer_surface.png
diff --git a/tests/baseline/test_plot_coefficients.png b/tests/baseline/test_plot_coefficients.png
diff --git a/tests/baseline/test_plot_normF_2d.png b/tests/baseline/test_plot_normF_2d.png
diff --git a/tests/baseline/test_plot_normF_section.png b/tests/baseline/test_plot_normF_section.png
diff --git a/tests/baseline/test_plot_poincare.png b/tests/baseline/test_plot_poincare.png
diff --git a/tests/baseline/test_section_F.png b/tests/baseline/test_section_F.png
diff --git a/tests/baseline/test_section_logF.png b/tests/baseline/test_section_logF.png
diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl
diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py
@@ -293,9 +293,10 @@ def test_limit_continuity(self):
             "iota_num_rr": {"atol": 5e-5},
             "grad(B)": {"rtol": 1e-4},
             "alpha_r (secular)": {"atol": 1e-4},
-            "grad(alpha) (secular)": {"atol": 1e-4},
+            "grad(alpha) (secular)": {"atol": 2e-4},
             "gbdrift (secular)": {"atol": 1e-4},
             "gbdrift (secular)/phi": {"atol": 1e-4},
+            "(psi_r/sqrt(g))_rr": {"rtol": 2e-5},
         }
         zero_map = dict.fromkeys(zero_limits, {"desired_at_axis": 0})
         kwargs = weaker_tolerance | zero_map

diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py
@@ -1589,13 +1589,12 @@ def test_bootstrap_optimization_comparison_qa():
         objective=objective,
         constraints=constraints,
         optimizer="proximal-lsq-exact",
-        maxiter=5,
-        gtol=1e-16,
+        maxiter=10,
         verbose=3,
     )
 
     # method 2
-    niters = 3
+    niters = 4
     for k in range(niters):
         eq2 = eq2.copy()
         data = eq2.compute("current Redl", grid)
@@ -1616,11 +1615,11 @@ def test_bootstrap_optimization_comparison_qa():
     data2 = eq2.compute(["<J*B> Redl", "<J*B>"], grid)
 
     np.testing.assert_allclose(
-        grid.compress(data1["<J*B>"]), grid.compress(data1["<J*B> Redl"]), rtol=2.1e-2
+        grid.compress(data1["<J*B>"]), grid.compress(data1["<J*B> Redl"]), rtol=2e-2
     )
     np.testing.assert_allclose(
-        grid.compress(data2["<J*B>"]), grid.compress(data2["<J*B> Redl"]), rtol=1.8e-2
+        grid.compress(data2["<J*B>"]), grid.compress(data2["<J*B> Redl"]), rtol=2e-2
     )
     np.testing.assert_allclose(
-        grid.compress(data1["<J*B>"]), grid.compress(data2["<J*B>"]), rtol=1.9e-2
+        grid.compress(data1["<J*B>"]), grid.compress(data2["<J*B>"]), rtol=2e-2
     )
diff --git a/tests/test_neoclassical.py b/tests/test_neoclassical.py
@@ -46,7 +46,7 @@ def test_fieldline_average():
     np.testing.assert_allclose(
         data["fieldline length"] / data["fieldline length/volume"],
         data["V_r(r)"] / (4 * np.pi**2),
-        rtol=1e-3,
+        rtol=2e-3,
     )
     assert np.all(data["fieldline length"] > 0)
     assert np.all(data["fieldline length/volume"] > 0)

diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py
@@ -373,7 +373,7 @@ def test_qh_boozer(self):
         # check that the objective returns the lowest amplitudes
         # 120 ~ smallest amplitudes BEFORE QH modes show up so that sorting both arrays
         # should have the same values up until then
-        np.testing.assert_allclose(f[idx_f][:120], B_mn[idx_B][:120])
+        np.testing.assert_allclose(f[idx_f][:120], B_mn[idx_B][:120], rtol=1e-6)
 
     @pytest.mark.unit
     def test_qh_boozer_multiple_surfaces(self):
@@ -2862,7 +2862,7 @@ def test_compute_scalar_resolution_others(self, objective):
             )
             obj.build(verbose=0)
             f[i] = obj.compute_scalar(obj.x())
-        np.testing.assert_allclose(f, f[-1], rtol=5e-2)
+        np.testing.assert_allclose(f, f[-1], rtol=6e-2)
 
     @pytest.mark.regression
     @pytest.mark.parametrize(

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -1416,6 +1416,6 @@ def test_optimize_coil_currents(DummyCoilSet):
     )
     # check that optimized coil currents changed by more than 15% from initial values
     np.testing.assert_array_less(
-        np.asarray(coils.current) * 0.15,
-        np.abs(np.asarray(coils_opt.current) - np.asarray(coils.current)),
+        np.asarray(coils.current).mean() * 0.15,
+        np.abs(np.asarray(coils_opt.current) - np.asarray(coils.current)).mean(),
     )
diff --git a/tests/test_stability_funs.py b/tests/test_stability_funs.py
@@ -160,7 +160,7 @@ def test(eq, vmec, rho_range=DEFAULT_RANGE, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL
     test(
         desc.examples.get("HELIOTRON"),
         ".//tests//inputs//wout_HELIOTRON.nc",
-        (0.01, 0.45),
+        (0.07, 0.45),
         rtol=1.75e-1,
     )
     test(
@@ -512,11 +512,11 @@ def test_ballooning_geometry(tmpdir_factory):
         )
         cvdrift_alt = -sign_psi * data["cvdrift"] * 2 * Bref * Lref**2 * np.sqrt(psi)
 
-        np.testing.assert_allclose(gds2, gds2_alt, rtol=4e-3)
+        np.testing.assert_allclose(gds2, gds2_alt, rtol=6e-3)
         np.testing.assert_allclose(gds22, gds22_alt)
         # gds21 is a zero crossing quantity, rtol won't work,
         # shifting the grid slightly can change rtol requirement significantly
-        np.testing.assert_allclose(gds21, gds21_alt, atol=1e-3)
+        np.testing.assert_allclose(gds21, gds21_alt, atol=2e-3)
         np.testing.assert_allclose(gbdrift, gbdrift_alt)
         # cvdrift is a zero crossing quantity, rtol won't work,
         # shifting the grid slightly can change rtol requirement significantly