Merge pull request hiddenSymmetries#393 from hiddenSymmetries/A_cyl

Implement magnetic potential in cylindrical coordinates

src/simsoptpp/magneticfield.h

@@ -106,6 +106,18 @@ class MagneticField {
             }
         }
 
+        virtual void _A_cyl_impl(Tensor2& A_cyl) {
+            Tensor2& A = this->A_ref();
+            Tensor2& rphiz = this->get_points_cyl_ref();
+            int npoints = A.shape(0);
+            for (int i = 0; i < npoints; ++i) {
+                double phi = rphiz(i, 1);
+                A_cyl(i, 0) = std::cos(phi)*A(i, 0) + std::sin(phi)*A(i, 1);
+                A_cyl(i, 1) = std::cos(phi)*A(i, 1) - std::sin(phi)*A(i, 0);
+                A_cyl(i, 2) = A(i, 2);
+            }
+        }
+
         virtual void _GradAbsB_cyl_impl(Tensor2& GradAbsB_cyl) {
             Tensor2& GradAbsB = this->GradAbsB_ref();
             Tensor2& rphiz = this->get_points_cyl_ref();
@@ -127,7 +139,7 @@ class MagneticField {
 
         CachedTensor<T, 2> points_cart;
         CachedTensor<T, 2> points_cyl;
-        CachedTensor<T, 2> data_B, data_A, data_GradAbsB, data_AbsB, data_Bcyl, data_GradAbsBcyl;
+        CachedTensor<T, 2> data_B, data_A, data_GradAbsB, data_AbsB, data_Bcyl, data_Acyl, data_GradAbsBcyl;
         CachedTensor<T, 3> data_dB, data_dA;
         CachedTensor<T, 4> data_ddB, data_ddA;
         int npoints;
@@ -148,6 +160,7 @@ class MagneticField {
             data_AbsB.invalidate_cache();
             data_GradAbsB.invalidate_cache();
             data_Bcyl.invalidate_cache();
+            data_Acyl.invalidate_cache();
             data_GradAbsBcyl.invalidate_cache();
         }
 
@@ -231,6 +244,13 @@ class MagneticField {
             return data_Bcyl.get_or_create_and_fill({npoints, 3}, [this](Tensor2& B) { return _B_cyl_impl(B);});
         }
 
+        Tensor2 A_cyl() {
+            return A_cyl_ref();
+        }
+
+        Tensor2& A_cyl_ref() {
+            return data_Acyl.get_or_create_and_fill({npoints, 3}, [this](Tensor2& A) { return _A_cyl_impl(A);});
+        }
 
         Tensor2 AbsB() {
             return AbsB_ref();

src/simsoptpp/python_magneticfield.cpp

@@ -37,6 +37,8 @@ template <typename T, typename S> void register_common_field_methods(S &c) {
      .def("B_cyl", py::overload_cast<>(&T::B_cyl), "Return a `(npoints, 3)` array containing the magnetic field (in cylindrical coordinates) (the order is :math:`(B_r, B_\\phi, B_z)`).")
      .def("B_cyl_ref", py::overload_cast<>(&T::B_cyl_ref), "As `B_cyl`, but returns a reference to the array (this array should be read only).")
      .def("A", py::overload_cast<>(&T::A), "Returns a `(npoints, 3)` array containing the magnetic potential (in cartesian coordinates). Denoting the indices by `i` and `l`, the result contains  `A_l(x_i)`.")
+     .def("A_cyl", py::overload_cast<>(&T::A_cyl), "Return a `(npoints, 3)` array containing the magnetic potential (in cylindrical coordinates) (the order is :math:`(A_r, A_\\phi, A_z)`).")
+     .def("A_cyl_ref", py::overload_cast<>(&T::A_cyl_ref), "As `A_cyl`, but returns a reference to the array (this array should be read only).")
      .def("dA_by_dX", py::overload_cast<>(&T::dA_by_dX), "Returns a `(npoints, 3, 3)` array containing the gradient of the magnetic potential (in cartesian coordinates). Denoting the indices by `i`, `j` and `l`, the result contains  `\\partial_j A_l(x_i)`.")
      .def("d2A_by_dXdX", py::overload_cast<>(&T::d2A_by_dXdX), "Returns a `(npoints, 3, 3)` array containing the hessian of the magnetic potential (in cartesian coordinates). Denoting the indices by `i`, `j`, `k` and `l`, the result contains  `\\partial_k\\partial_j  A_l(x_i)`.")
      .def("A_ref", py::overload_cast<>(&T::A_ref), "As `A`, but returns a reference to the array (this array should be read only).")

tests/field/test_magneticfields.py

@@ -846,6 +846,44 @@ def test_reiman_dBdX_taylortest(self):
             with self.subTest(idx=idx):
                 self.subtest_reiman_dBdX_taylortest(idx)
 
+    def test_cyl_versions(self):
+        R0test = 1.5
+        B0test = 0.8
+        B0 = ToroidalField(R0test, B0test)
+
+        curves, currents, ma = get_ncsx_data()
+        nfp = 3
+        coils = coils_via_symmetries(curves, currents, nfp, True)
+        bs = BiotSavart(coils)
+        btotal = bs + B0
+        rmin = 1.5
+        rmax = 1.7
+        phimin = 0
+        phimax = 2*np.pi/nfp
+        zmax = 0.1
+        N = 1000
+        points = np.random.uniform(size=(N, 3))
+        points[:, 0] = points[:, 0]*(rmax-rmin) + rmin
+        points[:, 1] = points[:, 1]*(nfp*phimax-phimin) + phimin
+        points[:, 2] = points[:, 2]*(2*zmax) - zmax
+        btotal.set_points_cyl(points)
+
+        dB = btotal.GradAbsB()
+        B = btotal.B()
+        A = btotal.A()
+        dB_cyl = btotal.GradAbsB_cyl()
+        B_cyl = btotal.B_cyl()
+        A_cyl = btotal.A_cyl()
+
+        for j in range(N):
+            phi = points[j, 1]
+            rotation = np.array([[np.cos(phi), np.sin(phi), 0],
+                                [-np.sin(phi), np.cos(phi), 0],
+                                [0, 0, 1]])
+            np.testing.assert_allclose(rotation @ B[j, :], B_cyl[j, :])
+            np.testing.assert_allclose(rotation @ dB[j, :], dB_cyl[j, :])
+            np.testing.assert_allclose(rotation @ A[j, :], A_cyl[j, :])
+
     def test_interpolated_field_close_with_symmetries(self):
         R0test = 1.5
         B0test = 0.8