gaussHermite.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Time-stamp: <2015-04-07 18:58:52 ycopin>

from __future__ import division, print_function

"""
Gauss-Hermite line profile.
"""

__author__ = "Yannick Copin <y.copin@ipnl.in2p3.fr>"

from astropy.modeling.models import PolynomialModel, Gaussian1D
from astropy.modeling.parameters import Parameter
from astropy.modeling.core import Fittable1DModel

import re

class Hermite1D(PolynomialModel):
    """
    1D (physicists') Hermite polynomial.

    Parameters
    ----------
    degree : int
        degree of the series
    domain : list or None
    window : list or None
        If None, it is set to [-1,1]
        Fitters will remap the domain to this window
    param_dim : int
        number of parameter sets
    **params : dict
        keyword: value pairs, representing parameter_name: value
    """

    inputs = ('x',)
    outputs = ('y',)

    def __init__(self, degree, domain=None, window=[-1, 1], n_models=None,
                 model_set_axis=None, name=None, meta=None, **params):
        self.domain = domain
        self.window = window
        super(Hermite1D, self).__init__(
            degree, n_models=n_models, model_set_axis=model_set_axis,
            name=name, meta=meta, **params)

    def prepare_inputs(self, x, **kwargs):
        inputs, format_info = \
                super(PolynomialModel, self).prepare_inputs(x, **kwargs)

        x = inputs[0]

        if self.domain is not None:
            x = poly_map_domain(x, self.domain, self.window)

        return (x,), format_info

    @classmethod
    def evaluate(cls, x, *coeffs):

        return cls.clenshaw(x, coeffs)

    def fit_deriv(self, x, *params):
        """
        Computes the Vandermonde matrix.

        Parameters
        ----------
        x : ndarray
            input
        params : throw away parameter
            parameter list returned by non-linear fitters

        Returns
        -------
        result : ndarray
            The Vandermonde matrix
        """

        x = np.array(x, dtype=np.float, copy=False, ndmin=1)
        v = np.empty((self.degree + 1,) + x.shape, dtype=x.dtype)
        v[0] = 1          # H_0
        if self.degree > 0:
            v[1] = 2 * x  # H_1
            for i in range(2, self.degree + 1):
                # Recurrence: H_n = 2 * x * H_{n-1} - 2 * (n - 1) * H_{n-2}
                v[i] = 2 * (x * v[i - 1] - (i - 1) * v[i - 2])

        return np.rollaxis(v, 0, v.ndim)

    @staticmethod
    def clenshaw(x, coeffs):

        # Clenshaw algo: H_n = alpha * H_{n-1} + beta * H_{n-2}
        alpha = 2 * x
        if len(coeffs) == 1:
            c0 = coeffs[0]
            c1 = 0
        elif len(coeffs) == 2:
            c0 = coeffs[0]
            c1 = coeffs[1]
        else:
            nd = len(coeffs)
            c0 = coeffs[-2]
            c1 = coeffs[-1]
            for i in range(3, len(coeffs) + 1):
                tmp = c0
                nd = nd - 1
                beta = -2 * (nd - 1)
                c0 = coeffs[-i] + c1 * beta
                c1 = tmp + c1 * alpha

        return c0 + c1 * alpha    # c0 * H_0 + c_1 * H_1


class GaussHermite(Fittable1DModel):

    _param_names = ()  # Will be filled at instanciation

    def __init__(self, order, *args, **kwargs):

        self._order = int(order)
        if self._order < 3:
            self._order = 0

        # Gaussian model
        self._gaussian = Gaussian1D()
        # Hermite series
        if self._order:
            self._hermite = Hermite1D(self._order)
        else:
            self._hermite = None

        self._param_names = self._generate_coeff_names()

        super(GaussHermite, self).__init__(*args, **kwargs)

    def _hi_order(self, name):
        # One could store the compiled regex, but it will crash the deepcopy:
        # "cannot deepcopy this pattern object"

        match = re.match('h(?P<order>\d+)', name)  # h3, h4, etc.
        order = int(match.groupdict()['order']) if match else 0

        return order

    def _generate_coeff_names(self):

        names = list(self._gaussian.param_names)  # Gaussian parameters
        names += [ 'h{}'.format(i)
                   for i in range(3, self._order + 1) ] # Hermite coeffs

        return tuple(names)

    def __getattr__(self, attr):

        if attr[0] == '_':
            super(GaussHermite, self).__getattr__(attr)
        elif attr in self._gaussian.param_names:
            return self._gaussian.__getattribute__(attr)
        elif self._order and self._hi_order(attr) >= 3:
            return self._hermite.__getattr__(attr.replace('h', 'c'))
        else:
            super(GaussHermite, self).__getattr__(attr)

    def __setattr__(self, attr, value):

        if attr[0] == '_':
            super(GaussHermite, self).__setattr__(attr, value)
        elif attr in self._gaussian.param_names:
            self._gaussian.__setattr__(attr, value)
        elif self._order and self._hi_order(attr) >= 3:
            self._hermite.__setattr__(attr.replace('h', 'c'), value)
        else:
            super(GaussHermite, self).__setattr__(attr, value)

    @property
    def param_names(self):

        return self._param_names

    def evaluate(self, x, *params):

        a, m, s = params[:3]                      # amplitude, mean, stddev
        f = self._gaussian.evaluate(x, a, m, s)
        if self._order:
            f *= (1 + self._hermite.evaluate((x - m)/s, 0, 0, 0, *params[3:]))

        return f



if __name__ == '__main__':

    import numpy as N
    import matplotlib.pyplot as P
    import astropy.modeling as AM

    if True:
        g0 = GaussHermite(4, amplitude=1, mean=0, stddev=1, h3=0.005, h4=0.005)
    else:
        g0 = AM.models.Gaussian1D(amplitude=1, mean=0, stddev=1)

    x = N.linspace(-10, 10, 201)
    y = g0(x)

    fitter = AM.fitting.SimplexLSQFitter()
    g0.amplitude = 0.8
    g0.mean = -0.5
    g0.stddev = 1.2
    g = fitter(g0, x, y)
    print(g)

    fig = P.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.plot(x, y, label='Input')
    ax.plot(x, g0(x), ls='--', label='Guess')
    ax.plot(x, g(x), ls=':', label='Fit')

    ax.legend()

    P.show()

@nden That's a good idea about using tied parameters to do this. It's too bad though that they only work in fitting. That's still something that we need to consider changing since I know it causes regular confusion.

Also yeah I don't know why the Scale and Shift models aren't already fittable = True by default. That makes no sense.