Made initial progress in a c++ function that calculates all the point…

…s for the flux calculation. Plan is to then set the Bfield as usual, then call another c++ function to do the integration.

src/simsopt/geo/psc_grid.py

@@ -240,8 +240,10 @@ def geo_setup_between_toroidal_surfaces(
             contig(psc_grid.xyz_inner), 
             contig(psc_grid.xyz_outer))
         inds = np.ravel(np.logical_not(np.all(psc_grid.grid_xyz == 0.0, axis=-1)))
-        psc_grid.grid_xyz = psc_grid.grid_xyz[inds, :]
+        psc_grid.grid_xyz = np.array(psc_grid.grid_xyz[inds, :], dtype=float)
         psc_grid.num_psc = psc_grid.grid_xyz.shape[0]
+        psc_grid.rho = np.linspace(0, psc_grid.R, N, endpoint=False)
+
         # psc_grid.pm_phi = np.arctan2(psc_grid.grid_xyz[:, 1], psc_grid.grid_xyz[:, 0])
         pointsToVTK('psc_grid',
                     contig(psc_grid.grid_xyz[:, 0]),
@@ -306,7 +308,7 @@ def geo_setup_manual(
         from simsopt.field import Current, coils_via_symmetries
 
         psc_grid = cls()
-        psc_grid.grid_xyz = points
+        psc_grid.grid_xyz = np.array(points, dtype=float)
         psc_grid.R = R
         psc_grid.a = a
         psc_grid.alphas = alphas
@@ -317,6 +319,7 @@ def geo_setup_manual(
         N = 50
         psc_grid.phi = np.linspace(0, 2 * np.pi, N, endpoint=False)
         psc_grid.dphi = psc_grid.phi[1] - psc_grid.phi[0]
+        psc_grid.rho = np.linspace(0, R, N, endpoint=False)
 
         Bn = kwargs.pop("Bn", np.zeros((1, 3)))
         Bn = np.array(Bn)
@@ -578,15 +581,29 @@ def least_squares(self, kappas):
 
     def setup_orientations(self, alphas, deltas):
         contig = np.ascontiguousarray
+        self.coil_normals = np.array(
+            [np.cos(alphas) * np.sin(deltas),
+              -np.sin(alphas),
+              np.cos(alphas) * np.cos(deltas)]
+        ).T
         t1 = time.time()
-        self.fluxes(alphas, deltas)
-        # sopp.TF_fluxes(
-        #     contig(self.grid_xyz), 
-        #     contig(alphas),
-        #     contig(deltas), 
-        #     contig(self.rho), 
-        #     contig(self.phi), 
-        #     self.I, Array& normal, double R);
+        # self.fluxes(alphas, deltas)
+        # flux_grid = np.zeros((len(alphas), 50, 50, 3))
+        # alphas = np.array(alphas, dtype=float)
+        # deltas = np.array(deltas, dtype=float)
+        # self.phi = np.array(self.phi, dtype=float)
+        # self.rho = np.array(self.rho, dtype=float)
+        # self.coil_normals = np.array(self.coil_normals, dtype=float)
+        # self.grid_xyz = np.array(self.grid_xyz, dtype=float)
+        print(self.grid_xyz.shape, alphas.shape, deltas.shape, self.rho.shape, self.phi.shape, self.coil_normals.shape)
+        flux_grid = sopp.flux_xyz(
+            contig(self.grid_xyz), 
+            contig(self.grid_xyz),
+            contig(deltas), 
+            contig(self.rho), 
+            contig(self.phi), 
+            contig(self.coil_normals)
+        )
         t2 = time.time()
         print('Fluxes time = ', t2 - t1)
         t1 = time.time()

src/simsoptpp/psc.cpp

@@ -1,4 +1,5 @@
 #include "psc.h"
+#include <cstdio>
 
 // Calculate the inductance matrix needed for the PSC forward problem
 Array L_matrix(Array& points, Array& alphas, Array& deltas, Array& phi, double R)
@@ -62,7 +63,8 @@ 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& I, Array& normal, double R)
+// Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, Array& I, Array& normal, double R)
+Array flux_xyz(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, 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)
@@ -75,8 +77,8 @@ Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array&
           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(I.layout() != xt::layout_type::row_major)
-          throw std::runtime_error("I needs to be in row-major storage order");
+//     if(I.layout() != xt::layout_type::row_major)
+//           throw std::runtime_error("I 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");
 
@@ -86,10 +88,11 @@ Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array&
     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});
