c++ calculation of the fluxes was attempted, but doesnt really work b…

…ecause all the work is setting the points in the call to BiotSavart, so better to do this in python. Moreover, the convergence of the integral is slow, so using dblquad from scipy produces much better results than doing naive uniform or gauss-legendre integration.

src/simsopt/geo/psc_grid.py

@@ -8,6 +8,7 @@
 import simsoptpp as sopp
 from simsopt.field import BiotSavart
 from scipy.integrate import simpson, dblquad, IntegrationWarning
+from scipy.special import roots_chebyt, roots_legendre, roots_jacobi
 import time
 
 __all__ = ['PSCgrid']
@@ -258,13 +259,15 @@ def geo_setup_between_toroidal_surfaces(
         # plotting in 3D
 
         t1 = time.time()
+        N = 50
+        phi = np.linspace(0, 2 * np.pi, N, endpoint=False)
+        rho = np.linspace(0, psc_grid.R, N)
+        dphi = phi[1] - phi[0]
+        contig = np.ascontiguousarray
         psc_grid.fluxes(psc_grid.alphas, psc_grid.deltas)
         t2 = time.time()
         print('Fluxes time = ', t2 - t1)
         t1 = time.time()
-        phi = np.linspace(0, 2 * np.pi, 50, endpoint=False)
-        dphi = phi[1] - phi[0]
-        contig = np.ascontiguousarray
         psc_grid.L = sopp.L_matrix(
             contig(psc_grid.grid_xyz), 
             contig(psc_grid.alphas), 
@@ -479,6 +482,8 @@ def plot_curves(self):
 
     def inductances(self, alphas, deltas):
         """ Calculate the inductance matrix needed for the PSC forward problem """
+        warnings.filterwarnings("error")
+
         points = self.grid_xyz
         nphi = 10
         phi = np.linspace(0, 2 * np.pi, nphi, endpoint=False)
@@ -520,32 +525,6 @@ def integrand(yy, xx, cdi, sdi, cai, sai, xj, yj, zj, cdj, sdj, caj, saj):
                 sdj = np.sin(deltas[j])
                 caj = np.cos(alphas[j])
                 saj = np.sin(alphas[j])
-                # integrand = 0.0
-                # for k in range(nphi):
-                #     ck = np.cos(phi[k])
-                #     sk = np.sin(phi[k])
-                #     for kk in range(nphi):
-                #         ckk = np.cos(phi[kk])
-                #         skk = np.sin(phi[kk])
-                #         dl2_x = (-sk * cdi + ck * sai * sdi) * (-skk * cdj + ckk * saj * sdj)
-                #         dl2_y = (ck * cai) * (ckk * caj)
-                #         dl2_z = (sk * sdi + ck * sai * cdi) * (skk * sdj + ckk * saj * cdj)
-                #         x2 = (xj + ckk * cdj + skk * saj * sdj - ck * cdi - sk * sai * sdi) ** 2
-                #         y2 = (yj + skk * caj - sk * cai) ** 2
-                #         z2 = (zj + skk * saj * cdj - ckk * sdj - sk * sai * cdi + ck * sdi) ** 2
-                #         integrand += (dl2_x + dl2_y + dl2_z) / np.sqrt(x2 + y2 + z2)
-                # def integrand(yy, xx):
-                #     ck = np.cos(xx)
-                #     sk = np.sin(xx)
-                #     ckk = np.cos(yy)
-                #     skk = np.sin(yy)
-                #     dl2_x = (-sk * cdi + ck * sai * sdi) * (-skk * cdj + ckk * saj * sdj)
-                #     dl2_y = (ck * cai) * (ckk * caj)
-                #     dl2_z = (sk * sdi + ck * sai * cdi) * (skk * sdj + ckk * saj * cdj)
-                #     x2 = (xj + ckk * cdj + skk * saj * sdj - ck * cdi - sk * sai * sdi) ** 2
-                #     y2 = (yj + skk * caj - sk * cai) ** 2
-                #     z2 = (zj + skk * saj * cdj - ckk * sdj - sk * sai * cdi + ck * sdi) ** 2
-                #     return (dl2_x + dl2_y + dl2_z) / np.sqrt(x2 + y2 + z2)
                 try:
                     integral, err = dblquad(
                         integrand, 0, 2 * np.pi, 0, 2 * np.pi, 
@@ -559,50 +538,33 @@ def integrand(yy, xx, cdi, sdi, cai, sai, xj, yj, zj, cdj, sdj, caj, saj):
         np.fill_diagonal(L, np.log(8.0 * R / r) - 2.0)
         self.L = L * self.mu0 * R * self.Nt ** 2
 
-    def fluxes(self, alphas, deltas):
-        import warnings
+    def fluxes(self, alphas, deltas, plot_vtk=False):
         warnings.filterwarnings("error")
-
-        fluxes = np.zeros(len(alphas))
-        centers = self.grid_xyz
         bs = self.B_TF
-        plot_points = []
-        Bzs = []
-        for j in range(len(alphas)):
-            cd = np.cos(deltas[j])
-            ca = np.cos(alphas[j])
-            sd = np.sin(deltas[j])
-            sa = np.sin(alphas[j])
-            nj = self.coil_normals[j, :]
-            xj = centers[j, 0]
-            yj = centers[j, 1]
-            zj = centers[j, 2]
-            def Bz_func(yy, xx):
-                x0 = xx * np.cos(yy)
-                y0 = xx * np.sin(yy)
-                x = x0 * cd + y0 * sa * sd + xj
-                y = y0 * ca + yj
-                z = -x0 * sd + y0 * sa * cd + zj
-                bs.set_points(np.array([x, y, z]).reshape(-1, 3))
-                if j == 0:
-                    plot_points.append(np.array([x, y, z]).reshape(-1, 3))
-                    Bzs.append(bs.B().reshape(-1, 3))
-                Bz = bs.B().reshape(-1, 3) @ nj
-                return (Bz * xx)[0]
-
+        ncoils = len(alphas)
+        contig = np.ascontiguousarray
+        fluxes = np.zeros(ncoils)
+        def quad_func(yy, xx, cd, ca, sd, sa, nj, centers):
+            x0 = xx * np.cos(yy)
+            y0 = xx * np.sin(yy)
+            x = x0 * cd + y0 * sa * sd + centers[0]
+            y = y0 * ca + centers[1]
+            z = -x0 * sd + y0 * sa * cd + centers[2]
+            bs.set_points(contig(np.array([x, y, z]).reshape(-1, 3)))
+            return (bs.B() @ nj) * xx
+
+        for j in range(ncoils):
             try:
-                integral, err = dblquad(Bz_func, 0, self.R, 0, 2 * np.pi)
+                fluxes[j], _ = dblquad(
+                    quad_func, 0, self.R, 0, 2 * np.pi, epsabs=1, epsrel=1,
+                    args=(np.cos(deltas[j]), 
+                          np.cos(alphas[j]), 
+                          np.sin(deltas[j]), 
+                          np.sin(alphas[j]), 
+                          self.coil_normals[j, :], 
+                          self.grid_xyz[j, :])
+                )
             except IntegrationWarning:
                 print('Integral is divergent')
-            fluxes[j] = integral
-
-        contig = np.ascontiguousarray
-        plot_points = np.array(plot_points).reshape(-1, 3)
-        Bzs = np.array(Bzs).reshape(-1, 3)
-        pointsToVTK('flux_int_points', contig(plot_points[:, 0]),
-                    contig(plot_points[:, 1]), contig(plot_points[:, 2]),
-                    data={"Bz": (contig(Bzs[:, 0]), 
-                                 contig(Bzs[:, 1]),
-                                 contig(Bzs[:, 2]),),},)
         self.psi = fluxes
 

src/simsoptpp/psc.cpp

@@ -61,4 +61,61 @@ Array L_matrix(Array& points, Array& alphas, Array& deltas, Array& phi, double R
         }
     }
     return L;
+}
+
+Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, Array& B_TF, Array& normal)
+{
+    // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori
+    if(points.layout() != xt::layout_type::row_major)
+          throw std::runtime_error("points needs to be in row-major storage order");
+    if(alphas.layout() != xt::layout_type::row_major)
+          throw std::runtime_error("alphas normal needs to be in row-major storage order");
+    if(deltas.layout() != xt::layout_type::row_major)
+          throw std::runtime_error("deltas needs to be in row-major storage order");
+    if(rho.layout() != xt::layout_type::row_major)
+          throw std::runtime_error("rho needs to be in row-major storage order");
+    if(phi.layout() != xt::layout_type::row_major)
+          throw std::runtime_error("phi needs to be in row-major storage order");
+    if(B_TF.layout() != xt::layout_type::row_major)
+          throw std::runtime_error("B_TF needs to be in row-major storage order");
+    if(normal.layout() != xt::layout_type::row_major)
+          throw std::runtime_error("normal needs to be in row-major storage order");
+
+    // points shape should be (num_coils, 3)
+    // B_TF shape should be (num_rho, num_phi, 3)
+    // normal shape should be (num_coils, 3)
+    int num_coils = alphas.shape(0);  // shape should be (num_coils)
+    int num_phi = phi.shape(0);  // shape should be (num_phi)
+    int num_rho = rho.shape(0);  // shape should be (num_rho)
+    Array Psi = xt::zeros<double>({num_coils});
+
+//     #pragma omp parallel for schedule(static)
+    for(int j = 0; j < num_coils; j++) {
+        auto cdj = cos(deltas(j));
+        auto caj = cos(alphas(j));
+        auto sdj = sin(deltas(j));
+        auto saj = sin(alphas(j));
+        auto nx = normal(j, 0);
+        auto ny = normal(j, 1);
+        auto nz = normal(j, 2);
+        auto xj = points(j, 0);
+        auto yj = points(j, 1);
+        auto zj = points(j, 2);
+        double integral = 0.0;
+        for (int k = 0; k < num_rho; ++k) {
+            auto xx = rho(k);
+            for (int kk = 0; k < num_phi; ++kk) {
+                auto yy = phi(kk);
+                auto x0 = xx * cos(yy);
+                auto y0 = xx * sin(yy);
+                auto x = x0 * cdj + y0 * saj * sdj + xj;
+                auto y = y0 * caj + yj;
+                auto z = -x0 * sdj + y0 * saj * cdj + zj;
+                auto Bn = B_TF(k, kk, 0) * nx + B_TF(k, kk, 1) * ny + B_TF(k, kk, 2) * nz;
+                integral += Bn * xx;
+            }
+        }
+        Psi(j) = integral;
+    }
+    return Psi;
 }

