#!/usr/bin/env python
# coding: utf-8

# # Example - Burst Variance Analysis
# 
# *This notebook is part of smFRET burst analysis software [FRETBursts](http://opensmfs.github.io/FRETBursts/).*

# > This notebook shows how to implement Burst Variance Analysis (BVA) ([Torella 2011](http://dx.doi.org/10.1016/j.bpj.2011.01.066)) using FRETBursts.
# 
# > For a complete tutorial on burst analysis see 
# > [FRETBursts - us-ALEX smFRET burst analysis](FRETBursts - us-ALEX smFRET burst analysis.ipynb).

# In[1]:


from fretbursts import *
sns = init_notebook(apionly=True)


# In[2]:


# Tweak here matplotlib style
import matplotlib as mpl
mpl.rcParams['font.sans-serif'].insert(0, 'Arial')
mpl.rcParams['font.size'] = 12
get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")


# # Load Data

# In[3]:


url = 'http://files.figshare.com/2182601/0023uLRpitc_NTP_20dT_0.5GndCl.hdf5'
download_file(url, save_dir='./data')

full_fname = "./data/0023uLRpitc_NTP_20dT_0.5GndCl.hdf5"
d = loader.photon_hdf5(full_fname)
loader.alex_apply_period(d)
d.calc_bg(bg.exp_fit, time_s=50.1, tail_min_us='auto', F_bg=1.7)


# In[4]:


d.burst_search(m=10, computefret=False, ph_sel=Ph_sel(Dex='DAem'))
d.calc_fret(count_ph=True, corrections=False)


# In[5]:


ds = d.select_bursts(select_bursts.naa, th1=30, computefret=False)
ds1 = ds.select_bursts(select_bursts.size, th1=30, computefret=False)
ds_FRET = ds1.select_bursts(select_bursts.S, S1=0.25, S2=0.85, computefret=False)


# In[6]:


dx=ds_FRET
alex_jointplot(dx);


# # Burst Variance Analysis

# We define a function to compute $s_E$:

# In[7]:


def bva_sigma_E(n, bursts, DexAem_mask, out=None):
    """
    Perform BVA analysis computing std.dev. of E for sub-bursts in each burst.
    
    Split each burst in n-photons chunks (sub-bursts), compute E for each sub-burst,
    then compute std.dev. of E across the sub-bursts.

    For details on BVA see:

    - Torella et al. (2011) Biophys. J. doi.org/10.1016/j.bpj.2011.01.066
    - Ingargiola et al. (2016) bioRxiv, doi.org/10.1101/039198

    Arguments:
        n (int): number of photons in each sub-burst
        bursts (Bursts object): burst-data object with indexes relative 
            to the Dex photon stream.
        DexAem_mask (bool array): mask of A-emitted photons during D-excitation 
            periods. It is a boolean array indexing the array of Dex timestamps 
            (`Ph_sel(Dex='DAem')`).
        out (None or list): append the result to the passed list. If None,
            creates a new list. This is useful to accumulate data from
            different spots in a single list.

    Returns:
        E_sub_std (1D array): contains for each burst, the standard deviation of 
        sub-bursts FRET efficiency. Same length of input argument `bursts`.
    """
    E_sub_std = [] if out is None else out
    
    for burst in bursts:
        E_sub_bursts = []
        startlist = range(burst.istart, burst.istop + 2 - n, n)
        stoplist = [i + n for i in startlist]
        for start, stop in zip(startlist, stoplist):
            A_D = DexAem_mask[start:stop].sum()
            assert stop - start == n
            E = A_D / n
            E_sub_bursts.append(E)
        E_sub_std.append(np.std(E_sub_bursts))
        
    return E_sub_std


# Next we prepare the data for BVA:

# In[8]:


ph_d = ds_FRET.get_ph_times(ph_sel=Ph_sel(Dex='DAem'))
bursts = ds_FRET.mburst[0]
bursts_d = bursts.recompute_index_reduce(ph_d)


# In[9]:


Dex_mask = ds_FRET.get_ph_mask(ph_sel=Ph_sel(Dex='DAem'))   
DexAem_mask = ds_FRET.get_ph_mask(ph_sel=Ph_sel(Dex='Aem')) 
DexAem_mask_d = DexAem_mask[Dex_mask]


# and call the `bva_sigma_E` function:

# In[10]:


n = 7
E_sub_std = bva_sigma_E(n, bursts_d, DexAem_mask_d)


# Finally, we make a KDE plot of the 2D distribution `E_sub_std` versus the burst FRET efficiency:

# In[11]:


plt.figure(figsize=(4.5, 4.5))
x = np.arange(0,1.01,0.01)
y = np.sqrt((x*(1-x))/n)
plt.plot(x, y, lw=2, color='k', ls='--')
im = sns.kdeplot(data={'E':ds_FRET.E[0], 'sigma':np.asfarray(E_sub_std)}, x='E', y='sigma', 
                 fill=True, cmap='Spectral_r', thresh=0.05, levels=20)
plt.xlim(0,1)
plt.ylim(0,np.sqrt(0.5**2/7)*2)
plt.xlabel('E', fontsize=16)
plt.ylabel(r'$\sigma_i$', fontsize=16);
plt.text(0.05, 0.95, 'BVA', va='top', fontsize=22, transform=plt.gca().transAxes)
plt.text(0.95, 0.95, '# Bursts: %d' % ds_FRET.num_bursts, 
         va='top', ha='right', transform=plt.gca().transAxes)
plt.savefig('BVA.png', bbox_inches='tight', dpi=200, transparent=False)


# In[12]:


x, y = ds_FRET.E[0], np.asfarray(E_sub_std)
hist_kws = dict(edgecolor='k', linewidth=0.2,
                facecolor=sns.color_palette('Spectral_r', 100)[10])

g = sns.JointGrid(x=x, y=y, ratio=3)
g.plot_joint(sns.kdeplot, cmap='Spectral_r', fill=True, thresh=0.05, levels=20)
g.ax_marg_x.hist(x, bins=np.arange(-0.2, 1.2, 0.025), **hist_kws)
g.ax_marg_y.hist(y, bins=50, orientation="horizontal", **hist_kws)

x1 = np.arange(0,1.01,0.01)
y1 = np.sqrt((x1*(1-x1))/n)
plt.plot(x1, y1, lw=2, color='k', ls='--')

g.ax_joint.set_xlim(0,1)
g.ax_joint.set_ylim(0,np.sqrt(0.5**2/7)*2)
g.ax_joint.set_xlabel('E', fontsize=16)
g.ax_joint.set_ylabel(r'$\sigma_i$', fontsize=16);
g.ax_joint.text(0.05, 0.95, 'BVA', va='top', fontsize=22, transform=g.ax_joint.transAxes)
g.ax_joint.text(0.95, 0.95, '# Bursts: %d' % ds_FRET.num_bursts, 
         va='top', ha='right', transform=g.ax_joint.transAxes)
plt.savefig('BVA_joint.png', bbox_inches='tight', dpi=200, transparent=False)