Add files

functions.py

@@ -0,0 +1,304 @@
+import numpy
+import scipy
+
+import MDAnalysis as mda
+import matplotlib.pylab as plt
+
+
+def perform_search_time(XYZ_UNI, COFF, INIT, END, STRIDE, BOX_const=False):
+    beads = XYZ_UNI.select_atoms('all')
+    cont_list = list()
+    #loop over traj
+    for i,ts in enumerate(XYZ_UNI.trajectory[INIT:END:STRIDE]):
+        if BOX_const == True:
+            if i==0:
+                nsearch = mda.lib.NeighborSearch.AtomNeighborSearch(beads, box=XYZ_UNI.dimensions)
+        else:
+            nsearch = mda.lib.NeighborSearch.AtomNeighborSearch(beads, box=XYZ_UNI.dimensions)
+        cont_list.append([nsearch.search(i, COFF, level='A')  for i in beads])
+    return cont_list
+
+def xyz_to_uni(xyz, box=None, traj_name_output=None):
+
+    # reading input xyz and box dimensions
+    xyz_u = mda.Universe(xyz)
+    n_monomers = len(xyz_u.atoms)
+    n_frames = len(xyz_u.trajectory)
+
+    if not isinstance(box, str):
+        xyz_box_sigle = mda.lib.mdamath.triclinic_box(box[0],box[1],box[2])[:3]
+        xyz_box = numpy.array([xyz_box_sigle for i in range(n_frames)])
+        angles = list(mda.lib.mdamath.triclinic_box(box[0],box[1],box[2])[3:])
+    else:
+        xyz_box = numpy.loadtxt(box)
+        angles = [90.,90.,90.]
+
+    # new Universe init
+    new_u = mda.Universe.empty(n_atoms=n_monomers,
+                               n_residues=n_monomers,
+                               n_segments=1,
+                               atom_resindex=numpy.arange(n_monomers),
+                               residue_segindex=[1]*n_monomers,
+                               trajectory=True)
+
+    # init new empty traj
+    new_traj = numpy.empty((n_frames, n_monomers, 3))
+
+    # loop over empty frame to insert the new frames
+    for i,ts in enumerate(xyz_u.trajectory):
+        empty_frame = numpy.empty((n_monomers, 3))
+        for j,pos in enumerate(xyz_u.atoms.positions):
+            empty_frame[j] = pos
+        new_traj[i] = empty_frame
+
+    # add the dimension of the box to the new traj
+    new_u.load_new(new_traj, format=mda.coordinates.memory.MemoryReader)
+    for s,snap in enumerate(new_u.trajectory):
+        box_dim_tmp = numpy.pad(xyz_box[s], (0, 3), 'constant') + numpy.array([0.,0.,0.]+angles)
+        new_u.trajectory[snap.frame].dimensions = box_dim_tmp
+
+    return new_u
+
+
+def local_dynamics(list_sum):
+    particle = [i for i in range(numpy.shape(list_sum)[1])]
+    ncont_tot = list()
+    nn_tot = list()
+    num_tot = list()
+    den_tot  = list()
+    for p in particle:
+        ncont = list()
+        nn = list()
+        num = list()
+        den = list()
+        for frame in range(len(list_sum)):
+            if frame == 0:
+                ncont.append(0)
+                nn.append(0)
+            else:
+                # se il set di primi vicini cambia totalmente, l'intersezione è lunga 1 ovvero la bead self
+                # vale anche se il numero di primi vicini prima e dopo cambia
+                if len(list(set(list_sum[frame-1][p]) & set(list_sum[frame][p])))==1:
+                    # se non ho NN lens è 0
+                    if len(list(set(list_sum[frame-1][p])))==1 and len(set(list_sum[frame][p]))==1:
+                        ncont.append(0)
+                        nn.append(0)
+                        num.append(0)
+                        den.append(0)
+                    # se ho NN lo metto 1
+                    else:
+                        ncont.append(1)
+                        nn.append(len(list_sum[frame][p])-1)
+                        num.append(1)
+                        den.append(len(list_sum[frame-1][p])-1+len(list_sum[frame][p])-1)    
+                else:
+                    # contrario dell'intersezione fra vicini al frame f-1 e al frame f
+                    c_diff = set(list_sum[frame-1][p]).symmetric_difference(set(list_sum[frame][p]))
+                    ncont.append(len(c_diff)/(len(list_sum[frame-1][p])-1+len(list_sum[frame][p])-1))
+                    nn.append(len(list_sum[frame][p])-1)
+                    num.append(len(c_diff))
+                    den.append(len(list_sum[frame-1][p])-1+len(list_sum[frame][p])-1)
+        num_tot.append(num)
+        den_tot.append(den)
+        ncont_tot.append(ncont)
+        nn_tot.append(nn)
+    return ncont_tot, nn_tot, num_tot, den_tot
+
+def check(value, b_chunk):
+    if b_chunk[0] <= value < b_chunk[1]:
+        return True
+    return False
+
+
+def savgol_filter_mod(ncont_tot,polyorder,window,plot=True, ylim=None, xticks=None, xticks_l=None,yticks=None, yticks_l=None, xunit='$\mu$', windows_study=[10,50,100,150],polyorder_study=[2,4,6]):
+    ncont_rolling = list()
+    particle = [i for i in range(numpy.shape(ncont_tot)[0])]
+    for p in particle:
+        savgol_2_10 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[0],polyorder=polyorder_study[0])
+        savgol_2_50 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[1],polyorder=polyorder_study[0])
+        savgol_2_100 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[2],polyorder=polyorder_study[0])
+        savgol_2_150 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[3],polyorder=polyorder_study[0])
+
+        savgol_4_10 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[0],polyorder=polyorder_study[1])
+        savgol_4_50 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[1],polyorder=polyorder_study[1])
+        savgol_4_100 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[2],polyorder=polyorder_study[1])
+        savgol_4_150 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[3],polyorder=polyorder_study[1])
+
+
+        savgol_6_10 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[0],polyorder=polyorder_study[2])
+        savgol_6_50 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[1],polyorder=polyorder_study[2])
+        savgol_6_100 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[2],polyorder=polyorder_study[2])    
+        savgol_6_150 = scipy.signal.savgol_filter(ncont_tot[p], window_length=windows_study[3],polyorder=polyorder_study[2])
+
+        savgol = scipy.signal.savgol_filter(ncont_tot[p], window_length=window,polyorder=polyorder)
+        ncont_rolling.append(savgol[int(window/2):-int(window/2)])
+
+        if p%100==0 and plot:
+            fig, ax = plt.subplots(1,4, figsize=(16,4), dpi=500)
+            fig.suptitle(r'Bead ID '+str(p), size=20)
+            ax[0].set_title("window: " +str(windows_study[0]), fontsize=20)
+            ax[0].plot(savgol_2_10, c='green', label='poly order '+str(polyorder_study[0]))
+            ax[0].plot(savgol_4_10, c='blue', label='poly order '+str(polyorder_study[1]))
+            ax[0].plot(savgol_6_10, c='red', label='poly order '+str(polyorder_study[2]))
+            ax[0].plot(ncont_tot[p], c='gray',lw=0.5, alpha=0.8)
+            ax[0].set_ylim(ylim)
+            ax[0].set_ylabel(r'$\delta_b^{\tau}$', size=20)
+            ax[0].set_xlabel(r't ['+xunit+'s]', size=20)
+            ax[0].set_xticks(xticks, fontsize=20)
+            ax[0].set_xticklabels(xticks_l, fontsize=20)
+            ax[0].set_yticks(yticks, fontsize=20)
+            ax[0].set_yticklabels(yticks_l, fontsize=20)
+            ax[0].legend()
+
+            ax[1].set_title("window: " +str(windows_study[1]), fontsize=20)
+            ax[1].plot(savgol_2_50, c='green', label='poly order '+str(polyorder_study[0]))
+            ax[1].plot(savgol_4_50, c='blue', label='poly order '+str(polyorder_study[1]))
+            ax[1].plot(savgol_6_50, c='red', label='poly order '+str(polyorder_study[2]))
+            ax[1].plot(ncont_tot[p], c='gray',lw=0.1, alpha=0.8)
+            ax[1].set_ylim(ylim)
+            ax[1].set_xlabel(r't ['+xunit+'s]', size=20)
+            ax[1].set_xticks(xticks, fontsize=20)
+            ax[1].set_xticklabels(xticks_l, fontsize=20)
+            ax[1].set_yticks([], fontsize=20)
+            ax[1].set_yticklabels([], fontsize=20)
+            ax[1].legend()
+
+            ax[2].set_title("window: " +str(windows_study[2]), fontsize=20)        
+            ax[2].plot(savgol_2_100, c='green', label='poly order '+str(polyorder_study[0]))
+            ax[2].plot(savgol_4_100, c='blue',label='poly order '+str(polyorder_study[1]))
+            ax[2].plot(savgol_6_100, c='red', label='poly order '+str(polyorder_study[2]))
+            ax[2].plot(ncont_tot[p], c='gray',lw=0.1, alpha=0.8)
+            ax[2].set_ylim(ylim)
+            ax[2].set_xlabel(r't ['+xunit+'s]', size=20)
+            ax[2].set_xticks(xticks, fontsize=20)
+            ax[2].set_xticklabels(xticks_l, fontsize=20)
+            ax[2].set_yticks([], fontsize=20)
+            ax[2].set_yticklabels([], fontsize=20)
+            ax[2].legend()
+
+            ax[3].set_title("window: " +str(windows_study[3]), fontsize=20)
+            ax[3].plot(savgol_2_150, c='green', label='poly order '+str(polyorder_study[0]))
+            ax[3].plot(savgol_4_150, c='blue', label='poly order '+str(polyorder_study[1]))
+            ax[3].plot(savgol_6_150, c='red', label='poly order '+str(polyorder_study[2]))
+            ax[3].plot(ncont_tot[p], c='gray',lw=0.1, alpha=0.8)
+            ax[3].set_ylim(ylim)
+            ax[3].set_xlabel(r't ['+xunit+'s]', size=20)
+            ax[3].set_xticks(xticks, fontsize=20)
+            ax[3].set_xticklabels(xticks_l, fontsize=20)
+            ax[3].set_yticks([], fontsize=20)
+            ax[3].set_yticklabels([], fontsize=20)        
+            ax[3].legend()
+
+    #         for coff in range(len(dynamic_coff)):
+    #             ax[0].hlines(dynamic_coff[coff], INIT/STRIDE, END/STRIDE, colors='black', linewidth=2, linestyle=':')
+    #             ax[1].hlines(dynamic_coff[coff], INIT/STRIDE, END/STRIDE, colors='black', linewidth=2, linestyle=':')
+    #             ax[2].hlines(dynamic_coff[coff], INIT/STRIDE, END/STRIDE, colors='black', linewidth=2, linestyle=':')
+    #             ax[3].hlines(dynamic_coff[coff], INIT/STRIDE, END/STRIDE, colors='black', linewidth=2, linestyle=':')
+    #             ax[3].text(END/STRIDE+END/STRIDE*0.066,dynamic_coff[coff],'c'+str(coff),fontsize=18)        
+            plt.tight_layout()
+    return ncont_rolling
+
+# Add properties based on properties list
+def add_properties(prop_dict):
+    prop = {"name" : [p[0] for p in prop_dict["Properties"]],\
+            "type" : [p[1] for p in prop_dict["Properties"]],\
+            "col" : [p[2] for p in prop_dict["Properties"]]}
+
+    properties_l = list()
+    for p in range(len(prop["name"])):
+        properties_l.append(str(prop["name"][p])+\
+                        ":"+str(prop["type"][p])+\
+                        ":"+str(prop["col"][p]))
+    return properties_l
+
+# Add field in xyz comment line
+def xyz_extender(prop_dict, i):
+    ext_prop = 'Lattice="'+str(' '.join([str(item) for sublist in prop_dict["Lattice"][i] for item in sublist]))+\
+            '" Origin="'+str(' '.join([str(item) for item in prop_dict["Origin"]]))+\
+                '" Properties='+str(':'.join(add_properties(prop_dict)))
+    return ext_prop
+
+def arrowed_spines(
+        ax,
+        x_width_fraction=0.05,
+        x_height_fraction=0.05,
+        lw=None,
+        ohg=0.3,
+        locations=('bottom right', 'left up'),
+        **arrow_kwargs
+):
+
+    # set/override some default plotting parameters if required
+    arrow_kwargs.setdefault('overhang', ohg)
+    arrow_kwargs.setdefault('clip_on', False)
+    arrow_kwargs.update({'length_includes_head': True})
+
+    # axis line width
+    if lw is None:
+        # FIXME: does this still work if the left spine has been deleted?
+        lw = ax.spines['left'].get_linewidth()
+
+    annots = {}
+
+    xmin, xmax = ax.get_xlim()
+    ymin, ymax = ax.get_ylim()
+
+    # get width and height of axes object to compute
+    # matching arrowhead length and width
+    fig = ax.get_figure()
+    dps = fig.dpi_scale_trans.inverted()
+    bbox = ax.get_window_extent().transformed(dps)
+    width, height = bbox.width, bbox.height
+
+    # manual arrowhead width and length
+    hw = x_width_fraction * (ymax-ymin)
+    hl = x_height_fraction * (xmax-xmin)
+
+    # compute matching arrowhead length and width
+    yhw = hw/(ymax-ymin)*(xmax-xmin)* height/width
+    yhl = hl/(xmax-xmin)*(ymax-ymin)* width/height
+
+    # draw x and y axis
+    for loc_str in locations:
+        side, direction = loc_str.split(' ')
+        assert side in {'top', 'bottom', 'left', 'right'}, "Unsupported side"
+        assert direction in {'up', 'down', 'left', 'right'}, "Unsupported direction"
+
+        if side in {'bottom', 'top'}:
+            if direction in {'up', 'down'}:
+                raise ValueError("Only left/right arrows supported on the bottom and top")
+
+            dy = 0
+            head_width = hw
+            head_length = hl
+
+            y = ymin if side == 'bottom' else ymax
+
+            if direction == 'right':
+                x = xmin
+                dx = xmax - xmin
+            else:
+                x = xmax
+                dx = xmin - xmax
+
+        else:
+            if direction in {'left', 'right'}:
+                raise ValueError("Only up/downarrows supported on the left and right")
+            dx = 0
+            head_width = yhw
+            head_length = yhl
+
+            x = xmin if side == 'left' else xmax
+
+            if direction == 'up':
+                y = ymin
+                dy = ymax - ymin
+            else:
+                y = ymax
+                dy = ymin - ymax
+
+
+        annots[loc_str] = ax.arrow(x, y, dx, dy, fc='k', ec='k', lw = lw,
+                 head_width=head_width, head_length=head_length, **arrow_kwargs)
+
+    return annots