inprep_s4_Analysis.py

# -*- coding: utf-8 -*-
"""
Created on Mon Aug  8 16:00:38 2022

@author: Pablo
"""


import os
import cv2
import pickle
import numpy as np
from scipy.optimize import curve_fit
from tqdm import tqdm
from datetime import datetime
from matplotlib import pyplot as plt
from sklearn.mixture import GaussianMixture

from pytracker.video_utils import MultiWormTracker
import pytracker.video_utils as vutils


# ... Filenames ... 
DIR_NAME          = './videos/Carla_EC/Carla_N2_EC_2211101415_002'
BLOB_FILENAME     = os.path.join( DIR_NAME, 'video_data_blobs.pkl')
BLOB_REF_FILENAME = os.path.join( DIR_NAME, 'video_reference_contour.pkl')
TRAJ_FILENAME     = os.path.join( DIR_NAME, 'trajectories.pkl')
NPZ_FILENAME      = os.path.join( DIR_NAME, 'trajectories.pkl.npz')
IMG_FILENAME      = os.path.join( DIR_NAME, 'trajectories.png')
BKGD_FILENAME     = os.path.join( DIR_NAME, 'video_fondo.png')
ROIS_FILENAME     = os.path.join( DIR_NAME, 'rois.pkl')


#... Parameters
AVG_WDW = 3
MIN_LENGTH = 2500 # in frames
MIN_DIST   = 5 # in mm
XY_SCALE = 36 # px/mm
T_SCALE  = 2 #fps

def moving_average(x, w):
    return np.convolve(x, np.ones(w), 'valid') / w

def draw_circle( p0 , r0, *args, **kwargs ):
    t = np.linspace( -np.pi, np.pi, 100)
    x = p0[0] + r0*np.cos(t)
    y = p0[1] + r0*np.sin(t)
    return plt.plot( x,y,*args, **kwargs)

def rotation_matrix(theta):
    return np.array((( np.cos(theta), - np.sin(theta)), ( np.sin(theta),  np.cos(theta))))

def generate_brownian( N , steps, c, p0, s0 , speed=[1,]):
    random_part = 2 * np.random.rand( N , steps , 2 ) - 1
    X = np.zeros_like( random_part ) + p0
    
    for t in tqdm(np.arange(1,steps-1,1)):
        v_final = speed[ np.random.randint(len(speed))]
        v_source = s0-X[:,t-1,:]
        v_notnorm = random_part[:, t,:] + c*v_source
        v_norm = v_final * v_notnorm / np.linalg.norm(v_notnorm, axis=1, keepdims=True)
        X[:,t+1,:] = X[:,t,:] + v_norm
    return X

def chemotaxis_index( X ):
    pos = np.sum( X[:,:,0]>0, axis=0) + 1e-90
    neg = np.sum( X[:,:,0]<0, axis=0) + 1e-90
    return (pos - neg)/(pos+neg)


def discard_nan( array):
    idx = np.nonzero( np.logical_not (np.isnan(array) ) )
    return array[idx]

def gaussian_pdf(x, mu, s):
    return np.exp(-(x - mu)**2/(2*s**2))/(np.sqrt(2*np.pi)*s)

def time_to_goal( X , p0 , min_dist ):
    d_to_goal = np.sqrt( (X[:,:,0]- p0[0])**2 + (X[:,:,1]-p0[1])**2)
    t_to_goal = np.ones( (X.shape[0],))

    for worm_jj in range( X.shape[0] ):
        tini = np.argwhere( X[worm_jj,:,0] >-9999).min()
        tfin = np.argwhere(d_to_goal[worm_jj,:] < min_dist)
        if len(tfin)>0:
            t_to_goal[worm_jj] = (tfin.min() - tini) 
        else:
            t_to_goal[worm_jj] = X.shape[1]
    return t_to_goal






# Load and rescale data
data = np.load( NPZ_FILENAME)['data']
data = data / XY_SCALE
# ... compute number of worms and frames
n_worms, n_frames , _ = data.shape
data_t = np.array( range(n_frames)) / T_SCALE / 60

# ... repeat for rois positions
with open( ROIS_FILENAME, 'rb') as f:
    ROIS = pickle.load( f) 
    plate = ROIS[0] / XY_SCALE
    chunk = ROIS[1]/ XY_SCALE
    

