1# Licensed under a 3-clause BSD style license - see LICENSE.rst
2"""
3Power law model variants
4"""
5# pylint: disable=invalid-name
6import numpy as np
7
8from astropy.units import Quantity
9from .core import Fittable1DModel
10from .parameters import Parameter, InputParameterError
11
12
13__all__ = ['PowerLaw1D', 'BrokenPowerLaw1D', 'SmoothlyBrokenPowerLaw1D',
14           'ExponentialCutoffPowerLaw1D', 'LogParabola1D']
15
16
17class PowerLaw1D(Fittable1DModel):
18    """
19    One dimensional power law model.
20
21    Parameters
22    ----------
23    amplitude : float
24        Model amplitude at the reference point
25    x_0 : float
26        Reference point
27    alpha : float
28        Power law index
29
30    See Also
31    --------
32    BrokenPowerLaw1D, ExponentialCutoffPowerLaw1D, LogParabola1D
33
34    Notes
35    -----
36    Model formula (with :math:`A` for ``amplitude`` and :math:`\\alpha` for ``alpha``):
37
38        .. math:: f(x) = A (x / x_0) ^ {-\\alpha}
39
40    """
41
42    amplitude = Parameter(default=1, description="Peak value at the reference point")
43    x_0 = Parameter(default=1, description="Reference point")
44    alpha = Parameter(default=1, description="Power law index")
45
46    @staticmethod
47    def evaluate(x, amplitude, x_0, alpha):
48        """One dimensional power law model function"""
49        xx = x / x_0
50        return amplitude * xx ** (-alpha)
51
52    @staticmethod
53    def fit_deriv(x, amplitude, x_0, alpha):
54        """One dimensional power law derivative with respect to parameters"""
55
56        xx = x / x_0
57
58        d_amplitude = xx ** (-alpha)
59        d_x_0 = amplitude * alpha * d_amplitude / x_0
60        d_alpha = -amplitude * d_amplitude * np.log(xx)
61
62        return [d_amplitude, d_x_0, d_alpha]
63
64    @property
65    def input_units(self):
66        if self.x_0.unit is None:
67            return None
68        return {self.inputs[0]: self.x_0.unit}
69
70    def _parameter_units_for_data_units(self, inputs_unit, outputs_unit):
71        return {'x_0': inputs_unit[self.inputs[0]],
72                'amplitude': outputs_unit[self.outputs[0]]}
73
74
75class BrokenPowerLaw1D(Fittable1DModel):
76    """
77    One dimensional power law model with a break.
78
79    Parameters
80    ----------
81    amplitude : float
82        Model amplitude at the break point.
83    x_break : float
84        Break point.
85    alpha_1 : float
86        Power law index for x < x_break.
87    alpha_2 : float
88        Power law index for x > x_break.
89
90    See Also
91    --------
92    PowerLaw1D, ExponentialCutoffPowerLaw1D, LogParabola1D
93
94    Notes
95    -----
96    Model formula (with :math:`A` for ``amplitude`` and :math:`\\alpha_1`
97    for ``alpha_1`` and :math:`\\alpha_2` for ``alpha_2``):
98
99        .. math::
100
101            f(x) = \\left \\{
102                     \\begin{array}{ll}
103                       A (x / x_{break}) ^ {-\\alpha_1} & : x < x_{break} \\\\
104                       A (x / x_{break}) ^ {-\\alpha_2} & :  x > x_{break} \\\\
105                     \\end{array}
106                   \\right.
107    """
108
109    amplitude = Parameter(default=1, description="Peak value at break point")
110    x_break = Parameter(default=1, description="Break point")
111    alpha_1 = Parameter(default=1, description="Power law index before break point")
112    alpha_2 = Parameter(default=1, description="Power law index after break point")
113
114    @staticmethod
115    def evaluate(x, amplitude, x_break, alpha_1, alpha_2):
116        """One dimensional broken power law model function"""
117
118        alpha = np.where(x < x_break, alpha_1, alpha_2)
119        xx = x / x_break
120        return amplitude * xx ** (-alpha)
121
122    @staticmethod
123    def fit_deriv(x, amplitude, x_break, alpha_1, alpha_2):
124        """One dimensional broken power law derivative with respect to parameters"""
125
126        alpha = np.where(x < x_break, alpha_1, alpha_2)
127        xx = x / x_break
128
129        d_amplitude = xx ** (-alpha)
130        d_x_break = amplitude * alpha * d_amplitude / x_break
131        d_alpha = -amplitude * d_amplitude * np.log(xx)
132        d_alpha_1 = np.where(x < x_break, d_alpha, 0)
133        d_alpha_2 = np.where(x >= x_break, d_alpha, 0)
134
135        return [d_amplitude, d_x_break, d_alpha_1, d_alpha_2]
136
137    @property
138    def input_units(self):
139        if self.x_break.unit is None:
140            return None
141        return {self.inputs[0]: self.x_break.unit}
142
143    def _parameter_units_for_data_units(self, inputs_unit, outputs_unit):
144        return {'x_break': inputs_unit[self.inputs[0]],
145                'amplitude': outputs_unit[self.outputs[0]]}
146
147
148class SmoothlyBrokenPowerLaw1D(Fittable1DModel):
149    """One dimensional smoothly broken power law model.
150
151    Parameters
152    ----------
153    amplitude : float
154        Model amplitude at the break point.
155    x_break : float
156        Break point.
157    alpha_1 : float
158        Power law index for ``x << x_break``.
159    alpha_2 : float
160        Power law index for ``x >> x_break``.
161    delta : float
162        Smoothness parameter.
163
164    See Also
165    --------
166    BrokenPowerLaw1D
167
168    Notes
169    -----
170    Model formula (with :math:`A` for ``amplitude``, :math:`x_b` for
171    ``x_break``, :math:`\\alpha_1` for ``alpha_1``,
172    :math:`\\alpha_2` for ``alpha_2`` and :math:`\\Delta` for
173    ``delta``):
174
175        .. math::
176
177            f(x) = A \\left( \\frac{x}{x_b} \\right) ^ {-\\alpha_1}
178                   \\left\\{
179                      \\frac{1}{2}
180                      \\left[
181                        1 + \\left( \\frac{x}{x_b}\\right)^{1 / \\Delta}
182                      \\right]
183                   \\right\\}^{(\\alpha_1 - \\alpha_2) \\Delta}
184
185
186    The change of slope occurs between the values :math:`x_1`
187    and :math:`x_2` such that:
188
189        .. math::
190            \\log_{10} \\frac{x_2}{x_b} = \\log_{10} \\frac{x_b}{x_1}
191            \\sim \\Delta
192
193
194    At values :math:`x \\lesssim x_1` and :math:`x \\gtrsim x_2` the
195    model is approximately a simple power law with index
196    :math:`\\alpha_1` and :math:`\\alpha_2` respectively.  The two
197    power laws are smoothly joined at values :math:`x_1 < x < x_2`,
198    hence the :math:`\\Delta` parameter sets the "smoothness" of the
199    slope change.
200
201    The ``delta`` parameter is bounded to values greater than 1e-3
202    (corresponding to :math:`x_2 / x_1 \\gtrsim 1.002`) to avoid
203    overflow errors.
204
205    The ``amplitude`` parameter is bounded to positive values since
206    this model is typically used to represent positive quantities.
207
208
209    Examples
210    --------
211    .. plot::
212        :include-source:
213
214        import numpy as np
215        import matplotlib.pyplot as plt
216        from astropy.modeling import models
217
218        x = np.logspace(0.7, 2.3, 500)
219        f = models.SmoothlyBrokenPowerLaw1D(amplitude=1, x_break=20,
220                                            alpha_1=-2, alpha_2=2)
221
222        plt.figure()
223        plt.title("amplitude=1, x_break=20, alpha_1=-2, alpha_2=2")
224
225        f.delta = 0.5
226        plt.loglog(x, f(x), '--', label='delta=0.5')
227
228        f.delta = 0.3
229        plt.loglog(x, f(x), '-.', label='delta=0.3')
230
231        f.delta = 0.1
232        plt.loglog(x, f(x), label='delta=0.1')
233
234        plt.axis([x.min(), x.max(), 0.1, 1.1])
235        plt.legend(loc='lower center')
236        plt.grid(True)
237        plt.show()
238
239    """
240
241    amplitude = Parameter(default=1, min=0, description="Peak value at break point")
242    x_break = Parameter(default=1, description="Break point")
243    alpha_1 = Parameter(default=-2, description="Power law index before break point")
244    alpha_2 = Parameter(default=2, description="Power law index after break point")
245    delta = Parameter(default=1, min=1.e-3, description="Smoothness Parameter")
246
247    @amplitude.validator
248    def amplitude(self, value):
249        if np.any(value <= 0):
250            raise InputParameterError(
251                "amplitude parameter must be > 0")
252
253    @delta.validator
254    def delta(self, value):
255        if np.any(value < 0.001):
256            raise InputParameterError(
257                "delta parameter must be >= 0.001")
258
259    @staticmethod
260    def evaluate(x, amplitude, x_break, alpha_1, alpha_2, delta):
261        """One dimensional smoothly broken power law model function"""
262
263        # Pre-calculate `x/x_b`
264        xx = x / x_break
265
266        # Initialize the return value
267        f = np.zeros_like(xx, subok=False)
268
269        if isinstance(amplitude, Quantity):
270            return_unit = amplitude.unit
271            amplitude = amplitude.value
272        else:
273            return_unit = None
274
275        # The quantity `t = (x / x_b)^(1 / delta)` can become quite
276        # large.  To avoid overflow errors we will start by calculating
277        # its natural logarithm:
278        logt = np.log(xx) / delta
279
280        # When `t >> 1` or `t << 1` we don't actually need to compute
281        # the `t` value since the main formula (see docstring) can be
282        # significantly simplified by neglecting `1` or `t`
283        # respectively.  In the following we will check whether `t` is
284        # much greater, much smaller, or comparable to 1 by comparing
285        # the `logt` value with an appropriate threshold.
286        threshold = 30  # corresponding to exp(30) ~ 1e13
287        i = logt > threshold
288        if i.max():
289            # In this case the main formula reduces to a simple power
290            # law with index `alpha_2`.
291            f[i] = amplitude * xx[i] ** (-alpha_2) \
292                   / (2. ** ((alpha_1 - alpha_2) * delta))
293
294        i = logt < -threshold
295        if i.max():
296            # In this case the main formula reduces to a simple power
297            # law with index `alpha_1`.
298            f[i] = amplitude * xx[i] ** (-alpha_1) \
299                   / (2. ** ((alpha_1 - alpha_2) * delta))
300
301        i = np.abs(logt) <= threshold
302        if i.max():
303            # In this case the `t` value is "comparable" to 1, hence we
304            # we will evaluate the whole formula.
305            t = np.exp(logt[i])
306            r = (1. + t) / 2.
307            f[i] = amplitude * xx[i] ** (-alpha_1) \
308                   * r ** ((alpha_1 - alpha_2) * delta)
309
310        if return_unit:
311            return Quantity(f, unit=return_unit, copy=False)
312        return f
313
314    @staticmethod
315    def fit_deriv(x, amplitude, x_break, alpha_1, alpha_2, delta):
316        """One dimensional smoothly broken power law derivative with respect
317           to parameters"""
318
319        # Pre-calculate `x_b` and `x/x_b` and `logt` (see comments in
320        # SmoothlyBrokenPowerLaw1D.evaluate)
321        xx = x / x_break
322        logt = np.log(xx) / delta
323
324        # Initialize the return values
325        f = np.zeros_like(xx)
326        d_amplitude = np.zeros_like(xx)
327        d_x_break = np.zeros_like(xx)
328        d_alpha_1 = np.zeros_like(xx)
329        d_alpha_2 = np.zeros_like(xx)
330        d_delta = np.zeros_like(xx)
331
332        threshold = 30  # (see comments in SmoothlyBrokenPowerLaw1D.evaluate)
333        i = logt > threshold
334        if i.max():
335            f[i] = amplitude * xx[i] ** (-alpha_2) \
336                   / (2. ** ((alpha_1 - alpha_2) * delta))
337
338            d_amplitude[i] = f[i] / amplitude
339            d_x_break[i] = f[i] * alpha_2 / x_break
340            d_alpha_1[i] = f[i] * (-delta * np.log(2))
341            d_alpha_2[i] = f[i] * (-np.log(xx[i]) + delta * np.log(2))
342            d_delta[i] = f[i] * (-(alpha_1 - alpha_2) * np.log(2))
343
344        i = logt < -threshold
345        if i.max():
346            f[i] = amplitude * xx[i] ** (-alpha_1) \
347                   / (2. ** ((alpha_1 - alpha_2) * delta))
348
349            d_amplitude[i] = f[i] / amplitude
350            d_x_break[i] = f[i] * alpha_1 / x_break
351            d_alpha_1[i] = f[i] * (-np.log(xx[i]) - delta * np.log(2))
352            d_alpha_2[i] = f[i] * delta * np.log(2)
353            d_delta[i] = f[i] * (-(alpha_1 - alpha_2) * np.log(2))
354
355        i = np.abs(logt) <= threshold
356        if i.max():
357            t = np.exp(logt[i])
358            r = (1. + t) / 2.
359            f[i] = amplitude * xx[i] ** (-alpha_1) \
360                   * r ** ((alpha_1 - alpha_2) * delta)
361
362            d_amplitude[i] = f[i] / amplitude
363            d_x_break[i] = f[i] * (alpha_1 - (alpha_1 - alpha_2) * t / 2. / r) / x_break
364            d_alpha_1[i] = f[i] * (-np.log(xx[i]) + delta * np.log(r))
365            d_alpha_2[i] = f[i] * (-delta * np.log(r))
366            d_delta[i] = f[i] * (alpha_1 - alpha_2) \
367                         * (np.log(r) - t / (1. + t) / delta * np.log(xx[i]))
368
369        return [d_amplitude, d_x_break, d_alpha_1, d_alpha_2, d_delta]
370
371    @property
372    def input_units(self):
373        if self.x_break.unit is None:
374            return None
375        return {self.inputs[0]: self.x_break.unit}
376
377    def _parameter_units_for_data_units(self, inputs_unit, outputs_unit):
378        return {'x_break': inputs_unit[self.inputs[0]],
379                'amplitude': outputs_unit[self.outputs[0]]}
380
381
382class ExponentialCutoffPowerLaw1D(Fittable1DModel):
383    """
384    One dimensional power law model with an exponential cutoff.
385
386    Parameters
387    ----------
388    amplitude : float
389        Model amplitude
390    x_0 : float
391        Reference point
392    alpha : float
393        Power law index
394    x_cutoff : float
395        Cutoff point
396
397    See Also
398    --------
399    PowerLaw1D, BrokenPowerLaw1D, LogParabola1D
400
401    Notes
402    -----
403    Model formula (with :math:`A` for ``amplitude`` and :math:`\\alpha` for ``alpha``):
404
405        .. math:: f(x) = A (x / x_0) ^ {-\\alpha} \\exp (-x / x_{cutoff})
406
407    """
408
409    amplitude = Parameter(default=1, description="Peak value of model")
410    x_0 = Parameter(default=1, description="Reference point")
411    alpha = Parameter(default=1, description="Power law index")
412    x_cutoff = Parameter(default=1, description="Cutoff point")
413
414    @staticmethod
415    def evaluate(x, amplitude, x_0, alpha, x_cutoff):
416        """One dimensional exponential cutoff power law model function"""
417
418        xx = x / x_0
419        return amplitude * xx ** (-alpha) * np.exp(-x / x_cutoff)
420
421    @staticmethod
422    def fit_deriv(x, amplitude, x_0, alpha, x_cutoff):
423        """One dimensional exponential cutoff power law derivative with respect to parameters"""
424
425        xx = x / x_0
426        xc = x / x_cutoff
427
428        d_amplitude = xx ** (-alpha) * np.exp(-xc)
429        d_x_0 = alpha * amplitude * d_amplitude / x_0
430        d_alpha = -amplitude * d_amplitude * np.log(xx)
431        d_x_cutoff = amplitude * x * d_amplitude / x_cutoff ** 2
432
433        return [d_amplitude, d_x_0, d_alpha, d_x_cutoff]
434
435    @property
436    def input_units(self):
437        if self.x_0.unit is None:
438            return None
439        return {self.inputs[0]: self.x_0.unit}
440
441    def _parameter_units_for_data_units(self, inputs_unit, outputs_unit):
442        return {'x_0': inputs_unit[self.inputs[0]],
443                'x_cutoff': inputs_unit[self.inputs[0]],
444                'amplitude': outputs_unit[self.outputs[0]]}
445
446
447class LogParabola1D(Fittable1DModel):
448    """
449    One dimensional log parabola model (sometimes called curved power law).
450
451    Parameters
452    ----------
453    amplitude : float
454        Model amplitude
455    x_0 : float
456        Reference point
457    alpha : float
458        Power law index
459    beta : float
460        Power law curvature
461
462    See Also
463    --------
464    PowerLaw1D, BrokenPowerLaw1D, ExponentialCutoffPowerLaw1D
465
466    Notes
467    -----
468    Model formula (with :math:`A` for ``amplitude`` and :math:`\\alpha` for ``alpha`` and :math:`\\beta` for ``beta``):
469
470        .. math:: f(x) = A \\left(\\frac{x}{x_{0}}\\right)^{- \\alpha - \\beta \\log{\\left (\\frac{x}{x_{0}} \\right )}}
471
472    """
473
474    amplitude = Parameter(default=1, description="Peak value of model")
475    x_0 = Parameter(default=1, description="Reference point")
476    alpha = Parameter(default=1, description="Power law index")
477    beta = Parameter(default=0, description="Power law curvature")
478
479    @staticmethod
480    def evaluate(x, amplitude, x_0, alpha, beta):
481        """One dimensional log parabola model function"""
482
483        xx = x / x_0
484        exponent = -alpha - beta * np.log(xx)
485        return amplitude * xx ** exponent
486
487    @staticmethod
488    def fit_deriv(x, amplitude, x_0, alpha, beta):
489        """One dimensional log parabola derivative with respect to parameters"""
490
491        xx = x / x_0
492        log_xx = np.log(xx)
493        exponent = -alpha - beta * log_xx
494
495        d_amplitude = xx ** exponent
496        d_beta = -amplitude * d_amplitude * log_xx ** 2
497        d_x_0 = amplitude * d_amplitude * (beta * log_xx / x_0 - exponent / x_0)
498        d_alpha = -amplitude * d_amplitude * log_xx
499        return [d_amplitude, d_x_0, d_alpha, d_beta]
500
501    @property
502    def input_units(self):
503        if self.x_0.unit is None:
504            return None
505        return {self.inputs[0]: self.x_0.unit}
506
507    def _parameter_units_for_data_units(self, inputs_unit, outputs_unit):
508        return {'x_0': inputs_unit[self.inputs[0]],
509                'amplitude': outputs_unit[self.outputs[0]]}
510