"Shadow" algorithm shouldn't actually be needed because B-field calcu…

…lation uses B_direct, does not actually use the A matrix. However, dipole grid somehow gets better objective with ExactField than exact optimization does, which seems like a bug.
hiddenSymmetries · Nov 10, 2024 · 00f6fff · 00f6fff
1 parent 56a2a2a
commit 00f6fff
Show file tree

Hide file tree

Showing 8 changed files with 171 additions and 148 deletions.
diff --git a/examples/2_Intermediate/dipole_QA.py b/examples/2_Intermediate/dipole_QA.py
@@ -21,7 +21,7 @@
     ntheta = nphi
     dr = 0.05  # cylindrical bricks with radial extent 5 cm
 else:
-    nphi = 16  # nphi = ntheta >= 64 needed for accurate full-resolution runs
+    nphi = 32  # 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?
@@ -86,20 +86,15 @@
 pm_opt = PermanentMagnetGrid.geo_setup_between_toroidal_surfaces(
     s, Bnormal, s_inner, s_outer, **kwargs_geo
 )
-kwargs_geo = {"Nx": Nx}  # will get popped out, so I think kwargs needs to be reset like this
-pm_shadow = ExactMagnetGrid.geo_setup_between_toroidal_surfaces(
-    s, Bnormal, s_inner, s_outer, **kwargs_geo
-)
 
 # Optimize the permanent magnets. This actually solves
 kwargs = initialize_default_kwargs('GPMO')
-nIter_max = 5000
+nIter_max = 10000
 # algorithm = 'baseline'
-algorithm = 'baseline_shadow'
+algorithm = 'baseline'
 nHistory = 20
 kwargs['K'] = nIter_max
 kwargs['nhistory'] = nHistory
-kwargs['A_shadow'] = pm_shadow.A_obj
 print('kwargs = ',kwargs)
 t1 = time.time()
 R2_history, Bn_history, m_history = GPMO(pm_opt, algorithm, **kwargs)
@@ -111,6 +106,8 @@
 dipoles = pm_opt.m.reshape(pm_opt.ndipoles, 3)
 print('Volume of permanent magnets is = ', np.sum(np.sqrt(np.sum(dipoles ** 2, axis=-1))) / M_max)
 print('sum(|m_i|)', np.sum(np.sqrt(np.sum(dipoles ** 2, axis=-1))))
+print('dip_Ms = ',pm_opt.m)
+
 b_dipole = DipoleField(
     pm_opt.dipole_grid_xyz,
     pm_opt.m,
@@ -146,20 +143,20 @@
 
 # Print optimized f_B and other metrics
 print('B_field shape = ',b_dipole.B().shape)
-print('B_field = ',b_dipole.B())
+# print('B_field = ',b_dipole.B())
 f_B_sf = SquaredFlux(s_plot, b_dipole, -Bnormal).J()
 
-print('s_plot = ',s_plot)
-print('b_dipole = ',b_dipole)
-print('- Bnorm = ',-Bnormal)
+# print('s_plot = ',s_plot)
+# print('b_dipole = ',b_dipole)
+# print('- Bnorm = ',-Bnormal)
 
 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)
 
 
 # field for cubic magents in dipole optimization positions
