Add Vector Potential Calculation to Coil classes and Most MagneticField Classes

desc/magnetic_fields/_core.py

@@ -239,6 +239,31 @@ def __call__(self, grid, params=None, basis="rpz"):
         """Compute magnetic field at a set of points."""
         return self.compute_magnetic_field(grid, params, basis)
 
+    @abstractmethod
+    def compute_magnetic_vector_potential(
+        self, coords, params=None, basis="rpz", source_grid=None
+    ):
+        """Compute magnetic vector potential at a set of points.
+
+        Parameters
+        ----------
+        coords : array-like shape(n,3)
+            Nodes to evaluate vector potential at in [R,phi,Z] or [X,Y,Z] coordinates.
+        params : dict or array-like of dict, optional
+            Dictionary of optimizable parameters, eg field.params_dict.
+        basis : {"rpz", "xyz"}
+            Basis for input coordinates and returned magnetic vector potential.
+        source_grid : Grid, int or None or array-like, optional
+            Grid used to discretize MagneticField object if calculating A from
+            Biot-Savart. Should NOT include endpoint at 2pi.
+
+        Returns
+        -------
+        A : ndarray, shape(N,3)
+            magnetic vector potential at specified points
+
+        """
+
     def compute_Bnormal(
         self,
         surface,
@@ -734,6 +759,33 @@ def compute_magnetic_field(
             coords, params, basis, source_grid
         )
 
+    def compute_magnetic_vector_potential(
+        self, coords, params=None, basis="rpz", source_grid=None
+    ):
+        """Compute magnetic vector potential at a set of points.
+
+        Parameters
+        ----------
+        coords : array-like shape(n,3)
+            Nodes to evaluate vector potential at in [R,phi,Z] or [X,Y,Z] coordinates.
+        params : dict or array-like of dict, optional
+            Dictionary of optimizable parameters, eg field.params_dict.
+        basis : {"rpz", "xyz"}
+            Basis for input coordinates and returned magnetic vector potential.
+        source_grid : Grid, int or None or array-like, optional
+            Grid used to discretize MagneticField object if calculating A from
+            Biot-Savart. Should NOT include endpoint at 2pi.
+
+        Returns
+        -------
+        A : ndarray, shape(N,3)
+            scaled magnetic vector potential at specified points
+
+        """
+        return self._scale * self._field.compute_magnetic_vector_potential(
+            coords, params, basis, source_grid
+        )
+
 
 class SumMagneticField(_MagneticField, MutableSequence, OptimizableCollection):
     """Sum of two or more magnetic field sources.
@@ -799,6 +851,50 @@ def compute_magnetic_field(
 
         return B
 
+    def compute_magnetic_vector_potential(
+        self, coords, params=None, basis="rpz", source_grid=None
+    ):
+        """Compute magnetic vector potential at a set of points.
+
+        Parameters
+        ----------
+        coords : array-like shape(n,3)
+            Nodes to evaluate vector potential at in [R,phi,Z] or [X,Y,Z] coordinates.
+        params : dict or array-like of dict, optional
+            Dictionary of optimizable parameters, eg field.params_dict.
+        basis : {"rpz", "xyz"}
+            Basis for input coordinates and returned magnetic vector potential.
+        source_grid : Grid, int or None or array-like, optional
+            Grid used to discretize MagneticField object if calculating A from
+            Biot-Savart. Should NOT include endpoint at 2pi.
+
+        Returns
+        -------
+        A : ndarray, shape(N,3)
+            scaled magnetic vector potential at specified points
+
+        """
+        if params is None:
+            params = [None] * len(self._fields)
+        if isinstance(params, dict):
+            params = [params]
+        if source_grid is None:
+            source_grid = [None] * len(self._fields)
+        if not isinstance(source_grid, (list, tuple)):
+            source_grid = [source_grid]
+        if len(source_grid) != len(self._fields):
+            # ensure that if source_grid is shorter, that it is simply repeated so that
+            # zip does not terminate early
+            source_grid = source_grid * len(self._fields)
+
+        A = 0
+        for i, (field, g) in enumerate(zip(self._fields, source_grid)):
+            A += field.compute_magnetic_vector_potential(
+                coords, params[i % len(params)], basis, source_grid=g
+            )
+
+        return A
+
     # dunder methods required by MutableSequence
     def __getitem__(self, i):
         return self._fields[i]
@@ -904,6 +1000,43 @@ def compute_magnetic_field(
 
         return B
 
+    def compute_magnetic_vector_potential(
+        self, coords, params=None, basis="rpz", source_grid=None
+    ):
+        """Compute magnetic vector potential at a set of points.
+
+        Parameters
+        ----------
+        coords : array-like shape(n,3)
+            Nodes to evaluate vector potential at in [R,phi,Z] or [X,Y,Z] coordinates.
+        params : dict or array-like of dict, optional
+            Dict of values for R0 and B0.
+        basis : {"rpz", "xyz"}
+            Basis for input coordinates and returned magnetic vector potential.
+        source_grid : Grid, int or None or array-like, optional
+            Unused by this MagneticField class.
+
+        Returns
+        -------
+        A : ndarray, shape(N,3)
+            magnetic vector potential at specified points
+
+        """
+        params = setdefault(params, {})
+        B0 = params.get("B0", self.B0)
+        R0 = params.get("R0", self.R0)
+
+        assert basis.lower() in ["rpz", "xyz"]
+        coords = jnp.atleast_2d(jnp.asarray(coords))
+        if basis == "xyz":
+            coords = xyz2rpz(coords)
+        az = -B0 * R0 * jnp.log(coords[:, 0])
+        arp = jnp.zeros_like(az)
+        A = jnp.array([arp, arp, az]).T
+        # b/c it only has a nonzero z component, no need
+        # to switch bases back if xyz is given
+        return A
+
 
 class VerticalMagneticField(_MagneticField, Optimizable):
     """Uniform magnetic field purely in the vertical (Z) direction.
@@ -955,18 +1088,52 @@ def compute_magnetic_field(
         params = setdefault(params, {})
         B0 = params.get("B0", self.B0)
 
-        assert basis.lower() in ["rpz", "xyz"]
         coords = jnp.atleast_2d(jnp.asarray(coords))
-        if basis == "xyz":
-            coords = xyz2rpz(coords)
         bz = B0 * jnp.ones_like(coords[:, 2])
         brp = jnp.zeros_like(bz)
         B = jnp.array([brp, brp, bz]).T
-        if basis == "xyz":
-            B = rpz2xyz_vec(B, phi=coords[:, 1])
+        # b/c it only has a nonzero z component, no need
+        # to switch bases back if xyz is given
 
         return B
 
+    def compute_magnetic_vector_potential(
+        self, coords, params=None, basis="rpz", source_grid=None
+    ):
+        """Compute magnetic vector potential at a set of points.
+
+        Parameters
+        ----------
+        coords : array-like shape(n,3)
+            Nodes to evaluate vector potential at in [R,phi,Z] or [X,Y,Z] coordinates.
+        params : dict or array-like of dict, optional
+            Dict of values for B0.
+        basis : {"rpz", "xyz"}
+            Basis for input coordinates and returned magnetic vector potential.
+        source_grid : Grid, int or None or array-like, optional
+            Unused by this MagneticField class.
+
+        Returns
+        -------
+        A : ndarray, shape(N,3)
+            magnetic vector potential at specified points
+
+        """
+        params = setdefault(params, {})
+        B0 = params.get("B0", self.B0)
+
+        assert basis.lower() in ["rpz", "xyz"]
+        coords = jnp.atleast_2d(jnp.asarray(coords))
+        if basis == "xyz":
+            coords = xyz2rpz(coords)
+        axy = (B0 / 2) * jnp.ones_like(coords[:, 2])
+        az = jnp.zeros_like(axy)
+        A = jnp.array([axy, -axy, az]).T
+        if basis == "xyz":
+            A = rpz2xyz_vec(A, phi=coords[:, 1])
+
+        return A
+
 
 class PoloidalMagneticField(_MagneticField, Optimizable):
     """Pure poloidal magnetic field (ie in theta direction).
@@ -1074,6 +1241,33 @@ def compute_magnetic_field(
 
         return B
 
+    def compute_magnetic_vector_potential(
+        self, coords, params=None, basis="rpz", source_grid=None
+    ):
+        """Compute magnetic vector potential at a set of points.
+
+        Parameters
+        ----------
+        coords : array-like shape(n,3)
+            Nodes to evaluate vector potential at in [R,phi,Z] or [X,Y,Z] coordinates.
+        params : dict or array-like of dict, optional
+            Dict of values for B0.
+        basis : {"rpz", "xyz"}
+            Basis for input coordinates and returned magnetic vector potential.
+        source_grid : Grid, int or None or array-like, optional
+            Unused by this MagneticField class.
+
+        Returns
+        -------
+        A : ndarray, shape(N,3)
+            magnetic vector potential at specified points
+
+        """
+        raise NotImplementedError(
+            "PoloidalMagneticField has nonzero divergence, therefore it can't be "
+            "represented with a vector potential."
+        )
+
 
 class SplineMagneticField(_MagneticField, Optimizable):
     """Magnetic field from precomputed values on a grid.

tests/test_magnetic_fields.py

@@ -104,17 +104,37 @@ def test_combined_fields(self):
         assert scaled_field.B0 == 2
         assert scaled_field.scale == 3.1
         np.testing.assert_allclose(scaled_field([1.0, 0, 0]), np.array([[0, 6.2, 0]]))
+        np.testing.assert_allclose(
+            scaled_field.compute_magnetic_vector_potential([2.0, 0, 0]),
+            np.array([[0, 0, -3.1 * 2 * 1 * np.log(2)]]),
+        )
+
         scaled_field.R0 = 1.3
         scaled_field.scale = 1.0
         np.testing.assert_allclose(scaled_field([1.3, 0, 0]), np.array([[0, 2, 0]]))
+        np.testing.assert_allclose(
+            scaled_field.compute_magnetic_vector_potential([2.0, 0, 0]),
+            np.array([[0, 0, -2 * 1.3 * np.log(2)]]),
+        )
         assert scaled_field.optimizable_params == ["B0", "R0", "scale"]
         assert hasattr(scaled_field, "B0")
 
         sum_field = vfield + pfield + tfield
+        sum_field_tv = vfield + tfield  # to test A since pfield does not have A
         assert len(sum_field) == 3
+        assert len(sum_field_tv) == 2
+
         np.testing.assert_allclose(
             sum_field([1.3, 0, 0.0]), [[0.0, 2, 3.2 + 2 * 1.2 * 0.3]]
         )
+        R = 1.3
+        tfield_A = np.array([[0, 0, -2 * 1.3 * np.log(R)]])
+        vfield_A = np.array([[vfield.B0, -vfield.B0, 0]]) / 2
+        np.testing.assert_allclose(
+            sum_field_tv.compute_magnetic_vector_potential([R, 0, 0.0], basis="xyz"),
+            tfield_A + vfield_A,
+        )
+
         assert sum_field.optimizable_params == [
             ["B0"],
             ["B0", "R0", "iota"],
@@ -337,7 +357,6 @@ def test_current_potential_vector_potential(self):
         )
         # test the loop integral of A around a curve encompassing the torus
         # against the analytic result for flux in an ideal toroidal solenoid
-        ## expression for flux inside of toroidal solenoid of radius a
         prefactors = mu_0 * G / 2 / jnp.pi
         correct_flux = -2 * np.pi * prefactors * (np.sqrt(R0**2 - a**2) - R0)