-    double R2 = R * R;
-    double fac = 4e-7 * R2;
-    using namespace boost::math;
+//     Array Psi = xt::zeros<double>({num_coils});
+    Array XYZs = xt::zeros<double>({num_coils, num_rho, num_phi, 3});
+//     double R2 = R * R;
+//     double fac = 4e-7 * R2;
+//     using namespace boost::math;
 
 //     #pragma omp parallel for schedule(static)
     for(int j = 0; j < num_coils; j++) {
@@ -103,38 +106,43 @@ Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array&
         auto xj = points(j, 0);
         auto yj = points(j, 1);
         auto zj = points(j, 2);
-        double integral = 0.0;
+//         double integral = 0.0;
         for (int k = 0; k < num_rho; ++k) {
             auto xx = rho(k);
-            for (int kk = 0; k < num_phi; ++kk) {
+            for (int kk = 0; kk < num_phi; ++kk) {
+//                 printf("%d %d %d\n", j, k, 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 rho2 = (x - xj) * (x - xj) + (y - yj) * (y - yj);
-                auto r2 = rho2 + (z - zj) * (z - zj);
-                auto rho = sqrt(rho2);
-                auto R2_r2 = R2 + r2;
-                auto gamma2 = R2_r2 + 2 * R * rho;
-                auto beta2 = R2_r2 - 2 * R * rho;
-                auto beta = sqrt(beta2);
-                auto k2 = 1 - gamma2 / beta2;
-                auto beta_gamma2 = beta * gamma2;
-                auto k = sqrt(k2);
-                auto ellipe = ellint_2(k);
-                auto ellipk = ellint_1(k);
-                auto Eplus = R2_r2 * ellipe - gamma2 * ellipk;
-                auto Eminus = (R2 - r2) * ellipe + gamma2 * ellipk;
-                auto Bx = (x - xj) * (z - zj) * Eplus / (rho2 * beta_gamma2);
-                auto By = (y - yj) * (z - zj) * Eplus / (rho2 * beta_gamma2);
-                auto Bz = Eminus / beta_gamma2;
-                auto Bn = Bx * nx + By * ny + Bz * nz;
-                integral += Bn * xx;
+//                 auto rho2 = (x - xj) * (x - xj) + (y - yj) * (y - yj);
+//                 auto r2 = rho2 + (z - zj) * (z - zj);
+//                 auto rho = sqrt(rho2);
+//                 auto R2_r2 = R2 + r2;
+//                 auto gamma2 = R2_r2 + 2 * R * rho;
+//                 auto beta2 = R2_r2 - 2 * R * rho;
+//                 auto beta = sqrt(beta2);
+//                 auto k2 = 1 - gamma2 / beta2;
+//                 auto beta_gamma2 = beta * gamma2;
+//                 auto k = sqrt(k2);
+//                 auto ellipe = ellint_2(k);
+//                 auto ellipk = ellint_1(k);
+//                 auto Eplus = R2_r2 * ellipe - gamma2 * ellipk;
+//                 auto Eminus = (R2 - r2) * ellipe + gamma2 * ellipk;
+//                 auto Bx = (x - xj) * (z - zj) * Eplus / (rho2 * beta_gamma2);
+//                 auto By = (y - yj) * (z - zj) * Eplus / (rho2 * beta_gamma2);
+//                 auto Bz = Eminus / beta_gamma2;
+//                 auto Bn = Bx * nx + By * ny + Bz * nz;
+//                 integral += Bn * xx;
+                XYZs(j, k, kk, 0) = x;
+                XYZs(j, k, kk, 1) = y;
+                XYZs(j, k, kk, 2) = z;
             }
         }
-        Psi(j) = integral * fac * I(j);
+//         Psi(j) = integral * fac * I(j);
     }
-    return Psi;
+//     return Psi;
+    return XYZs;
 }

src/simsoptpp/psc.h

@@ -4,9 +4,9 @@
 // #include <tr1/cmath>
 // #define BOOST_CONFIG_SUPPRESS_OUTDATED_MESSAGE
 // #include <boost/lambda/lambda.hpp>
-#include <boost/math/special_functions/ellint_1.hpp>  
-#include <boost/math/special_functions/ellint_2.hpp>  
-#include <boost/math/special_functions/ellint_3.hpp>
+// #include <boost/math/special_functions/ellint_1.hpp>  
+// #include <boost/math/special_functions/ellint_2.hpp>  
+// #include <boost/math/special_functions/ellint_3.hpp>
 #include <cmath>
 #include <tuple>  // c++ tuples
 #include <string> // for string class
@@ -16,6 +16,7 @@ typedef xt::pyarray<double> Array;
 
 Array L_matrix(Array& points, Array& alphas, Array& deltas, Array& phi, double R);
 // 
-Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, Array& I, Array& normal, double R);
+// Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, Array& I, Array& normal, double R);
+Array flux_xyz(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, Array& normal);
 // 
 // Array A_matrix(Array& plasma_surface_normals, Array& B, Array& alphas, Array& deltas);

src/simsoptpp/python.cpp

@@ -61,7 +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);
+    m.def("flux_xyz" , &flux_xyz);
 
     // Functions below are implemented for permanent magnet optimization
     m.def("dipole_field_B" , &dipole_field_B);