Skip to content
Snippets Groups Projects
simulate.py 8.43 KiB
import os
import tempfile
from .arbdmodel.coords import readArbdCoords, readAvgArbdCoords
import shutil
from . import get_resource_path

import numpy as np

## TODO: implement replicas, initial conditions specified through some restart, and a custom simulation schedule

def _remove_bonds_beyond(model, cutoff):
    bonds = model.get_bonds()
    for seg in model.segments:
        bonds.extend(seg.get_bonds())
    new_bonds = []
    for bond in bonds:
        n1,n2,b,ex = bond
        r2 = np.sum( (n1.position - n2.position)**2 )
        if r2 > cutoff**2:
            for l in [model] + model.children:
                try:
                    l.bonds.remove(bond)
                except:
                    ...

def minimize_and_simulate_oxdna( model, 
                                 output_name,
                                 directory = None,
                                 num_min_steps = 5e2,
                                 num_steps = 1e7,
                                 output_period = None,
                                 **oxdna_args
                             ):

    d_orig = os.getcwd()
    try:
        if directory is not None:
            if not os.path.exists(directory):
                os.makedirs(directory)
            os.chdir(directory)

        model.generate_oxdna_model()

        top = "{}-oxdna.top".format(output_name)
        model._write_oxdna_topology(top)

        min_args = {k:v for k,v in oxdna_args.items() if k != "sim_type"}

        top,conf = model.simulate_oxdna(output_name = "{}-oxdna-min".format(output_name),
                                        topology = top,
                                        sim_type = "min",
                                        num_steps = num_min_steps,
                                        print_conf_interval = 100,
                                        print_energy_every = 10,
                                        **min_args)

        model.simulate_oxdna(output_name = "{}-oxdna".format(output_name),
                             topology = top,
                             configuration = conf,
                             num_steps = num_steps,
                             print_conf_interval = output_period,
                             print_energy_every = output_period,
                             **oxdna_args)
    finally:
        os.chdir(d_orig)
        

def multiresolution_simulation( model, output_name, 
                                job_id=None,
                                gpu = 0,
                                directory=None,
                                minimization_steps = 0,
                                coarse_steps = 5e7,
                                fine_steps = 5e7,
                                minimization_output_period = 1e5,
                                coarse_output_period = 1e5,
                                fine_output_period = 1e5,
                                coarse_local_twist = False,
                                fix_linking_number = False,
                                bond_cutoff = 0,
                                backbone_scale=1.0,
                                oxdna_steps = None,
                                oxdna_output_period = None,
                                run_enrg_md = False
                            ):

    ## Round steps up to nearest multiple of output_period, plus 1
    if minimization_steps > 0:
        raise NotImplementedError
        minimization_steps = ((minimization_steps+minimization_output_period-1)//minimization_output_period)*minimization_output_period+1
    coarse_steps = ((coarse_steps+coarse_output_period-1)//coarse_output_period)*coarse_output_period+1
    fine_steps = ((fine_steps+fine_output_period-1)//fine_output_period)
    if fine_steps == 1: fine_steps += 1
    fine_steps = fine_steps*fine_output_period+1 

    ret = None
    d_orig = os.getcwd()
    try:
        if directory is None:
            tmp = job_id if job_id is not None else output_name
            directory = tempfile.mkdtemp(prefix='origami-%s-' % tmp, dir='/dev/shm/')
        elif not os.path.exists(directory):
            os.makedirs(directory)
        os.chdir(directory)

        output_directory = "output"

        simulation_step = 0

        def run_step(num_steps, simargs, do_coordinate_average=False):
            """ Function runs one step of the simulation protocol process """
            nonlocal simulation_step

            output_prefix = "{}-{}".format(output_name, simulation_step)
            full_output_prefix = "{}/{}".format(output_directory,output_prefix)

            model.simulate( output_name = output_prefix, num_steps = num_steps, **simargs )
            
            if do_coordinate_average:
                try:
                    coordinates = readAvgArbdCoords('%s.psf' % output_prefix,'%s.pdb' % output_prefix, '%s.dcd' % full_output_prefix, rmsdThreshold=1)
                except:
                    coordinates = readAvgArbdCoords('%s.psf' % output_prefix,'%s.pdb' % output_prefix, '%s.0.dcd' % full_output_prefix, rmsdThreshold=1)

            else:
                try:
                    coordinates = readArbdCoords('%s.restart' % full_output_prefix)
                except:
                    coordinates = readArbdCoords('%s.0.restart' % full_output_prefix)
                    
            model.update_splines(coordinates)
            simulation_step += 1

        coarse_model_args = dict( max_basepairs_per_bead = 5,
                                  max_nucleotides_per_bead = 5,
                                  local_twist = coarse_local_twist,
                                  escapable_twist = fix_linking_number )
        
        # TODO: add the output directory in to the readArbdCoords functions or make it an attribute of the model object 

        """ Prepare coarse model for minimization and coarse simulation """
        model.clear_beads()
        model.generate_bead_model( **coarse_model_args )
        
        if bond_cutoff > 0:
            _remove_bonds_beyond(model, bond_cutoff)

        """ Minimization """
        if minimization_steps > 0:
            simargs = dict(timestep=200e-6, output_period=minimization_output_period, gpu=gpu)
            run_step(num_steps=minimization_steps, simargs=simargs)
            model.clear_beads()
            model.generate_bead_model( **coarse_model_args )

        """ Coarse simulation """
        simargs = dict(timestep=150e-6 if coarse_local_twist else 200e-6, output_period=coarse_output_period, gpu=gpu)
        if bond_cutoff > 0 and minimization_steps <= 0:
            run_step(num_steps=0.25*coarse_steps, simargs=simargs)
            model.clear_beads()
            model.generate_bead_model( **coarse_model_args )
            run_step(num_steps=0.75*coarse_steps, simargs=simargs)
        else:
            run_step(num_steps=coarse_steps, simargs=simargs)

        """ Fine simulation """ 
        simargs = dict(timestep=40e-6, output_period=fine_output_period, gpu=gpu)
        if not fix_linking_number:
            model.clear_beads()
            model.generate_bead_model( 1, 1, local_twist=True, escapable_twist=True )
            run_step(num_steps=fine_steps, simargs=simargs, do_coordinate_average=(fine_steps > 2*fine_output_period+1))

        """ Freeze twist """
        model.clear_beads()
        model.generate_bead_model( 1, 1, local_twist=True, escapable_twist=False )
        run_step(num_steps=fine_steps, simargs=simargs, do_coordinate_average=(fine_steps > 2*fine_output_period+1))


        if oxdna_steps is not None:
            minimize_and_simulate_oxdna( model, 
                                         output_name = output_name,
                                         num_min_steps = 1e3,
                                         num_steps = oxdna_steps,
                                         output_period = oxdna_output_period )
        else:
            """ Atomic simulation """
            model.generate_atomic_model(scale=backbone_scale)

            output_prefix = "{}-{}".format(output_name, simulation_step)
            # full_output_prefix = "%s/%s" % (output_directory,output_prefix)

            model.write_atomic_ENM( output_prefix )
            try:
                shutil.copytree( get_resource_path("charmm36.nbfix"), "charmm36.nbfix" )
            except FileExistsError:
                ...

            model.atomic_simulate( output_name = output_prefix, dry_run = not run_enrg_md )



        ret = directory
        
    except:
        raise
    finally:
        os.chdir(d_orig)
        
    return ret