Photoactivation Project - Intensity Profile Analysis Framework

profile.py

from __future__ import division
import statsmodels.api as sm
import matplotlib.pyplot as plt
from lmfit import Model
import numpy as np
from numpy import exp, convolve
import pandas as pd
from pandas import Series, DataFrame


def gaussian(x, sigma, center):
    return 1/(np.sqrt(2*np.pi) * sigma) * exp(-(x-center)**2/(2*sigma**2))


def fit_line(x, y):
    """Return slope, intercept of best fit line."""
    # intercept is forced to zero
    model = sm.OLS(y, x, missing='drop') # ignores entires where x or y is NaN
    fit = model.fit()
    return fit
    
    
class ConvolvedGaussianModel(Model):
    def __init__(self, kernel, margin=50, **kwargs):
        assert margin < len(kernel) // 2 and margin > 0, "User is a doofis."
        self.kernel = kernel
        self.margin = margin
        
        def func(x, sigma, center, constant, norm):
            # Note: This uses kernel, defined just above.
            
            result = convolve(norm * gaussian(x, sigma, center), self.kernel, 'same')
            return result[self.margin:-self.margin] + constant
        
        super(ConvolvedGaussianModel, self).__init__(func, **kwargs)
        
    def guess(self, profile):
        self.set_param_hint('sigma' , value=10,
                            min=0.001, max=100)
        self.set_param_hint('center', value=len(profile) / 2, min=0, max=len(profile))
        self.set_param_hint('constant', value=0)
        self.set_param_hint('norm', value=1, min=0, vary=True)
        self.has_initial_guess = True
    
    def fit(self, profile, **kwargs):
        profile = profile.copy()
        m = self.margin  # for brevity
        x = np.arange(m, 512 - m)

        if not self.param_hints:
            self.guess(profile[m:-m])
            
        fit_result = super(ConvolvedGaussianModel, self).fit(profile[m:-m], x=x, **kwargs)

        return fit_result

    
class Profile(object):
    def __init__(self, frame, calibration, background=0, axis=0):
        self.background = background
        self.profile = frame.sum(axis) - background
        self.frame_no = frame.frame_no
        self.calibration = calibration
    
    @property
    def raw_profile(self):
        """Return a copy of this profile before any
        corrections were applied.
        
        Do we ever need this? I dunno."""
        return self.profile + self.background
    
    def plot(self, ax=None, **kws):
        "Convenience method that just plots the profile"
        if ax is None:
            fig, ax = plt.subplots()
        ax.plot(self.profile, **kws)
        
    def as_kernel(self):
        """Return a copy of this profile but with the
        baseline brightness set to zero."""
        return self.profile # - self.profile.min()
    
    def map_onto(self, other):
        """Perform a best fit mapping this profile onto another.
        Return a Mapping object that contains information about
        the fit and the mapping in general."""

        return GaussianMapping(self, other)
    
    
class GaussianMapping(object):
    def __init__(self, profile1, profile2, margin=50):
        """Map profile1 onto profile2 through convolution with 
        a Gaussian."""
        assert profile1.calibration == profile2.calibration  # sanity check
        self.calibration = profile1.calibration
        cg_model = ConvolvedGaussianModel(profile1.as_kernel(), margin=margin)
        self.fit_result = cg_model.fit(profile2.profile)
        self.x = np.arange(margin, margin + len(self.fit_result.best_fit))
        # Keep private copies to discourage accidentally modifying in place.
        self._profile1 = profile1
        self._profile2 = profile2
        
    @property
    def kernel(self):
        values = self.fit_result.values
        return Series(gaussian(self.x, values['sigma'], values['center']))
    
    @property
    def sigma(self):
        "Return sigma in physical units"
        return self.fit_result.values['sigma'] * self.calibration
    
    def plot(self, ax=None, **kws):
        """Plot profile1, profile2,
        and profile1 mapped onto profile2."""
        if ax is None:
            fig, ax = plt.subplots()
        fn1 = self._profile1.frame_no
        fn2 = self._profile2.frame_no
        self._profile1.plot(ax=ax, label=fn1, **kws)
        self._profile2.plot(ax=ax, label=fn2, **kws)
        ax.plot(self.x, self.fit_result.best_fit, label='{0} mapped onto {1}'.format(fn1, fn2), **kws)
        ax.set(xlabel=r'$x$ [px]', ylabel='intensity')
        ax.legend()
        return ax