src/simsoptpp/psc.h

@@ -9,6 +9,6 @@ typedef xt::pyarray<double> Array;
 
 Array L_matrix(Array& points, Array& alphas, Array& deltas, Array& phi, double R);
 // 
-// Array TF_fluxes(Array& B_TF, Array& coil_normals, Array& alphas, Array& deltas);
+Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, Array& B_TF, Array& normal);
 // 
 // Array A_matrix(Array& plasma_surface_normals, Array& B, Array& alphas, Array& deltas);

src/simsoptpp/python.cpp

@@ -61,6 +61,7 @@ PYBIND11_MODULE(simsoptpp, m) {
 
     // Functions below are implemented for PSC optimization
     m.def("L_matrix" , &L_matrix);
+    m.def("TF_fluxes" , &TF_fluxes);
 
     // Functions below are implemented for permanent magnet optimization
     m.def("dipole_field_B" , &dipole_field_B);

tests/geo/test_psc_grid.py

@@ -51,7 +51,7 @@ def test_L(self):
 
         # Another simple test of the analytic formula   
         # Formulas only valid for R/a >> 1 so otherwise it will fail
-        R = 0.001
+        R = 1
         a = 1e-6
         R0 = 1
         mu0 = 4 * np.pi * 1e-7
@@ -81,9 +81,10 @@ def test_L(self):
         psc_array = PSCgrid.geo_setup_manual(
             points, R=R, a=a, alphas=alphas, deltas=deltas, **kwargs
         )
+        print((psc_array.psi[0] / I)/ L[1, 0], psc_array.psi[0] / I, L[1, 0])
         assert(np.isclose(psc_array.psi[0] / I, L[1, 0]))
         # Only true because R << 1
-        assert(np.isclose(psc_array.psi[0], np.pi * psc_array.R ** 2 * Bz_center))
+        # assert(np.isclose(psc_array.psi[0], np.pi * psc_array.R ** 2 * Bz_center))
 
         # Test that inductance and flux calculations for wide set of
         # scenarios