Added GeneralizedLorentz1D and SmoothBrokenPowerLaw classes to stingray.simulator.models.

docs/changes/889.feature.rst

@@ -0,0 +1 @@
+Added two new analytical model classes, `GeneralizedLorentz1D` and `SmoothBrokenPowerLaw`, to `stingray.simulator.models`.

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):

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)