Replace np.product with np.prod; replace more asserts with np.testing…

….assert_allclose for extra diagnostics info
hiddenSymmetries · Jan 13, 2025 · f57abea · f57abea
1 parent 35553ae
commit f57abea
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/tests/core/test_optimizable.py b/tests/core/test_optimizable.py
@@ -97,7 +97,7 @@ def sum(self):
         return np.sum(self.local_full_x)
 
     def product(self):
-        return np.product(self.local_full_x)
+        return np.prod(self.local_full_x)
 
     return_fn_map = {'sum': sum, 'prod': product}
 

diff --git a/tests/geo/test_pm_grid.py b/tests/geo/test_pm_grid.py
@@ -221,12 +221,12 @@ def test_Bn(self):
             # check <Bn>
             B_opt = np.mean(np.abs(pm_opt.A_obj.dot(pm_opt.m) - pm_opt.b_obj) * np.sqrt(Ngrid / Nnorms))
             B_dipole_field = np.mean(np.abs(np.sum((bs.B() + b_dipole.B()).reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)))
-            assert np.isclose(B_opt, B_dipole_field)
+            np.testing.assert_allclose(B_opt, B_dipole_field)
 
             # check integral Bn^2
             f_B_Am = 0.5 * np.linalg.norm(pm_opt.A_obj.dot(pm_opt.m) - pm_opt.b_obj, ord=2) ** 2
             f_B = SquaredFlux(s, b_dipole, -Bn).J()
-            assert np.isclose(f_B, f_B_Am)
+            np.testing.assert_allclose(f_B, f_B_Am)
 
             # Create PM class with cylindrical bricks
             Bn = np.sum(bs.B().reshape(nphi, ntheta, 3) * s.unitnormal(), axis=-1)
@@ -235,11 +235,11 @@ def test_Bn(self):
             mmax_new = pm_opt.m_maxima / 2.0
             kwargs_geo = {"dr": 0.15, "coordinate_flag": "cylindrical", "m_maxima": mmax_new}
             pm_opt = PermanentMagnetGrid.geo_setup_between_toroidal_surfaces(s, Bn, s1, s2, **kwargs_geo)
-            assert np.allclose(pm_opt.m_maxima, mmax_new)
+            np.testing.assert_allclose(pm_opt.m_maxima, mmax_new)
             mmax_new = pm_opt.m_maxima[-1] / 2.0
             kwargs_geo = {"dr": 0.15, "coordinate_flag": "cylindrical", "m_maxima": mmax_new}
             pm_opt = PermanentMagnetGrid.geo_setup_between_toroidal_surfaces(s, Bn, s1, s2, **kwargs_geo)
-            assert np.allclose(pm_opt.m_maxima, mmax_new)
+            np.testing.assert_allclose(pm_opt.m_maxima, mmax_new)
             pm_opt = PermanentMagnetGrid.geo_setup_between_toroidal_surfaces(s, Bn, s1, s2)
             _, _, _, = relax_and_split(pm_opt)
             b_dipole = DipoleField(pm_opt.dipole_grid_xyz, pm_opt.m_proxy, nfp=s.nfp,
@@ -250,17 +250,17 @@ def test_Bn(self):
         Nnorms = np.ravel(np.sqrt(np.sum(s.normal() ** 2, axis=-1)))
         Ngrid = nphi * ntheta
         Bn_Am = (pm_opt.A_obj.dot(pm_opt.m) - pm_opt.b_obj) * np.sqrt(Ngrid / Nnorms) 
-        assert np.allclose(Bn_Am.reshape(nphi, ntheta), np.sum((bs.B() + b_dipole.B()).reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2))
+        np.testing.assert_allclose(Bn_Am.reshape(nphi, ntheta), np.sum((bs.B() + b_dipole.B()).reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2), atol=1e-15)
 
         # check <Bn>
         B_opt = np.mean(np.abs(pm_opt.A_obj.dot(pm_opt.m) - pm_opt.b_obj) * np.sqrt(Ngrid / Nnorms))
         B_dipole_field = np.mean(np.abs(np.sum((bs.B() + b_dipole.B()).reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)))
-        assert np.isclose(B_opt, B_dipole_field)
+        np.testing.assert_allclose(B_opt, B_dipole_field)
 
         # check integral Bn^2
         f_B_Am = 0.5 * np.linalg.norm(pm_opt.A_obj.dot(pm_opt.m) - pm_opt.b_obj, ord=2) ** 2
         f_B = SquaredFlux(s, b_dipole, -Bn).J()
-        assert np.isclose(f_B, f_B_Am)
+        np.testing.assert_allclose(f_B, f_B_Am)
 
     def test_grid_chopping(self):
         """