Source code for cytocalc.csmsimulation

#! /usr/bin/env python3
# A CSMSimulation is a collection of CSMFrames and the
# simulation parameters (read from configurations.cym)

import cytocalc.csmframe
import freud
import warnings
import numpy as np
import pandas as pd
from scipy.stats import linregress
from bisect import insort
from typing import List



class CSMSimulation:
    """
    Contains a set of CSMFrames that form a simulation.
    Also contains the simulation parameters.
    """

    def __init__(self):
        self.frames = [] # list containing cytosim frames
        self.params = {} # dict containing simulation



    def add_frame(self, csmframe):
        """ Adds a CSMFrame to the simulation
        at the correct time (frame1.time < frame2.time)"""
        insort(self.frames, csmframe, key=lambda x: x.time)




    def get_np_object(self, start_frame='0', end_frame=None):
        """ Constructs a numpy object from csm frame data
        of shape (n_frames,n_particles,3) """
        np_frames = [frame._npobj for frame in self.frames[start_frame:end_frame]]
        return np.stack(np_frames)




    def contraction_rate(self, start_frame=None, end_frame=None, r_sq = 0.9):
        """ Obtain the contraction rate (rate of decrease in radius of network)
        between start_frame and end_frame """
        time = []
        radius = []
        for frame in self.frames[start_frame:end_frame]:
            time.append(frame.time)
            radius.append(frame.compute_radius())
        fit = linregress(time, radius)
        # warn if R^2 value is low
        if (fit.rvalue)**2 <= r_sq:
            warnings.warn('R-square value is too low!')
        return fit.slope




    def compute_msd(self, start_frame : int = 0, end_frame : int = -1):
        """
        Compute Mean Squared Displacement (direct) of objects.

        :param start_frame: int, frame number for starting MSD calculation, defaults to the first frame
        :param end_frame: int, frame at which MSD calculation ends, defaults to the last (-1) frame
        :return: dict, with time difference as key and corresponding MSD as value
        """
        init_frame = self.frames[start_frame]
        init_disp = init_frame._npobj

        msd_vs_time = {}
        for frame in self.frames[start_frame+1:end_frame]:
            time = round(frame.time - init_frame.time,4) # time diff between start and curr frame
            disp = frame._npobj - init_disp
            msd = np.mean(np.sum(disp**2, axis=1))
            msd_vs_time[time] = msd
        return msd_vs_time




    def compute_time_avg_msd(self,
                             start_frame : int = 0,
                             end_frame : int = -1,
                             box_size : List[float] = []):
        """
        Computes a 'time averaged MSD' i.e using all distinct displacements to extract more data.
        See [1] for futher details

        [1] https://en.wikipedia.org/wiki/Mean_squared_displacement#Definition_of_MSD_for_time_lags

        :param start_frame: int, frame number for starting MSD calculation, defaults to the first frame
        :param end_frame: int, frame at which MSD calculation ends, defaults to the last (-1) frame
        :param box_size: List[float], dimensions of box (assumes periodicty), ignore if non-periodic
        :return: dict, with time difference as key and corresponding MSD as value
        """
        init_frame = self.frames[start_frame]

        sim_array = self.get_np_object(start_frame,end_frame)

        # freud isn't thread_safe, breaks if we multithread elsewhere
        # so we disable multithreading in freud
        freud.parallel.NumThreads(1)

        msd_vs_DeltaT = {}

        if not box_size:
            # no box size specified, assume non-periodic
            t_avg_msd = freud.msd.MSD().compute(sim_array).msd
        else:
            box = freud.Box.from_box(box_size, dimensions=3) # our position data is always in 3D
            t_avg_msd = freud.msd.MSD(box).compute(sim_array).msd

        # return timestamps along with MSD
        for count,frame in enumerate(self.frames[start_frame:end_frame]):
            time = frame.time - init_frame.time # time diff between start and curr frame
            msd_vs_DeltaT[time] = t_avg_msd[count]

        return msd_vs_DeltaT




    def bound_couple_vs_time(self, start_frame : int = 0, end_frame: int = -1):
        """
        Get the number of bound hands as a function of time

        :param start_frame: int, frame number for starting calculation, defaults to the first frame
        :param end_frame: int, frame at which calculation ends, defaults to the last (-1) frame
        :return: dict, with time as key and corresponding bound crosslinker count `s value
        """
        AA_vs_time = {}
        for frame in self.frames[start_frame:end_frame]:
            AA_vs_time[frame.time] = int(frame.data['AA'].iloc[0])
        return AA_vs_time