qdsfit/libyaff.py

__author__ = 'markusro'

from numpy import *
import multiprocessing

from PyQt4.QtCore import pyqtSignal, QObject

import scipy.special
import scipy.integrate

#define norm1(a,b)  exp(gammln(b/a))/(a*pow(b/a,b/a))
#define glntau1(x,t,a,b)  exp( -exp( (log(x)-log(t))*a )*b/a)  *  pow(x/t,b)


def filon(oms, unsorted_x,unsorted_y):
    x = unsorted_x[unsorted_x.argsort()]
    y = unsorted_y[unsorted_x.argsort()]
    amps = zeros(oms.size, dtype='complex')
    for i,om in enumerate(oms):
        amps[i]  =        sum(diff(y)/diff(x)*(cos(om*x[1:])-cos(om*x[:-1])))/om**2
        amps[i] += 1j*( - sum(diff(y)/diff(x)*(sin(om*x[1:])-sin(om*x[:-1])))/om**2)
    #-y[0]/om
    return amps

class PhiT:
    def __init__(self, dist_x, dist_y):
        self.dist_x = dist_x
        self.dist_y = dist_y
    def __call__(self, ts):
        phis = zeros(ts.size)
        for i,t in enumerate(ts):
            kern = self.dist_y*exp(-t/self.dist_x)
            phis[i] = scipy.integrate.simps(kern, log(self.dist_x))
        return phis