# Compute speeds
v = np.sqrt( np.diff( data[:,:,0], axis=1)**2 + np.diff( data[:,:,1], axis=1)**2 )
w = np.arctan2( np.diff( data[:,:,1], axis=1), np.diff( data[:,:,0], axis=1) )
# w = np.arctan( np.diff( data[:,:,1], axis=1) / np.diff( data[:,:,0], axis=1) )

#... do the running average on coordinates and speeds
data_xy = np.nan*np.zeros( (n_worms, n_frames-AVG_WDW+1, 2))
data_vw = np.nan*np.zeros( (n_worms, n_frames-AVG_WDW, 2) )

# ... to every worm
for jj in range( data.shape[0] ):
    data_vw[jj,:,0] = moving_average( v[jj,:], AVG_WDW) * T_SCALE
    data_vw[jj,:,1] = moving_average( w[jj,:], AVG_WDW) * T_SCALE

    # ... to every coordinate
    for kk in range(data.shape[2]):
        data_xy[jj,:,kk] = moving_average( data[jj,:,kk], AVG_WDW)     





# Select trajectories with lengths less than MIN_LENGTH
# plt.hist( traj_length, bins=np.linspace(0,1000,30) )
traj_length = np.sum( data_xy[:,:,0]>0, axis=1)
I_long = np.argwhere( traj_length > MIN_LENGTH )[:,0]


# Select trajectories that start close to the chunk
# plt.hist( dst_start, bins=np.linspace(0,20,30) )
# idx_start = (data_xy[:,:,0]>0).argmax(axis=1) #... find at which idx each trajectory starts
# pos_start = np.array( [ data_xy[_, idx_start[_],: ] for _ in range(n_worms) ] )
# dst_start = np.sqrt( (pos_start[:,0]-chunk[0])**2 + (pos_start[:,1]-chunk[1])**2 )
# I_close = np.argwhere( dst_start > MIN_DIST )[:,0]

# Select trajectories that start inside the plate
# plt.hist( dst_start, bins=np.linspace(0,20,30) )
# idx_within = (data_xy[:,:,0]>0).argmax(axis=1) #... find at which idx each trajectory starts
# pos_within = np.array( [ data_xy[_, idx_within[_],: ] for _ in range(n_worms) ] )
# dst_within = np.sqrt( (pos_within[:,0]-plate[0])**2 + (pos_within[:,1]-plate[1])**2 )
# I_within  = np.argwhere( dst_within < 0.85*plate[2] )[:,0]


plt.figure( figsize=(15,5) )
plt.subplot(1,2,1)
plt.hist( data_vw[:,:,0].flatten(), bins=np.linspace(0,1,20), density=True)
plt.xlabel('Speed |v| (mm/s)')

plt.subplot(1,2,2)
plt.hist( data_vw[:,:,1].flatten(), bins=20, density=True)
plt.xlabel('Ang. speed $\omega$ (1/s)')



#%% 
######## FIGURE 1 - PANEL A ########
fondo = vutils.load_background(BKGD_FILENAME)
color = plt.cm.hsv(np.linspace(0, 1, n_worms))

plt.figure( figsize=(6,6), dpi=600 )
#... draw background image
plt.imshow(  1.2*fondo , cmap='gray', vmax=255, vmin=0)
#... draw ROIS
draw_circle( [XY_SCALE*chunk[0], XY_SCALE*chunk[1]], XY_SCALE*MIN_DIST ,'r', lw=1, alpha=0.8)
draw_circle( [XY_SCALE*plate[0], XY_SCALE*plate[1]], XY_SCALE*1*plate[2] ,'r', lw=1, alpha=0.8)
plt.plot( XY_SCALE*plate[0], XY_SCALE*plate[1], '+r', ms=5, zorder=1000)
#... draw all trajectories in black
plt.plot( XY_SCALE*data_xy[:,:,0].T, XY_SCALE*data_xy[:,:,1].T, lw=0.5, c='k');
#... draw only good trajectories in colorines
# I_good = [ idx for idx in range(n_worms) if idx in I_long and idx in I_close and idx in I_within ]
I_good = I_long
rr = np.random.permutation(n_worms)
for ii, idx in enumerate( I_good ):
    plt.plot( XY_SCALE*data_xy[idx,:,0].T, XY_SCALE*data_xy[idx,:,1].T, lw=0.75, c=0.9*color[rr[ii],:])

plt.xticks( ticks=plt.xticks()[0], labels='')
plt.yticks( ticks=plt.yticks()[0], labels='')
plt.xlim(0, fondo.shape[1])
plt.ylim(0, fondo.shape[0])
plt.gca().invert_yaxis()