-b_test = ExactField(
+b_comp = ExactField(
     pm_opt.dipole_grid_xyz,
     pm_opt.m,
     pm_opt.dims,
@@ -169,27 +166,29 @@
     m_maxima = pm_opt.m_maxima
 )
 
-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('dip_grid = ',pm_opt.dipole_grid_xyz)
+
+b_comp.set_points(s_plot.gamma().reshape((-1, 3)))
+b_comp._toVTK(out_dir / "comp_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 order to get a correct fB, have to have an exact grid that compows 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)
+make_Bnormal_plots(b_comp, s_plot, out_dir, "only_m__comp_optimized")
+s_plot.to_vtk(out_dir / "m_comp_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_Field_comp shape = ',b_comp.B().shape)
+# print('B_Field_comp = ',b_comp.B())
+f_Btest_sf = SquaredFlux(s_plot, b_comp, -Bnormal).J()
 
-print('b_test = ',b_test)
+print('b_comp = ',b_comp)
 
-print('f_testB = ', f_Btest_sf)
+print('f_B_comp = ', f_Btest_sf)
 
 # b_dipole = DipoleField(
 #     pm_opt.dipole_grid_xyz,

diff --git a/examples/2_Intermediate/exact_CSX.py b/examples/2_Intermediate/exact_CSX.py
@@ -23,12 +23,12 @@
     ntheta = nphi
     dx = 0.05  # bricks with radial extent 5 cm
 else:
-    nphi = 64  # nphi = ntheta >= 64 needed for accurate full-resolution runs
+    nphi = 32  # nphi = ntheta >= 64 needed for accurate full-resolution runs
     ntheta = nphi
-    Nx = 30 # bricks with radial extent ??? cm
+    Nx = 60 # bricks with radial extent ??? cm
 
 coff = 0.1  # PM grid starts offset ~ 10 cm from the plasma surface
-poff = 0.1  # PM grid end offset ~ 15 cm from the plasma surface
+poff = 0.12  # PM grid end offset ~ 15 cm from the plasma surface
 
 TEST_DIR = "../../tests/test_files/"
 # configuration_name = "CSX_5.0_WPs"
@@ -134,9 +134,9 @@
 # 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'
+# algorithm = 'ArbVec_backtracking'
 nBacktracking = 50
 nAdjacent = 10
 thresh_angle = np.pi  # / np.sqrt(2)
@@ -168,7 +168,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,

diff --git a/examples/2_Intermediate/exact_QA.py b/examples/2_Intermediate/exact_QA.py
@@ -21,15 +21,15 @@
     ntheta = nphi
     dx = 0.05  # bricks with radial extent 5 cm
 else:
-    nphi = 16  # nphi = ntheta >= 64 needed for accurate full-resolution runs
+    nphi = 32  # 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
+max_nMagnets = 10000
 
 # Read in the plas/ma equilibrium file
 TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve()
@@ -45,7 +45,7 @@
 # Make the output directory
 out_str = f"exact_QA_nphi{nphi}_maxIter{max_nMagnets}"
 out_dir = Path(out_str)
-out_dir.mkdir(parents=True, exist_ok=False)
+out_dir.mkdir(parents=True, exist_ok=True)
 
 # initialize the coils
 base_curves, curves, coils = initialize_coils('qa', TEST_DIR, s, out_dir)
@@ -87,6 +87,8 @@
     s, Bnormal, s_inner, s_outer, **kwargs_geo
 )
 
+print(pm_opt.phiThetas)
+
 # Optimize the permanent magnets. This actually solves
 kwargs = initialize_default_kwargs('GPMO')
 # nIter_max = 50000
@@ -108,7 +110,7 @@
     # Below line required for the backtracking to be backwards 
     # compatible with the PermanentMagnetGrid class
     # pm_opt.coordinate_flag = 'cartesian'  
-nHistory = 100
+nHistory = 20
 kwargs['nhistory'] = nHistory
 t1 = time.time()
 R2_history, Bn_history, m_history = GPMO(pm_opt, algorithm, **kwargs)

diff --git a/src/simsopt/geo/permanent_magnet_grid.py b/src/simsopt/geo/permanent_magnet_grid.py
@@ -390,6 +390,7 @@ def geo_setup_between_toroidal_surfaces(
             contig(pm_grid.xyz_outer))
         inds = np.ravel(np.logical_not(np.all(pm_grid.dipole_grid_xyz == 0.0, axis=-1)))
         pm_grid.dipole_grid_xyz = pm_grid.dipole_grid_xyz[inds, :]
+        print('DIPOLE GRID LOOKS LIKE ',pm_grid.dipole_grid_xyz)
         pm_grid.ndipoles = pm_grid.dipole_grid_xyz.shape[0]
         pm_grid.pm_phi = np.arctan2(pm_grid.dipole_grid_xyz[:, 1], pm_grid.dipole_grid_xyz[:, 0])
         if coordinate_flag == 'cylindrical':
@@ -434,7 +435,6 @@ 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):
@@ -904,6 +904,7 @@ def geo_setup_between_toroidal_surfaces(
             contig(pm_grid.xyz_outer))
         inds = np.ravel(np.logical_not(np.all(pm_grid.pm_grid_xyz == 0.0, axis=-1)))
         pm_grid.pm_grid_xyz = pm_grid.pm_grid_xyz[inds, :]
+        print('GRID LOOKS LIKE',pm_grid.pm_grid_xyz)
         pm_grid.ndipoles = pm_grid.pm_grid_xyz.shape[0]
         pm_grid.pm_phi = np.arctan2(pm_grid.pm_grid_xyz[:, 1], pm_grid.pm_grid_xyz[:, 0])
         cell_vol = pm_grid.dx * pm_grid.dy * pm_grid.dz * np.ones(pm_grid.ndipoles)     

diff --git a/src/simsopt/solve/permanent_magnet_optimization.py b/src/simsopt/solve/permanent_magnet_optimization.py
@@ -376,11 +376,12 @@ def GPMO(pm_opt, algorithm='baseline', **kwargs):
     mmax_vec = contig(np.array([mmax, mmax, mmax]).T.reshape(pm_opt.ndipoles * 3))
     A_obj = pm_opt.A_obj * mmax_vec
 
-    if algorithm == 'baseline_shadow' and 'A_shadow' not in kwargs:
-        raise KeyError('must provide a shadow grid to use baseline_shadow algorithm')
+    # if algorithm == 'baseline_shadow' and 'A_shadow' not in kwargs:
+    #     raise KeyError('must provide a shadow grid to use baseline_shadow algorithm')
 
-    shadow_grid = kwargs.pop("A_shadow", None) # how can I be sure locations are even the same and that they are indexed in the same way?
-    A_shad = shadow_grid * mmax_vec
+    # shadow_grid = kwargs.pop("A_shadow", None) # how can I be sure locations are even the same and that they are indexed in the same way?
+    # if shadow_grid != None:
+    #     A_shad = shadow_grid * mmax_vec
     print('KWARGS = ',kwargs)
 
     if (algorithm != 'baseline' and algorithm != 'mutual_coherence' and algorithm != 'ArbVec' and algorithm != 'baseline_shadow') and 'dipole_grid_xyz' not in kwargs:
@@ -413,15 +414,15 @@ def GPMO(pm_opt, algorithm='baseline', **kwargs):
             normal_norms=Nnorms,
             **kwargs
         )
-    elif algorithm == 'baseline_shadow':
-        algorithm_history, Bn_history, m_history, m = sopp.GPMO_baseline_shadow(
-            A_obj=contig(A_obj.T),
-            b_obj=contig(pm_opt.b_obj),
-            mmax=np.sqrt(reg_l2)*mmax_vec,
-            normal_norms=Nnorms,
-            A_shadow = contig(A_shad.T)
-            **kwargs
-        )
+    # elif algorithm == 'baseline_shadow':
+    #     algorithm_history, Bn_history, m_history, shadowAlg_history, BnShadow_history, m = sopp.GPMO_baseline_shadow(
+    #         A_obj=contig(A_obj.T),
+    #         b_obj=contig(pm_opt.b_obj),
+    #         mmax=np.sqrt(reg_l2)*mmax_vec,
+    #         normal_norms=Nnorms,
+    #         A_shadow = contig(A_shad.T),
+    #         **kwargs
+    #     )
     elif algorithm == 'ArbVec':  # GPMO with arbitrary polarization vectors
         algorithm_history, Bn_history, m_history, m = sopp.GPMO_ArbVec(
             A_obj=contig(A_obj.T),
@@ -488,8 +489,18 @@ def GPMO(pm_opt, algorithm='baseline', **kwargs):
         m_history[:, :, i] = m_history[:, :, i] * (mmax_vec.reshape(pm_opt.ndipoles, 3))
     errors = algorithm_history[algorithm_history != 0]
     Bn_errors = Bn_history[Bn_history != 0]
+
+    # if algorithm == 'baseline_shadow':
+    #     shad_errors = shadowAlg_history[shadowAlg_history != 0]
+    #     BnShad_errors = BnShadow_history[BnShadow_history != 0]
 
     # note m = m_proxy for GPMO because this is not using relax-and-split
     pm_opt.m = np.ravel(m)
+    print('m according to GPMO = ',pm_opt.m)
     pm_opt.m_proxy = pm_opt.m
+
+    # if algorithm == 'baseline_shadow':
+    #     return errors, Bn_errors, shad_errors, BnShad_errors, m_history
+
+    # else:
     return errors, Bn_errors, m_history