Support numpy 2.0 (#1319)

PlasmaControl · Nov 15, 2024 · 35ec831 · 35ec831
2 parents cc45766 + d45f5d3
commit 35ec831
Show file tree

Hide file tree

Showing 74 changed files with 31 additions and 21 deletions.
diff --git a/.github/workflows/benchmark.yml b/.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'

diff --git a/.github/workflows/cache_dependencies.yml b/.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

diff --git a/.github/workflows/jax_tests.yml b/.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

diff --git a/.github/workflows/notebook_tests.yml b/.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'

diff --git a/.github/workflows/regression_tests.yml b/.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'

diff --git a/.github/workflows/unit_tests.yml b/.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'

diff --git a/.github/workflows/weekly_tests.yml b/.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:

diff --git a/desc/basis.py b/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)

diff --git a/desc/magnetic_fields/_core.py b/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))

diff --git a/devtools/dev-requirements_conda.yml b/devtools/dev-requirements_conda.yml
@@ -7,7 +7,7 @@ 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
   - termcolor

diff --git a/requirements.txt b/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
diff --git a/requirements_conda.yml b/requirements_conda.yml
@@ -7,7 +7,7 @@ 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
   - termcolor

diff --git a/tests/baseline/test_1d_elongation.png b/tests/baseline/test_1d_elongation.png
diff --git a/tests/baseline/test_1d_iota.png b/tests/baseline/test_1d_iota.png
diff --git a/tests/baseline/test_1d_iota_radial.png b/tests/baseline/test_1d_iota_radial.png
diff --git a/tests/baseline/test_1d_logpsi.png b/tests/baseline/test_1d_logpsi.png
diff --git a/tests/baseline/test_1d_p.png b/tests/baseline/test_1d_p.png
diff --git a/tests/baseline/test_2d_g_rz.png b/tests/baseline/test_2d_g_rz.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_2d_plot_Bn.png b/tests/baseline/test_2d_plot_Bn.png
diff --git a/tests/baseline/test_Redl_figures_2_3.png b/tests/baseline/test_Redl_figures_2_3.png
diff --git a/tests/baseline/test_Redl_figures_4_5.png b/tests/baseline/test_Redl_figures_4_5.png
diff --git a/tests/baseline/test_Redl_sfincs_QA.png b/tests/baseline/test_Redl_sfincs_QA.png
diff --git a/tests/baseline/test_Redl_sfincs_QH.png b/tests/baseline/test_Redl_sfincs_QH.png
diff --git a/tests/baseline/test_Redl_sfincs_tokamak_benchmark.png b/tests/baseline/test_Redl_sfincs_tokamak_benchmark.png
diff --git a/tests/baseline/test_binormal_drift_bounce1d.png b/tests/baseline/test_binormal_drift_bounce1d.png
diff --git a/tests/baseline/test_bounce1d_checks.png b/tests/baseline/test_bounce1d_checks.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_fsa_I.png b/tests/baseline/test_fsa_I.png
diff --git a/tests/baseline/test_plot_1d_curve.png b/tests/baseline/test_plot_1d_curve.png
diff --git a/tests/baseline/test_plot_1d_surface.png b/tests/baseline/test_plot_1d_surface.png
diff --git a/tests/baseline/test_plot_2d_surface.png b/tests/baseline/test_plot_2d_surface.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_basis_doublefourierseries.png b/tests/baseline/test_plot_basis_doublefourierseries.png
diff --git a/tests/baseline/test_plot_basis_fourierseries.png b/tests/baseline/test_plot_basis_fourierseries.png
diff --git a/tests/baseline/test_plot_basis_fourierzernike.png b/tests/baseline/test_plot_basis_fourierzernike.png
diff --git a/tests/baseline/test_plot_basis_powerseries.png b/tests/baseline/test_plot_basis_powerseries.png
diff --git a/tests/baseline/test_plot_boozer_modes.png b/tests/baseline/test_plot_boozer_modes.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_max.png b/tests/baseline/test_plot_boozer_modes_max.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_boundaries.png b/tests/baseline/test_plot_boundaries.png
diff --git a/tests/baseline/test_plot_boundary.png b/tests/baseline/test_plot_boundary.png
diff --git a/tests/baseline/test_plot_boundary_surface.png b/tests/baseline/test_plot_boundary_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_comparison.png b/tests/baseline/test_plot_comparison.png
diff --git a/tests/baseline/test_plot_comparison_different_NFPs.png b/tests/baseline/test_plot_comparison_different_NFPs.png
diff --git a/tests/baseline/test_plot_comparison_no_theta.png b/tests/baseline/test_plot_comparison_no_theta.png
diff --git a/tests/baseline/test_plot_con_basis.png b/tests/baseline/test_plot_con_basis.png
diff --git a/tests/baseline/test_plot_cov_basis.png b/tests/baseline/test_plot_cov_basis.png
diff --git a/tests/baseline/test_plot_grid_cheb1.png b/tests/baseline/test_plot_grid_cheb1.png
diff --git a/tests/baseline/test_plot_grid_cheb2.png b/tests/baseline/test_plot_grid_cheb2.png
diff --git a/tests/baseline/test_plot_grid_jacobi.png b/tests/baseline/test_plot_grid_jacobi.png
diff --git a/tests/baseline/test_plot_grid_linear.png b/tests/baseline/test_plot_grid_linear.png
diff --git a/tests/baseline/test_plot_grid_ocs.png b/tests/baseline/test_plot_grid_ocs.png
diff --git a/tests/baseline/test_plot_grid_quad.png b/tests/baseline/test_plot_grid_quad.png
diff --git a/tests/baseline/test_plot_logo.png b/tests/baseline/test_plot_logo.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_omnigenous_field.png b/tests/baseline/test_plot_omnigenous_field.png
diff --git a/tests/baseline/test_plot_poincare.png b/tests/baseline/test_plot_poincare.png
diff --git a/tests/baseline/test_plot_qs_error.png b/tests/baseline/test_plot_qs_error.png
diff --git a/tests/baseline/test_plot_surfaces.png b/tests/baseline/test_plot_surfaces.png
diff --git a/tests/baseline/test_plot_surfaces_HELIOTRON.png b/tests/baseline/test_plot_surfaces_HELIOTRON.png
diff --git a/tests/baseline/test_plot_surfaces_no_theta.png b/tests/baseline/test_plot_surfaces_no_theta.png
diff --git a/tests/baseline/test_plot_vmec_comparison.png b/tests/baseline/test_plot_vmec_comparison.png
diff --git a/tests/baseline/test_qh_optimization3.png b/tests/baseline/test_qh_optimization3.png
diff --git a/tests/baseline/test_section_F.png b/tests/baseline/test_section_F.png
diff --git a/tests/baseline/test_section_J.png b/tests/baseline/test_section_J.png
diff --git a/tests/baseline/test_section_logF.png b/tests/baseline/test_section_logF.png
diff --git a/tests/baseline/test_trapped_fraction_Kim.png b/tests/baseline/test_trapped_fraction_Kim.png
diff --git a/tests/test_stability_funs.py b/tests/test_stability_funs.py
@@ -64,7 +64,8 @@ def get_vmec_data(path, quantity):
 
     """
     f = Dataset(path)
-    rho = np.sqrt(f.variables["phi"] / np.array(f.variables["phi"])[-1])
+    phi = np.array(f.variables["phi"])
+    rho = np.sqrt(phi / phi[-1])
     q = np.array(f.variables[quantity])
     f.close()
     return rho, q