def _init_system(self, mol, forces, nreplicas, T):
system = System(mol.numAtoms, nreplicas, self.dtype, self.device)
system.set_positions(mol.coords)
system.set_box(mol.box)
system.set_velocities(
maxwell_boltzmann(forces.par.masses, T=T, replicas=nreplicas))
return system
def setup(args):
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
device = torch.device(args.device)
if args.topology is not None:
mol = Molecule(args.topology)
elif args.structure is not None:
mol = Molecule(args.structure)
mol.box = np.array(
[mol.crystalinfo['a'], mol.crystalinfo['b'],
mol.crystalinfo['c']]).reshape(3, 1).astype(np.float32)
if args.coordinates is not None:
mol.read(args.coordinates)
if args.extended_system is not None:
mol.read(args.extended_system)
precision = precisionmap[args.precision]
print("Force terms: ", args.forceterms)
ff = ForceField.create(mol, args.forcefield)
parameters = Parameters(ff,
mol,
args.forceterms,
precision=precision,
device=device)
external = None
if args.external is not None:
externalmodule = importlib.import_module(args.external["module"])
embeddings = torch.tensor(args.external["embeddings"]).repeat(
args.replicas, 1)
external = externalmodule.External(args.external["file"], embeddings,
device)
system = System(mol.numAtoms, args.replicas, precision, device)
system.set_positions(mol.coords)
system.set_box(mol.box)
system.set_velocities(
maxwell_boltzmann(parameters.masses, args.temperature, args.replicas))
forces = Forces(parameters,
terms=args.forceterms,
external=external,
cutoff=args.cutoff,
rfa=args.rfa,
switch_dist=args.switch_dist)
return mol, system, forces
def getTorchMDSystem(mol, device, precision):
from torchmd.systems import System
system = System(mol.numAtoms, 1, precision, device)
system.set_positions(mol.coords)
system.set_box(mol.box)
return system
def test_replicas(self):
from moleculekit.molecule import Molecule
from torchmd.systems import System
import os
n_replicas = 2
testdir = os.path.join("test-data", "prod_alanine_dipeptide_amber")
mol = Molecule(os.path.join(testdir, "structure.prmtop"))
mol.read(os.path.join(testdir, "input.coor"))
struct = parmed.load_file(os.path.join(testdir, "structure.prmtop"))
prm = parmed.amber.AmberParameterSet().from_structure(struct)
terms = [
"bonds",
"angles",
"dihedrals",
"impropers",
"1-4",
"electrostatics",
"lj",
]
cutoff = 9
switch_dist = 7.5
rfa = True
precision = torch.double
device = "cpu"
ff = ForceField.create(mol, prm)
parameters = Parameters(ff, mol, precision=precision, device=device)
system = System(mol.numAtoms, n_replicas, precision, device)
system.set_positions(mol.coords)
system.set_box(mol.box)
forces = Forces(
parameters,
terms=terms,
cutoff=cutoff,
switch_dist=switch_dist,
rfa=rfa,
)
Epot = forces.compute(
system.pos.detach().requires_grad_(True),
system.box,
system.forces,
returnDetails=False,
explicit_forces=False,
)
assert len(Epot) == 2
Epot = [ee.detach().numpy() for ee in Epot]
assert np.abs(Epot[0] + 1722.3567) < 1e-4 and np.abs(Epot[1] +
1722.3567) < 1e-4
ommforceterms = ["lennardjones", "bond", "angle", "electrostatic"]
mol_org = mol.copy()
datastore = {}
for forceterm, ommforceterm in zip(forceterms, ommforceterms):
forceterm = [
forceterm,
]
ommforceterm = [
ommforceterm,
]
mol = mol_org.copy()
print("Force terms: ", forceterm)
system = System(atom_pos, atom_vel, box_full, atom_forces, precision,
device)
forces = Forces(
parameters,
forceterm,
external=None,
cutoff=7.3,
rfa=True,
)
Epot = forces.compute(system.pos,
system.box,
system.forces,
returnDetails=True)[0]
myforces = system.forces.cpu().numpy()[0]
prm = parmed.charmm.CharmmParameterSet("./tests/water/parameters.prm")
struct = parmed.charmm.CharmmPsfFile("./tests/water/structure.psf")
def make_deltaforces(
coords_npz,
forces_npz,
delta_forces_npz,
forcefield,
psf,
exclusions=("bonds"),
device="cpu",
forceterms=["Bonds", "Angles", "RepulsionCG"],
):
device = torch.device(device)
precision = torch.double
replicas = 1
mol = Molecule(psf)
atom_types = mol.atomtype
natoms = mol.numAtoms
coords = np.load(coords_npz)
### test if any bonds are longer than 10A
print("Check for broken coords.")
broken_frames = []
for bond in tqdm(mol.bonds):
crds = coords[:, bond, :]
crds_dif = crds[:, 0, :] - crds[:, 1, :]
dists = np.linalg.norm(crds_dif, axis=1)
broken = dists > 10.0
broken_frames.append(broken)
broken_frames = np.stack(broken_frames)
broken_frames = broken_frames.any(axis=0)
if broken_frames.any():
print("Removing broken coords with distances larger than 10A.")
coords_good = coords[~broken_frames, :, :] # remove broken frames
np.save(coords_npz.replace(".", "_fix."), coords_good)
broken_coords = coords[broken_frames, :, :]
np.save(coords_npz.replace(".", "_broken."), broken_coords)
coords = coords_good
else:
print("No broken frames")
all_forces = np.load(forces_npz)[~broken_frames, :, :]
coords = torch.tensor(coords, dtype=precision).to(device)
atom_vel = torch.zeros(replicas, natoms, 3)
atom_forces = torch.zeros(natoms, 3, replicas)
atom_pos = torch.zeros(natoms, 3, replicas)
box_full = torch.zeros(3, replicas)
ff = ForceField.create(mol, forcefield)
parameters = Parameters(ff,
mol,
forceterms,
precision=precision,
device=device)
system = System(natoms, replicas, precision, device)
system.set_positions(atom_pos)
system.set_box(box_full)
system.set_velocities(atom_vel)
forces = Forces(parameters, terms=forceterms, exclusions=exclusions)
print("Producing delta forces")
prior_forces = []
for co in tqdm(coords):
Epot = forces.compute(co.reshape([1, natoms, 3]), system.box,
system.forces)
fr = (system.forces.detach().cpu().numpy().astype(np.float32).reshape(
[natoms, 3]))
prior_forces.append(fr)
prior_forces = np.array(prior_forces)
delta_forces = all_forces - prior_forces
np.save(delta_forces_npz, delta_forces)