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
box_full = torch.zeros((replicas, 3, 3)).type(precision)
for r in range(box.shape[0]):
box_full[r][torch.eye(3).bool()] = torch.tensor(box[r]).type(precision)
natoms = len(atom_types)
bonds = mol.bonds.astype(int).copy()
angles = mol.angles.astype(int).copy()
print("Force terms: ", args.forceterms)
ff = ForceField.create(mol, args.forcefield)
parameters = Parameters(ff, mol)
# forcefield = ForceField(args.forcefield, precision=precision)
# parameters = forcefield.create(atom_types, bonds=bonds, angles=angles)
atom_vel = maxwell_boltzmann(parameters.masses, args.temperature, replicas)
atom_forces = torch.zeros(replicas, natoms, 3).to(device).type(precision)
forceterms = ["LJ", "Bonds", "Angles", "Electrostatics"]
ommforceterms = ["lennardjones", "bond", "angle", "electrostatic"]
mol_org = mol.copy()
datastore = {}
for forceterm, ommforceterm in zip(forceterms, ommforceterms):
forceterm = [
forceterm,
]
ommforceterm = [
ommforceterm,
]