simulate.py 9.08 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
                                 num_steps = 1e7,
                                 output_period = None,
32
                                 gpu = 0,
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
                                 **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)

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

50
51
52
53
54
55
        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,
56
                                        gpu = gpu,
57
                                        **min_args)
58

59
60
61
62
63
64
        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,
65
                                        gpu = gpu,
66
                                        **oxdna_args)
67
68
    finally:
        os.chdir(d_orig)
69
70
71
72
73

    try:
        return top,conf
    except:
        pass
74

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

97
    ## Round steps up to nearest multiple of output_period, plus 1
98
99
100
    if minimization_steps > 0:
        raise NotImplementedError
        minimization_steps = ((minimization_steps+minimization_output_period-1)//minimization_output_period)*minimization_output_period+1
101
    coarse_steps = ((coarse_steps+coarse_output_period-1)//coarse_output_period)*coarse_output_period+1
cmaffeo2's avatar
cmaffeo2 committed
102
103
104
    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 
105

106
107
108
109
    ret = None
    d_orig = os.getcwd()
    try:
        if directory is None:
110
111
            tmp = job_id if job_id is not None else output_name
            directory = tempfile.mkdtemp(prefix='origami-%s-' % tmp, dir='/dev/shm/')
112
113
114
115
116
        elif not os.path.exists(directory):
            os.makedirs(directory)
        os.chdir(directory)

        output_directory = "output"
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

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

151
152
153
        """ Prepare coarse model for minimization and coarse simulation """
        model.clear_beads()
        model.generate_bead_model( **coarse_model_args )
154
155
156
        
        if bond_cutoff > 0:
            _remove_bonds_beyond(model, bond_cutoff)
157

158
159
        """ Minimization """
        if minimization_steps > 0:
cmaffeo2's avatar
cmaffeo2 committed
160
161
            simargs = dict(timestep=200e-6, output_period=minimization_output_period, write_pqr=write_pqr,
                           gpu=gpu, arbd=arbd)
162
163
164
            run_step(num_steps=minimization_steps, simargs=simargs)
            model.clear_beads()
            model.generate_bead_model( **coarse_model_args )
165

166
        """ Coarse simulation """
cmaffeo2's avatar
cmaffeo2 committed
167
168
        simargs = dict(timestep=100e-6 if coarse_local_twist else 200e-6, output_period=coarse_output_period,
                       write_pqr=write_pqr, gpu=gpu, arbd=arbd)
169
170
        if bond_cutoff > 0 and minimization_steps <= 0:
            run_step(num_steps=0.25*coarse_steps, simargs=simargs)
171
172
            model.clear_beads()
            model.generate_bead_model( **coarse_model_args )
173
            run_step(num_steps=0.75*coarse_steps, simargs=simargs)
174
        else:
175
            run_step(num_steps=coarse_steps, simargs=simargs)
176
177

        """ Fine simulation """ 
cmaffeo2's avatar
cmaffeo2 committed
178
        simargs = dict(timestep=40e-6, output_period=fine_output_period, write_pqr=write_pqr, gpu=gpu, arbd=arbd)
179
180
181
182
        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))
183
184

        """ Freeze twist """
185
186
187
        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))
188

189

190
191
192
193
194
        if oxdna_steps is not None:
            minimize_and_simulate_oxdna( model, 
                                         output_name = output_name,
                                         num_min_steps = 1e3,
                                         num_steps = oxdna_steps,
195
196
                                         output_period = oxdna_output_period,
                                         gpu = gpu)
197
198
199
200
201
202
203
204
205
206
207
        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:
208
                pass
209

cmaffeo2's avatar
cmaffeo2 committed
210
            model.atomic_simulate( output_name = output_prefix, num_steps=enrg_md_steps, write_pqr=write_pqr, dry_run = not run_enrg_md )
cmaffeo2's avatar
cmaffeo2 committed
211
212


213
214
215
216
217
218
219
220
221

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