RAU_MasterStudy/thesis/old/openmm_help/openmm_run.py
2025-06-04 20:04:29 +03:00

258 lines
8.2 KiB
Python

"""
Generated by CHARMM-GUI (http://www.charmm-gui.org)
openmm_run.py
This program is OpenMM running scripts written in python.
Correspondance: jul316@lehigh.edu or wonpil@lehigh.edu
Last update: February 5, 2025
"""
from __future__ import print_function
import argparse
import sys
import os
from omm_readinputs import *
from omm_readparams import *
from omm_vfswitch import *
from omm_barostat import *
from omm_restraints import *
from omm_rewrap import *
from openmm import *
from openmm.app import *
from openmm.unit import *
parser = argparse.ArgumentParser()
parser.add_argument(
"--platform", nargs=1, help="OpenMM platform (default: CUDA or OpenCL)"
)
parser.add_argument("-i", dest="inpfile", help="Input parameter file", required=True)
parser.add_argument("-p", dest="topfile", help="Input topology file", required=True)
parser.add_argument("-c", dest="crdfile", help="Input coordinate file", required=True)
parser.add_argument(
"-t", dest="toppar", help="Input CHARMM-GUI toppar stream file (optional)"
)
parser.add_argument(
"-b", dest="sysinfo", help="Input CHARMM-GUI sysinfo stream file (optional)"
)
parser.add_argument(
"-ff",
dest="fftype",
help="Input force field type (default: CHARMM)",
default="CHARMM",
)
parser.add_argument(
"-icrst", metavar="RSTFILE", dest="icrst", help="Input CHARMM RST file (optional)"
)
parser.add_argument(
"-irst", metavar="RSTFILE", dest="irst", help="Input restart file (optional)"
)
parser.add_argument(
"-ichk", metavar="CHKFILE", dest="ichk", help="Input checkpoint file (optional)"
)
parser.add_argument(
"-opdb", metavar="PDBFILE", dest="opdb", help="Output PDB file (optional)"
)
parser.add_argument(
"-orst", metavar="RSTFILE", dest="orst", help="Output restart file (optional)"
)
parser.add_argument(
"-ochk", metavar="CHKFILE", dest="ochk", help="Output checkpoint file (optional)"
)
parser.add_argument(
"-odcd", metavar="DCDFILE", dest="odcd", help="Output trajectory file (optional)"
)
parser.add_argument(
"-rewrap",
dest="rewrap",
help="Re-wrap the coordinates in a molecular basis (optional)",
action="store_true",
default=False,
)
args = parser.parse_args()
# Load parameters
print("Loading parameters")
inputs = read_inputs(args.inpfile)
top = read_top(args.topfile, args.fftype.upper())
crd = read_crd(args.crdfile, args.fftype.upper())
if args.fftype.upper() == "CHARMM":
params = read_params(args.toppar)
top = read_box(top, args.sysinfo) if args.sysinfo else gen_box(top, crd)
# Build system
nboptions = dict(
nonbondedMethod=inputs.coulomb,
nonbondedCutoff=inputs.r_off * nanometers,
constraints=inputs.cons,
ewaldErrorTolerance=inputs.ewald_Tol,
)
if inputs.vdw == "Switch":
nboptions["switchDistance"] = inputs.r_on * nanometers
if inputs.vdw == "LJPME":
nboptions["nonbondedMethod"] = LJPME
if inputs.implicitSolvent:
nboptions["implicitSolvent"] = inputs.implicitSolvent
nboptions["implicitSolventSaltConc"] = inputs.implicit_salt * (moles / liter)
nboptions["temperature"] = inputs.temp * kelvin
nboptions["soluteDielectric"] = inputs.solut_diele
nboptions["solventDielectric"] = inputs.solve_diele
nboptions["gbsaModel"] = inputs.gbsamodel
if args.fftype.upper() == "CHARMM":
system = top.createSystem(params, **nboptions)
elif args.fftype.upper() == "AMBER":
system = top.createSystem(**nboptions)
if inputs.vdw == "Force-switch":
system = vfswitch(system, top, inputs)
if inputs.lj_lrc == "yes":
for force in system.getForces():
if isinstance(force, NonbondedForce):
force.setUseDispersionCorrection(True)
if (
isinstance(force, CustomNonbondedForce)
and force.getNumTabulatedFunctions() != 1
):
force.setUseLongRangeCorrection(True)
if inputs.e14scale != 1.0:
for force in system.getForces():
if isinstance(force, NonbondedForce):
nonbonded = force
break
for i in range(nonbonded.getNumExceptions()):
atom1, atom2, chg, sig, eps = nonbonded.getExceptionParameters(i)
nonbonded.setExceptionParameters(
i, atom1, atom2, chg * inputs.e14scale, sig, eps
)
if inputs.pcouple == "yes":
system = barostat(system, inputs)
if inputs.rest == "yes":
system = restraints(system, crd, inputs)
integrator = LangevinIntegrator(
inputs.temp * kelvin, inputs.fric_coeff / picosecond, inputs.dt * picoseconds
)
# Set platform
DEFAULT_PLATFORMS = "CUDA", "OpenCL", "CPU"
enabled_platforms = [
Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())
]
if args.platform:
if not args.platform[0] in enabled_platforms:
print(
"Unable to find OpenMM platform '{}'; exiting".format(args.platform[0]),
file=sys.stderr,
)
sys.exit(1)
platform = Platform.getPlatformByName(args.platform[0])
else:
for platform in DEFAULT_PLATFORMS:
if platform in enabled_platforms:
platform = Platform.getPlatformByName(platform)
break
if isinstance(platform, str):
print(
"Unable to find any OpenMM platform; exiting".format(args.platform[0]),
file=sys.stderr,
)
sys.exit(1)
print("Using platform:", platform.getName())
prop = dict(CudaPrecision="single") if platform.getName() == "CUDA" else dict()
# Build simulation context
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:
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")
print(simulation.context.getState(getEnergy=True).getPotentialEnergy())
# Energy minimization
if inputs.mini_nstep > 0:
print("\nEnergy minimization: %s steps" % inputs.mini_nstep)
simulation.minimizeEnergy(
tolerance=inputs.mini_Tol * kilojoule / mole / nanometers,
maxIterations=inputs.mini_nstep,
)
print(simulation.context.getState(getEnergy=True).getPotentialEnergy())
# Generate initial velocities
if inputs.gen_vel == "yes":
print("\nGenerate initial velocities")
if inputs.gen_seed:
simulation.context.setVelocitiesToTemperature(inputs.gen_temp, inputs.gen_seed)
else:
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",
)
)
# 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(top.topology, crd, open(args.opdb, "w"))