Added new GPMO algorithm that includes "shadow" grid. Working out bugs

examples/2_Intermediate/dipole_QA.py

@@ -7,7 +7,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from simsopt.field import BiotSavart, DipoleField, ExactField
-from simsopt.geo import PermanentMagnetGrid, SurfaceRZFourier
+from simsopt.geo import PermanentMagnetGrid, SurfaceRZFourier, ExactMagnetGrid
 from simsopt.objectives import SquaredFlux
 from simsopt.solve import GPMO
 from simsopt.util import in_github_actions
@@ -86,14 +86,20 @@
 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
-algorithm = 'baseline'
+# algorithm = 'baseline'
+algorithm = 'baseline_shadow'
 nHistory = 20
 kwargs['K'] = nIter_max
 kwargs['nhistory'] = nHistory
+kwargs['shadow_grid'] = pm_shadow
 t1 = time.time()
 R2_history, Bn_history, m_history = GPMO(pm_opt, algorithm, **kwargs)
 

src/simsopt/solve/permanent_magnet_optimization.py

@@ -376,7 +376,13 @@ 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' and algorithm != 'mutual_coherence' and algorithm != 'ArbVec') and 'dipole_grid_xyz' not in kwargs:
+    if algorithm == 'baseline_shadow' and 'shadow_grid' not in kwargs:
+        raise KeyError('must provide shadow_grid to use baseline_shadow algorithm')
+
+    shadow_grid = kwargs.pop("shadow_grid", None) # how can I be sure locations are even the same and that they are indexed in the same way?
+    A_shad = shadow_grid.A_obj * mmax_vec
+
+    if (algorithm != 'baseline' and algorithm != 'mutual_coherence' and algorithm != 'ArbVec' and algorithm != 'baseline_shadow') and 'dipole_grid_xyz' not in kwargs:
         raise ValueError('GPMO variants require dipole_grid_xyz to be defined.')
 
     # Set the L2 regularization if it is included in the kwargs 
@@ -406,6 +412,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 == 'ArbVec':  # GPMO with arbitrary polarization vectors
         algorithm_history, Bn_history, m_history, m = sopp.GPMO_ArbVec(
             A_obj=contig(A_obj.T),

src/simsoptpp/permanent_magnet_optimization.cpp

@@ -1329,3 +1329,114 @@ std::tuple<Array, Array, Array, Array> GPMO_baseline(Array& A_obj, Array& b_obj,
     }
     return std::make_tuple(objective_history, Bn_history, m_history, x);
 }
+
+// uses a "shadow" matrix that takes magnets in all the same locations as the active one used by GPMO. Allows for direct comparison
+// between dipole and exact optimizations
+std::tuple<Array, Array, Array, Array, Array> GPMO_baseline_shadow(Array& A_obj, Array& b_obj, Array& mmax, Array& normal_norms, Array& A_shadow, int K, bool verbose, int nhistory, int single_direction) 
+{
+    if (A_obj.shape() != A_shadow.shape()) {
+        std::invalid_argument("A dipole and A shadow are not the samae shape");
+    }
+
+    int ngrid = A_obj.shape(1);
+    int N = int(A_obj.shape(0) / 3);
+    int N3 = 3 * N;
+    int print_iter = 0;
+
+    Array x = xt::zeros<double>({N, 3});
+
+    // record the history of the algorithm iterations
+    Array m_history = xt::zeros<double>({N, 3, nhistory + 1});
+    Array objective_history = xt::zeros<double>({nhistory + 1});
+    Array Bn_history = xt::zeros<double>({nhistory + 1});
+
+    Array objShad_hist = xt::zeros<double>({nhistory + 1});
+    Array BnShad_hist = xt::zeros<double>({nhistory + 1});
+
+    // print out the names of the error columns
+    if (verbose)
+        printf("Iteration ... |Am - b|^2 ... lam*|m|^2\n");
+
+    // initialize Gamma_complement with all indices available
+    Array Gamma_complement = xt::ones<bool>({N, 3});
+
+    // initialize least-square values to large numbers    
+    vector<double> R2s(6 * N, 1e50);
+    vector<int> skj(K);
+    vector<int> skjj(K);
+    vector<double> sign_fac(K);
+
+    double* R2s_ptr = &(R2s[0]);
+    double* Aij_ptr = &(A_obj(0, 0));
+    double* Gamma_ptr = &(Gamma_complement(0, 0));
+    double* Aij_shad_ptr = &(A_shadow(0, 0));
+
+    // initialize running matrix-vector product
+    Array Aij_mj_sum = -b_obj;
+    double mmax_sum = 0.0;
+    double* Aij_mj_ptr = &(Aij_mj_sum(0));
+    double* normal_norms_ptr = &(normal_norms(0));
+    double* mmax_ptr = &(mmax(0));
+
+    Array Aij_shad_mj_sum = -b_obj;
+    double* Aij_shad_mj_ptr = &(Aij_shad_mj_sum);
+
+    // if using a single direction, increase j by 3 each iteration
+    int j_update = 1;
+    if (single_direction >= 0) j_update = 3;
+
+    // Main loop over the optimization iterations
+    for (int k = 0; k < K; ++k) {
+#pragma omp parallel for schedule(static)
+	for (int j = std::max(0, single_direction); j < N3; j += j_update) {
+
+	    // Check all the allowed dipole positions
+	    if (Gamma_ptr[j]) {
+		double R2 = 0.0;
+		double R2minus = 0.0;
+		int nj = ngrid * j;
+
+		// Compute contribution of jth dipole component, either with +- orientation
+		for(int i = 0; i < ngrid; ++i) {
+		    R2 += (Aij_mj_ptr[i] + Aij_ptr[i + nj]) * (Aij_mj_ptr[i] + Aij_ptr[i + nj]);
+		    R2minus += (Aij_mj_ptr[i] - Aij_ptr[i + nj]) * (Aij_mj_ptr[i] - Aij_ptr[i + nj]); 
+		}
+		R2s_ptr[j] = R2 + (mmax_ptr[j] * mmax_ptr[j]);
+		R2s_ptr[j + N3] = R2minus + (mmax_ptr[j] * mmax_ptr[j]);
+	    }
+	}
+
+	// find the dipole that most minimizes the least-squares term
+        skj[k] = int(std::distance(R2s.begin(), std::min_element(R2s.begin(), R2s.end())));
+	if (skj[k] >= N3) {
+	    skj[k] -= N3;
+	    sign_fac[k] = -1.0;
+	}
+	else {
+            sign_fac[k] = 1.0;
+	}
+	skjj[k] = (skj[k] % 3); 
+	skj[k] = int(skj[k] / 3.0);
+        x(skj[k], skjj[k]) = sign_fac[k];
+
+	// Add binary magnet and get rid of the magnet (all three components)
+        // from the complement of Gamma 
+	int skj_inds = (3 * skj[k] + skjj[k]) * ngrid;
+#pragma omp parallel for schedule(static)
+	for(int i = 0; i < ngrid; ++i) {
+            Aij_mj_ptr[i] += sign_fac[k] * Aij_ptr[i + skj_inds];
+            Aij_shad_mj_ptr[i] += sign_fac[k] * Aij_shad_ptr[i + skj_inds];
+	}
+        for (int j = 0; j < 3; ++j) {
+            Gamma_complement(skj[k], j) = false;
+	    R2s[3 * skj[k] + j] = 1e50;
+	    R2s[N3 + 3 * skj[k] + j] = 1e50;
+        }
+
+	if (verbose && (((k % int(K / nhistory)) == 0) || k == 0 || k == K - 1)) {
+            print_GPMO(k, ngrid, print_iter, x, Aij_mj_ptr, objective_history, Bn_history, m_history, mmax_sum, normal_norms_ptr);
+            print_GPMO(k, ngrid, print_iter, x, Aij_shad_mj_ptr, objShad_hist, BnShad_hist, m_history, mmax_sum, normal_norms_ptr);
+	}
+    }
+    return std::make_tuple(objective_history, Bn_history, m_history, onjShad_hist, BnShad_hist, x);
+}

