Skip to content

Commit

Permalink
adapting the openmm_run skript
Browse files Browse the repository at this point in the history
  • Loading branch information
@JohannesKarwou
JohannesKarwou committed Nov 7, 2023
1 parent 9c38526 commit d461e10
Show file tree
Hide file tree
Showing 2 changed files with 93 additions and 201 deletions.
277 changes: 93 additions & 184 deletions transformato/bin/openmm_run.py
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,6 @@
"""
Generated by CHARMM-GUI (http://www.charmm-gui.org)
openmm_run.py
This program is OpenMM running scripts written in python.
Correspondance: [email protected] or [email protected]
Last update: March 29, 2017
File created in analogy to CHARMM-GUI (http://www.charmm-gui.org)
Last update: November, 2023
"""

from __future__ import print_function
Expand All @@ -17,173 +11,108 @@
from omm_readinputs import *
from omm_readparams import *
from omm_vfswitch import *
from omm_barostat import *
from omm_restraints import *
from omm_rewrap import *

import openmm.unit as unit
from openmm.unit import *
from openmm import *
from openmm.app import *

parser = argparse.ArgumentParser()
parser.add_argument("-i", dest="inpfile", help="Input parameter file", required=True)
parser.add_argument("-p", dest="psffile", help="Input CHARMM PSF file", required=True)
parser.add_argument("-c", dest="crdfile", help="Input CHARMM CRD file", required=True)
parser.add_argument(
"-t", dest="toppar", help="Input CHARMM-GUI toppar stream file", required=True
)
parser.add_argument(
"-b",
dest="sysinfo",
help="Input CHARMM-GUI sysinfo stream file (optional)",
default=None,
)
parser.add_argument(
"-icrst",
metavar="RSTFILE",
dest="icrst",
help="Input CHARMM RST file (optional)",
default=None,
)
parser.add_argument(
"-irst",
metavar="RSTFILE",
dest="irst",
help="Input restart file (optional)",
default=None,
)
parser.add_argument(
"-ichk",
metavar="CHKFILE",
dest="ichk",
help="Input checkpoint file (optional)",
default=None,
)
parser.add_argument(
"-opdb",
metavar="PDBFILE",
dest="opdb",
help="Output PDB file (optional)",
default=None,
)
parser.add_argument(
"-orst",
metavar="RSTFILE",
dest="orst",
help="Output restart file (optional)",
default=None,
)
parser.add_argument(
"-ochk",
metavar="CHKFILE",
dest="ochk",
help="Output checkpoint file (optional)",
default=None,
)
parser.add_argument(
"-odcd",
metavar="DCDFILE",
dest="odcd",
help="Output trajectory file (optional)",
default=None,
)
parser.add_argument(
"-rewrap",
dest="rewrap",
help="Re-wrap the coordinates in a molecular basis (optional)",
action="store_true",
default=False,
"-odcd", metavar="DCDFILE", dest="odcd", help="Output trajectory file (optional)"
)
args = parser.parse_args()

# Load parameters
print("Loading parameters")
inputs = read_inputs(args.inpfile)
params = read_params(args.toppar)
psf = CharmmPsfFile(args.psffile)
crd = read_crd(args.crdfile)
if args.sysinfo:
psf = read_box(psf, args.sysinfo)
env = "waterbox"
inputs = read_inputs(f"lig_in_{env}.inp")

if os.path.isfile(f"lig_in_{env}.parm7"):
top = AmberPrmtopFile(f"lig_in_{env}.parm7")
crd = AmberInpcrdFile(f"lig_in_{env}.rst7")
fftype = "amber"
else:
psf = gen_box(psf, crd)
params = read_params("toppar.str")
top = CharmmPsfFile(f"lig_in_{env}.psf")
crd = read_crd(f"lig_in_{env}.crd")
top = gen_box(top, crd)

# Build system
nboptions = dict(
nonbondedMethod=inputs.coulomb,
nonbondedCutoff=inputs.r_off * nanometers,
constraints=inputs.cons,
ewaldErrorTolerance=inputs.ewald_Tol,
)
if env == "waterbox":
nboptions = dict(
nonbondedMethod=inputs.coulomb,
nonbondedCutoff=inputs.r_off * unit.nanometers,
constraints=inputs.cons,
ewaldErrorTolerance=inputs.ewald_Tol,
)
elif env == "vacuum":
nboptions = dict(
nonbondedMethod=NoCutoff,
constraints=inputs.cons,
)