######## FIGURE 3 - PANEL A ########    
# Time series of linear and angular speed
idx = np.argmax( np.sum( data_xy[:,:,0]>-9999, axis=1 ))
# tini, tfin = np.argwhere( np.diff( 1.0*(data_vw[idx,:,0]>-9999)) != 0).flatten()/T_SCALE/60
tini, tfin = 0, n_frames/T_SCALE/60
dt = data_vw.shape[1]
col_gray= np.ones((3,))*0.6


fondo = vutils.load_background(BKGD_FILENAME)
color = plt.cm.hsv(np.linspace(0, 1, n_worms))

plt.figure( figsize=(3,3), dpi=600 )
plt.imshow(  1.2*fondo , cmap='gray', vmax=255, vmin=0)
plt.plot( XY_SCALE*data_xy[idx,:,0].T, XY_SCALE*data_xy[idx,:,1].T, lw=0.5, c='r');
plt.xticks( plt.xticks()[0], labels='')
plt.yticks( plt.yticks()[0], labels='')
plt.xlim(0, fondo.shape[1])
plt.ylim(0, fondo.shape[0])


plt.figure( figsize=(6,4), dpi=600 )
plt.subplot(2,1,1)
plt.plot( data_t[:dt], data_vw[idx,:,0] , color='k', lw=0.5)
plt.ylabel('|v| (mm/s)')
plt.xlim(tini, tfin)
ylims1 = plt.ylim()

plt.subplot(2,1,2)
plt.plot( data_t[:dt], data_vw[idx,:,1] , color= 'k', lw=0.5)
plt.xlabel('Time (min)')
plt.ylabel('$\omega$ (rad/s)')
plt.xlim(tini, tfin)
ylims2 = plt.ylim()

plt.tight_layout()


######## FIGURE 3 - PANEL A2 ########    
# Histograms of instantaneous linear and angular speeds
plt.figure( figsize=(3,4), dpi=600 )
plt.subplot(2,1,1)
bbins = np.linspace( np.min(ylims1), np.max(ylims1), 20 )
plt.hist( data_vw[:,:,0].flatten() ,bins=bbins, orientation='horizontal', color=col_good, density=True, lw=1.5, histtype='step')
plt.hist( data_vw[idx,:,0] ,bins=20, orientation='horizontal', color='k', alpha=0.75, density=True)
plt.yticks( plt.yticks()[0], labels='')
plt.ylim(ylims1)

plt.subplot(2,1,2)
bbins = np.linspace( np.min(ylims2), np.max(ylims2), 20 )
plt.hist( data_vw[:,:,1].flatten() ,bins=bbins, orientation='horizontal', color=col_good, density=True, lw=1.5, histtype='step')
plt.hist( data_vw[idx,:,1] ,bins=20, orientation='horizontal', color='k', alpha=0.75, density=True)
plt.yticks( plt.yticks()[0], labels='')
plt.ylim(ylims2)
plt.xlabel('Probability densities')

plt.legend(('All traj.','Only $w^1$'))

plt.tight_layout()



######## FIGURE 3 - PANEL B ########    
# Individuality matrix
v_ = np.linspace( 0 , 0.2 , 14 )
w_ = np.linspace(-2*np.pi, 2*np.pi, 12 )
Dt = 30*T_SCALE # in frames, typically 30s, 60frames

extent = [ v_[0], v_[-1],w_[0], w_[-1] ]

Indiv = []
for ii in range(n_worms):
    Indiv.append( np.histogram2d( data_vw[ii,:,0], data_vw[ii,:,1], bins=[v_,w_], density=True)[0] )
Indiv= np.array(Indiv)

M_indiv = np.reshape( Indiv, (n_worms, (len(v_)-1)*(len(w_)-1) ) )

# plt.colorbar(label='Probability density')

plt.figure( figsize=(2,2), dpi=600)
plt.imshow( Indiv[idx,:,:].T, aspect='auto',  vmin=0, vmax=5, cmap='inferno_r')
plt.xlabel('|v| (mm/s)')
plt.ylabel('$\omega$ (rad/s)')


plt.figure( figsize=(16,10), dpi=100)
for jj in range(17):
    plt.subplot(4,5,jj+1)
    plt.imshow( Indiv[jj,:,:].T, aspect='auto',  vmin=0, vmax=5, cmap='inferno_r')
    plt.xti
    plt.xlabel('|v| (mm/s)')
    plt.ylabel('$\omega$ (rad/s)')
