Fit data like spectra or something to a sum of gaussians.

mog.py

"""
mog.py

Mixture of gaussians. For fitting data that looks like Gaussians.

author: Colin Clement
date: 2016-2-29

"""

import numpy as np
from scipy.optimize import leastsq


class Mog(object):
    def __init__(self, L, N_g, params = None, 
                 spacing = None, **kwargs):
        """
        input:
            L: (int) length of data array
            N_g: (int) number of gaussians in model
            (optional)
            params: (N_g, 3)-shaped array for initial parameters
                    2nd index is [amplitude, mean, standard deviation]
            spacing: (float) if you want the space between points to not
                    be 1
        """
        self.L = L
        self.N_g = N_g
        self.d = spacing if spacing is not None else 1.

        self.x = np.arange(self.L)
        self.mixture = np.zeros(L)
        self.d_params = np.zeros((self.L, self.N_g*3))
        self.params = params
        if self.params is None:
            #Initialize A, x, sx to some random values
            means = np.cumsum(np.random.rand(self.N_g))
            self.params = np.vstack([np.random.rand(self.N_g)/self.d,
                                     np.random.rand(self.N_g)*self.L*self.d,
                                     self.d*3*(np.random.rand(self.N_g)+1)]).T
            self.params[:,0] /= self.params[:,0].sum()
        self.logparams = self.getlogp(self.params)
        self.update(self.logparams)

    def gaussian(self, params):
        """
        Gaussians model. Log parameters for amplitude and
        standard deviation enforce positivity
        input:
            params: (N_g, 3)-shaped array
        """
        #Log params for a, sx to enforce positivity
        x, sx = params[:,1], np.exp(params[:,2])
        Dx = self.x[:,None] - x[None,:]
        g = np.exp(- Dx**2/(2*sx[None,:]**2))
        g = np.exp(params[:,0])[None,:]*g/np.sqrt(2*np.pi*sx[None,:]**2)
        return g.sum(1).ravel()
 
    def d_gaussian(self, params):
        """
        Gaussians model. Log parameters for amplitude and
        standard deviation enforce positivity
        input:
            params: (N_g, 3)-shaped array
        """
        a = np.exp(params[:,0])
        x, sx = params[:,1], np.exp(params[:,2])
        sx2 = sx*sx
        Dx = self.x[:,None]-x[None,:]
        Dx2 = Dx*Dx
        g = np.exp(- Dx2/(2*sx2[None,:]))
        g = g*(a/np.sqrt(2*np.pi*sx2))[None,:]
        da = g #a = exp(alpha) 
        dx = g*Dx/sx[None,:]**2
        dsx = g*(2*Dx2-sx2[None,:])/(2*sx2[None,:])
        for i, d in enumerate([da, dx, dsx]):
            self.d_params[:,i::3] = d.reshape(-1, self.N_g)
        return self.d_params

    def getp(self, logp):
        """
        Convert log parameters to parameters
        """
        a, x, sx = np.exp(logp[:,0]), logp[:,1], np.exp(logp[:,2])
        return np.vstack([a, x, sx]).T

    def getlogp(self, p):
        """
        Convert parameters to log parameters
        """
        a, x, sx = np.log(p[:,0]), p[:,1], np.log(p[:,2])
        return np.vstack([a, x, sx]).T

    def update(self, params):
        """
        update state of self.mixture with params
        """
        self.mixture = self.gaussian(params.reshape(self.N_g, -1))

    def residual(self, params, data):
        """
        Compute difference between model and data.
        input:
            params: (N_g, 3)-shaped array
            data: (L)-shaped array of data
        returns:
            residuals : (L)-shaped array
        """
        self.update(params)
        return self.mixture - data.ravel()

    def jacobian(self, params, data):
        """
        Compute jacobian of the residuals w.r.t. model parameters.
        Does not actually depend on data, but leastsq likes consistent
        function signatures.   
        input:
            params: (N_g, 3)-shaped array
            data: (L)-shaped array of data
        returns:
            jacobian: (L, 3*N_g)-shaped array
        """
        return self.d_gaussian(params.reshape(self.N_g, -1))

    def fit(self, data, p0 = None, maxfev = 200):
        """
        Fit a Mog model to data by minimizing squared error.
        input:
            data:  (L)-shaped array
            (optional)
            p0: (N_g,3)-shaped array of parameter initial guesses.
            maxfev: (int) limit on leastsq so you don't wait too long
        returns:
            best fit parameters (N_g, 3)-shaped array.
        NOTE:
            self.opt stores full output of leastsq. Examine in case of
            weird behavior.
        """
        p0 = p0 if p0 is not None else self.params
        logp0 = self.getlogp(p0)
        self.opt = leastsq(self.residual, logp0.ravel(),
                           Dfun = self.jacobian, args = (data,),
                           full_output = 1, maxfev = maxfev)
        self.logparams = self.opt[0].reshape(self.N_g, -1)
        self.params = self.getp(self.logparams)
        return self.params