def test_write():
clp = ForceField.open(cwd + '/files/CLP.ff', cwd + '/files/CLPol-alpha.ff')
tmp = os.path.join(tmpdir, 'out-CLPol.zfp')
Zfp.save_to(clp, tmp)
assert filecmp.cmp(tmp, cwd + '/files/baselines/out-CLPol.zfp')
il = ForceField.open(cwd + '/files/TEAM_IL.ppf')
tmp = os.path.join(tmpdir, 'out-TEAM_IL.zfp')
Zfp.save_to(il, tmp)
assert filecmp.cmp(tmp, cwd + '/files/baselines/out-TEAM_IL.zfp')
spce = ForceField.open(cwd + '/files/SPCE.ppf')
tmp = os.path.join(tmpdir, 'out-SPCE.zfp')
spce.write(tmp)
assert filecmp.cmp(tmp, cwd + '/files/baselines/out-SPCE.zfp')
def test_guess_connectivity():
top = Topology.open(cwd + '/files/MoS2-13x8-layer1.xyz')
top.cell.set_box([4.109, 4.380, 1.230])
ff = ForceField.open(cwd + '/files/MoS2.ff')
top.guess_connectivity_from_ff(ff, angle_tolerance=15, pbc='xy')
assert top.n_bond == 1248
assert top.n_angle == 3120
assert top.n_dihedral == 0
assert top.n_improper == 0
def test_transfer_bonded_terms():
top = Topology.open(cwd + '/../topology/files/c_3oh.msd')
ff = ForceField.open(cwd + '/../topology/files/c_3oh.ppf')
top.assign_charge_from_ff(ff, transfer_bci_terms=True)
system = System(top, ff, transfer_bonded_terms=True)
angle = next(a for a in system.topology.angles if a.name == 'C2-C3-H8')
assert ff.get_eqt_for_angle(angle) == ('c_3', 'c_3o', 'h_1')
aterm = system.angle_terms[id(angle)]
assert aterm.name == 'c_3,c_3,h_1'
def test_charmm():
ff = ForceField.open(cwd + '/files/10-benzene.ppf')
top = Topology.open(cwd + '/files/10-benzene.lmp', improper_center=3)
top.assign_charge_from_ff(ff)
system = System(top, ff)
tmppdb = os.path.join(tmpdir, 'conf.pdb')
tmppsf = os.path.join(tmpdir, 'topol.psf')
tmpprm = os.path.join(tmpdir, 'ff.prm')
system.export_charmm(pdb_out=tmppdb, psf_out=tmppsf, prm_out=tmpprm)
assert filecmp.cmp(tmppdb, cwd + '/files/baselines/conf.pdb')
assert filecmp.cmp(tmppsf, cwd + '/files/baselines/topol.psf')
assert filecmp.cmp(tmpprm, cwd + '/files/baselines/ff.prm')
def test_write_swm4():
ff = ForceField.open(cwd + '/../forcefield/files/SWM4-NDP.zfp')
mol = Topology.open(cwd + '/files/TIP3P.zmat').molecules[0]
mol.generate_virtual_sites(ff)
mol.generate_drude_particles(ff)
mol.assign_charge_from_ff(ff)
top = Topology([mol], numbers=[10])
tmp = os.path.join(tmpdir, 'swm4-out.psf')
top.write(tmp)
assert filecmp.cmp(tmp, cwd + '/files/baselines/swm4-out.psf')
def test_gmx():
ff = ForceField.open(cwd + '/files/10-benzene.ppf')
top = Topology.open(cwd + '/files/10-benzene.lmp', improper_center=3)
top.assign_charge_from_ff(ff)
system = System(top, ff)
tmpgro = os.path.join(tmpdir, 'conf.gro')
tmptop = os.path.join(tmpdir, 'topol.top')
tmpmdp = os.path.join(tmpdir, 'grompp.mdp')
system.export_gromacs(gro_out=tmpgro, top_out=tmptop, mdp_out=tmpmdp)
assert filecmp.cmp(tmpgro, cwd + '/files/baselines/conf.gro')
assert filecmp.cmp(tmptop, cwd + '/files/baselines/topol.top')
assert filecmp.cmp(tmpmdp, cwd + '/files/baselines/grompp.mdp')
def test_lmp_drude():
ff = ForceField.open(cwd + '/../forcefield/files/CLP.ff',
cwd + '/../forcefield/files/CLPol-alpha.ff')
top = Topology.open(cwd + '/files/5-Im21-BF4-drude.lmp')
top.generate_angle_dihedral_improper()
top.generate_drude_particles(ff)
top.assign_charge_from_ff(ff)
system = System(top, ff)
tmpdata = os.path.join(tmpdir, 'data-drude.lmp')
tmpin = os.path.join(tmpdir, 'in-drude.lmp')
system.export_lammps(data_out=tmpdata, in_out=tmpin)
assert filecmp.cmp(tmpdata, cwd + '/files/baselines/data-drude.lmp')
assert filecmp.cmp(tmpin, cwd + '/files/baselines/in-drude.lmp')
def test_eqt_vdw():
ff = ForceField.open(cwd + '/files/c_3ad.ppf')
top = Topology.open(cwd + '/files/c_3ad.msd')
top.assign_charge_from_ff(ff)
system = System(top, ff)
context = mm.Context(system.to_omm_system(), *get_omm_integrator_platform())
context.setPositions(top.positions)
print_energy_terms(context)
pe = context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(kcal_mol)
assert pytest.approx(pe, rel=0.001) == 46.55
def test_team_vacuum():
ff = ForceField.open(cwd + '/files/10-benzene.ppf')
top = Topology.open(cwd + '/files/10-benzene.lmp', improper_center=3)
top.assign_charge_from_ff(ff)
system = System(top, ff, cell=UnitCell([0, 0, 0]))
context = mm.Context(system.to_omm_system(), *get_omm_integrator_platform())
context.setPositions(top.positions)
print_energy_terms(context)
pe = context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(kcal_mol)
assert pytest.approx(pe, rel=0.001) == 490.8
def test_sdk():
ff = ForceField.open(cwd + '/../forcefield/files/SPICA_v1.zfp')
top = Topology.open(cwd + '/files/10-SDS-20-W.lmp')
for atom in top.atoms:
atom.charge /= 80 ** 0.5
system = System(top, ff)
context = mm.Context(system.to_omm_system(), *get_omm_integrator_platform())
context.setPositions(top.positions)
print_energy_terms(context)
pe = context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(kcal_mol)
assert pytest.approx(pe, rel=0.001) == 186.0
def test_assign_charge():
top = Topology.open(cwd + '/files/c_3oh.msd')
ff = ForceField.open(cwd + '/files/c_3oh.ppf')
top.assign_charge_from_ff(ff)
charges = [atom.charge for atom in top.atoms]
assert pytest.approx(charges, abs=1E-6) == [
-0.32, -0.16, 0.4791, -0.45, 0.16, 0.16, 0.16, -0.0291
]
top.assign_charge_from_ff(ff, transfer_bci_terms=True)
charges = [atom.charge for atom in top.atoms]
assert pytest.approx(charges, abs=1E-6) == [
-0.32, -0.1954, 0.5145, -0.45, 0.16, 0.16, 0.16, -0.0291
]
def test_vdw_shift():
ff = ForceField.open(cwd + '/files/10-benzene.ppf')
ff.vdw_long_range = ff.VDW_LONGRANGE_SHIFT
top = Topology.open(cwd + '/files/10-benzene.lmp', improper_center=3)
# top.assign_charge_from_ff(ff)
system = System(top, ff)
context = mm.Context(system.to_omm_system(), *get_omm_integrator_platform())
context.setPositions(top.positions)
print_energy_terms(context)
pe = context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(kcal_mol)
assert pytest.approx(pe, rel=0.001) == 494.3
def test_gmx_drude():
ff = ForceField.open(cwd + '/../forcefield/files/CLP.ff',
cwd + '/../forcefield/files/CLPol-alpha.ff')
top = Topology.open(cwd + '/files/5-Im21-BF4-drude.lmp')
top.generate_angle_dihedral_improper()
top.generate_drude_particles(ff)
top.assign_charge_from_ff(ff)
system = System(top, ff)
tmpgro = os.path.join(tmpdir, 'conf-drude.gro')
tmptop = os.path.join(tmpdir, 'topol-drude.top')
tmpmdp = os.path.join(tmpdir, 'grompp-drude.mdp')
system.export_gromacs(gro_out=tmpgro, top_out=tmptop, mdp_out=tmpmdp)
assert filecmp.cmp(tmpgro, cwd + '/files/baselines/conf-drude.gro')
assert filecmp.cmp(tmptop, cwd + '/files/baselines/topol-drude.top')
assert filecmp.cmp(tmpmdp, cwd + '/files/baselines/grompp-drude.mdp')
def test_tip4p():
top = Topology.open(os.path.join(cwd, 'files/TIP3P.zmat'))
mol = top.molecules[0]
ff = ForceField.open(os.path.join(cwd, '../forcefield/files/TIP4P.zfp'))
mol.generate_virtual_sites(ff)
mol.assign_charge_from_ff(ff)
avsite = top.atoms[-1]
assert avsite.name == 'VS1'
assert avsite.type == 'MW'
assert pytest.approx(avsite.position,
abs=1E-6) == [0.0091824, -0.011861, 0.]
assert pytest.approx([atom.charge for atom in top.atoms],
abs=1E-6) == [0, 0.52, 0.52, -1.04]
def test_gmx_tip4p():
mol = Topology.open(cwd + '/../topology/files/TIP3P.zmat').molecules[0]
ff = ForceField.open(cwd + '/../forcefield/files/TIP4P.zfp')
mol.generate_virtual_sites(ff)
mol.assign_charge_from_ff(ff)
top = Topology([mol], numbers=[100], cell=UnitCell([3, 3, 3]))
top.set_positions(Topology.open(cwd + '/files/100-TIP4P.pdb').positions)
system = System(top, ff)
tmpgro = os.path.join(tmpdir, 'conf-vsite.gro')
tmptop = os.path.join(tmpdir, 'topol-vsite.top')
tmpmdp = os.path.join(tmpdir, 'grompp-vsite.mdp')
system.export_gromacs(gro_out=tmpgro, top_out=tmptop, mdp_out=tmpmdp)
assert filecmp.cmp(tmpgro, cwd + '/files/baselines/conf-vsite.gro')
assert filecmp.cmp(tmptop, cwd + '/files/baselines/topol-vsite.top')
assert filecmp.cmp(tmpmdp, cwd + '/files/baselines/grompp-vsite.mdp')
def test_drude():
ff = ForceField.open(cwd + '/../forcefield/files/CLP.ff', cwd + '/../forcefield/files/CLPol-alpha.ff')
top = Topology.open(cwd + '/files/5-Im21-BF4-drude.lmp')
top.generate_angle_dihedral_improper()
# top.remove_drude_particles()
# top.generate_drude_particles(ff)
top.assign_charge_from_ff(ff)
system = System(top, ff)
context = mm.Context(system.to_omm_system(), *get_omm_integrator_platform())
context.setPositions(top.positions)
print_energy_terms(context)
pe = context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(kcal_mol)
assert pytest.approx(pe, rel=0.001) == 100.8
def main():
args = read_args()
ff = ForceField.open(*args.forcefield)
if args.scale:
scaler = PaduaLJScaler(args.scale)
scaler.scale(ff)
mols = read_mols(args.files, ff)
system = make_system(
mols,
args.numbers,
ff,
box=args.box,
virtual_sites=args.virtuals,
drudes=args.drudes,
)
system.export_gromacs(gro_out="conf.gro", top_out=None, mdp_out=None)
system.export_charmm(psf_out="topol.psf", prm_out="ff.prm", pdb_out=None)
def test_tip4p():
mol = Topology.open(cwd + '/../topology/files/TIP3P.zmat').molecules[0]
ff = ForceField.open(cwd + '/../forcefield/files/TIP4P.zfp')
mol.generate_virtual_sites(ff)
mol.assign_charge_from_ff(ff)
top = Topology([mol], numbers=[100], cell=UnitCell([3, 3, 3]))
top.set_positions(Topology.open(cwd + '/files/100-TIP4P.pdb').positions)
system = System(top, ff)
context = mm.Context(system.to_omm_system(), *get_omm_integrator_platform())
context.setPositions(top.positions)
print_energy_terms(context)
pe = context.getState(getEnergy=True).getPotentialEnergy().value_in_unit(kJ_mol)
assert pytest.approx(pe, rel=0.001) == 383.796 # results of GROMACS
def test_read():
ff = ForceField.open(cwd + '/files/TEAM_IL.zfp')
assert ff.vdw_cutoff == 1.2
assert ff.vdw_long_range == 'correct'
assert ff.scale_14_vdw == 0.5
assert ff.lj_mixing_rule == ForceField.LJ_MIXING_LB
assert pytest.approx(ff.scale_14_coulomb, abs=1E-4) == 0.8333
atype = ff.atom_types['s_1-']
assert atype.name == 's_1-'
assert atype.charge == -1
assert atype.eqt_imp_s == 's_1'
assert atype.eqt_imp_c == 's_1-'
term = ff.bci_terms['b_4-,c_2nb']
assert term.type1 == 'b_4-'
assert term.type2 == 'c_2nb'
assert term.value == 0.262
term = ff.bci_terms['b_4-,f_1']
assert term.type1 == 'b_4-'
assert term.type2 == 'f_1'
assert term.value == 0.4215
term = ff.dihedral_terms['o_1,s_4o,n_2-,s_4']
assert type(term) == PeriodicDihedralTerm
assert term.type1 == 'o_1'
assert term.type2 == 's_4o'
assert term.type3 == 'n_2-'
assert term.type4 == 's_4'
assert term.is_opls_convention
assert pytest.approx(term.get_opls_parameters(),
abs=1E-4) == [12.9809, 7.9709, 1.6389, 0]
term = ff.improper_terms['c_3a,c_3a,c_3a,h_1']
assert type(term) == OplsImproperTerm
assert term.type1 == 'c_3a'
assert term.type2 == 'c_3a'
assert term.type3 == 'c_3a'
assert term.type4 == 'h_1'
assert pytest.approx(term.k, abs=1E-4) == 7.1270
raise Exception('--typer is required for SMILES input')
top = Topology.open(inp)
molecules += top.molecules * n
top = Topology(molecules)
top.remove_drude_particles()
top.remove_virtual_sites()
if args.trj is not None:
frame = Trajectory.read_frame_from_file(args.trj, -1)
if frame.cell.volume != 0:
top.cell.set_box(frame.cell.vectors)
if args.box is not None:
top.cell.set_box(args.box)
ff = ForceField.open(*args.forcefield)
if args.nodrude:
ff.polarizable_terms.clear()
ff.virtual_site_terms.clear()
if args.ljscale is not None:
scaler = PaduaLJScaler(args.ljscale)
scaler.scale(ff)
if args.scaleeps != 1.0 or args.scalesig != 1.0:
for vdw in list(ff.vdw_terms.values()) + list(ff.pairwise_vdw_terms.values()):
if vdw.type1 in args.scaleignoreatom or vdw.type2 in args.scaleignoreatom:
continue
if args.scaleeps != 1.0:
vdw.epsilon *= args.scaleeps
vdw.comments.append('eps*%.3f' % args.scaleeps)
if args.scalesig != 1.0:
c_4h3 [CX4;H3] 0
HierarchicalTree
h_1
c_4
c_4h2
c_4h3
'''
typer = ZftTyper(io.StringIO(definition))
butane = Molecule.from_smiles('CCCC butane')
typer.type_molecule(butane)
# Load force field parameters from ZFP file
ff = ForceField.open('alkane.zfp')
# Initialize a topology with periodic boundary condition
# For now, it contains only one butane molecule
top = Topology([butane], cell=UnitCell([3, 3, 3]))
# Assign atomic charges based on the charge parameters in the force field
top.assign_charge_from_ff(ff)
# Call Packmol to build a configuration containing 100 butane molecules
packmol = Packmol(r'/path/of/packmol')
top.scale_with_packmol([100], packmol=packmol)
# Associate force field parameters to the topology
# And then export input files for simulation engines
system = System(top, ff)
def test_read():
ff = ForceField.open(cwd + '/files/TEAM_IL.ppf', cwd + '/files/SPCE.ppf')
assert len(ff.atom_types) == 65
c_4pp = ff.atom_types.get('c_4pp')
assert c_4pp.charge == 0
assert c_4pp.eqt_vdw == c_4pp.name
assert c_4pp.eqt_q_inc == c_4pp.name
assert c_4pp.eqt_bond == 'c_4'
assert c_4pp.eqt_ang_c == 'c_4'
assert c_4pp.eqt_ang_s == 'c_4'
assert c_4pp.eqt_dih_c == 'c_4'
assert c_4pp.eqt_dih_s == 'c_4'
assert c_4pp.eqt_imp_c == 'c_4'
assert c_4pp.eqt_imp_s == 'c_4'
assert c_4pp.version == '0.25'
c_3_ = ff.atom_types.get('c_3-')
assert c_3_.charge == -1
assert c_3_.version == '0.1'
o_1_t = ff.atom_types.get('o_1-t')
assert o_1_t.charge == -0.3333
binc = ff.bci_terms.get('c_35an2,h_1')
assert binc.type1 == 'c_35an2'
assert binc.type2 == 'h_1'
assert binc.version == '0.16'
assert binc.value == -0.22
vdw = ff.vdw_terms.get('b_4-,b_4-')
assert vdw.type1 == 'b_4-'
assert vdw.type2 == 'b_4-'
assert pytest.approx(vdw.epsilon / 4.184, abs=1E-6) == 0.05
assert pytest.approx(vdw.sigma * 2**(1 / 6) * 10, abs=1E-6) == 3.8
assert vdw.version == '0.36'
vdw = ff.pairwise_vdw_terms.get('o_2,o_2w')
assert vdw.type1 == 'o_2'
assert vdw.type2 == 'o_2w'
assert pytest.approx(vdw.epsilon / 4.184, abs=1E-6) == 0.22936
assert pytest.approx(vdw.sigma * 2**(1 / 6) * 10, abs=1E-6) == 3.44671
assert vdw.version == None
bond = ff.bond_terms.get('c_2n,n_2-')
assert bond.type1 == 'c_2n'
assert bond.type2 == 'n_2-'
assert pytest.approx(bond.length * 10, abs=1E-6) == 1.276
assert pytest.approx(bond.k / 4.184 / 100, abs=1E-6) == 487.36
assert bond.version == '0.3'
angle = ff.angle_terms.get('c_3a,c_3a,h_1')
assert angle.type1 == 'c_3a'
assert angle.type2 == 'c_3a'
assert angle.type3 == 'h_1'
assert pytest.approx(angle.theta, abs=1E-6) == 118.916
assert pytest.approx(angle.k / 4.184, abs=1E-6) == 43.8769
assert angle.version == '0.1'
dihedral = ff.dihedral_terms.get('*,c_4,c_4,*')
assert dihedral.type1 == '*'
assert dihedral.type2 == 'c_4'
assert dihedral.type3 == 'c_4'
assert dihedral.type4 == '*'
para1 = dihedral.parameters[0]
assert pytest.approx(para1.k / 4.184, abs=1E-6) == 0.6214
assert para1.n == 1
assert para1.phi == 0.0
para2 = dihedral.parameters[1]
assert pytest.approx(para2.k / 4.184, abs=1E-6) == 0.0507
assert para2.n == 2
assert para2.phi == 180.0
para0 = dihedral.parameters[2]
assert pytest.approx(para0.k / 4.184, abs=1E-6) == 0.0688
assert para0.n == 3
assert para0.phi == 0.0
assert dihedral.version == '0.12'
improper = ff.improper_terms.get('c_3o,o_1,*,*')
assert pytest.approx(improper.k / 4.184, abs=1E-6) == 18
assert improper.version == '0.21'
assert len(ff.polarizable_terms) == 0
def test_write():
ff = ForceField.open(cwd + '/files/CLP.ff', cwd + '/files/CLPol-alpha.ff')
tmp = os.path.join(tmpdir, 'out-CLP.ppf')
Ppf.save_to(ff, tmp)
assert filecmp.cmp(tmp, cwd + '/files/baselines/out-CLP.ppf')
import sys
import argparse
from mstools.forcefield import ForceField
parser = argparse.ArgumentParser()
parser.add_argument('input', nargs='+', type=str, help='force field files')
parser.add_argument('-o',
'--output',
required=True,
type=str,
help='output file')
args = parser.parse_args()
ff = ForceField.open(*args.input)
print('\nForce field info:\n'
'%6i atom types\n'
'%6i self vdW terms\n'
'%6i pairwise vdW terms\n'
'%6i charge increment terms\n'
'%6i bond terms\n'
'%6i angle terms\n'
'%6i dihedral terms\n'
'%6i improper terms\n'
'%6i polarizable terms' %
(len(ff.atom_types), len(ff.vdw_terms), len(ff.pairwise_vdw_terms),
len(ff.bci_terms), len(ff.bond_terms), len(
ff.angle_terms), len(ff.dihedral_terms), len(
ff.improper_terms), len(ff.polarizable_terms)))
def msd2dgl_ff(mol_list, ff_file):
top = Topology(mol_list)
ff = ForceField.open(ff_file)
top.assign_charge_from_ff(ff, transfer_bci_terms=True)
system = System(top,
ff,
transfer_bonded_terms=True,
suppress_pbc_warning=True)
top, ff = system.topology, system.ff
graph_list = []
feats_node_list = []
feats_bond_list = []
feats_angle_list = []
feats_dihedral_list = []
for mol in top.molecules:
feats_node = np.zeros((mol.n_atom, 4))
feats_bond = np.zeros((mol.n_bond * 2, 2)) # bidirectional
feats_angle = np.zeros((mol.n_angle * 2, 2)) # bidirectional
feats_dihedral = np.zeros((mol.n_dihedral * 2, 3)) # bidirectional
for i, atom in enumerate(mol.atoms):
atype = ff.atom_types[atom.type]
vdw = ff.get_vdw_term(atype, atype)
feats_node[i] = vdw.sigma, vdw.epsilon, atom.charge, 0
for i, bond in enumerate(mol.bonds):
term = system.bond_terms[id(bond)]
inv_k = 1e5 / term.k if not term.fixed else 0
feats_bond[i] = feats_bond[i +
mol.n_bond] = term.length * 10, inv_k
for i, angle in enumerate(mol.angles):
term = system.angle_terms[id(angle)]
inv_k = 100 / term.k if not term.fixed else 0
feats_angle[i] = feats_angle[i +
mol.n_angle] = term.theta / 100, inv_k
for i, dihedral in enumerate(mol.dihedrals):
term = system.dihedral_terms[id(dihedral)]
k1, k2, k3, k4 = term.get_opls_parameters()
feats_dihedral[i] = feats_dihedral[
i + mol.n_dihedral] = k1 / 10, k2 / 10, k3 / 10
for i, improper in enumerate(mol.impropers):
term = system.improper_terms[id(improper)]
feats_node[improper.atom1.id_in_mol][-1] = term.k / 10
feats_node_list.append(feats_node)
feats_bond_list.append(feats_bond)
feats_angle_list.append(feats_angle)
feats_dihedral_list.append(feats_dihedral)
bonds = [(bond.atom1.id_in_mol, bond.atom2.id_in_mol)
for bond in mol.bonds]
edges = list(zip(*bonds))
u = torch.tensor(edges[0] + edges[1]) # bidirectional
v = torch.tensor(edges[1] + edges[0])
graph_data = {('atom', 'bond', 'atom'): (u, v)}
angles = [(angle.atom1.id_in_mol, angle.atom3.id_in_mol)
for angle in mol.angles]
edges = list(zip(*angles))
u = torch.tensor(edges[0] + edges[1]) # bidirectional
v = torch.tensor(edges[1] + edges[0])
graph_data.update({('atom', 'angle', 'atom'): (u, v)})
dihedrals = [(dihedral.atom1.id_in_mol, dihedral.atom4.id_in_mol)
for dihedral in mol.dihedrals]
edges = list(zip(*dihedrals))
u = torch.tensor(edges[0] + edges[1]) # bidirectional
v = torch.tensor(edges[1] + edges[0])
graph_data.update({('atom', 'dihedral', 'atom'): (u, v)})
graph = dgl.heterograph(graph_data)
graph_list.append(graph)
return graph_list, feats_node_list, {
'bond': feats_bond_list,
'angle': feats_angle_list,
'dihedral': feats_dihedral_list
}
def mol2dgl_ff_pairs(mol_list, ff_file, dist_list, distinguish_pairs=False):
top = Topology(mol_list)
ff = ForceField.open(ff_file)
top.assign_charge_from_ff(ff, transfer_bci_terms=True)
system = System(top,
ff,
transfer_bonded_terms=True,
suppress_pbc_warning=True)
top, ff = system.topology, system.ff
graph_list = []
feats_node_list = []
feats_p12_list = []
feats_p13_list = []
feats_p14_list = []
feats_edge_length = len(dist_list[0])
for i, (mol, df) in enumerate(zip(top.molecules, dist_list)):
if i % 1000 == 0:
print('Processing molecule %i / %i' % (i, len(dist_list)))
pairs12, pairs13, pairs14 = mol.get_12_13_14_pairs()
edges = list(zip(*[(p[0].id_in_mol, p[1].id_in_mol) for p in pairs12]))
u12, v12 = edges[0] + edges[1], edges[1] + edges[0]
edges = list(zip(*[(p[0].id_in_mol, p[1].id_in_mol) for p in pairs13]))
u13, v13 = tuple(), tuple()
if len(pairs13) > 0:
u13, v13 = edges[0] + edges[1], edges[1] + edges[0]
edges = list(zip(*[(p[0].id_in_mol, p[1].id_in_mol) for p in pairs14]))
u14, v14 = tuple(), tuple()
if len(pairs14) > 0:
u14, v14 = edges[0] + edges[1], edges[1] + edges[0]
if distinguish_pairs:
graph = dgl.heterograph({
('atom', 'pair12', 'atom'): (u12, v12),
('atom', 'pair13', 'atom'): (u13, v13),
('atom', 'pair14', 'atom'): (u14, v14),
})
else:
uv_self = tuple(range(mol.n_atom))
graph = dgl.heterograph({
('atom', 'pair', 'atom'):
(u12 + u13 + u14 + uv_self, v12 + v13 + v14 + uv_self)
})
graph_list.append(graph)
feats_node = np.zeros((mol.n_atom, 3))
for i, atom in enumerate(mol.atoms):
atype = ff.atom_types[atom.type]
vdw = ff.get_vdw_term(atype, atype)
feats_node[i] = vdw.sigma, vdw.epsilon, atom.charge
feats_node_list.append(feats_node)
feats_p12 = np.zeros(
(len(pairs12) * 2, feats_edge_length)) # bidirectional
feats_p13 = np.zeros(
(len(pairs13) * 2, feats_edge_length)) # bidirectional
feats_p14 = np.zeros(
(len(pairs14) * 2, feats_edge_length)) # bidirectional
for i, pair in enumerate(pairs12):
key = '%s-%s' % (pair[0].name, pair[1].name)
feats_p12[i][:] = feats_p12[i +
len(pairs12)][:] = df[key].values / 10
for i, pair in enumerate(pairs13):
key = '%s-%s' % (pair[0].name, pair[1].name)
feats_p13[i][:] = feats_p13[i +
len(pairs13)][:] = df[key].values / 10
for i, pair in enumerate(pairs14):
key = '%s-%s' % (pair[0].name, pair[1].name)
feats_p14[i][:] = feats_p14[i +
len(pairs14)][:] = df[key].values / 10
feats_p12_list.append(feats_p12)
feats_p13_list.append(feats_p13)
feats_p14_list.append(feats_p14)
if distinguish_pairs:
feats_edges = {
'pair12': [feats[:, 2:13] for feats in feats_p12_list],
'pair13': [feats[:, 7:18] for feats in feats_p13_list],
'pair14': [feats[:, 12:23] for feats in feats_p14_list],
}
else:
feats_self_list = [
np.zeros((mol.n_atom, feats_edge_length)) for mol in top.molecules
]
feats_edges = {
'pair': [
np.concatenate(x)[:, 2:23]
for x in zip(feats_p12_list, feats_p13_list, feats_p14_list,
feats_self_list)
]
}
return graph_list, feats_node_list, feats_edges