simulate.py 8.55 KB
Newer Older
1
import os
2
import tempfile
3
from .arbdmodel.coords import readArbdCoords, readAvgArbdCoords
4
5
import shutil
from . import get_resource_path
6

7
8
import numpy as np

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

11
12
def _remove_bonds_beyond(model, cutoff):
    bonds = model.get_bonds()
13
14
    for seg in model.segments:
        bonds.extend(seg.get_bonds())
15
16
17
18
19
20
21
22
23
24
    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:
                    ...
25
26
27
28

def minimize_and_simulate_oxdna( model, 
                                 output_name,
                                 directory = None,
29
                                 num_min_steps = 5e2,
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
                                 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)

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

49
50
51
52
53
54
        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,
55
                                        **min_args)
56

57
58
59
60
61
62
63
        top,conf = 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)
64
65
    finally:
        os.chdir(d_orig)
66
67
68
69
70

    try:
        return top,conf
    except:
        pass
71

72
73
def multiresolution_simulation( model, output_name, 
                                job_id=None,
74
75
                                gpu = 0,
                                directory=None,
76
                                minimization_steps = 0,
77
78
                                coarse_steps = 5e7,
                                fine_steps = 5e7,
79
                                minimization_output_period = 1e5,
80
81
                                coarse_output_period = 1e5,
                                fine_output_period = 1e5,
82
                                coarse_local_twist = False,
83
                                fix_linking_number = False,
84
                                bond_cutoff = 0,
85
                                backbone_scale=1.0,
86
87
                                oxdna_steps = None,
                                oxdna_output_period = None,
cmaffeo2's avatar
cmaffeo2 committed
88
                                run_enrg_md = False
89
                            ):
90

91
    ## Round steps up to nearest multiple of output_period, plus 1
92
93
94
    if minimization_steps > 0:
        raise NotImplementedError
        minimization_steps = ((minimization_steps+minimization_output_period-1)//minimization_output_period)*minimization_output_period+1
95
    coarse_steps = ((coarse_steps+coarse_output_period-1)//coarse_output_period)*coarse_output_period+1
cmaffeo2's avatar
cmaffeo2 committed
96
97
98
    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 
99

100
101
102
103
    ret = None
    d_orig = os.getcwd()
    try:
        if directory is None:
104
105
            tmp = job_id if job_id is not None else output_name
            directory = tempfile.mkdtemp(prefix='origami-%s-' % tmp, dir='/dev/shm/')
106
107
108
109
110
        elif not os.path.exists(directory):
            os.makedirs(directory)
        os.chdir(directory)

        output_directory = "output"
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

        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,
141
                                  escapable_twist = fix_linking_number )
142
143
144
        
        # TODO: add the output directory in to the readArbdCoords functions or make it an attribute of the model object 

145
146
147
        """ Prepare coarse model for minimization and coarse simulation """
        model.clear_beads()
        model.generate_bead_model( **coarse_model_args )
148
149
150
        
        if bond_cutoff > 0:
            _remove_bonds_beyond(model, bond_cutoff)
151

152
153
154
155
156
157
        """ 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 )
158

159
        """ Coarse simulation """
160
        simargs = dict(timestep=150e-6 if coarse_local_twist else 200e-6, output_period=coarse_output_period, gpu=gpu)
161
162
        if bond_cutoff > 0 and minimization_steps <= 0:
            run_step(num_steps=0.25*coarse_steps, simargs=simargs)
163
164
            model.clear_beads()
            model.generate_bead_model( **coarse_model_args )
165
            run_step(num_steps=0.75*coarse_steps, simargs=simargs)
166
        else:
167
            run_step(num_steps=coarse_steps, simargs=simargs)
168
169

        """ Fine simulation """ 
170
        simargs = dict(timestep=40e-6, output_period=fine_output_period, gpu=gpu)
171
172
173
174
        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))
175
176

        """ Freeze twist """
177
178
179
        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))
180

181

182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
        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:
                ...
200

cmaffeo2's avatar
cmaffeo2 committed
201
202
203
            model.atomic_simulate( output_name = output_prefix, dry_run = not run_enrg_md )


204
205
206
207
208
209
210
211
212

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