Did a refactorization of how the psc_array interacts with the BiotSav…

…art object. Before, there was a psc_array associated with every PSCCurve. Now there is a PSC_BiotSavart class that has one associated psc_array, and correctly feeds things to the curves and currents. Unit tests using this class with CurvePlanarFourier looks good so far, next step to sub in the PSCCurve class.

examples/2_Intermediate/stage_two_optimization_planar_coils.py

@@ -99,6 +99,10 @@
 # Create the initial coils:
 base_curves = create_equally_spaced_planar_curves(ncoils, s.nfp, stellsym=True, R0=R0, R1=R1, order=order)
 for i in range(len(base_curves)):
+    base_curves[i].set('x' + str(2 * order + 1), np.random.rand(1) - 0.5)
+    base_curves[i].set('x' + str(2 * order + 2), np.random.rand(1) - 0.5)
+    base_curves[i].set('x' + str(2 * order + 3), np.random.rand(1) - 0.5)
+    base_curves[i].set('x' + str(2 * order + 4), np.random.rand(1) - 0.5)
     for j in range(2 * order + 1):
         base_curves[i].fix('x' + str(j))
     base_curves[i].fix('x' + str(2 * order + 5))

src/simsopt/field/biotsavart.py

@@ -5,7 +5,7 @@
 from .._core.json import GSONDecoder
 from .._core.derivative import Derivative
 
-__all__ = ['BiotSavart']
+__all__ = ['BiotSavart', 'PSC_BiotSavart']
 
 
 class BiotSavart(sopp.BiotSavart, MagneticField):
@@ -296,3 +296,76 @@ def from_dict(cls, d, serial_objs_dict, recon_objs):
         bs = cls(coils)
         bs.set_points_cart(xyz)
         return bs
+
+class PSC_BiotSavart(BiotSavart):
+    """
+    """
+    def __init__(self, psc_array):
+        from simsopt.field import coils_via_symmetries, Current, psc_coils_via_symmetries
+        self.psc_array = psc_array
+        self.npsc = psc_array.num_psc
+        curves = [self.psc_array.curves[i] for i in range(self.npsc)]
+        currents = [Current(self.psc_array.I[i] * 1e-5) * 1e5 for i in range(self.npsc)]
+        self.curves = curves
+        [currents[i].fix_all() for i in range(self.npsc)]
+        self.currents = currents
+        coils = psc_coils_via_symmetries(self.curves, self.currents, psc_array.nfp, psc_array.stellsym)
+        BiotSavart.__init__(self, coils)
+
+    def B_vjp(self, v):
+        r"""
+        Assume the field was evaluated at points :math:`\mathbf{x}_i, i\in \{1, \ldots, n\}` and denote the value of the field at those points by
+        :math:`\{\mathbf{B}_i\}_{i=1}^n`.
+        These values depend on the shape of the coils, i.e. on the dofs :math:`\mathbf{c}_k` of each coil.
+        This function returns the vector Jacobian product of this dependency, i.e.
+
+        .. math::
+
+            \{ \sum_{i=1}^{n} \mathbf{v}_i \cdot \partial_{\mathbf{c}_k} \mathbf{B}_i \}_k.
+
+        """
+        from simsopt.geo.curveplanarfourier import PSCCurve
+
+        coils = self._coils
+        gammas = [coil.curve.gamma() for coil in coils]
+        gammadashs = [coil.curve.gammadash() for coil in coils]
+        currents = [coil.current.get_value() for coil in coils]
+        res_gamma = [np.zeros_like(gamma) for gamma in gammas]
+        res_gammadash = [np.zeros_like(gammadash) for gammadash in gammadashs]
+
+        points = self.get_points_cart_ref()
+        sopp.biot_savart_vjp_graph(points, gammas, gammadashs, currents, v,
+                                   res_gamma, res_gammadash, [], [], [])
+        dB_by_dcoilcurrents = self.dB_by_dcoilcurrents()
+        res_current = [np.sum(v * dB_by_dcoilcurrents[i]) for i in range(len(dB_by_dcoilcurrents))]
+
+        curve_flags = [isinstance(coil.curve, PSCCurve) for coil in coils]
+        if np.any(curve_flags):
+            order = coils[0].curve.order 
+            ndofs = 2 * order + 8
+            dofs = np.array([self.curves[i].get_dofs() for i in range(self.npsc)])
+            dofs = dofs[:, 2 * order + 1:2 * order + 5]
+            normalization = np.sqrt(np.sum(dofs ** 2, axis=-1))
+            dofs = dofs / normalization[:, None]
+            alphas1 = np.arctan2(2 * (dofs[:, 0] * dofs[:, 1] + dofs[:, 2] * dofs[:, 3]), 
+                                1 - 2.0 * (dofs[:, 1] ** 2 + dofs[:, 2] ** 2))
+            deltas1 = -np.pi / 2.0 + 2.0 * np.arctan2(
+                np.sqrt(1.0 + 2 * (dofs[:, 0] * dofs[:, 2] - dofs[:, 1] * dofs[:, 3])), 
+                np.sqrt(1.0 - 2 * (dofs[:, 0] * dofs[:, 2] - dofs[:, 1] * dofs[:, 3])))
+            self.psc_array.setup_orientations(alphas1, deltas1)
+            self.psc_array.update_psi()
+            # self.psc_array.setup_currents_and_fields()
+            self.psc_array.psi_deriv()
+            dI = np.zeros((len(coils), ndofs))
+            q = 0
+            for fp in range(self.psc_array.nfp):
+                for stell in self.psc_array.stell_list:
+                    for i in range(self.npsc):
+                        dI[i, :] += coils[i + q * self.npsc].curve.dkappa_dcoef_vjp(
+                            [res_current[i + q * self.npsc]], self.psc_array.dpsi) * stell
+                    q += 1
+            Linv = self.psc_array.L_inv
+            dI = - Linv @ dI
+            return sum([coils[i].vjp(res_gamma[i], res_gammadash[i], dI[i, :]) for i in range(len(coils))])
+        else:
+            return sum([coils[i].vjp(res_gamma[i], res_gammadash[i], res_current[i]) for i in range(len(coils))])

src/simsopt/field/coil.py

@@ -13,7 +13,7 @@
            'Current', 'coils_via_symmetries',
            'load_coils_from_makegrid_file',
            'apply_symmetries_to_currents', 'apply_symmetries_to_curves',
-           'apply_symmetries_to_psc_curves',
+           'apply_symmetries_to_psc_curves', 'psc_coils_via_symmetries',
            'coils_to_makegrid', 'coils_to_focus']
 
 
@@ -60,14 +60,9 @@ def __init__(self, curve, current):
         Optimizable.__init__(self, depends_on=[curve, current])
 
     def vjp(self, v_gamma, v_gammadash, v_current):
-        # print('other derivs = ', v_gamma.shape, v_gammadash.shape, self.curve.dgamma_by_dcoeff_vjp_impl(v_gamma).shape, self.curve.dgammadash_by_dcoeff_vjp_impl(v_gammadash).shape)
-        # print(self.curve.dgamma_by_dcoeff_vjp(v_gamma),
-        #     self.curve.dgammadash_by_dcoeff_vjp(v_gammadash),
-        #     Derivative({self.curve: v_current}))
         return self.curve.dgamma_by_dcoeff_vjp(v_gamma) \
             + self.curve.dgammadash_by_dcoeff_vjp(v_gammadash) \
-            + Derivative({self.curve: v_current}) #\
-            # + self.curve.psc_current_contribution_vjp(v_current)
+            # + Derivative({self.curve: v_current})
 
     def plot(self, **kwargs):
         """
@@ -249,6 +244,20 @@ def coils_via_symmetries(curves, currents, nfp, stellsym):
     coils = [Coil(curv, curr) for (curv, curr) in zip(curves, currents)]
     return coils
 
+def psc_coils_via_symmetries(curves, currents, nfp, stellsym):
+    """
+    Take a list of ``n`` curves and return ``n * nfp * (1+int(stellsym))``
+    ``Coil`` objects obtained by applying rotations and flipping corresponding
+    to ``nfp`` fold rotational symmetry and optionally stellarator symmetry.
+    """
+
+    assert len(curves) == len(currents)
+    curves = apply_symmetries_to_psc_curves(curves, nfp, stellsym)
+    # [currents[i].fix_all() for i in range(len(currents))]
+    currents = apply_symmetries_to_currents(currents, nfp, stellsym)
+    coils = [Coil(curv, curr) for (curv, curr) in zip(curves, currents)]
+    return coils
+
 
 def load_coils_from_makegrid_file(filename, order, ppp=20):
     """