plt.colorbar(label='Probability density')
plt.suptitle( DIR_NAME.split('/')[-1])
plt.tight_layout()



plt.figure( figsize=(1,3.5) , dpi=600)
plt.imshow( M_indiv[idx:(idx+1),:].T, aspect='auto',vmin=0, vmax=5, cmap='inferno_r')
plt.xticks( ticks=[])
plt.ylabel('Features')

plt.tight_layout()



import umap
from scipy.spatial import ConvexHull

DRA = umap.UMAP(metric='euclidean',n_neighbors=15,random_state=9999, verbose=False, min_dist=0.1)
Xnew = DRA.fit_transform( M_indiv[:,:] )
hull = ConvexHull(Xnew)

plt.figure( figsize=(5,4), dpi=600)
plt.plot( Xnew[:,0], Xnew[:,1], '.k', ms=5, zorder=0)
plt.scatter( Xnew[I_good,0], Xnew[I_good,1], 
            c = np.nanmean( data_vw[I_good,:,0], axis=1),  
            cmap = 'turbo',
            linewidths=0.5,
            edgecolors='k')
plt.xlabel('UMAP #1')
plt.ylabel('UMAP #2')
plt.colorbar( label='Mean speed' )
for simplex in hull.simplices:
    plt.plot(Xnew[simplex, 0], Xnew[simplex, 1], '--',color=np.ones((3,))*0.75, zorder=0)

plt.tight_layout()





# New points
x_new = np.linspace(1,11,4)
y_new = np.linspace(7,12,4)
test_pts = np.array( [ [xx,yy] for xx in x_new for yy in y_new] )

new_points = DRA.inverse_transform(test_pts)

ii = 15
A = np.reshape( new_points[ii, :], ( (len(v_)-1),(len(w_)-1)  ) )
plt.imshow(A, aspect='auto',vmin=0, vmax=5, cmap='inferno_r')







Mndiv = []
Mwhere= []
t_fin = -1

while t_fin<data_vw.shape[1]:
    t_ini = t_ini+1
    t_fin = t_ini+Dt
    _M_submat =[]
    _Mw_sub   = np.zeros((n_worms))
    
    for ii in range(n_worms):
        _v = data_vw[ii,t_ini:t_fin,0]
        _w = data_vw[ii,t_ini:t_fin,1]
        
        _is_valid = not np.all(np.isnan(_v))
    
        if _is_valid:
            
            _M = np.histogram2d( _v, _w, bins=[v_,w_], density=True)[0]
            _M_submat.append( _M  )
            _Mw_sub[ii] = 1
        else:
            _M_submat.append( np.zeros_like(Indiv[0,:,:])*np.nan )
            _Mw_sub[ii] = 0
    
    Mndiv.append(_M_submat)
    Mwhere.append( _Mw_sub )
    
Mndiv = np.array(Mndiv)
Mwhere= np.array(Mwhere)

print(Mndiv.shape)


plt.figure( figsize=(5,4), dpi=600)
plt.imshow( Indiv[idx,:,:].T, aspect='auto', extent=extent, vmin=0, vmax=5, cmap='inferno_r')
plt.xlabel(' |v|')
plt.ylabel('$\omega$')
plt.colorbar(label='Probability density')

plt.tight_layout()


# plt.figure( figsize=(8,4), dpi=600)
# for jj



# plt.xticks(ticks=range(len(w_)-1), labels='')
# plt.yticks(ticks=range(len(v_)-1), labels='')




# plt.subplot(2,2,4)
# bbins = np.linspace(0,15,20)
# plt.hist( np.clip( t_to_goal, bbins[0], bbins[-1]) , bins = bbins , color= col_all )
# plt.hist( t_to_goal[I_good] , bins = bbins , color= col_good )
# plt.xlabel('Time to goal')
# plt.ylabel('Number of\ntrajectories')

# plt.tight_layout()


# H=plt.hist( t_to_goal, bins=40)

# ys = H[0]
# xs = H[1][0] + np.cumsum(np.diff(H[1]))
# xq = np.linspace( xs.min(), xs.max(), 100)
# p0 = [1, -1]
# params, cv = curve_fit(monoExp, xs, ys, p0)

# plt.figure( figsize=(5,4), dpi=300 )
# plt.bar(xs,ys)
# plt.plot( xq, monoExp(xq, *params) ,'r')
# plt.ylim(1e-3, 1.05*ys.max() )
# plt.xlabel('Time to NaCl (min)')
# plt.ylabel('Probability density')

