pyLoopSage 0.1.1__tar.gz → 1.0.1__tar.gz

{pyloopsage-0.1.1 → pyloopsage-1.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pyLoopSage
-Version: 0.1.1
+Version: 1.0.1
 Summary: An energy-based stochastic model of loop extrusion in chromatin.
 Home-page: https://github.com/SFGLab/pyLoopSage
 Author: Sebastian Korsak

pyloopsage-1.0.1/loopsage/__init__.py

@@ -0,0 +1,4 @@
+# Importing specific functions or classes from submodules
+from .run import main
+
+__version__ = "1.0.1"

pyloopsage-1.0.1/loopsage/args_definition.py

@@ -0,0 +1,160 @@
+import datetime
+import re
+from dataclasses import dataclass
+from math import pi
+from typing import Union
+import argparse
+import openmm as mm
+from openmm.unit import Quantity
+
+@dataclass
+class Arg(object):
+    name: str
+    help: str
+    type: type
+    default: Union[str, float, int, bool, Quantity, None]
+    val: Union[str, float, int, bool, Quantity, None]
+
+# Define custom type to parse list from string
+def parse_list(s):
+    try:
+        return [int(x.strip()) for x in s.strip('[]').split(',')]
+    except ValueError:
+        raise argparse.ArgumentTypeError("Invalid list format. Must be a comma-separated list of integers.")
+
+class ListOfArgs(list):
+    quantity_regexp = re.compile(r'(?P<value>[-+]?\d+(?:\.\d+)?) ?(?P<unit>\w+)')
+
+    def get_arg(self, name: str) -> Arg:
+        """Stupid arg search in list of args"""
+        name = name.upper()
+        for i in self:
+            if i.name == name:
+                return i
+        raise ValueError(f"No such arg: {name}")
+
+    def __getattr__(self, item):
+        return self.get_arg(item).val
+
+    def parse_args(self):
+        parser = argparse.ArgumentParser()
+        for arg in self.arg_list:
+            parser.add_argument(arg['name'], help=arg['help'], type=arg.get('type', str), default=arg.get('default', ''), val=arg.get('val', ''))
+
+        args = parser.parse_args()
+        parsed_args = {arg['name']: getattr(args, arg['name']) for arg in self.arg_list}
+        return parsed_args
+
+    def parse_quantity(self, val: str) -> Union[Quantity, None]:
+        if val == '':
+            return None
+        match_obj = self.quantity_regexp.match(val)
+        value, unit = match_obj.groups()
+        try:
+            unit = getattr(mm.unit, unit)
+        except AttributeError:
+            raise ValueError(f"I Can't recognise unit {unit} in expression {val}. Example of valid quantity: 12.3 femtosecond.")
+        return Quantity(value=float(value), unit=unit)
+
+    def to_python(self):
+        """Casts string args to ints, floats, bool..."""
+        for i in self:
+            if i.val == '':
+                i.val = None
+            elif i.name == "HR_K_PARAM":  # Workaround for complex unit
+                i.val = Quantity(float(i.val), mm.unit.kilojoule_per_mole / mm.unit.nanometer ** 2)
+            elif i.type == str:
+                continue
+            elif i.type == int:
+                i.val = int(i.val)
+            elif i.type == float:
+                i.val = float(i.val)
+            elif i.type == bool:
+                if i.val.lower() in ['true', '1', 'y', 'yes']:
+                    i.val = True
+                elif i.val.lower() in ['false', '0', 'n', 'no']:
+                    i.val = False
+                else:
+                    raise ValueError(f"Can't convert {i.val} into bool type.")
+            elif i.type == Quantity:
+                try:
+                    i.val = self.parse_quantity(i.val)
+                except AttributeError:
+                    raise ValueError(f"Can't parse: {i.name} = {i.val}")
+            else:
+                raise ValueError(f"Can't parse: {i.name} = {i.val}")
+
+    def get_complete_config(self) -> str:
+        w = "####################\n"
+        w += "#   LoopSage Model   #\n"
+        w += "####################\n\n"
+        w += "# This is automatically generated config file.\n"
+        w += f"# Generated at: {datetime.datetime.now().isoformat()}\n\n"
+        w += "# Notes:\n"
+        w += "# Some fields require units. Units are represented as objects from mm.units module.\n"
+        w += "# Simple units are parsed directly. For example: \n"
+        w += "# HR_R0_PARAM = 0.2 nanometer\n"
+        w += "# But more complex units does not have any more sophisticated parser written, and will fail.'\n"
+        w += "# In such cases the unit is fixed (and noted in comment), so please convert complex units manually if needed.\n"
+        w += "# <float> and <int> types does not require any unit. Quantity require unit.\n\n"
+        w += "# Default values does not mean valid value. In many places it's only a empty field that need to be filled.\n\n"
+
+        w += '[Main]'
+        for i in self:
+            w += f'; {i.help}, type: {i.type.__name__}, default: {i.default}\n'
+            if i.val is None:
+                w += f'{i.name} = \n\n'
+            else:
+                if i.type == Quantity:
+                    # noinspection PyProtectedMember
+                    w += f'{i.name} = {i.val._value} {i.val.unit.get_name()}\n\n'
+                else:
+                    w += f'{i.name} = {i.val}\n\n'
+        w = w[:-2]
+        return w
+
+    def write_config_file(self):
+        auto_config_filename = 'config_auto.ini'
+        with open(auto_config_filename, 'w') as f:
+            f.write(self.get_complete_config())
+        print(f"Automatically generated config file saved in {auto_config_filename}")
+
+available_platforms = [mm.Platform.getPlatform(i).getName() for i in range(mm.Platform.getNumPlatforms())]
+
+args = ListOfArgs([
+    # Platform settings
+    Arg('PLATFORM', help=f"name of the platform. Available choices: {' '.join(available_platforms)}", type=str, default='CPU', val='CPU'),
+    Arg('DEVICE', help="device index for CUDA or OpenCL (count from 0)", type=str, default='', val=''),
+    
+    # Input data
+    Arg('N_BEADS', help="Number of Simulation Beads.", type=int, default='', val=''),
+    Arg('BEDPE_PATH', help="A .bedpe file path with loops. It is required.", type=str, default='', val=''),
+    Arg('OUT_PATH', help="Output folder name.", type=str, default='../results', val='../results'),
+    Arg('REGION_START', help="Starting region coordinate.", type=int, default='', val=''),
+    Arg('REGION_END', help="Ending region coordinate.", type=int, default='', val=''),
+    Arg('CHROM', help="Chromosome that corresponds the the modelling region of interest (in case that you do not want to model the whole genome).", type=str, default='', val=''),
+    
+    # Stochastic Simulation parameters
+    Arg('LEF_RW', help="True in case that you would like to make cohesins slide as random walk, instead of sliding only in one direction.", type=bool, default='True', val='True'),
+    Arg('N_STEPS', help="Number of Monte Carlo steps.", type=int, default='', val=''),
+    Arg('N_LEF', help="Number of loop extrusion factors (condensins and cohesins).", type=int, default='', val=''),
+    Arg('MC_STEP', help="Monte Carlo frequency.", type=int, default='200', val='200'),
+    Arg('BURNIN', help="Burnin-period (steps that are considered before equillibrium).", type=int, default='1000', val='1000'),
+    Arg('T_INIT', help="Initial Temperature of the Stochastic Model.", type=float, default='1.8', val='1.8'),
+    Arg('T_FINAL', help="Final Temperature of the Stochastic Model.", type=float, default='0.01', val='0.01'),
+    Arg('METHOD', help="Stochastic modelling method. It can be Metropolis or Simulated Annealing.", type=str, default='Annealing', val='Annealing'),
+    Arg('FOLDING_COEFF', help="Folding coefficient.", type=float, default='1.0', val='1.0'),
+    Arg('CROSS_COEFF', help="LEF crossing coefficient.", type=float, default='1.0', val='1.0'),
+    Arg('BIND_COEFF', help="CTCF binding coefficient.", type=float, default='1.0', val='1.0'),
+
+    # Molecular Dynamic Properties
+    Arg('INITIAL_STRUCTURE_TYPE', help="you can choose between: rw, confined_rw, self_avoiding_rw, helix, circle, spiral, sphere.", type=str, default='rw', val='rw'),
+    Arg('SIMULATION_TYPE', help="It can be either EM (multiple energy minimizations) or MD (one energy minimization and then run molecular dynamics).", type=str, default='', val=''),
+    Arg('FORCEFIELD_PATH', help="Path to XML file with forcefield.", type=str, default='forcefield/classic_sm_ff.xml', val='forcefield/classic_sm_ff.xml'),
+    Arg('ANGLE_FF_STRENGTH', help="Angle force strength.", type=float, default='200.0', val='200.0'),
+    Arg('LE_FF_LENGTH', help="Equillibrium distance of loop forces.", type=float, default='0.0', val='0.0'),
+    Arg('LE_FF_STRENGTH', help="Equillibrium distance of loop forces.", type=float, default='300000.0', val='300000.0'),
+    Arg('EV_FF_STRENGTH', help="Excluded-volume strength.", type=float, default='10.0', val='10.0'),
+    Arg('TOLERANCE', help="Tolerance that works as stopping condition for energy minimization.", type=float, default='0.001', val='0.001'),
+    Arg('SIM_STEP', help="This is the amount of simulation steps that are perform each time that we change the loop forces. If this number is too high, the simulation is slow, if is too low it may not have enough time to adapt the structure to the new constraints.", type=int, default='100', val='100'),
+])

pyloopsage-1.0.1/loopsage/em.py

@@ -0,0 +1,131 @@
+import copy
+import time
+import numpy as np
+import openmm as mm
+import openmm.unit as u
+from tqdm import tqdm
+from sys import stdout
+from mdtraj.reporters import HDF5Reporter
+from scipy import ndimage
+from openmm.app import PDBFile, PDBxFile, ForceField, Simulation, PDBReporter, PDBxReporter, DCDReporter, StateDataReporter, CharmmPsfFile
+from .utils import *
+from .initial_structures import *
+
+class EM_LE:
+    def __init__(self,M,N,N_beads,burnin,MC_step,path,platform,angle_ff_strength=200,le_distance=0.1,le_ff_strength=300000.0,ev_ff_strength=10.0,tolerance=0.001):
+        '''
+        M, N (np arrays): Position matrix of two legs of cohesin m,n. 
+                          Rows represent  loops/cohesins and columns represent time
+        N_beads (int): The number of beads of initial structure.
+        step (int): sampling rate
+        path (int): the path where the simulation will save structures etc.
+        '''
+        self.M, self.N = M, N
+        self.N_coh, self.N_steps = M.shape
+        self.N_beads, self.step, self.burnin = N_beads, MC_step, burnin//MC_step
+        self.path = path
+        self.platform = platform
+        self.angle_ff_strength = angle_ff_strength
+        self.le_distance = le_distance
+        self.le_ff_strength = le_ff_strength
+        self.ev_ff_strength = ev_ff_strength
+        self.tolerance = tolerance
+    
+    def run_pipeline(self,write_files=False,plots=False):
+        '''
+        This is the basic function that runs the molecular simulation pipeline.
+
+        Input parameters:
+        run_MD (bool): True if user wants to run molecular simulation (not only energy minimization).
+        sim_step (int): the simulation step of Langevin integrator.
+        write_files (bool): True if the user wants to save the structures that determine the simulation ensemble.
+        plots (bool): True if the user wants to see the output average heatmaps.
+        '''
+        # Define initial structure
+        print('Building initial structure...')
+        points = compute_init_struct(self.N_beads,mode='rw')
+        write_mmcif(points,self.path+'/LE_init_struct.cif')
+        generate_psf(self.N_beads,self.path+'/other/LE_init_struct.psf')
+        print('Done brother ;D\n')
+
+        counter = 0
+        sum_heat = np.zeros((self.N_beads,self.N_beads))
+        print('Running energy minimizations...')
+        for i in tqdm(range(self.burnin,self.N_steps)):
+            # Define System
+            pdb = PDBxFile(self.path+'/LE_init_struct.cif')
+            forcefield = ForceField('forcefields/classic_sm_ff.xml')
+            self.system = forcefield.createSystem(pdb.topology, nonbondedCutoff=1*u.nanometer)
+            integrator = mm.LangevinIntegrator(310, 0.05, 100 * mm.unit.femtosecond)
+
+            # Add forces
+            ms,ns = self.M[:,i], self.N[:,i]
+            self.add_forcefield(ms,ns)
+
+            # Minimize energy
+            platform = mm.Platform.getPlatformByName(self.platform)
+            simulation = Simulation(pdb.topology, self.system, integrator, platform)
+            simulation.context.setPositions(pdb.positions)
+            simulation.minimizeEnergy(tolerance=self.tolerance)
+            self.state = simulation.context.getState(getPositions=True)
+            PDBxFile.writeFile(pdb.topology, self.state.getPositions(), open(self.path+f'/ensemble/EMLE_{i-self.burnin+1}.cif', 'w'))
+            save_path = self.path+f'/heatmaps/heat_{i-self.burnin+1}.svg' if write_files else None
+            sum_heat+=get_heatmap(self.state.getPositions(),save_path=save_path,save=write_files)
+            counter+=1
+        print('Energy minimizations done :D\n')
+
+        self.avg_heat = sum_heat/counter
+        np.save(self.path+f'/other/avg_heatmap.npy',self.avg_heat)
+        if plots:
+            figure(figsize=(10, 10))
+            plt.imshow(self.avg_heat,cmap="Reds",vmax=1)
+            plt.colorbar()
+            plt.savefig(self.path+f'/plots/avg_heatmap.svg',format='svg',dpi=500)
+            plt.savefig(self.path+f'/plots/avg_heatmap.pdf',format='pdf',dpi=500)
+            # plt.colorbar()
+            plt.close()
+            
+            return self.avg_heat
+
+    def add_forcefield(self,ms,ns):
+        '''
+        Here is the definition of the forcefield.
+
+        There are the following energies:
+        - ev force: repelling LJ-like forcefield
+        - harmonic bond force: to connect adjacent beads.
+        - angle force: for polymer stiffness.
+        - LE forces: this is a list of force objects. Each object corresponds to a different cohesin. It is needed to define a force for each time step.
+        '''
+        # Leonard-Jones potential for excluded volume
+        self.ev_force = mm.CustomNonbondedForce('epsilon*((sigma1+sigma2)/r)^6')
+        self.ev_force.addGlobalParameter('epsilon', defaultValue=self.ev_ff_strength)
+        self.ev_force.addPerParticleParameter('sigma')
+        self.system.addForce(self.ev_force)
+        for i in range(self.system.getNumParticles()):
+            self.ev_force.addParticle([0.1])
+
+        # Harmonic bond borce between succesive beads
+        self.bond_force = mm.HarmonicBondForce()
+        self.system.addForce(self.bond_force)
+        for i in range(self.system.getNumParticles() - 1):
+            self.bond_force.addBond(i, i + 1, 0.1, 300000.0)
+
+        # Harmonic angle force between successive beads so as to make chromatin rigid
+        self.angle_force = mm.HarmonicAngleForce()
+        self.system.addForce(self.angle_force)
+        for i in range(self.system.getNumParticles() - 2):
+            self.angle_force.addAngle(i, i + 1, i + 2, np.pi, self.angle_ff_strength)
+        
+        # LE force that connects cohesin restraints
+        self.LE_force = mm.HarmonicBondForce()
+        self.system.addForce(self.LE_force)
+        for nn in range(self.N_coh):
+            self.LE_force.addBond(ms[nn], ns[nn], self.le_distance, self.le_ff_strength)
+        
+def main():
+    # A potential example
+    M = np.load('/home/skorsak/Dropbox/LoopSage/files/region_[48100000,48700000]_chr3/Annealing_Nbeads500_ncoh50/Ms.npy')
+    N = np.load('/home/skorsak/Dropbox/LoopSage/files/region_[48100000,48700000]_chr3/Annealing_Nbeads500_ncoh50/Ns.npy')
+    md = EM_LE(4*M,4*N,2000,5,1)
+    md.run_pipeline(write_files=False,plots=True,sim_step=100)

pyloopsage-1.0.1/loopsage/initial_structures.py

@@ -0,0 +1,137 @@
+import numpy as np
+from tqdm import tqdm
+
+def dist(p1: np.ndarray, p2: np.ndarray) -> float:
+    """Mierzy dystans w przestrzeni R^3"""
+    x1, y1, z1 = p1
+    x2, y2, z2 = p2
+    return ((x1 - x2) ** 2 + (y1 - y2) ** 2 + (z1 - z2) ** 2) ** 0.5  # faster than np.linalg.norm
+
+def random_versor() -> np.ndarray:
+    """Losuje wersor"""
+    x = np.random.uniform(-1, 1)
+    y = np.random.uniform(-1, 1)
+    z = np.random.uniform(-1, 1)
+    d = (x ** 2 + y ** 2 + z ** 2) ** 0.5
+    return np.array([x / d, y / d, z / d])
+
+def self_avoiding_random_walk(n: int, step: float = 1.0, bead_radius: float = 0.5, epsilon: float = 0.001, two_dimensions=False) -> np.ndarray:
+    potential_new_step = [0, 0, 0]
+    while True:
+        points = [np.array([0, 0, 0])]
+        for _ in tqdm(range(n - 1)):
+            step_is_ok = False
+            trials = 0
+            while not step_is_ok and trials < 1000:
+                potential_new_step = points[-1] + step * random_versor()
+                if two_dimensions:
+                    potential_new_step[2] = 0
+                for j in points:
+                    d = dist(j, potential_new_step)
+                    if d < 2 * bead_radius - epsilon:
+                        trials += 1
+                        break
+                else:
+                    step_is_ok = True
+            points.append(potential_new_step)
+        points = np.array(points)
+        return points
+
+def polymer_circle(n: int, z_stretch: float = 1.0, radius: float = 5.0) -> np.ndarray:
+    points = []
+    angle_increment = 360 / float(n)
+    radius = 1 / (2 * np.sin(np.radians(angle_increment) / 2.)) if radius==None else radius
+    z_stretch = z_stretch / n
+    z = 0
+    for i in range(n):
+        x = radius * np.cos(angle_increment * i * np.pi / 180)
+        y = radius * np.sin(angle_increment * i * np.pi / 180)
+        if z_stretch != 0:
+            z += z_stretch
+        points.append((x, y, z))
+    points = np.array(points)
+    return points
+
+def helix_structure(N_beads, radius=1, pitch=2):
+    theta = np.linspace(0, 4 * np.pi, N_beads)  # 2 full turns
+    x = radius * np.cos(theta)
+    y = radius * np.sin(theta)
+    z = np.linspace(0, pitch * N_beads, N_beads)
+    V = np.column_stack((x, y, z))
+    return V
+
+def spiral_structure(N_beads, initial_radius=1, pitch=1, growth_factor=0.05):
+    theta = np.linspace(0, 4 * np.pi, N_beads)
+    radius = initial_radius + growth_factor * np.arange(N_beads)
+    x = radius * np.cos(theta)
+    y = radius * np.sin(theta)
+    z = np.linspace(0, pitch * N_beads, N_beads)
+    V = np.column_stack((x, y, z))
+    return V
+
+def sphere_surface_structure(N_beads, radius=1):
+    phi = np.random.uniform(0, 2 * np.pi, N_beads)
+    costheta = np.random.uniform(-1, 1, N_beads)
+    u = np.random.uniform(0, 1, N_beads)
+    
+    theta = np.arccos(costheta)
+    r = radius * u ** (1/3)
+
+    x = r * np.sin(theta) * np.cos(phi)
+    y = r * np.sin(theta) * np.sin(phi)
+    z = r * np.cos(theta)
+    
+    V = np.column_stack((x, y, z))
+    return V
+
+def confined_random_walk(N_beads, box_size=5):
+    V = np.zeros((N_beads, 3))
+    for i in range(1, N_beads):
+        step = np.random.choice([-1, 1], size=3)  # Random step in x, y, z
+        V[i] = V[i-1] + step
+        # Keep the points within a confined box (folding in)
+        V[i] = np.clip(V[i], -box_size, box_size)
+    return V
+
+def trefoil_knot_structure(N_beads, scale=5):
+    t = np.linspace(0, 2 * np.pi, N_beads)
+    x = scale * (np.sin(t) + 2 * np.sin(2 * t))
+    y = scale * (np.cos(t) - 2 * np.cos(2 * t))
+    z = -scale * np.sin(3 * t)
+    
+    V = np.column_stack((x, y, z))
+    return V
+
+def random_walk_structure(N_beads, step_size=1):
+    # Initialize the structure array
+    V = np.zeros((N_beads, 3))
+    
+    # Loop over each bead, starting from the second one
+    for i in range(1, N_beads):
+        # Randomly pick a direction for each step
+        step_direction = np.random.normal(size=3)
+        step_direction /= np.linalg.norm(step_direction)  # Normalize to make it unit length
+        
+        # Move the current bead from the last bead by a fixed step size
+        V[i] = V[i-1] + step_size * step_direction
+    
+    return V
+
+def compute_init_struct(N_beads,mode='rw'):
+    match mode:
+        case 'rw':
+            return random_walk_structure(N_beads)
+        case 'confined_rw':
+            return confined_random_walk(N_beads)
+        case 'self_avoiding_rw':
+            return self_avoiding_random_walk(N_beads)
+        case 'circle':
+            return polymer_circle(N_beads)
+        case 'helix':
+            return helix_structure(N_beads)
+        case 'spiral':
+            return spiral_structure(N_beads)
+        case 'sphere':
+            return sphere_surface_structure(N_beads)
+        case _:
+            return IndentationError('Invalid option for initial structure.')

pyloopsage-1.0.1/loopsage/md.py

@@ -0,0 +1,175 @@
+#########################################################################
+########### CREATOR: SEBASTIAN KORSAK, WARSAW 2022 ######################
+#########################################################################
+
+import copy
+import time
+import numpy as np
+import openmm as mm
+import openmm.unit as u
+from tqdm import tqdm
+from sys import stdout
+from mdtraj.reporters import HDF5Reporter
+from scipy import ndimage
+from openmm.app import PDBFile, PDBxFile, ForceField, Simulation, PDBReporter, PDBxReporter, DCDReporter, StateDataReporter, CharmmPsfFile
+from .utils import *
+from .initial_structures import *
+
+class MD_LE:
+    def __init__(self,M,N,N_beads,burnin,MC_step,path,platform,angle_ff_strength=200,le_distance=0.1,le_ff_strength=300000.0,ev_ff_strength=10.0,tolerance=0.001):
+        '''
+        M, N (np arrays): Position matrix of two legs of cohesin m,n. 
+                          Rows represent  loops/cohesins and columns represent time
+        N_beads (int): The number of beads of initial structure.
+        step (int): sampling rate
+        path (int): the path where the simulation will save structures etc.
+        '''
+        self.M, self.N = M, N
+        self.N_coh, self.N_steps = M.shape
+        self.N_beads, self.step, self.burnin = N_beads, MC_step, burnin//MC_step
+        self.path = path
+        self.platform = platform
+        self.angle_ff_strength = angle_ff_strength
+        self.le_distance = le_distance
+        self.le_ff_strength = le_ff_strength
+        self.ev_ff_strength = ev_ff_strength
+        self.tolerance = tolerance
+    
+    def run_pipeline(self,run_MD=True,sim_step=100,write_files=False,plots=False):
+        '''
+        This is the basic function that runs the molecular simulation pipeline.
+
+        Input parameters:
+        run_MD (bool): True if user wants to run molecular simulation (not only energy minimization).
+        sim_step (int): the simulation step of Langevin integrator.
+        write_files (bool): True if the user wants to save the structures that determine the simulation ensemble.
+        plots (bool): True if the user wants to see the output average heatmaps.
+        '''
+        # Parameters
+        self.angle_ff_strength=200
+        self.le_distance=0.1
+        self.le_ff_strength=300000.0
+        self.tolerance=0.001
+
+        # Define initial structure
+        print('Building initial structure...')
+        points = compute_init_struct(self.N_beads,mode='rw')
+        write_mmcif(points,self.path+'/LE_init_struct.cif')
+        generate_psf(self.N_beads,self.path+'/other/LE_init_struct.psf')
+        print('Done brother ;D\n')
+
+        # Define System
+        pdb = PDBxFile(self.path+'/LE_init_struct.cif')
+        forcefield = ForceField('forcefields/classic_sm_ff.xml')
+        self.system = forcefield.createSystem(pdb.topology, nonbondedCutoff=1*u.nanometer)
+        integrator = mm.LangevinIntegrator(310, 0.05, 100 * mm.unit.femtosecond)
+
+        # Add forces
+        print('Adding forces...')
+        self.add_forcefield()
+        print('Forces added ;)\n')
+
+        # Minimize energy
+        print('Minimizing energy...')
+        platform = mm.Platform.getPlatformByName(self.platform)
+        self.simulation = Simulation(pdb.topology, self.system, integrator, platform)
+        self.simulation.reporters.append(StateDataReporter(stdout, (self.N_steps*sim_step)//100, step=True, totalEnergy=True, potentialEnergy=True, temperature=True))
+        self.simulation.reporters.append(DCDReporter(self.path+'/other/stochastic_LE.dcd', 5))
+        self.simulation.context.setPositions(pdb.positions)
+        current_platform = self.simulation.context.getPlatform()
+        print(f"Simulation will run on platform: {current_platform.getName()}")
+        self.simulation.minimizeEnergy(tolerance=self.tolerance)
+        print('Energy minimization done :D\n')
+
+        # Run molecular dynamics simulation
+        if run_MD:
+            print('Running molecular dynamics (wait for 100 steps)...')
+            start = time.time()
+            heats = list()
+            for i in range(1,self.N_steps):
+                self.change_loop(i)
+                self.simulation.step(sim_step)
+                if i>=self.burnin:
+                    self.state = self.simulation.context.getState(getPositions=True)
+                    if write_files: PDBxFile.writeFile(pdb.topology, self.state.getPositions(), open(self.path+f'/ensemble/MDLE_{i-self.burnin+1}.cif', 'w'))
+                    save_path = self.path+f'/heatmaps/heat_{i-self.burnin+1}.svg' if write_files else None
+                    heats.append(get_heatmap(self.state.getPositions(),save_path=save_path,save=write_files))
+            end = time.time()
+            elapsed = end - start
+
+            print(f'Everything is done! Simulation finished succesfully!\nMD finished in {elapsed/60:.2f} minutes.\n')
+
+            self.avg_heat = np.average(heats,axis=0)
+            self.std_heat = np.std(heats,axis=0)
+            np.save(self.path+f'/other/avg_heatmap.npy',self.avg_heat)
+            np.save(self.path+f'/other/std_heatmap.npy',self.std_heat)
+            if plots:
+                self.plot_heat(self.avg_heat,f'/plots/avg_heatmap.svg')
+                self.plot_heat(self.std_heat,f'/plots/std_heatmap.svg')
+            return self.avg_heat
+    
+    def change_loop(self,i):
+        force_idx = self.system.getNumForces()-1
+        self.system.removeForce(force_idx)
+        self.add_loops(i)
+        self.simulation.context.reinitialize(preserveState=True)
+        self.LE_force.updateParametersInContext(self.simulation.context)
+
+    def add_evforce(self):
+        'Leonard-Jones potential for excluded volume'
+        self.ev_force = mm.CustomNonbondedForce('epsilon*((sigma1+sigma2)/(r+r_small))^6')
+        self.ev_force.addGlobalParameter('epsilon', defaultValue=self.ev_ff_strength)
+        self.ev_force.addGlobalParameter('r_small', defaultValue=0.01)
+        self.ev_force.addPerParticleParameter('sigma')
+        for i in range(self.N_beads):
+            self.ev_force.addParticle([0.1])
+        self.system.addForce(self.ev_force)
+
+    def add_bonds(self):
+        'Harmonic bond borce between succesive beads'
+        self.bond_force = mm.HarmonicBondForce()
+        for i in range(self.N_beads - 1):
+            self.bond_force.addBond(i, i + 1, 0.1, 3e5)
+        self.system.addForce(self.bond_force)
+    
+    def add_stiffness(self):
+        'Harmonic angle force between successive beads so as to make chromatin rigid'
+        self.angle_force = mm.HarmonicAngleForce()
+        for i in range(self.N_beads - 2):
+            self.angle_force.addAngle(i, i + 1, i + 2, np.pi, self.angle_ff_strength)
+        self.system.addForce(self.angle_force)
+    
+    def add_loops(self,i=0):
+        'LE force that connects cohesin restraints'
+        self.LE_force = mm.HarmonicBondForce()
+        for nn in range(self.N_coh):
+            self.LE_force.addBond(self.M[nn,i], self.N[nn,i], self.le_distance, self.le_ff_strength)
+        self.system.addForce(self.LE_force)
+
+    def add_forcefield(self):
+        '''
+        Here is the definition of the forcefield.
+
+        There are the following energies:
+        - ev force: repelling LJ-like forcefield
+        - harmonic bond force: to connect adjacent beads.
+        - angle force: for polymer stiffness.
+        - loop forces: this is a list of force objects. Each object corresponds to a different cohesin. It is needed to define a force for each time step.
+        '''
+        self.add_evforce()
+        self.add_bonds()
+        self.add_stiffness()
+        self.add_loops()
+
+    def plot_heat(self,img,file_name):
+        figure(figsize=(10, 10))
+        plt.imshow(img,cmap="Reds",vmax=1)
+        plt.savefig(self.path+file_name,format='svg',dpi=500)
+        plt.close()
+
+def main():
+    # A potential example
+    M = np.load('/home/skorsak/Dropbox/LoopSage/files/region_[48100000,48700000]_chr3/Annealing_Nbeads500_ncoh50/Ms.npy')
+    N = np.load('/home/skorsak/Dropbox/LoopSage/files/region_[48100000,48700000]_chr3/Annealing_Nbeads500_ncoh50/Ns.npy')
+    md = MD_LE(4*M,4*N,2000,5,1)
+    md.run_pipeline(write_files=False,plots=True,sim_step=100)