Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

TwoPowerProfile #984

Merged
merged 12 commits into from
Apr 16, 2024
104 changes: 103 additions & 1 deletion desc/profiles.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -687,6 +693,102 @@
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(
ddudt marked this conversation as resolved.
Show resolved Hide resolved
self._params[1] < 1,
UserWarning,
"Derivatives of this profile will be infinite at rho=0!",
)
warnif(
self._params[2] < 1,
UserWarning,
"Derivatives of this profile will be infinite at rho=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)}.")

Check warning on line 750 in desc/profiles.py

View check run for this annotation

Codecov / codecov/patch

desc/profiles.py#L750

Added line #L750 was not covered by tests

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)

Check warning on line 774 in desc/profiles.py

View check run for this annotation

Codecov / codecov/patch

desc/profiles.py#L774

Added line #L774 was not covered by tests
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!")

Check warning on line 788 in desc/profiles.py

View check run for this annotation

Codecov / codecov/patch

desc/profiles.py#L788

Added line #L788 was not covered by tests
ddudt marked this conversation as resolved.
Show resolved Hide resolved
return f


class SplineProfile(_Profile):
"""Profile represented by a piecewise cubic spline.

Expand Down
1 change: 1 addition & 0 deletions docs/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,7 @@ Profiles
:template: class.rst

desc.profiles.PowerSeriesProfile
desc.profiles.TwoPowerProfile
desc.profiles.SplineProfile
desc.profiles.MTanhProfile
desc.profiles.ScaledProfile
Expand Down
21 changes: 21 additions & 0 deletions tests/test_profiles.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
MTanhProfile,
PowerSeriesProfile,
SplineProfile,
TwoPowerProfile,
)

from .utils import area_difference, compute_coords
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -356,17 +373,21 @@ 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])

@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)
Expand Down
Loading