TwoPowerProfile

desc/profiles.py

@@ -12,7 +12,12 @@
 from desc.derivatives import Derivative
 from desc.grid import Grid, _Grid
 from desc.io import IOAble
-from desc.utils import combination_permutation, copy_coeffs, multinomial_coefficients
+from desc.utils import (
+    combination_permutation,
+    copy_coeffs,
+    errorif,
+    multinomial_coefficients,
+)
 
 
 class _Profile(IOAble, ABC):
@@ -687,6 +692,88 @@
         return cls(params, sym=sym, name=name)
 
 
+class PowerLawProfile(_Profile):
+    """Profile represented by a power law.
+
+    f(x) = a[0]*(1 - x**a[1])**a[2]
+
+    Parameters
+    ----------
+    params: array-like
+        Coefficients of the power law series. Must be an array of size 3.
+        Assumed to be zero if not specified.
+    name : str
+        Name of the profile.
+
+    """
+
+    _io_attrs_ = _Profile._io_attrs_
+
+    def __init__(self, params=None, name=""):
+        super().__init__(name)
+
+        if params is None:
+            params = [0, 0, 0]
+        self._params = np.atleast_1d(params)
+
+        errorif(
+            self._params.size != 3, ValueError, "params must be an array of size 3."
+        )
+
+    @property
+    def params(self):
+        """ndarray: Parameter values."""
+        return self._params
+
+    @params.setter
+    def params(self, new):
+        new = jnp.atleast_1d(jnp.asarray(new))
+        if new.size == 3:
+            self._params = jnp.asarray(new)
+        else:
+            raise ValueError(f"params should be an array of size 3, got {len(new)}.")
+
+    def compute(self, grid, params=None, dr=0, dt=0, dz=0):
+        """Compute values of profile at specified nodes.
+
+        Parameters
+        ----------
+        grid : Grid
+            Locations to compute values at.
+        params : array-like
+            Power law coefficients to use. Must be an array of size 3.
+            If not given, uses the values given by the params attribute.
+        dr, dt, dz : int
+            Derivative order in rho, theta, zeta.
+
+        Returns
+        -------
+        values : ndarray
+            Values of the profile or its derivative at the points specified.
+
+        """
+        if params is None:
+            params = self.params
+        if (dt != 0) or (dz != 0):
+            return jnp.zeros(grid.num_nodes)
+        a, b, c = params
+        r = grid.nodes[:, 0]
+        if dr == 0:
+            f = a * (1 - r**b) ** c
+        # df/dr = inf at rho = 1 for all dr > 0
+        elif dr == 1:
+            f = r ** (b - 1) * self.compute(grid, params=[-a * b * c, b, c - 1])
+        elif dr == 2:
+            f = (
+                r ** (b - 2)
+                * ((b * c - 1) * r**b - b + 1)
+                * self.compute(grid, params=[a * b * c, b, c - 2])
+            )
+        else:
+            raise NotImplementedError("dr > 2 not implemented for PowerLawProfile!")
+        return f
+
+
 class SplineProfile(_Profile):
     """Profile represented by a piecewise cubic spline.
 

tests/test_profiles.py

@@ -16,6 +16,7 @@
 from desc.profiles import (
     FourierZernikeProfile,
     MTanhProfile,
+    PowerLawProfile,
     PowerSeriesProfile,
     SplineProfile,
 )
@@ -89,6 +90,15 @@ def test_close_values(self):
         np.testing.assert_allclose(mp(x, dt=1), 0)
         np.testing.assert_allclose(mp(x, dz=1), 0)
 
+        pp = PowerSeriesProfile([0.5, 0, -1, 0, 0.5])
+        lp = PowerLawProfile([0.5, 2, 2])
+        np.testing.assert_allclose(pp(x), lp(x))
+        np.testing.assert_allclose(pp(x, dr=1), lp(x, dr=1))
+        np.testing.assert_allclose(pp(x, dr=2), lp(x, dr=2))
+        np.testing.assert_allclose(
+            pp.params, lp.to_powerseries(order=4, sym=True).params
+        )
+
     @pytest.mark.unit
     def test_PowerSeriesProfile_even_sym(self):
         """Test that even symmetry is enforced properly in PowerSeriesProfile."""
@@ -165,10 +175,12 @@ def test_SplineProfile_methods(self):
     def test_repr(self):
         """Test string representation of profile classes."""
         pp = PowerSeriesProfile(modes=np.array([0, 2, 4]), params=np.array([1, -2, 1]))
+        lp = PowerLawProfile([1, 2, 1])
         sp = pp.to_spline()
         mp = pp.to_mtanh(order=4, ftol=1e-4, xtol=1e-4)
         zp = pp.to_fourierzernike()
         assert "PowerSeriesProfile" in str(pp)
+        assert "PowerLawProfile" in str(lp)
         assert "SplineProfile" in str(sp)
         assert "MTanhProfile" in str(mp)
         assert "FourierZernikeProfile" in str(zp)
@@ -195,6 +207,11 @@ def test_get_set(self):
         sp.params = sp.params + 1
         np.testing.assert_allclose(sp.params, 1 + pp(sp._knots))
 
+        lp = PowerLawProfile([1, 2, 1 / 2])
+        np.testing.assert_allclose(lp.params, [1, 2, 1 / 2])
+        lp.params = [2, 1, 1 / 3]
+        np.testing.assert_allclose(lp.params, [2, 1, 1 / 3])
+
         zp = FourierZernikeProfile([1 - 1 / 3 - 1 / 6, -1 / 2, 1 / 6])
         assert zp.get_params(2, 0, 0) == -1 / 2
         zp.set_params(2, 0, 0, 1 / 4)
@@ -356,17 +373,21 @@ def test_profile_errors(self):
         with pytest.raises(ValueError):
             a = sp * pp
             a.params = sp.params
+        with pytest.raises(ValueError):
+            _ = PowerLawProfile([1, 2, 3, 4])
 
     @pytest.mark.unit
     def test_default_profiles(self):
         """Test that default profiles are just zeros."""
         pp = PowerSeriesProfile()
+        lp = PowerLawProfile()
         sp = SplineProfile()
         mp = MTanhProfile()
         zp = FourierZernikeProfile()
 
         x = np.linspace(0, 1, 10)
         np.testing.assert_allclose(pp(x), 0)
+        np.testing.assert_allclose(lp(x), 0)
         np.testing.assert_allclose(sp(x), 0)
         np.testing.assert_allclose(mp(x), 0)
         np.testing.assert_allclose(zp(x), 0)