Added some initial functionality to perform minimization, but the cal…

…l to calculate the fluxes is wayyyy to slow, so need to speed this up. Tried writing up a version of this in c++ but realized I actually calculated the fluxes from the PSCs, which is not what is needed. We need to have a fast calculation of the fluxes from the fixed TF coils.

examples/2_Intermediate/psc_example.py

@@ -11,21 +11,22 @@
 from simsopt.objectives import SquaredFlux
 from simsopt.util import in_github_actions
 from simsopt.util.permanent_magnet_helper_functions import *
+import time
 
 
 # Set some parameters -- if doing CI, lower the resolution
 if in_github_actions:
     nphi = 4  # nphi = ntheta >= 64 needed for accurate full-resolution runs
     ntheta = nphi
 else:
-    nphi = 64  # nphi = ntheta >= 64 needed for accurate full-resolution runs
-    ntheta = 64
+    nphi = 32  # nphi = ntheta >= 64 needed for accurate full-resolution runs
+    ntheta = 32
     # Make higher resolution surface for plotting Bnormal
     qphi = 2 * nphi
     quadpoints_phi = np.linspace(0, 1, qphi, endpoint=True)
     quadpoints_theta = np.linspace(0, 1, ntheta, endpoint=True)
 
-coff = 0.25  # PM grid starts offset ~ 10 cm from the plasma surface
+coff = 0.2  # PM grid starts offset ~ 10 cm from the plasma surface
 poff = 0.05  # PM grid end offset ~ 15 cm from the plasma surface
 input_name = 'input.LandremanPaul2021_QA_lowres'
 
@@ -34,12 +35,29 @@
 surface_filename = TEST_DIR / input_name
 range_param = 'full torus'
 s = SurfaceRZFourier.from_vmec_input(surface_filename, range=range_param, nphi=qphi, ntheta=ntheta)
+print('s.R = ', s.rc[0, 0])
+
+# input_name = 'tests/test_files/input.circular_tokamak'
+# s_inner = SurfaceRZFourier(
+#     stellsym=False,
+#     quadpoints_phi=s.quadpoints_phi, 
+#     quadpoints_theta=s.quadpoints_theta
+# )
+# s_inner.set_rc(1, 0, 0.1)
+# s_inner.set_zs(1, 0, 0.1)
+# s_outer = SurfaceRZFourier(
+#     stellsym=False,
+#     quadpoints_phi=s.quadpoints_phi, 
+#     quadpoints_theta=s.quadpoints_theta
+# )
+# s_outer.set_rc(1, 0, 0.2)
+# s_outer.set_zs(1, 0, 0.2)
 s_inner = SurfaceRZFourier.from_vmec_input(surface_filename, range=range_param, nphi=qphi, ntheta=ntheta)
 s_outer = SurfaceRZFourier.from_vmec_input(surface_filename, range=range_param, nphi=qphi, ntheta=ntheta)
 
 # Make the inner and outer surfaces by extending the plasma surface
-s_inner.extend_via_projected_normal(poff)
-s_outer.extend_via_projected_normal(poff + coff)
+s_inner.extend_via_normal(poff)
+s_outer.extend_via_normal(poff + coff)
 
 # Make the output directory
 out_dir = Path("psc_output")
@@ -65,7 +83,7 @@
 make_Bnormal_plots(bs, s_plot, out_dir, "biot_savart_initial")
 
 # optimize the currents in the TF coils
-bs = coil_optimization(s, bs, base_curves, curves, out_dir)
+# bs = coil_optimization(s, bs, base_curves, curves, out_dir)
 curves_to_vtk(curves, out_dir / "TF_coils", close=True)
 bs.set_points(s.gamma().reshape((-1, 3)))
 Bnormal = np.sum(bs.B().reshape((qphi, ntheta, 3)) * s.unitnormal(), axis=2)
@@ -75,12 +93,30 @@
 calculate_on_axis_B(bs, s)
 
 # Finally, initialize the psc class
-kwargs_geo = {"Nx": 12}  
+kwargs_geo = {"Nx": 10}  
+# note Bnormal below is additional Bnormal (not from the TF coils or the PSCs)
 psc_array = PSCgrid.geo_setup_between_toroidal_surfaces(
-    s, Bnormal, bs, s_inner, s_outer,  **kwargs_geo
+    s, np.zeros(Bnormal.shape), bs, s_inner, s_outer,  **kwargs_geo
 )
 
+# print('Currents = ', psc_array.I)
 make_Bnormal_plots(psc_array.B_PSC, s_plot, out_dir, "biot_savart_PSC_initial")
 make_Bnormal_plots(bs + psc_array.B_PSC, s_plot, out_dir, "PSC_and_TF_initial")
 
+# Try and optimize with scipy
+t1 = time.time()
+x0 = np.ravel(np.array([psc_array.alphas, psc_array.deltas]))
+psc_array.least_squares(x0)
+t2 = time.time()
+print('Time for call to least-squares = ', t2 - t1)
+# psc_array.least_squares(x0 + np.random.rand(len(x0)))
+
+# from scipy.optimize import minimize
+# print('beginning optimization: ')
+# options = {"disp": True, "maxiter": 2}
+# x0 = np.ravel(np.array([psc_array.alphas, psc_array.deltas]))
+# # print(x0)
+# x_opt = minimize(psc_array.least_squares, x0, options=options)
+# print(x_opt)
+
 # plt.show()

src/simsopt/geo/curve.py

@@ -845,7 +845,7 @@ def wrap(data):
     else:
         coil_data = np.zeros(data.shape)
         for i in range(len(scalar_data)):
-            coil_data[i * ppl[i]: (i + 1) * ppl[i]] = scalar_data[i]
+            coil_data[i * ppl[i]: (i + 1) * ppl[i]] = abs(scalar_data[i])
         coil_data = np.ascontiguousarray(coil_data)
         polyLinesToVTK(str(filename), x, y, z, pointsPerLine=ppl, pointData={'idx': data, 'I': coil_data})
 

src/simsopt/geo/psc_grid.py

@@ -68,8 +68,8 @@ def _setup_uniform_grid(self):
         self.dz = 2 * z_max / (Nz - 1)
         print(Nx, Ny, Nz, self.dx, self.dy, self.dz)
         Nmin = min(self.dx, min(self.dy, self.dz))