if inputs.vdw == "Switch" and env != "vacuum":
nboptions["switchDistance"] = inputs.r_on * unit.nanometers
if fftype == "amber":
system = top.createSystem(**nboptions)
else:
system = top.createSystem(params, **nboptions)

if inputs.vdw == "Switch":
nboptions["switchDistance"] = inputs.r_on * nanometers
system = psf.createSystem(params, **nboptions)
if inputs.vdw == "Force-switch":
system = vfswitch(system, psf, inputs)

system = barostat(system, inputs)
if inputs.rest == "yes":
system = restraints(system, crd, inputs)
if inputs.vdw == "Force-switch" and fftype != "amber" and env != "vacuum":
system = vfswitch(system, top, inputs)

barostat = MonteCarloBarostat(inputs.p_ref * bar, inputs.temp * kelvin)
system.addForce(barostat)

integrator = LangevinIntegrator(
inputs.temp * kelvin, inputs.fric_coeff / picosecond, inputs.dt * picoseconds
inputs.temp * kelvin, 1 / unit.picosecond, inputs.dt * unit.picoseconds
)

# Set platform
platform = Platform.getPlatformByName("CUDA")
prop = dict()
prop = dict(CudaPrecision="mixed")

# Check if restraints.yaml exists - if it does, system uses restraints

pdbpath = args.inpfile.replace(".inp", ".pdb")
# pdbpath = args.inpfile.replace(".inp", ".pdb")

if os.path.exists("./restraints.yaml") and "complex" in pdbpath:
import transformato.restraints as tfrs
import yaml
# if os.path.exists("./restraints.yaml") and "complex" in pdbpath:
# import transformato.restraints as tfrs
# import yaml

print("Found restraints.yaml - applying restraints")
# Load tiny restraints config
with open("./restraints.yaml", "r") as stream:
try:
configuration = yaml.safe_load(stream)
except yaml.YAMLError as exc:
print(exc)
# print("Found restraints.yaml - applying restraints")
# # Load tiny restraints config
# with open("./restraints.yaml", "r") as stream:
# try:
# configuration = yaml.safe_load(stream)
# except yaml.YAMLError as exc:
# print(exc)

cc_names = configuration["system"]["structure"]["ccs"]
# cc_names = configuration["system"]["structure"]["ccs"]

# Add forces via transformato.restraints
# # Add forces via transformato.restraints

if not os.path.exists(pdbpath):
raise FileNotFoundError(
f"Couldnt find {pdbpath} necessary for Restraint Analysis"
)
# if not os.path.exists(pdbpath):
# raise FileNotFoundError(
# f"Couldnt find {pdbpath} necessary for Restraint Analysis"
# )

restraintList = tfrs.create_restraints_from_config(configuration, pdbpath)
# restraintList = tfrs.create_restraints_from_config(configuration, pdbpath)

for restraint in restraintList:
restraint.createForce(cc_names)
restraint.applyForce(system)
# for restraint in restraintList:
# restraint.createForce(cc_names)
# restraint.applyForce(system)


# Build simulation context
simulation = Simulation(psf.topology, system, integrator, platform, prop)
simulation = Simulation(top.topology, system, integrator, platform, prop)
simulation.context.setPositions(crd.positions)
if args.icrst:
charmm_rst = read_charmm_rst(args.icrst)
simulation.context.setPositions(charmm_rst.positions)
simulation.context.setVelocities(charmm_rst.velocities)
simulation.context.setPeriodicBoxVectors(
charmm_rst.box[0], charmm_rst.box[1], charmm_rst.box[2]
)
if args.irst:
with open(args.irst, "r") as f:
if os.path.isfile(f"lig_in_{env}.irst"):
with open(f"lig_in_{env}.irst", "r") as f:
simulation.context.setState(XmlSerializer.deserialize(f.read()))
if args.ichk:
with open(args.ichk, "rb") as f:
simulation.context.loadCheckpoint(f.read())

# Re-wrap
if args.rewrap:
simulation = rewrap(simulation)