# plt.text( 10, np.max( plt.ylim())*0.90, r"y ~ exp( t / $\tau$ )" )
# plt.text( 10, np.max( plt.ylim())*0.83, r"$\tau$ = %1.3f $\pm$ %1.3f" % (params[1], np.sqrt( cv[1,1]))  )

# # determine quality of the fit
# squaredDiffs = np.square(ys - monoExp(xs, *params) )
# squaredDiffsFromMean = np.square(ys - np.mean(ys))
# rSquared = 1 - np.sum(squaredDiffs) / np.sum(squaredDiffsFromMean)
# print(f"R² = {rSquared}")
# print(params)




# ####
# # Compute the 2D histograms of (v,w)
# v_ = np.linspace( 0 , 0.2 , 11 )
# w_ = np.linspace(-2*np.pi, 2*np.pi, 11 )
# t_ = 20

# indiv = np.zeros((n_worms, len(v_)-1, len(w_)-1 ) )
# #partial = np.zeros( (n_worms, n_times, len(v_)-1, len(w_)-1) )

# for worm_jj in range( n_worms ):
#     indiv[ worm_jj ] = np.histogram2d( data_vw[worm_jj,:,0], data_vw[worm_jj,:,1], bins=[v_,w_], density=True)[0]


# plt.figure( figsize=(16,16), dpi=300)
# for jj in range(5):
#     for kk in range(5):
#         idx = kk+5*jj
#         plt.subplot(5,5, idx+1 )
#         plt.imshow(indiv[idx,:,:], aspect='auto')
#         plt.xticks(ticks=range(len(w_)-1), labels='')
#         plt.yticks(ticks=range(len(v_)-1), labels='')
# plt.tight_layout()



# v_indiv = [ indiv[_,:,:].flatten() for _ in range(n_worms) ]
# v_indiv = np.array( v_indiv )


# from sklearn.manifold import TSNE
# # plt.figure(dpi=300)
# # for jj in range(9):
# #     perp = jj*2+1
# #     X_embedded = TSNE(n_components=2, learning_rate='auto',
# #                       init='random', perplexity=perp).fit_transform( v_indiv)
    

# #     plt.subplot(3,3,jj+1)    
# #     plt.scatter(X_embedded[:,0], X_embedded[:,1], c= np.nanmean(data_vw[:,:,0], axis=1))
# #     plt.xticks( ticks= (), labels='')
# #     plt.yticks( ticks= (), labels='')
# #     plt.colorbar()
# #     plt.title(perp)
# # plt.tight_layout()


# plt.figure(dpi=300)
# X_embedded = TSNE(n_components=2, learning_rate='auto',
#                   init='pca', perplexity=10).fit_transform( v_indiv)
# plt.scatter(X_embedded[:,0], X_embedded[:,1], c= np.nanmean(data_vw[:,:,0], axis=1), s=28-t_to_goal)
# plt.xticks( ticks= (), labels='')
# plt.yticks( ticks= (), labels='')
# plt.colorbar(label=' $v_0$ (mm/s)')










# # # Distance traveled by each worm, and mean speed of each one
# # plt.subplot(1,2,1)
# # plt.hist( np.nansum( data_vw[:,:,0], axis=1) , bins=10)
# # plt.xlabel('Distance traveled (mm)')

# # plt.subplot(1,2,2)
# # # # plt.hist( np.nanmean( data_vw[:,:,0], axis=1 ))
# # # # plt.hist( np.nanmax( data_vw[:,:,0], axis=1 ), bins=np.linspace(0,1,20), alpha=0.5, zorder=0)
# # # # plt.xlabel('Speed (mm/s)')
# # # # plt.legend(('Mean','Max'))

# # # idx=3
# # plt.plot( data_xy[:,:,0].T, data_xy[:,:,1].T,'.-', ms=1, lw=0.5)
# # plt.scatter( chunk[0], chunk[1], s=1000, alpha=0.5)
# # plt.scatter( chunk[0], chunk[1], s=10)
# # plt.scatter( plate[0], plate[1], s=50000, c='r',alpha=0.2)
# # plt.scatter( plate[0], plate[1], s=10,c='r')
# # plt.scatter( x_goal, y_goal, s=10,c='b')
# # plt.xlim(0,60)
# # plt.ylim(0,60)