-        self.R = Nmin / 4.0
-        self.a = self.R / 10.0
+        self.R = Nmin / 2.0
+        self.a = self.R / 100.0
         print('Major radius of the coils is R = ', self.R)
         print('Coils are spaced so that every coil of radius R '
               ' is at least 2R away from the next coil'
@@ -202,6 +202,13 @@ def geo_setup_between_toroidal_surfaces(
         psc_grid.plasma_boundary = plasma_boundary.to_RZFourier()
         psc_grid.nphi = len(psc_grid.plasma_boundary.quadpoints_phi)
         psc_grid.ntheta = len(psc_grid.plasma_boundary.quadpoints_theta)
+        Ngrid = psc_grid.nphi * psc_grid.ntheta
+        Nnorms = np.ravel(np.sqrt(np.sum(psc_grid.plasma_boundary.normal() ** 2, axis=-1)))
+        psc_grid.grid_normalization = np.sqrt(Nnorms / Ngrid)
+        # integration over phi for L
+        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]
         Nx = kwargs.pop("Nx", 10)
         Ny = Nx  # kwargs.pop("Ny", 10)
         Nz = Nx  # kwargs.pop("Nz", 10)
@@ -244,46 +251,42 @@ def geo_setup_between_toroidal_surfaces(
 
         # PSC coil geometry determined by its center point in grid_xyz
         # and its alpha and delta angles, which we initialize randomly here.
-        psc_grid.alphas = np.random.rand(psc_grid.num_psc) * 2 * np.pi
-        psc_grid.deltas = np.random.rand(psc_grid.num_psc) * 2 * np.pi
-
-        psc_grid.coil_normals = np.array(
-            [np.cos(psc_grid.alphas) * np.sin(psc_grid.deltas),
-             -np.sin(psc_grid.alphas),
-             np.cos(psc_grid.alphas) * np.cos(psc_grid.deltas)]
-        ).T
+
+        # Randomly initialize the coil orientations
+        # psc_grid.alphas = np.random.rand(psc_grid.num_psc) * 2 * np.pi
+        # psc_grid.deltas = np.random.rand(psc_grid.num_psc) * 2 * np.pi
+        # psc_grid.coil_normals = np.array(
+        #     [np.cos(psc_grid.alphas) * np.sin(psc_grid.deltas),
+        #      -np.sin(psc_grid.alphas),
+        #      np.cos(psc_grid.alphas) * np.cos(psc_grid.deltas)]
+        # ).T
+
+        # Initialize the coil orientations parallel to the B field. 
+        B_TF.set_points(psc_grid.grid_xyz)
+        B = B_TF.B()
+        print(B.shape)
+        psc_grid.coil_normals = (B.T / np.sqrt(np.sum(B ** 2, axis=-1))).T
+        psc_grid.alphas = np.arcsin(-psc_grid.coil_normals[:, 1])
+        minus_yinds = np.ravel(np.where(psc_grid.grid_xyz[:, 1] > 0.0))
+        psc_grid.alphas[minus_yinds] = -psc_grid.alphas[minus_yinds]
+        psc_grid.deltas = np.arccos(psc_grid.coil_normals[:, 2] / np.cos(psc_grid.alphas))
+        # determine the alphas and deltas from these normal vectors
+        print(psc_grid.coil_normals.shape, B.shape)
+
         t2 = time.time()
         print('Initialize grid time = ', t2 - t1)
 
         # Initialize curve objects corresponding to each PSC coil for 
         # 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()
-        psc_grid.L = sopp.L_matrix(
-            contig(psc_grid.grid_xyz), 
-            contig(psc_grid.alphas), 
-            contig(psc_grid.deltas),
-            contig(phi),
-            psc_grid.R,
-        ) * dphi ** 2 / (4 * np.pi)
-        print(psc_grid.L)
-        psc_grid.L = (psc_grid.L + psc_grid.L.T)
-        np.fill_diagonal(psc_grid.L, np.log(8.0 * psc_grid.R / psc_grid.a) - 2.0)
-        psc_grid.L = psc_grid.L * psc_grid.mu0 * psc_grid.R * psc_grid.Nt ** 2
+        psc_grid.setup_orientations(psc_grid.alphas, psc_grid.deltas)
         t2 = time.time()
-        print('Inductances time = ', t2 - t1)
-        psc_grid.setup_curves()
-        psc_grid.setup_currents_and_fields()
+        print('Geo setup time = ', t2 - t1)
         psc_grid.plot_curves()
+        psc_grid.b_vector()
+        kappas = np.ravel(np.array([psc_grid.alphas, psc_grid.deltas]))
+        psc_grid.A_opt = psc_grid.least_squares(kappas)
 
         # psc_grid._optimization_setup()
         return psc_grid
@@ -311,6 +314,9 @@ def geo_setup_manual(
         Nt = kwargs.pop("Nt", 1)
         psc_grid.Nt = Nt
         psc_grid.num_psc = psc_grid.grid_xyz.shape[0]
+        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]
 
         Bn = kwargs.pop("Bn", np.zeros((1, 3)))
         Bn = np.array(Bn)
@@ -373,23 +379,7 @@ def geo_setup_manual(
 
         # Initialize curve objects corresponding to each PSC coil for 
         # plotting in 3D
-        psc_grid.fluxes(psc_grid.alphas, psc_grid.deltas)
-        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), 
-            contig(psc_grid.deltas),
-            contig(phi),
-            psc_grid.R,
-        ) * dphi ** 2 / (4 * np.pi)
-        psc_grid.L = (psc_grid.L + psc_grid.L.T)
-        np.fill_diagonal(psc_grid.L, np.log(8.0 * R / psc_grid.a) - 2.0)
-        psc_grid.L *= psc_grid.mu0 * psc_grid.R * psc_grid.Nt ** 2
-        # psc_grid.inductances(psc_grid.alphas, psc_grid.deltas)
-        psc_grid.setup_curves()
-        psc_grid.setup_currents_and_fields()
+        psc_grid.setup_orientations(psc_grid.alphas, psc_grid.deltas)
         psc_grid.plot_curves()
 
         return psc_grid
@@ -567,4 +557,57 @@ def quad_func(yy, xx, cd, ca, sd, sa, nj, centers):
             except IntegrationWarning:
                 print('Integral is divergent')
         self.psi = fluxes
+
+    def b_vector(self):
+        Bn_target = self.Bn.reshape(-1)
+        self.B_TF.set_points(self.plasma_boundary.gamma().reshape(-1, 3))
+        Bn_TF = np.sum(
+            self.B_TF.B().reshape(-1, 3) * self.plasma_boundary.unitnormal(
+                ).reshape(-1, 3), axis=-1
+        )
+        self.b_opt = (Bn_TF - Bn_target) * self.grid_normalization
+
+    def least_squares(self, kappas):
+        alphas = kappas[:len(kappas) // 2]
+        deltas = kappas[len(kappas) // 2:]
+        self.setup_orientations(alphas, deltas)
+        BdotN2 = np.sum((self.Bn_PSC.reshape(-1) * self.grid_normalization + self.b_opt) ** 2)
+        outstr = f"||Bn||^2 = {BdotN2:.2e}"
+        print(outstr)
+        return BdotN2
+
+    def setup_orientations(self, alphas, deltas):
+        contig = np.ascontiguousarray
+        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);
+        t2 = time.time()
+        print('Fluxes time = ', t2 - t1)
+        t1 = time.time()
+        self.L = sopp.L_matrix(
+            contig(self.grid_xyz), 
+            contig(alphas), 
+            contig(deltas),
+            contig(self.phi),
+            self.R,
+        ) * self.dphi ** 2 / (4 * np.pi)
+        self.L = (self.L + self.L.T)
+        np.fill_diagonal(self.L, np.log(8.0 * self.R / self.a) - 2.0)
+        self.L = self.L * self.mu0 * self.R * self.Nt ** 2
+        t2 = time.time()
+        print('Inductances time = ', t2 - t1)
+        t1 = time.time()
+        self.setup_curves()
+        t2 = time.time()
+        print('Setup curves time = ', t2 - t1)
+        t1 = time.time()
+        self.setup_currents_and_fields()
+        t2 = time.time()
+        print('Setup fields time = ', t2 - t1)
 

src/simsoptpp/psc.cpp

@@ -1,5 +1,4 @@
 #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)
@@ -63,7 +62,7 @@ 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)
+Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, Array& I, Array& normal, double R)
 {
     // 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)
@@ -76,18 +75,21 @@ 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(B_TF.layout() != xt::layout_type::row_major)
-          throw std::runtime_error("B_TF 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");
 
     // points shape should be (num_coils, 3)
-    // B_TF shape should be (num_rho, num_phi, 3)
+    // I shape should be (num_coils, 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});
+    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++) {
@@ -111,11 +113,28 @@ Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array&
                 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;
+                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;
             }
         }
-        Psi(j) = integral;
+        Psi(j) = integral * fac * I(j);
     }
     return Psi;
 }

src/simsoptpp/psc.h

@@ -1,5 +1,12 @@
 #pragma once
 
+// #define __STDCPP_WANT_MATH_SPEC_FUNCS__ 1
+// #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 <cmath>
 #include <tuple>  // c++ tuples
 #include <string> // for string class
@@ -9,6 +16,6 @@ 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& B_TF, Array& normal);
+Array TF_fluxes(Array& points, Array& alphas, Array& deltas, Array& rho, Array& phi, Array& I, Array& normal, double R);
 // 
 // Array A_matrix(Array& plasma_surface_normals, Array& B, Array& alphas, Array& deltas);