Add NFP to ScalarPotentialField and its subclasses

CHANGELOG.md

@@ -13,6 +13,7 @@ New Features
     * use of both this and the ``QuadraticFlux`` objective allows for REGCOIL solutions to be obtained through the optimization framework, and combined with other objectives as well.
 - Changes local area weighting of Bn in QuadraticFlux objective to be the square root of the local area element (Note that any existing optimizations using this objective may need different weights to achieve the same result now.)
 - Adds a new tutorial showing how to use``REGCOIL`` features.
+- Adds an ``NFP`` attribute to ``ScalarPotentialField``, ``VectorPotentialField`` and ``DommaschkPotentialField``, to allow ``SplineMagneticField.from_field`` and ``MagneticField.save_mgrid`` to efficiently take advantage of the discrete toroidal symmetry of these fields, if present.
 
 
 Bug Fixes

desc/magnetic_fields/_core.py

@@ -1885,7 +1885,7 @@
 
     @classmethod
     def from_field(
-        cls, field, R, phi, Z, params=None, method="cubic", extrap=False, NFP=1
+        cls, field, R, phi, Z, params=None, method="cubic", extrap=False, NFP=None
     ):
         """Create a splined magnetic field from another field for faster evaluation.
 
@@ -1903,14 +1903,16 @@
         extrap : bool
             whether to extrapolate splines beyond specified R,phi,Z
         NFP : int, optional
-            Number of toroidal field periods.
+            Number of toroidal field periods.  If not provided, will default to 1 or
+        the provided field's NFP, if it has that attribute.
 
         """
         R, phi, Z = map(np.asarray, (R, phi, Z))
         rr, pp, zz = np.meshgrid(R, phi, Z, indexing="ij")
         shp = rr.shape
         coords = np.array([rr.flatten(), pp.flatten(), zz.flatten()]).T
         BR, BP, BZ = field.compute_magnetic_field(coords, params, basis="rpz").T
+        NFP = getattr(field, "_NFP", 1)
         try:
             AR, AP, AZ = field.compute_magnetic_vector_potential(
                 coords, params, basis="rpz"
@@ -1948,12 +1950,22 @@
         R,phi,Z are arrays of cylindrical coordinates.
     params : dict, optional
         default parameters to pass to potential function
+    NFP : int, optional
+        Whether the field has a discrete periodicity. This is only used when making
+        a ``SplineMagneticField`` from this field using its ``from_field`` method,
+        or when saving this field as an mgrid file using the ``save_mgrid`` method.
 
     """
 
-    def __init__(self, potential, params=None):
+    def __init__(self, potential, params=None, NFP=1):
         self._potential = potential
         self._params = params
+        self._NFP = NFP
+
+    @property
+    def NFP(self):
+        """int: Number of (toroidal) field periods."""
+        return self._NFP
 
     def compute_magnetic_field(
         self, coords, params=None, basis="rpz", source_grid=None, transforms=None
@@ -2042,12 +2054,22 @@
         R,phi,Z are arrays of cylindrical coordinates.
     params : dict, optional
         default parameters to pass to potential function
+    NFP : int, optional
+        Whether the field has a discrete periodicity. This is only used when making
+        a ``SplineMagneticField`` from this field using its ``from_field`` method,
+        or when saving this field as an mgrid file using the ``save_mgrid`` method.
 
     """
 
-    def __init__(self, potential, params=None):
+    def __init__(self, potential, params=None, NFP=1):
         self._potential = potential
         self._params = params
+        self._NFP = NFP
+
+    @property
+    def NFP(self):
+        """int: Number of (toroidal) field periods."""
+        return self._NFP
 
     def _compute_A_or_B(
         self,

desc/magnetic_fields/_dommaschk.py

@@ -38,6 +38,10 @@ class DommaschkPotentialField(ScalarPotentialField):
         d_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*N_m_l-1 terms
     B0: float
         scale strength of the magnetic field's 1/R portion
+    NFP : int, optional
+        Whether the field has a discrete periodicity. This is only used when making
+        a ``SplineMagneticField`` from this field using its ``from_field`` method,
+        or when saving this field as an mgrid file using the ``save_mgrid`` method.
 
     """
 
@@ -50,6 +54,7 @@ def __init__(
         c_arr=jnp.array([0.0]),
         d_arr=jnp.array([0.0]),
         B0=1.0,
+        NFP=1,
     ):
         ms = jnp.atleast_1d(jnp.asarray(ms))
         ls = jnp.atleast_1d(jnp.asarray(ls))
@@ -68,6 +73,11 @@ def __init__(
             jnp.isscalar(B0) or jnp.atleast_1d(B0).size == 1
         ), "B0 should be a scalar value!"
 
+        ms_over_NFP = ms / NFP
+        assert jnp.allclose(
+            ms_over_NFP, ms_over_NFP.astype(int)
+        ), "To enforce desired NFP, `ms` should be all integer multiples of NFP"
+
         params = {}
         params["ms"] = ms
         params["ls"] = ls
@@ -77,7 +87,7 @@ def __init__(
         params["d_arr"] = d_arr
         params["B0"] = B0
 
-        super().__init__(dommaschk_potential, params)
+        super().__init__(dommaschk_potential, params, NFP)
 
     @classmethod
     def fit_magnetic_field(  # noqa: C901

tests/test_magnetic_fields.py

@@ -1045,11 +1045,11 @@ def test_fourier_current_potential_field_coil_cut_warnings(self):
     @pytest.mark.unit
     def test_spline_field(self, tmpdir_factory):
         """Test accuracy of spline magnetic field."""
-        field1 = ScalarPotentialField(phi_lm, args)
+        field1 = ScalarPotentialField(phi_lm, args, NFP=5)
         R = np.linspace(0.5, 1.5, 20)
         Z = np.linspace(-1.5, 1.5, 20)
         p = np.linspace(0, 2 * np.pi / 5, 40)
-        field2 = SplineMagneticField.from_field(field1, R, p, Z, NFP=5)
+        field2 = SplineMagneticField.from_field(field1, R, p, Z)
 
         np.testing.assert_allclose(
             field1([1.0, 1.0, 1.0]), field2([1.0, 1.0, 1.0]), rtol=1e-2, atol=1e-2
@@ -1450,11 +1450,13 @@ def test_dommaschk_field_errors():
     b_arr = [1]
     c_arr = [1]
     d_arr = [1, 1]  # length is not equal to the rest
-    with pytest.raises(AssertionError):
+    with pytest.raises(AssertionError, match="size"):
         DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr)
-    d_arr = [1]
+    d_arr = [1]  # test with incorrect NFP
+    with pytest.raises(AssertionError, match="desired NFP"):
+        DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr, NFP=2)
     ms = [-1]  # negative mode number
-    with pytest.raises(AssertionError):
+    with pytest.raises(AssertionError, match=">= 0"):
         DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr)