diff --git a/desc/profiles.py b/desc/profiles.py index 65e413989f..2190359501 100644 --- a/desc/profiles.py +++ b/desc/profiles.py @@ -12,7 +12,13 @@ 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, + warnif, +) class _Profile(IOAble, ABC): @@ -687,6 +693,104 @@ def from_values(cls, x, y, order=6, rcond=None, w=None, sym="auto", name=""): return cls(params, sym=sym, name=name) +class TwoPowerProfile(_Profile): + """Profile represented by two powers. + + f(x) = a[0]*(1 - x**a[1])**a[2] + + Notes + ----- + df/dx = inf at x = 0 if a[1] < 1 + df/dx = inf at x = 1 if a[2] < dr + + Parameters + ---------- + params: array-like + Coefficients of the two power formula. Must be an array of size 3. + Default if not specified is [0, 1, 1]. + 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, 1, 1] + self._params = np.atleast_1d(params) + + errorif( + self._params.size != 3, ValueError, "params must be an array of size 3." + ) + warnif( + self._params[1] < 1, + UserWarning, + "Derivatives of this profile will be infinite at rho=0 " + + "because params[1] < 1.", + ) + warnif( + self._params[2] < 1, + UserWarning, + "Derivatives of this profile will be infinite at rho=1 " + + "because params[2] < 1.", + ) + + @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 + 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 TwoPowerProfile!") + return f + + class SplineProfile(_Profile): """Profile represented by a piecewise cubic spline. @@ -917,6 +1021,8 @@ def compute(self, grid, params=None, dr=0, dt=0, dz=0): """ if params is None: params = self.params + if dr > 2: + raise NotImplementedError("dr > 2 not implemented for MTanhProfile!") if dt != 0 or dz != 0: return jnp.zeros_like(grid.nodes[:, 0]) diff --git a/docs/api.rst b/docs/api.rst index 8d730601b0..7f5cb4c9ae 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -274,6 +274,7 @@ Profiles :template: class.rst desc.profiles.PowerSeriesProfile + desc.profiles.TwoPowerProfile desc.profiles.SplineProfile desc.profiles.MTanhProfile desc.profiles.ScaledProfile diff --git a/tests/test_profiles.py b/tests/test_profiles.py index f81d6e2d42..530fd5f765 100644 --- a/tests/test_profiles.py +++ b/tests/test_profiles.py @@ -18,6 +18,7 @@ MTanhProfile, PowerSeriesProfile, SplineProfile, + TwoPowerProfile, ) from .utils import area_difference, compute_coords @@ -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]) + tp = TwoPowerProfile([0.5, 2, 2]) + np.testing.assert_allclose(pp(x), tp(x)) + np.testing.assert_allclose(pp(x, dr=1), tp(x, dr=1)) + np.testing.assert_allclose(pp(x, dr=2), tp(x, dr=2)) + np.testing.assert_allclose( + pp.params, tp.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])) + tp = TwoPowerProfile([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 "TwoPowerProfile" in str(tp) 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)) + tp = TwoPowerProfile([1, 2, 3]) + np.testing.assert_allclose(tp.params, [1, 2, 3]) + tp.params = [1 / 2, 3 / 2, 4 / 3] + np.testing.assert_allclose(tp.params, [1 / 2, 3 / 2, 4 / 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) @@ -337,6 +354,8 @@ def test_profile_errors(self): sp = pp.to_spline() zp = pp.to_fourierzernike() mp = pp.to_mtanh(order=4, ftol=1e-4, xtol=1e-4) + tp = TwoPowerProfile(params=np.array([0.5, 2, 1.5])) + grid = LinearGrid(L=9) with pytest.raises(ValueError): zp.params = 4 @@ -356,17 +375,25 @@ def test_profile_errors(self): with pytest.raises(ValueError): a = sp * pp a.params = sp.params + with pytest.raises(ValueError): + _ = TwoPowerProfile([1, 2, 3, 4]) + with pytest.raises(NotImplementedError): + tp.compute(grid, dr=3) + with pytest.raises(NotImplementedError): + mp.compute(grid, dr=3) @pytest.mark.unit def test_default_profiles(self): """Test that default profiles are just zeros.""" pp = PowerSeriesProfile() + tp = TwoPowerProfile() sp = SplineProfile() mp = MTanhProfile() zp = FourierZernikeProfile() x = np.linspace(0, 1, 10) np.testing.assert_allclose(pp(x), 0) + np.testing.assert_allclose(tp(x), 0) np.testing.assert_allclose(sp(x), 0) np.testing.assert_allclose(mp(x), 0) np.testing.assert_allclose(zp(x), 0)