Added plots to dipole_QA.py that plot exact fields resulting from mag…

…nets in the positions selected by the dipole optimization. Will add a "shadow" exact grid that follows the dipole grid, placing magnets in all the same locations. Allows calculation of the actual objective for direct comparison to exact optimizer

examples/2_Intermediate/dipole_QA.py

@@ -6,7 +6,7 @@
 from pathlib import Path
 import numpy as np
 import matplotlib.pyplot as plt
-from simsopt.field import BiotSavart, DipoleField
+from simsopt.field import BiotSavart, DipoleField, ExactField
 from simsopt.geo import PermanentMagnetGrid, SurfaceRZFourier
 from simsopt.objectives import SquaredFlux
 from simsopt.solve import GPMO
@@ -21,7 +21,7 @@
     ntheta = nphi
     dr = 0.05  # cylindrical bricks with radial extent 5 cm
 else:
-    nphi = 32  # nphi = ntheta >= 64 needed for accurate full-resolution runs
+    nphi = 16  # nphi = ntheta >= 64 needed for accurate full-resolution runs
     ntheta = nphi
     # dr = 0.02  # cylindrical bricks with radial extent 2 cm
     # how do I manipulate the density of the dipole grid?
@@ -149,16 +149,51 @@
 print('f_B = ', f_B_sf)
 total_volume = np.sum(np.sqrt(np.sum(pm_opt.m.reshape(pm_opt.ndipoles, 3) ** 2, axis=-1))) * s.nfp * 2 * mu0 / B_max
 print('Total volume = ', total_volume)
-b_dipole = DipoleField(
+
+
+# field for cubic magents in dipole optimization positions
+b_test = ExactField(
     pm_opt.dipole_grid_xyz,
     pm_opt.m,
-    nfp=s.nfp,
-    coordinate_flag=pm_opt.coordinate_flag,
-    m_maxima=pm_opt.m_maxima
+    pm_opt.dims,
+    pm_opt.phiThetas,
+    nfp = s.nfp,
+    stellsym = s.stellsym,
+    m_maxima = pm_opt.m_maxima
 )
-b_dipole.set_points(s_plot.gamma().reshape((-1, 3)))
-dipoles = pm_opt.m.reshape(pm_opt.ndipoles, 3)
-num_nonzero_sparse = np.count_nonzero(np.sum(dipoles ** 2, axis=-1)) / pm_opt.ndipoles * 100
+
+b_test.set_points(s_plot.gamma().reshape((-1, 3)))
+b_test._toVTK(out_dir / "Test_Fields", pm_opt.dx, pm_opt.dy, pm_opt.dz)
+
+# Print optimized metrics
+# print("Total test fB = ",
+#       0.5 * np.sum((pm_opt.A_obj @ pm_opt.m - pm_opt.b_obj) ** 2))
+# in order to get a correct fB, have to have an exact grid that shadows the dipole optimzation, placing magnets
+# in all the same positions. Actually, do I need this even for the correct field?
+
+# For plotting Bn on the full torus surface at the end with just the dipole fields
+make_Bnormal_plots(b_test, s_plot, out_dir, "only_m_optimized")
+s_plot.to_vtk(out_dir / "mtest_optimized", extra_data=pointData)
+
+# Print optimized f_B and other metrics
+print('B_testField shape = ',b_test.B().shape)
+print('B_testField = ',b_test.B())
+f_Btest_sf = SquaredFlux(s_plot, b_test, -Bnormal).J()
+
+print('b_test = ',b_test)
+
+print('f_testB = ', f_Btest_sf)
+
+# b_dipole = DipoleField(
+#     pm_opt.dipole_grid_xyz,
+#     pm_opt.m,
+#     nfp=s.nfp,
+#     coordinate_flag=pm_opt.coordinate_flag,
+#     m_maxima=pm_opt.m_maxima
+# )
+# b_dipole.set_points(s_plot.gamma().reshape((-1, 3)))
+# dipoles = pm_opt.m.reshape(pm_opt.ndipoles, 3)
+# num_nonzero_sparse = np.count_nonzero(np.sum(dipoles ** 2, axis=-1)) / pm_opt.ndipoles * 100
 
 t_end = time.time()
 print('Total time = ', t_end - t_start)

examples/2_Intermediate/exact_QA.py

@@ -21,14 +21,16 @@
     ntheta = nphi
     dx = 0.05  # bricks with radial extent 5 cm
 else:
-    nphi = 64  # nphi = ntheta >= 64 needed for accurate full-resolution runs
+    nphi = 16  # nphi = ntheta >= 64 needed for accurate full-resolution runs
     ntheta = nphi
     Nx = 60 # bricks with radial extent ??? cm
 
 coff = 0.2  # PM grid starts offset ~ 10 cm from the plasma surface
 poff = 0.1  # PM grid end offset ~ 15 cm from the plasma surface
 input_name = 'input.LandremanPaul2021_QA_lowres'
 
+max_nMagnets = 5000
+
 # Read in the plas/ma equilibrium file
 TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve()
 surface_filename = TEST_DIR / input_name
@@ -41,7 +43,7 @@
 s_outer.extend_via_projected_normal(poff + coff)
 
 # Make the output directory
-out_str = "exact_QA_nphi{64}_maxIter{20k}"
+out_str = f"exact_QA_nphi{nphi}_maxIter{max_nMagnets}"
 out_dir = Path(out_str)
 out_dir.mkdir(parents=True, exist_ok=False)
 
@@ -88,7 +90,7 @@
 # Optimize the permanent magnets. This actually solves
 kwargs = initialize_default_kwargs('GPMO')
 # nIter_max = 50000
-max_nMagnets = 20000
+# max_nMagnets = 5000
 algorithm = 'baseline'
 # algorithm = 'ArbVec_backtracking'
 # nBacktracking = 200 
@@ -125,7 +127,7 @@
     pm_opt.pm_grid_xyz,
     pm_opt.m,
     pm_opt.dims,
-    pm_opt.get_phiThetas,
+    pm_opt.phiThetas,
     stellsym=s_plot.stellsym,
     nfp=s_plot.nfp,
     m_maxima=pm_opt.m_maxima,

src/simsopt/geo/permanent_magnet_grid.py

@@ -242,6 +242,10 @@ def geo_setup_from_famus(cls, plasma_boundary, Bn, famus_filename, **kwargs):
             mu0 = 4 * np.pi * 1e-7
             cell_vol = M0s * mu0 / B_max
             pm_grid.m_maxima = B_max * cell_vol[nonzero_inds] / mu0
+
+            cubrt_vol = np.cbrt(cell_vol[0])
+            pm_grid.dims = np.array([cubrt_vol, cubrt_vol, cubrt_vol])
+            print('DIPOLE DIMS = ',pm_grid.dims)
         else:
             if isinstance(m_maxima, float):
                 pm_grid.m_maxima = m_maxima * np.ones(pm_grid.ndipoles)
@@ -396,6 +400,9 @@ def geo_setup_between_toroidal_surfaces(
                     contig(pm_grid.dipole_grid_xyz[:, 0]),
                     contig(pm_grid.dipole_grid_xyz[:, 1]),
                     contig(pm_grid.dipole_grid_xyz[:, 2]))
+
+        pm_grid.dims = np.array([pm_grid.dx, pm_grid.dy, pm_grid.dz])
+        print('DIPOLE DIMS = ',pm_grid.dims)
         if m_maxima is None:
             B_max = 1.465  # value used in FAMUS runs for MUSE
             mu0 = 4 * np.pi * 1e-7
@@ -427,6 +434,9 @@ def geo_setup_between_toroidal_surfaces(
 
     def _optimization_setup(self):
 
+        # for now, set magnets all in the same direction
+        self.phiThetas = np.repeat(np.array([[0, 0]]), self.dipole_grid_xyz.shape[0], axis = 0)
+
         if self.Bn.shape != (self.nphi, self.ntheta):
             raise ValueError('Normal magnetic field surface data is incorrect shape.')
 
@@ -776,6 +786,7 @@ def geo_setup_from_famus(cls, plasma_boundary, Bn, famus_filename, **kwargs):
             # Assumes that the magnets are all the same shape, and are perfect cubes!
             cube_root_vol = np.cbrt(cell_vol[0]) # Assume FAMUS magnets are cubes
             pm_grid.dims = np.array([cube_root_vol, cube_root_vol, cube_root_vol])
+            print('MAG DIMS = ',pm_grid.dims)
             pm_grid.m_maxima = B_max * cell_vol[nonzero_inds] / mu0
         else:
             if isinstance(m_maxima, float):
@@ -902,6 +913,7 @@ def geo_setup_between_toroidal_surfaces(
                     contig(pm_grid.pm_grid_xyz[:, 2]))
 
         pm_grid.dims = np.array([pm_grid.dx, pm_grid.dy, pm_grid.dz])
+        print('MAG DIMS = ',pm_grid.dims)
         if m_maxima is None:
             B_max = 1.465  # value used in FAMUS runs for MUSE
             mu0 = 4 * np.pi * 1e-7
@@ -936,6 +948,7 @@ def geo_setup_between_toroidal_surfaces(
 
     def _optimization_setup(self):
 
+        # for now, set all magnets in same direction
         self.phiThetas = np.repeat(np.array([[0, 0]]), self.pm_grid_xyz.shape[0], axis = 0)
 
         if self.Bn.shape != (self.nphi, self.ntheta):
@@ -1095,18 +1108,19 @@ def rescale_for_opt(self, reg_l0, reg_l1, reg_l2, nu):
         return reg_l0, reg_l1, reg_l2, nu
 
     #transform grid_xyz to toroidal to get phiThetas of each grid point
-    def get_phiThetas(self):
-        grid = self.pm_grid_xyz
-        D = len(grid)
-        phiThetas = np.zeros((D,2))
+    # def get_phiThetas(self):
+    #     grid = self.pm_grid_xyz
+    #     D = len(grid)
+    #     phiThetas = np.zeros((D,2))
 
-        for d in range(D):
-            phi = np.arctan2(grid[d][1], grid[d][0])
-            theta = np.arctan2(grid[d][2], np.sqrt(grid[d][1]**2 + grid[d][0]**2) - self.R0)
+    #     for d in range(D):
+    #         phi = np.arctan2(grid[d][1], grid[d][0])
+    #         theta = np.arctan2(grid[d][2], np.sqrt(grid[d][1]**2 + grid[d][0]**2) - self.R0)
 
-            phiThetas[d] = np.array([phi, theta])
+    #         phiThetas[d] = np.array([phi, theta])
 
-        return phiThetas
+    #     return phiThetas
+    # assume pm_grid is a perfect non-twisted torus. works for MUSE
 
 def define_a_uniform_cartesian_grid_between_two_toroidal_surfaces(
     normal_inner,