src/simsopt/field/magneticfieldclasses.py

@@ -124,7 +124,6 @@ def from_dict(cls, d, serial_objs_dict, recon_objs):
         field.set_points_cart(xyz)
         return field
 
-
 class PoloidalField(MagneticField):
     '''
     Magnetic field purely in the poloidal direction, that is, in the

src/simsopt/geo/curve.py

@@ -617,16 +617,10 @@ class RotatedCurve(sopp.Curve, Curve):
     input a Curve, rotates it about the ``z`` axis by a toroidal angle
     ``phi``, and optionally completes a reflection when ``flip=True``.
     """
-
     def __init__(self, curve, phi, flip):
-        from simsopt.geo.curveplanarfourier import PSCCurve
-
         self.curve = curve
         sopp.Curve.__init__(self, curve.quadpoints)
-        # print(curve)
         Curve.__init__(self, depends_on=[curve])
-        # if isinstance(curve, PSCCurve):
-        #     exit()
         self._phi = phi
         self.rotmat = np.asarray(
             [[cos(phi), -sin(phi), 0],
@@ -816,53 +810,6 @@ def dgammadashdashdash_by_dcoeff_vjp(self, v):
     @property
     def flip(self):
         return True if self.rotmat[2][2] == -1 else False
-
-    def psc_current_contribution_vjp(self, v_current):
-        Linv_partial = self.curve._psc_array.L_inv[self._index, self._index]
-        # print(np.ravel(-Linv_partial * self.dkappa_dcoef_vjp(v_current)), v_current)
-        return Derivative({self: np.zeros(10)})
-        # Linv_partial = self.curve._psc_array.L_inv[self._index % self.curve.npsc, self._index % self.curve.npsc]
-        # return Derivative({self: np.ravel(-Linv_partial * self.dkappa_dcoef_vjp(v_current)) }) 
-
-    def dkappa_dcoef_vjp(self, v_current):
-        dofs = self.get_dofs()
-
-        # For the rotated coils, need to compute dalpha_prime / dcoef
-        if len(dofs) == 0:
-            dofs = self.curve.get_dofs()
-
-        dofs = dofs[2 * self.order + 1:2 * self.order + 5]  # don't need the coordinate variables
-        normalization = np.sqrt(np.sum(dofs ** 2))
-        dofs = dofs / normalization  # normalize the quaternion
-        w = dofs[0]
-        x = dofs[1]
-        y = dofs[2]
-        z = dofs[3]
-        dalpha_dw = (2 * x * (2 * (x ** 2 + y ** 2) - 1)) / \
-            (4 * (w * x + y * z) ** 2 + (1 - 2 * ( x ** 2 + y ** 2)) ** 2) / normalization
-        dalpha_dx = (w * (-0.5 - x ** 2 + y ** 2) - 2 * x * y * z) / \
-            (x ** 2 * (w ** 2 + 2 * y ** 2 - 1) + 2 * w * x * y * z + x ** 4 + y ** 4 + y ** 2 * (z ** 2 - 1) + 0.25) / normalization
-        dalpha_dy = (z * (-0.5 + x ** 2 - y ** 2) - 2 * x * y * w) / \
-            (x ** 2 * (w ** 2 + 2 * y ** 2 - 1) + 2 * w * x * y * z + x ** 4 + y ** 4 + y ** 2 * (z ** 2 - 1) + 0.25) / normalization
-        dalpha_dz = (2 * y * (2 * (x ** 2 + y ** 2) - 1)) / \
-            (4 * (w * x + y * z) ** 2 + (1 - 2 * ( x ** 2 + y ** 2)) ** 2) / normalization
-        ddelta_dw = -y / np.sqrt(-(w * y - x * z - 0.5) * (w * y - x * z + 0.5)) / normalization
-        ddelta_dx = z / np.sqrt(-(w * y - x * z - 0.5) * (w * y - x * z + 0.5)) / normalization
-        ddelta_dy = -w / np.sqrt(-(w * y - x * z - 0.5) * (w * y - x * z + 0.5)) / normalization
-        ddelta_dz = x / np.sqrt(-(w * y - x * z - 0.5) * (w * y - x * z + 0.5)) / normalization
-
-        dalpha = np.array([dalpha_dw, dalpha_dx, dalpha_dy, dalpha_dz])
-        ddelta = np.array([ddelta_dw, ddelta_dx, ddelta_dy, ddelta_dz])
-
-        dalpha = self.curve._psc_array.dpsi_daa[self._index] * dalpha + self.curve._psc_array.dpsi_dad[self._index] * ddelta
-        ddelta = self.curve._psc_array.dpsi_ddd[self._index] * ddelta + self.curve._psc_array.dpsi_dda[self._index] * dalpha
-
-        dalpha = np.hstack((np.zeros(2 * self.order + 1), dalpha))
-        dalpha = np.hstack((dalpha, np.zeros(3)))
-        ddelta = np.hstack((np.zeros(2 * self.order + 1), ddelta))
-        ddelta = np.hstack((ddelta, np.zeros(3)))
-        deriv = dalpha + ddelta
-        return deriv * v_current[0]
 
 def curves_to_vtk(curves, filename, close=False, scalar_data=None):
     """

src/simsopt/geo/curveplanarfourier.py

@@ -4,7 +4,7 @@
 from .curve import Curve, RotatedCurve
 from .._core.derivative import Derivative
 
-__all__ = ['CurvePlanarFourier', 'PSCCurve']  # , 'RotatedPSCCurve']
+__all__ = ['CurvePlanarFourier', 'PSCCurve']
 
 
 class CurvePlanarFourier(sopp.CurvePlanarFourier, Curve):
@@ -76,9 +76,8 @@ def set_dofs(self, dofs):
         sopp.CurvePlanarFourier.set_dofs(self, dofs)
 
 class PSCCurve(CurvePlanarFourier):
-    def __init__(self, quadpoints, order, nfp, stellsym, psc_array, index, dofs=None):
+    def __init__(self, quadpoints, order, nfp, stellsym, index, npsc, dofs=None):
         CurvePlanarFourier.__init__(self, quadpoints, order, nfp, stellsym, dofs=dofs)
-        self._psc_array = psc_array
         self._index = index
         names = self.local_dof_names
         if len(names) > 0: 
@@ -87,19 +86,12 @@ def __init__(self, quadpoints, order, nfp, stellsym, psc_array, index, dofs=None
             self.fix(names[2 * self.order + 5])
             self.fix(names[2 * self.order + 6])
             self.fix(names[2 * self.order + 7])
-        self.npsc = self._psc_array.num_psc
+        self.npsc = npsc
 
-    def dkappa_dcoef_vjp(self, v_current, index):
-        dofs = self.get_dofs()  # should already be only the orientation dofs
-
-        # For the rotated coils, need to compute dalpha_prime / dcoef
-        # if len(dofs) == 0:
-        #     dofs = self.curve.get_dofs()
-            # rotated curves need the orientation dofs of the rotated curve
-
-        dofs_unnormalized = dofs[2 * self.order + 1:2 * self.order + 5]  # don't need the coordinate variables
+    def dkappa_dcoef_vjp(self, v_current, dpsi):
+        dofs = self.get_dofs()  
+        dofs_unnormalized = dofs[2 * self.order + 1:2 * self.order + 5]
         normalization = np.sqrt(np.sum(dofs_unnormalized ** 2))
-        # print('dofs = ', dofs, normalization)
         dofs = dofs_unnormalized / normalization  # normalize the quaternion
         w = dofs[0]
         x = dofs[1]
@@ -118,9 +110,6 @@ def dkappa_dcoef_vjp(self, v_current, index):
         ddelta_dy = w / np.sqrt(-(w * y - x * z - 0.5) * (w * y - x * z + 0.5))
         ddelta_dz = -x / np.sqrt(-(w * y - x * z - 0.5) * (w * y - x * z + 0.5))
 
-        # dalpha = np.array([dalpha_dw, dalpha_dx, dalpha_dy, dalpha_dz])
-        # ddelta = np.array([ddelta_dw, ddelta_dx, ddelta_dy, ddelta_dz])
-
         # So far this is dalpha with respect to the normalized variables 
         # so need to convert to the original dof derivatives
         dnormalization = np.zeros((4, 4))
@@ -134,58 +123,12 @@ def dkappa_dcoef_vjp(self, v_current, index):
         # dalpha_total = self._psc_array.dpsi_full[index] * dalpha
         # ddelta_total = self._psc_array.dpsi_full[index + self.npsc] * ddelta
 
-        # dalpha_total = self._psc_array.dpsi_daa[self._index] * dalpha + self._psc_array.dpsi_dad[self._index] * ddelta
-        # ddelta_total = self._psc_array.dpsi_dda[self._index] * dalpha + self._psc_array.dpsi_ddd[self._index] * ddelta
-
-        # indices = np.arange(self._index, 2 * self.npsc * self._psc_array.symmetry, 2 * self.npsc)
-        # indices2 = np.arange(self._index + self.npsc, 2 * self.npsc * self._psc_array.symmetry, 2 * self.npsc)
-        # dalpha_total = self._psc_array.dpsi_full[indices] * dalpha
-        # ddelta_total = self._psc_array.dpsi_full[indices2] * ddelta
-
-        dalpha_total = self._psc_array.dpsi[self._index] * dalpha
-        ddelta_total = self._psc_array.dpsi[self._index + self.npsc] * ddelta
-
-        # if len(dofs) == 0:
-        # print(self._index)
-        # dalpha_total = np.zeros(dalpha.shape)
-        # ddelta_total = np.zeros(ddelta.shape)
-        # for i in range(self._psc_array.symmetry):
-        #     index = self._index + self.npsc * i
-        #     dalpha_total += self._psc_array.dpsi_daa[index] * dalpha + self._psc_array.dpsi_dad[index] * ddelta
-        #     ddelta_total += self._psc_array.dpsi_ddd[index] * ddelta + self._psc_array.dpsi_dda[index] * dalpha
+        dalpha_total = dpsi[self._index] * dalpha
+        ddelta_total = dpsi[self._index + self.npsc] * ddelta
 
         dalpha = np.hstack((np.zeros(2 * self.order + 1), dalpha_total))
         dalpha = np.hstack((dalpha, np.zeros(3)))
         ddelta = np.hstack((np.zeros(2 * self.order + 1), ddelta_total))
         ddelta = np.hstack((ddelta, np.zeros(3)))
         deriv = dalpha + ddelta
-        # print(deriv.shape)
-        # exit()
-        return deriv * v_current[0]  # Derivative({self: deriv})
-
-
-# class RotatedPSCCurve(RotatedCurve, PSCCurve):
-#     """
-#     RotatedCurve inherits from the Curve base class.  It takes an
-#     input a Curve, rotates it about the ``z`` axis by a toroidal angle
-#     ``phi``, and optionally completes a reflection when ``flip=True``.
-#     """
-
-#     def __init__(self, curve, phi, flip):
-#         # sopp.Curve.__init__(self, curve.quadpoints)
-#         # print(curve)
-#         # print(curve.quadpoints, curve.order, curve.nfp, curve.stellsym, curve._psc_array, curve._index, curve.get_dofs(), phi, flip)
-#         # print(self)
-#         # PSCCurve.__init__(self, curve.quadpoints, curve.order, curve.nfp, curve.stellsym, curve._psc_array, curve._index, curve.get_dofs())
-#         RotatedCurve.__init__(self, curve, phi, flip)
-#         exit()
-
-#         print(self, self.order, self.curve)
-#         # rotcurve.order = base_curves[i].order
-#         # rotcurve._psc_array = base_curves[i]._psc_array
-#         # rotcurve._index = base_curves[i]._index + base_curves[i].npsc * q
-#         # rotcurve.psc_current_contribution_vjp = base_curves[i].psc_current_contribution_vjp
-#         # rotcurve.dkappa_dcoef_vjp = base_curves[i].dkappa_dcoef_vjp
-#         #self.curve = curve
-#         #self.order = curve.order
-
+        return deriv * v_current[0]

src/simsopt/geo/psc_grid.py

@@ -861,12 +861,12 @@ def setup_curves(self):
         Convert the (alphas, deltas) angles into the actual circular coils
         that they represent. Also generates the symmetrized coils.
         """
-        from . import PSCCurve
+        from . import PSCCurve, CurvePlanarFourier
         from simsopt.field import apply_symmetries_to_psc_curves
 
         order = 1
         ncoils = self.num_psc
-        self.curves = [PSCCurve(order*self.ppp, order, nfp=1, stellsym=False, psc_array=self, index=i) for i in range(ncoils)]
+        self.curves = [CurvePlanarFourier(order*self.ppp, order, nfp=1, stellsym=False) for i in range(ncoils)]
         for ic in range(ncoils):
             alpha2 = self.alphas[ic] / 2.0
             delta2 = self.deltas[ic] / 2.0