# Calculate initial system energy
print("\nInitial system energy")
Expand All @@ -192,9 +121,7 @@
# Energy minimization
if inputs.mini_nstep > 0:
print("\nEnergy minimization: %s steps" % inputs.mini_nstep)
simulation.minimizeEnergy(
tolerance=inputs.mini_Tol * kilojoule / mole, maxIterations=inputs.mini_nstep
)
simulation.minimizeEnergy(maxIterations=inputs.mini_nstep)
print(simulation.context.getState(getEnergy=True).getPotentialEnergy())

# Generate initial velocities
Expand All @@ -206,48 +133,30 @@
simulation.context.setVelocitiesToTemperature(inputs.gen_temp)

# Production
if inputs.nstep > 0:
print("\nMD run: %s steps" % inputs.nstep)
if inputs.nstdcd > 0:
if not args.odcd:
args.odcd = "output.dcd"
simulation.reporters.append(DCDReporter(args.odcd, inputs.nstdcd))
simulation.reporters.append(
StateDataReporter(
sys.stdout,
inputs.nstout,
step=True,
time=True,
potentialEnergy=True,
temperature=True,
progress=True,
remainingTime=True,
speed=True,
totalSteps=inputs.nstep,
separator="\t",
)
print("\nMD run: %s steps" % inputs.nstep)
simulation.reporters.append(DCDReporter(args.odcd, inputs.nstdcd))
simulation.reporters.append(
StateDataReporter(
sys.stdout,
inputs.nstout,
step=True,
time=True,
potentialEnergy=True,
temperature=True,
progress=True,
remainingTime=True,
speed=True,
totalSteps=inputs.nstep,
separator="\t",
)
# Simulated annealing?
if inputs.annealing == "yes":
interval = inputs.interval
temp = inputs.temp_init
for i in range(inputs.nstep):
integrator.setTemperature(temp * kelvin)
simulation.step(1)
temp += interval
else:
simulation.step(inputs.nstep)

# Write restart file
if not (args.orst or args.ochk):
args.orst = "output.rst"
if args.orst:
state = simulation.context.getState(getPositions=True, getVelocities=True)
with open(args.orst, "w") as f:
f.write(XmlSerializer.serialize(state))
if args.ochk:
with open(args.ochk, "wb") as f:
f.write(simulation.context.createCheckpoint())
if args.opdb:
crd = simulation.context.getState(getPositions=True).getPositions()
PDBFile.writeFile(psf.topology, crd, open(args.opdb, "w"))
)

simulation.step(inputs.nstep)

# needed for later analysis
file_name = f"lig_in_{env}"
print(file_name)
with open(file_name + "_integrator.xml", "w") as outfile:
outfile.write(XmlSerializer.serialize(integrator))
with open(file_name + "_system.xml", "w") as outfile:
outfile.write(XmlSerializer.serialize(system))
17 changes: 0 additions & 17 deletions transformato/state.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -204,21 +204,6 @@ def write_state(

self.current_step += 1

def _add_serializer(self, file):
# adding serializer functions
with open(file, "a") as f:
f.write(
"""
# mw: adding xml serializer to the simulation script
file_name = str(args.psffile).replace('.psf', '')
print(file_name)
with open(file_name + '_integrator.xml','w') as outfile:
outfile.write(XmlSerializer.serialize(integrator))
with open(file_name + '_system.xml','w') as outfile:
outfile.write(XmlSerializer.serialize(system))
"""
)

def _get_simulations_parameters(self):
prms = {}
for key in self.configuration["simulation"]["parameters"]:
Expand Down Expand Up @@ -531,7 +516,6 @@ def _copy_omm_files(self, intermediate_state_file_path: str):
shutil.copyfile(omm_simulation_script_source, omm_simulation_script_target)
# add serialization
self._check_hmr(omm_simulation_script_target)
self._add_serializer(omm_simulation_script_target)
self._change_platform(omm_simulation_script_target)
check_switching_function(self.vdw_switch)

Expand All @@ -548,7 +532,6 @@ def _copy_omm_files(self, intermediate_state_file_path: str):
)
shutil.copyfile(omm_simulation_script_source, omm_simulation_script_target)
self._check_hmr(omm_simulation_script_target)
self._add_serializer(omm_simulation_script_target)
self._change_platform(omm_simulation_script_target)

def _check_hmr(self, file):
Expand Down

0 comments on commit d461e10

Please sign in to comment.