class Yaff(QObject):
    step_signal = pyqtSignal(list)
    def __init__(self, dist_type=0):
        super(Yaff,self).__init__()
        self._dist_tau = logspace(-12,6,512)
        self.dist_y = zeros(self.dist_tau.size)
        self._time_points = logspace(-10,5,512*2)

        yaff = [
            (self.yaff_gg, "GG"),
            (self.yaff_gge, "GGe"),
            (self.yaff_gb, "Gb"),
            (self.yaff, "GG + GGb"),
            (self.yaff_extended, "GGe + GGb"),
            ]
        self.loss, self.label = yaff[dist_type]
        self.params = 10


    @property
    def time_points(self):
        return self._time_points

    @time_points.setter
    def time_points(self, t_list):
        tmin, tmax, num = t_list
        tmi,tma = tmin, tmax
        self._time_points = logspace(log10(tmi), log10(tma), num)

    @property
    def dist_tau(self):
        return self._dist_tau

    @dist_tau.setter
    def dist_tau(self, tau_list):
        tmin, tmax, num = tau_list
        tmi,tma = tmin, tmax
        self._dist_tau = logspace(log10(tmi), log10(tma), num)


    def gg(self, p, tau):
        tau0, a, b = p
        """
        Generalized Gamma function
        """
        NGG = a * (b/a)**(b/a) / scipy.special.gamma(b/a)
        #g = exp(-b/a*exp( (log(tau)-log(tau0))*a))*(tau/tau0)**b
        g = exp(-b/a*exp( (log(tau)-log(tau0))*a))*(tau/tau0)**b
        return g*NGG

    def ggb(self, p, tau):
        tau0,a,b = p
        norm_ggb = a*(1+b)/pi *b**(b/(1+b)) * sin(pi*b/(1+b))
        g = (b*exp( (log(tau)-log(tau0)) *a) + exp( (log(tau)-log(tau0))*(-a*b)))**(-1)
        return norm_ggb*g

    def gg_hw(self, p ,tau):
        tau0, a, b, s,g  = p
        """
        Generalized Gamma function with HF wing.
        s gives the onset of the wing in units of tau0, i.e. 100 means 100*tau0 or two decades later
        g is the slope of the wing, smaller than b; if not the slope will be b
        """
        NGG = a * (b/a)**(b/a) / ( scipy.special.gamma(b/a) + s**(g-b)*(a/b)**((g-b)/a)* scipy.special.gamma(g/a) )
        g = exp(-b/a*exp((log(tau) - log(tau0))*a))*(tau/tau0)**b * (1 +  (tau*s/tau0)**(g-b))
        return g*NGG

    def dist_to_eps(omegas,dist_x,dist_y):
        epp = zeros(len(omegas), dtype='complex')
        for i,om in enumerate(omegas):
            kern_imag = dist_y*om*dist_x/(1+(om*dist_x)**2)
            kern_real = dist_y/(1+(om*dist_x)**2)
            epp[i] = scipy.integrate.simps(kern_real,log(dist_x))
            epp[i] += 1j*scipy.integrate.simps(kern_imag,log(dist_x))
        return epp

    # serial version
    def dist_to_relax_(self, ts,dist_x,dist_y):
        phi = zeros(len(ts))
        for i,t in enumerate(ts):
            phi[i] = self.phi_t(t,dist_x,dist_y)
        return phi

    # parallel version
    def dist_to_relax(self, ts,dist_x,dist_y):
        pool = multiprocessing.Pool()
        num_cpus = multiprocessing.cpu_count()
        ts_chunks = array_split(ts,num_cpus)
        result = pool.map(PhiT(dist_x,dist_y), ts_chunks )
        pool.close()
        pool.terminate()
        pool.join()
        return concatenate(result)


    def phi_t(self, t, dist_x, dist_y):
        kern = dist_y*exp(-t/dist_x)
        phi = scipy.integrate.simps(kern,log(dist_x))
        return phi

    def williams_watts(self, phi_1, phi_2, lambd):
        return phi_1 * ( (1-lambd) + lambd*phi_2)

    def yaff2(self, p, x):
        """
        GGe+G b ---> GG+GB - lambda determined by p
        """
        delta_eps, s_c, tau1,a1,b1, g,s, tau2, a2, b2 = p
        p_ggb = tau2,a2,b2
        p_gghw=tau1,a1,b1,g,s
        lambd = 2*g/pi*a1*(b1/a1)**(b1/a1)
        lambd /= scipy.special.gamma(b1/a1)*s**(b1-g) + (a1/b1)**((g-b1)/a1)*scipy.special.gamma(g/a1)
        print "LAMBDA:",lambd,a1,b1,g,s
        dist_ggb  = self.ggb(p_ggb, dist_t)
        phi_beta  = self.dist_to_relax(ts, dist_t, dist_ggb ).real
        phi_alpha = self.dist_to_relax(ts, dist_t, self.gg_hw(p_gghw, dist_t)).real
        phi_tot = (1-lambd) + lambd*phi_beta
        phi_tot *= phi_alpha
        epp = delta_eps*2*pi*x*filon(2*pi*x, ts, phi_tot ).real
        epp +=  s_c/(2*pi*x)
        return epp

    def yaff_extended(self, p,x, cb=None):
        """
        GG + GB
        """
        delta_eps, tau1, a1, b1, lambd, tau2, a2, b2, s, g = p

        dist_ggb = self.ggb(p=[tau2,a2,b2], tau=self.dist_tau)
        dist_gg  = self.gg_hw (p=[tau1,a1,b1,s,g], tau=self.dist_tau)

        #ts = logspace(-10,5,512*2)

        self.dist_y = dist_ggb
        phi_beta  = self.dist_to_relax(self.time_points, self.dist_tau, self.dist_y).real

        self.dist_y = dist_gg
        phi_alpha = self.dist_to_relax(self.time_points, self.dist_tau, self.dist_y).real

        # William-Watts-Ansatz
        phi_tot  = (1-lambd) + lambd*phi_beta
        phi_tot *= phi_alpha
        epp = delta_eps*2*pi*x*filon(2*pi*x, self.time_points, phi_tot)

        self.step_signal.emit(list(p))
        return epp

    def yaff(self, p,x, cb=None):
        """
        GG + GB
        """
        delta_eps, tau1, a1, b1, lambd, tau2, a2, b2, s, g = p

        dist_ggb = self.ggb(p=[tau2,a2,b2], tau=self.dist_tau)
        dist_gg  = self.gg (p=[tau1,a1,b1], tau=self.dist_tau)

        #ts = logspace(-10,5,512*2)

        self.dist_y = dist_ggb
        phi_beta  = self.dist_to_relax(self.time_points, self.dist_tau, self.dist_y).real

        self.dist_y = dist_gg
        phi_alpha = self.dist_to_relax(self.time_points, self.dist_tau, self.dist_y).real

        # William-Watts-Ansatz
        phi_tot  = (1-lambd) + lambd*phi_beta
        phi_tot *= phi_alpha
        epp = delta_eps*2*pi*x*filon(2*pi*x, self.time_points, phi_tot)

        self.step_signal.emit(list(p))
        return epp

    def yaff_gb(self, p,x, cb=None):
        """
        GG
        """
        delta_eps, tau1, a1, b1, lambd, tau2, a2, b2, s, g = p
        #ts = logspace(-10,5,512*2)
        dist_ggb   =  self.ggb(p=[tau2, a2, b2], tau=self.dist_tau)
        self.dist_y = dist_ggb
        phi_alpha = self.dist_to_relax(self.time_points, self.dist_tau, self.dist_y).real
        epp = delta_eps*2*pi*x*filon(2*pi*x, self.time_points, phi_alpha )
        self.step_signal.emit(list(p))
        return epp

    def yaff_gg(self, p,x, cb=None):
        """
        GG
        """
        delta_eps, tau1, a1, b1, lambd, tau2, a2, b2, s, g = p
        #ts = logspace(-10,5,512*2)
        dist_gg   =  self.gg(p=[tau1,a1,b1], tau=self.dist_tau)
        self.dist_y = dist_gg
        phi_alpha = self.dist_to_relax(self.time_points, self.dist_tau, self.dist_y).real
        epp = delta_eps*2*pi*x*filon(2*pi*x, self.time_points, phi_alpha )
        self.step_signal.emit(list(p))
        return epp

    def yaff_gge(self, p,x):
        """
        GG
        """
        delta_eps, tau1, a1, b1, lambd, tau2, a2, b2, s, g = p
        #ts = logspace(-10,5,512*2)
        dist_gg   =  self.gg_hw(p=[tau1, a1, b1, s, g], tau=self.dist_tau)
        self.dist_y = dist_gg
        phi_alpha = self.dist_to_relax(self.time_points, self.dist_tau, self.dist_y).real
        epp = delta_eps*2*pi*x*filon(2*pi*x, self.time_points, phi_alpha )
        self.step_signal.emit(list(p))
        return epp


def show_correlations(cov, ax):
    cor = zeros(cov.shape)
    for i in range(cor.shape[0]):
        for j in range(cor.shape[1]):
            if cov[i,i] == 0 or cov[j,j]==0:
                cor[i,j] = 0
            else:
                cor[i,j] = cov[i,j]/sqrt(cov[i,i]*cov[j,j])
    ax.matshow(abs(cor), cmap=cm.RdYlGn)
    ax.colorbar()