diff --git a/docs/changes/889.feature.rst b/docs/changes/889.feature.rst new file mode 100644 index 000000000..d18e31f4c --- /dev/null +++ b/docs/changes/889.feature.rst @@ -0,0 +1 @@ +Added two new analytical model classes, `GeneralizedLorentz1D` and `SmoothBrokenPowerLaw`, to `stingray.simulator.models`. diff --git a/stingray/simulator/models.py b/stingray/simulator/models.py index c209d5a29..233868329 100644 --- a/stingray/simulator/models.py +++ b/stingray/simulator/models.py @@ -1,12 +1,12 @@ -from astropy.modeling.models import custom_model +import numpy as np +from astropy.modeling import Fittable1DModel +from astropy.modeling.parameters import InputParameterError, Parameter -# TODO: Add Jacobian -@custom_model -def GeneralizedLorentz1D(x, x_0=1.0, fwhm=1.0, value=1.0, power_coeff=1.0): +class GeneralizedLorentz1D(Fittable1DModel): """ Generalized Lorentzian function, - implemented using astropy.modeling.models custom model + implemented using astropy.modeling.models Lorentz1D Parameters ---------- @@ -25,26 +25,100 @@ def GeneralizedLorentz1D(x, x_0=1.0, fwhm=1.0, value=1.0, power_coeff=1.0): power_coeff : float power coefficient [n] + Notes + ----- + Model formula (with :math:`V` for ``value``, :math:`x_0` for ``x_0``, + :math:`w` for ``fwhm``, and :math:`p` for ``power_coeff``): + + .. math:: + + f(x) = V \\cdot \\left( \\frac{w}{2} \\right)^p + \\cdot \\frac{1}{\\left( |x - x_0|^p + \\left( \\frac{w}{2} \\right)^p \\right)} + Returns ------- model: astropy.modeling.Model generalized Lorentzian psd model """ - assert power_coeff > 0.0, "The power coefficient should be greater than zero." - return ( - value - * (fwhm / 2) ** power_coeff - * 1.0 - / (abs(x - x_0) ** power_coeff + (fwhm / 2) ** power_coeff) - ) - - -# TODO: Add Jacobian -@custom_model -def SmoothBrokenPowerLaw(x, norm=1.0, gamma_low=1.0, gamma_high=1.0, break_freq=1.0): + + x_0 = Parameter(default=1.0, description="Peak central frequency") + fwhm = Parameter(default=1.0, description="FWHM of the peak (gamma)") + value = Parameter(default=1.0, description="Peak value at x=x0") + power_coeff = Parameter(default=1.0, description="Power coefficient [n]") + + def _power_coeff_validator(self, power_coeff): + if not np.any(power_coeff > 0): + raise InputParameterError("The power coefficient should be greater than zero.") + + power_coeff._validator = _power_coeff_validator + + @staticmethod + def evaluate(x, x_0, fwhm, value, power_coeff): + """ + Generalized Lorentzian function + """ + fwhm_pc = np.power(fwhm / 2, power_coeff) + return value * fwhm_pc * 1.0 / (np.power(np.abs(x - x_0), power_coeff) + fwhm_pc) + + @staticmethod + def fit_deriv(x, x_0, fwhm, value, power_coeff): + """ + Gaussian1D model function derivatives with respect + to parameters. + """ + fwhm_pc = np.power(fwhm / 2, power_coeff) + num = value * fwhm_pc + mod_x_pc = np.power(np.abs(x - x_0), power_coeff) + denom = mod_x_pc + fwhm_pc + denom_sq = np.power(denom, 2) + + d_x_0 = 1.0 * num / denom_sq * (power_coeff * mod_x_pc / np.abs(x - x_0)) * np.sign(x - x_0) + d_value = fwhm_pc / denom + + pre_compute = 1.0 / 2.0 * power_coeff * fwhm_pc / (fwhm / 2) + d_fwhm = 1.0 / denom_sq * (denom * (value * pre_compute) - num * pre_compute) + + d_power_coeff = ( + 1.0 + / denom_sq + * ( + denom * (value * np.log(fwhm / 2) * fwhm_pc) + - num * (np.log(abs(x - x_0)) * mod_x_pc + np.log(fwhm / 2) * fwhm_pc) + ) + ) + return np.array([d_x_0, d_value, d_fwhm, d_power_coeff]) + + def bounding_box(self, factor=25): + """Tuple defining the default ``bounding_box`` limits, + ``(x_low, x_high)``. + + Parameters + ---------- + factor : float + The multiple of FWHM used to define the limits. + Default is chosen to include most (99%) of the + area under the curve, while still showing the + central feature of interest. + + """ + x0 = self.x_0 + dx = factor * self.fwhm + + return (x0 - dx, x0 + dx) + + # NOTE: + # In astropy 4.3 'Parameter' object has no attribute 'input_unit', + # whereas newer versions of Astropy include this attribute. + + # TODO: + # Add 'input_units' and '_parameter_units_for_data_units' methods + # when we drop support for Astropy < 5.3. + + +class SmoothBrokenPowerLaw(Fittable1DModel): """ Smooth broken power law function, - implemented using astropy.modeling.models custom model + implemented using astropy.modeling.models SmoothlyBrokenPowerLaw1D Parameters ---------- @@ -63,14 +137,137 @@ def SmoothBrokenPowerLaw(x, norm=1.0, gamma_low=1.0, gamma_high=1.0, break_freq= break_freq: float break frequency + Notes + ----- + Model formula (with :math:`N` for ``norm``, :math:`x_b` for + ``break_freq``, :math:`\\gamma_1` for ``gamma_low``, + and :math:`\\gamma_2` for ``gamma_high``): + + .. math:: + + f(x) = N \\cdot x^{-\\gamma_1} + \\left( 1 + \\left( \\frac{x}{x_b} \\right)^2 \\right)^{\\frac{\\gamma_1 - \\gamma_2}{2}} + Returns ------- model: astropy.modeling.Model generalized smooth broken power law psd model """ - return ( - norm * x ** (-gamma_low) / (1.0 + (x / break_freq) ** 2) ** (-(gamma_low - gamma_high) / 2) - ) + + norm = Parameter(default=1.0, description="normalization frequency") + gamma_low = Parameter(default=-1.0, description="Power law index for f --> zero") + gamma_high = Parameter(default=1.0, description="Power law index for f --> infinity") + break_freq = Parameter(default=1.0, description="Break frequency") + + def _norm_validator(self, value): + if np.any(value <= 0): + raise InputParameterError("norm parameter must be > 0") + + norm._validator = _norm_validator + + @staticmethod + def evaluate(x, norm, gamma_low, gamma_high, break_freq): + """One dimensional smoothly broken power law model function.""" + # Pre-calculate `x/x_b` + xx = x / break_freq + + # Initialize the return value + f = np.zeros_like(x, subok=False) + + # The quantity `t = (x / x_b)^(1 / 2)` can become quite + # large. To avoid overflow errors we will start by calculating + # its natural logarithm: + logt = np.log(xx) / 2 + + # When `t >> 1` or `t << 1` we don't actually need to compute + # the `t` value since the main formula (see docstring) can be + # significantly simplified by neglecting `1` or `t` + # respectively. In the following we will check whether `t` is + # much greater, much smaller, or comparable to 1 by comparing + # the `logt` value with an appropriate threshold. + threshold = 30 # corresponding to exp(30) ~ 1e13 + i = logt > threshold + if i.max(): + f[i] = norm * np.power(x[i], -gamma_low) * np.power(xx[i], gamma_low - gamma_high) + + i = logt < -threshold + if i.max(): + f[i] = norm * np.power(x[i], -gamma_low) + + i = np.abs(logt) <= threshold + if i.max(): + # In this case the `t` value is "comparable" to 1, hence + # we will evaluate the whole formula. + f[i] = ( + norm + * np.power(x[i], -gamma_low) + * np.power(1.0 + np.power(xx[i], 2), (gamma_low - gamma_high) / 2) + ) + return f + + @staticmethod + def fit_deriv(x, norm, gamma_low, gamma_high, break_freq): + """One dimensional smoothly broken power law derivative with respect + to parameters. + """ + # Pre-calculate `x_b` and `x/x_b` and `logt` (see comments in + # SmoothBrokenPowerLaw.evaluate) + xx = x / break_freq + logt = np.log(xx) / 2 + + # Initialize the return values + f = np.zeros_like(x) + d_norm = np.zeros_like(x) + d_gamma_low = np.zeros_like(x) + d_gamma_high = np.zeros_like(x) + d_break_freq = np.zeros_like(x) + + threshold = 30 # (see comments in SmoothBrokenPowerLaw.evaluate) + i = logt > threshold + if i.max(): + f[i] = norm * np.power(x[i], -gamma_low) * np.power(xx[i], gamma_low - gamma_high) + + d_norm[i] = f[i] / norm + d_gamma_low[i] = f[i] * (-np.log(x[i]) + np.log(xx[i])) + d_gamma_high[i] = -f[i] * np.log(xx[i]) + d_break_freq[i] = f[i] * (gamma_high - gamma_low) / break_freq + + i = logt < -threshold + if i.max(): + f[i] = norm * np.power(x[i], -gamma_low) + + d_norm[i] = f[i] / norm + d_gamma_low[i] = -f[i] * np.log(x[i]) + d_gamma_high[i] = 0 + d_break_freq[i] = 0 + + i = np.abs(logt) <= threshold + if i.max(): + # In this case the `t` value is "comparable" to 1, hence we + # we will evaluate the whole formula. + f[i] = ( + norm + * np.power(x[i], -gamma_low) + * np.power(1.0 + np.power(xx[i], 2), (gamma_low - gamma_high) / 2) + ) + d_norm[i] = f[i] / norm + d_gamma_low[i] = f[i] * (-np.log(x[i]) + np.log(1.0 + np.power(xx[i], 2)) / 2) + d_gamma_high[i] = -f[i] * np.log(1.0 + np.power(xx[i], 2)) / 2 + d_break_freq[i] = ( + f[i] + * (np.power(x[i], 2) * (gamma_high - gamma_low)) + / (break_freq * (np.power(break_freq, 2) + np.power(x[i], 2))) + ) + + return np.array([d_norm, d_gamma_low, d_gamma_high, d_break_freq]) + + # NOTE: + # In astropy 4.3 'Parameter' object has no attribute 'input_unit', + # whereas newer versions of Astropy include this attribute. + + # TODO: + # Add 'input_units' and '_parameter_units_for_data_units' methods + # when we drop support for Astropy < 5.3. def generalized_lorentzian(x, p): diff --git a/stingray/simulator/tests/test_models.py b/stingray/simulator/tests/test_models.py new file mode 100644 index 000000000..db25c37ac --- /dev/null +++ b/stingray/simulator/tests/test_models.py @@ -0,0 +1,88 @@ +import numpy as np +import pytest +from numpy.testing import assert_allclose + +from astropy.modeling.parameters import InputParameterError +from astropy.modeling import fitting + +from stingray.simulator import models + + +class TestModel(object): + @classmethod + def setup_class(self): + self.lorentz1D = models.GeneralizedLorentz1D(x_0=3, fwhm=32, value=2.5, power_coeff=2) + self.smoothPowerlaw = models.SmoothBrokenPowerLaw( + norm=1, gamma_low=-2, gamma_high=2, break_freq=10 + ) + + def test_model_param(self): + lorentz1D = self.lorentz1D + smoothPowerlaw = self.smoothPowerlaw + + assert np.allclose(smoothPowerlaw.parameters, np.array([1, -2, 2, 10])) + assert np.allclose(lorentz1D.parameters, np.array([3.0, 32.0, 2.5, 2.0])) + + assert np.array_equal( + lorentz1D.param_names, np.array(["x_0", "fwhm", "value", "power_coeff"]) + ) + assert np.array_equal( + smoothPowerlaw.param_names, np.array(["norm", "gamma_low", "gamma_high", "break_freq"]) + ) + + def test_lorentz_power_coeff(self): + with pytest.raises( + InputParameterError, match="The power coefficient should be greater than zero." + ): + models.GeneralizedLorentz1D(x_0=2, fwhm=100, value=3, power_coeff=-1) + + def test_lorentz_bounding_box(self): + lorentz1D = self.lorentz1D + + assert np.allclose(lorentz1D.bounding_box(), [-797, 803]) + + @pytest.mark.parametrize( + "model, yy_func, params, x_lim", + [ + (models.SmoothBrokenPowerLaw, models.smoothbknpo, [1, -2, 2, 10], [0.01, 100]), + ( + models.GeneralizedLorentz1D, + models.generalized_lorentzian, + [3, 32, 2.5, 2], + [-10, 10], + ), + ], + ) + def test_model_evaluate(self, model, yy_func, params, x_lim): + model = model(*params) + xx = np.logspace(x_lim[0], x_lim[1], 100) + yy = model(xx) + yy_ref = yy_func(xx, params) + + assert_allclose(yy, yy_ref, rtol=0, atol=1e-8) + assert xx.shape == yy.shape == yy_ref.shape + + @pytest.mark.parametrize( + "model, x_lim", + [ + (models.SmoothBrokenPowerLaw(1, -2, 2, 10), [0.01, 70]), + (models.GeneralizedLorentz1D(3, 32, 2.5, 2), [-10, 10]), + ], + ) + def test_model_fitting(self, model, x_lim): + x = np.logspace(x_lim[0], x_lim[1], 100) + + model_with_deriv = model + model_no_deriv = model + + # add 10% noise to the amplitude + rng = np.random.default_rng(0) + rsn_rand_0 = rng.random(x.shape) + n = 0.1 * (rsn_rand_0 - 0.5) + + data = model_with_deriv(x) + n + fitter_with_deriv = fitting.LevMarLSQFitter() + new_model_with_deriv = fitter_with_deriv(model_with_deriv, x, data) + fitter_no_deriv = fitting.LevMarLSQFitter() + new_model_no_deriv = fitter_no_deriv(model_no_deriv, x, data, estimate_jacobian=True) + assert_allclose(new_model_with_deriv.parameters, new_model_no_deriv.parameters, atol=0.5)