src/simsoptpp/permanent_magnet_optimization.h

@@ -29,6 +29,7 @@ std::tuple<Array, Array, Array, Array, Array> GPMO_ArbVec_backtracking(
     Array& dipole_grid_xyz, int Nadjacent, double thresh_angle, 
     int max_nMagnets, Array& x_init);
 std::tuple<Array, Array, Array, Array> GPMO_baseline(Array& A_obj, Array& b_obj, Array&mmax, Array& normal_norms, int K, bool verbose, int nhistory, int single_direction);
+std::tuple<Array, Array, Array, Array, Array> GPMO_baseline_shadow(Array& A_obj, Array& b_obj, Array& mmax, Array& normal_norms, Array& A_shadow, int K, bool verbose, int nhistory, int single_direction);
 
 // helper functions for GPMO algorithm
 void print_GPMO(int k, int ngrid, int& print_iter, Array& x, double* Aij_mj_ptr, Array& objective_history, Array& Bn_history, Array& m_history, double mmax_sum, double* normal_norms_ptr); 

src/simsoptpp/python.cpp

@@ -87,6 +87,7 @@ PYBIND11_MODULE(simsoptpp, m) {
     m.def("GPMO_ArbVec", &GPMO_ArbVec, py::arg("A_obj"), py::arg("b_obj"), py::arg("mmax"), py::arg("normal_norms"), py::arg("pol_vectors"), py::arg("K") = 1000, py::arg("verbose") = false, py::arg("nhistory") = 100);
     m.def("GPMO_ArbVec_backtracking", &GPMO_ArbVec_backtracking, py::arg("A_obj"), py::arg("b_obj"), py::arg("mmax"), py::arg("normal_norms"), py::arg("pol_vectors"), py::arg("K") = 1000, py::arg("verbose") = false, py::arg("nhistory") = 100, py::arg("backtracking") = 100, py::arg("dipole_grid_xyz"), py::arg("Nadjacent") = 7, py::arg("thresh_angle") = 3.1415926535897931, py::arg("max_nMagnets"), py::arg("x_init"));
     m.def("GPMO_baseline", &GPMO_baseline, py::arg("A_obj"), py::arg("b_obj"), py::arg("mmax"), py::arg("normal_norms"), py::arg("K") = 1000, py::arg("verbose") = false, py::arg("nhistory") = 100, py::arg("single_direction") = -1);
+    m.def("GPMO_baseline_shadow", &GPMO_baseline_shadow, py::arg("A_obj"), py::arg("b_obj"), py::arg("mmax"), py::arg("normal_norms"), py::arg("A_shadow"), py::arg("K") = 1000, py::arg("verbose") = false, py::arg("nhistory") = 100, py::arg("single_direction") = -1);
 
     m.def("DommaschkB" , &DommaschkB);
     m.def("DommaschkdB", &DommaschkdB);