def addLigandCylinderBox(topology, positions, system, resname, dummies, radius, worker):
center, base, top_base = dummies
masses = []
coords = np.ndarray(shape=(0, 3))
ligand_atoms = []
for atom in topology.atoms():
if atom.residue.name == resname and atom.element.symbol != "H":
masses.append(atom.element.mass.value_in_unit(unit=unit.dalton))
coords = np.vstack((coords, positions[atom.index].value_in_unit(unit=unit.nanometer)))
ligand_atoms.append(atom)
masses = np.array(masses)
masses /= masses.sum()
mass_center = coords.astype('float64').T.dot(masses)
atomClosestToMassCenter = min(ligand_atoms, key=lambda x: np.linalg.norm(mass_center - positions[x.index].value_in_unit(unit=unit.nanometer)))
length = np.linalg.norm(positions[center].value_in_unit(unit.nanometers)-positions[base].value_in_unit(unit.nanometers))
if worker == 0:
utilities.print_unbuffered("Ligand atom selected to check distance to the box", atomClosestToMassCenter.residue.name, atomClosestToMassCenter.name, atomClosestToMassCenter.index)
ligand_atom = atomClosestToMassCenter.index
# forceFB = mm.CustomBondForce('step(r-r_l)*(k_box/2) * (r-r_l)^2')
forceFB = mm.CustomCompoundBondForce(3, '(step(r_par-2*r_l)+step(-r_par))*(k_box/2) * (r_par-r_l)^2; r_par=ax*dx+ay*dy+az*dz; ax=(x1-x2)/l; ay=(y1-y2)/l; az=(z1-z2)/l; dx=x3-x2; dy=y3-y2; dz=z3-z2; l=distance(p1, p2)')
forceFB.addPerBondParameter("k_box")
forceFB.addPerBondParameter("r_l")
#forceFB.addBond(center, ligand_atom, [5.0 * unit.kilocalories_per_mole / unit.angstroms ** 2, length*unit.nanometers])
forceFB.addBond([center, base, ligand_atom], [5.0 * unit.kilocalories_per_mole / unit.angstroms ** 2, length*unit.nanometers])
system.addForce(forceFB)
force_side = mm.CustomCompoundBondForce(3, 'step(r_normal-r0)*(k_box/2) * (r_normal-r0)^2; r_normal=sqrt((ay*dz-az*dy)^2+(az*dx-ax*dz)^2+(ax*dy-ay*dx)^2); ax=(x1-x2)/l; ay=(y1-y2)/l; az=(z1-z2)/l; dx=x3-x2; dy=y3-y2; dz=z3-z2; l=distance(p1, p2)')
force_side.addPerBondParameter("k_box")
force_side.addPerBondParameter("r0")
# the order of the particles involved are center of the cilinder, base and
# atom to constrain (ligand)
force_side.addBond([center, base, ligand_atom], [5.0 * unit.kilocalories_per_mole / unit.angstroms ** 2, radius*unit.angstroms])
system.addForce(force_side)
def get_Upair_force(Ucg, coeff, r_vals):
"""Create pairwise potentials as tabulated functions"""
N = Ucg.n_atoms
bcut = Ucg.bond_cutoff
Up_vals = Ucg.Upair_values(coeff, r_vals)
if len(Up_vals) > 1:
table_vals = np.zeros((len(Up_vals), len(r_vals)))
for i in range(len(Up_vals)):
table_vals[i,:] = Up_vals[i]
xvals = np.arange(len(Up_vals)).astype(float)
yvals = r_vals
Table_func = omm.Continuous2DFunction(len(xvals), len(yvals),
table_vals.flatten(), xvals[0], xvals[-1], yvals[0], yvals[-1])
Up_force = omm.CustomCompoundBondForce(2, "Table(p_idx, distance(p1, p2))")
Up_force.addPerBondParameter("p_idx")
Up_force.addTabulatedFunction("Table", Table_func)
else:
Table_func = omm.Continuous1DFunction(Up_vals[0], r_vals[0], r_vals[-1])
Up_force = omm.CustomCompoundBondForce(2, "Table(distance(p1, p2))")
Up_force.addTabulatedFunction("Table", Table_func)
if Ucg.pair_symmetry == "shared":
# all pairs share same interaction
for i in range(N - bcut):
for j in range(i + bcut, N):
Up_force.addBond((i, j))
elif Ucg.pair_symmetry == "seq_sep":
# pairs with same sequence sep share same interaction
for i in range(N - bcut):
for j in range(i + bcut, N):
p_idx = np.abs(j - i) - bcut
Up_force.addBond([i, j], p_idx)
elif Ucg.pair_symmetry == "unique":
# all pair have different interactions
p_idx = 0
for i in range(N - bcut):
for j in range(i + bcut, N):
Up_force.addBond((i, j), (p_idx))
p_idx += 1
else:
raise ValueError("pair_symmetry must be: shared, seq_sep, unique. Gave:" + str(Ucg.pair_symmetry))
return Up_force
def add_external_field(system: mm.System, args: ListOfArgs):
"""Add external forcefield for image-driven modelling purposes."""
print(' Adding external forcefield.')
size = os.stat(args.EF_PATH).st_size
print(f" Reading {args.EF_PATH} file ({sizeof_fmt(size)})...")
img = np.load(args.EF_PATH)
print(f" Array of shape {img.shape} loaded.")
print(f" Number of values: {img.size}")
print(f" Min: {np.min(img)}")
print(f" Max: {np.max(img)}")
if args.EF_NORMALIZE:
print(' [INFO] Field will be normalized to [0, -1]')
img = standardize_image(img)
print(f' [INFO] IMG min = {np.min(img)}, max = {np.max(img)}')
print(f' [INFO] Adding funnel like border to image')
mask_p = (img < -0.1)
mask_n = np.logical_not(mask_p)
img = add_funnel(img, mask_n)
print(" Creating a force based on density...")
voxel_size = np.array((args.EF_VOXEL_SIZE_X, args.EF_VOXEL_SIZE_Y, args.EF_VOXEL_SIZE_Z))
real_size = img.shape * voxel_size
density_fun_args = dict(
xsize=img.shape[2],
ysize=img.shape[1],
zsize=img.shape[0],
values=img.flatten().astype(np.float64),
xmin=0 * simtk.unit.angstrom - 0.5 * voxel_size[0],
ymin=0 * simtk.unit.angstrom - 0.5 * voxel_size[1],
zmin=0 * simtk.unit.angstrom - 0.5 * voxel_size[2],
xmax=(img.shape[0] - 1) * voxel_size[0] + 0.5 * voxel_size[0],
ymax=(img.shape[1] - 1) * voxel_size[1] + 0.5 * voxel_size[1],
zmax=(img.shape[2] - 1) * voxel_size[2] + 0.5 * voxel_size[2])
print(f' [INFO] Voxel size: ({args.EF_VOXEL_SIZE_X}, {args.EF_VOXEL_SIZE_Y}, {args.EF_VOXEL_SIZE_Z})')
print(f' [INFO] Real size (Shape * voxel size): ({real_size[0]}, {real_size[1]}, {real_size[2]})')
print(
f" [INFO] begin coords: ({density_fun_args['xmin']}, {density_fun_args['ymin']}, {density_fun_args['zmin']})")
print(
f" [INFO] end coords: ({density_fun_args['xmax']}, {density_fun_args['ymax']}, {density_fun_args['zmax']})")
center_x = (density_fun_args['xmax'] - density_fun_args['xmin']) / 2 + density_fun_args['xmin']
center_y = (density_fun_args['ymax'] - density_fun_args['ymin']) / 2 + density_fun_args['ymin']
center_z = (density_fun_args['zmax'] - density_fun_args['zmin']) / 2 + density_fun_args['zmin']
print(f" [INFO] Image central point: ({center_x}, {center_y}, {center_z}) ")
field_function = mm.Continuous3DFunction(**density_fun_args)
field_force = mm.CustomCompoundBondForce(1, 'ksi*fi(x1,y1,z1)')
field_force.addTabulatedFunction('fi', field_function)
field_force.addGlobalParameter('ksi', args.EF_SCALING_FACTOR)
print(" Adding force to the system...")
for i in range(system.getNumParticles()):
field_force.addBond([i], [])
system.addForce(field_force)
def setPositions(self, positions, extended_positions=None):
"""
Sets the positions of all particles and extended-space variables in
this context. If the latter are not provided, then suitable values
are automatically determined from the particle positions.
Parameters
----------
positions : list of openmm.Vec3
The positions of physical particles.
Keyword Args
------------
extended_positions : list of float or unit.Quantity
The positions of extended-space particles.
"""
a, b, _ = self.getState().getPeriodicBoxVectors()
nvars = len(self.variables)
particle_positions = _standardized(positions)
if extended_positions is None:
extra_positions = [openmm.Vec3(0, b.y*(i + 1)/(nvars + 2), 0) for i in range(nvars)]
minisystem = openmm.System()
expression = _get_energy_function(self.variables)
for i, v in enumerate(self.variables):
expression += f'; {v.id}={v._get_energy_function(index=i+1)}'
force = openmm.CustomCompoundBondForce(nvars, expression)
force.addBond(range(nvars), [])
for name, value in _get_parameters(self.variables).items():
force.addGlobalParameter(name, value)
force.addGlobalParameter('Lx', a.x)
for v in self.variables:
minisystem.addParticle(v._particle_mass(a.x))
for cv in v.colvars:
value = cv.evaluate(self.getSystem(), particle_positions + extra_positions)
force.addGlobalParameter(cv.id, value)
minisystem.addForce(force)
minicontext = openmm.Context(minisystem, openmm.CustomIntegrator(0),
openmm.Platform.getPlatformByName('Reference'))
minicontext.setPositions(extra_positions)
openmm.LocalEnergyMinimizer.minimize(minicontext, 1*unit.kilojoules_per_mole, 0)
ministate = minicontext.getState(getPositions=True)
extra_positions = ministate.getPositions().value_in_unit(unit.nanometers)
else:
extra_positions = [openmm.Vec3(x, b.y*(i + 1)/(nvars + 2), 0)
for i, x in enumerate(extended_positions)]
super().setPositions(particle_positions + extra_positions)
def get_restraint_force(self):
"""
Create a new copy of the receptor-ligand restraint force.
Returns
-------
force : simtk.openmm.CustomCompoundBondForce
A restraint force object
WARNING
-------
Because the domain of the `dihedral()` function in `CustomCompoundBondForce` is unknown, it is not currently clear whether the dihedral restraints will work.
We should replace those restraints with a negative log von Mises distribution: https://en.wikipedia.org/wiki/Von_Mises_distribution
where the von Mises parameter kappa would be computed from the desired standard deviation (kappa ~ sigma**(-2))
and the standard state correction would need to be modified
Notes
-----
There may be a better way to implement this restraint.
"""
energy_function = """
lambda_restraints * E;
E = (K_r/2)*(distance(p3,p4) - r_aA0)^2
+ (K_thetaA/2)*(angle(p2,p3,p4)-theta_A0)^2 + (K_thetaB/2)*(angle(p3,p4,p5)-theta_B0)^2
+ (K_phiA/2)*dphi_A^2 + (K_phiB/2)*dphi_B^2 + (K_phiC/2)*dphi_C^2;
dphi_A = dA - floor(dA/(2*pi)+0.5)*(2*pi); dA = dihedral(p1,p2,p3,p4) - phi_A0;
dphi_B = dB - floor(dB/(2*pi)+0.5)*(2*pi); dB = dihedral(p2,p3,p4,p5) - phi_B0;
dphi_C = dC - floor(dC/(2*pi)+0.5)*(2*pi); dC = dihedral(p3,p4,p5,p6) - phi_C0;
pi = %f;
""" % np.pi
# Add constant definitions to the energy function
for name in self._parameters:
energy_function += '%s = %f; ' % (name, self._parameters[name].value_in_unit_system(unit.md_unit_system))
# Create the force
nparticles = 6 # number of particles involved in restraint: p1 ... p6
force = openmm.CustomCompoundBondForce(nparticles, energy_function)
force.addGlobalParameter('lambda_restraints', 1.0)
force.addBond(self._restraint_atoms, [])
is_periodic = self._system.usesPeriodicBoundaryConditions()
try: # This was added in OpenMM 7.1
force.setUsesPeriodicBoundaryConditions(is_periodic)
except AttributeError:
pass
return force
def _hills_force(self):
num_bias_variables = len(self.bias_variables)
centers = [f'center{i+1}' for i in range(num_bias_variables)]
exponents = []
for v, center in zip(self.bias_variables, centers):
if v.periodic: # von Mises
factor = 2*np.pi/v._range
exponents.append(f'{1.0/(factor*v.sigma)**2}*(cos({factor}*({v.id}-{center}))-1)')
else: # Gauss
exponents.append(f'({-0.5/v.sigma**2})*({v.id}-{center})^2')
expression = f'height*exp({"+".join(exponents)})'
for i, v in enumerate(self.bias_variables):
expression += f';{v.id}={v._get_energy_function(i+1)}'
force = openmm.CustomCompoundBondForce(num_bias_variables, expression)
force.addPerBondParameter('height')
for center in centers:
force.addPerBondParameter(center)
force.addGlobalParameter('Lx', 0)
return force
def __init__(self, variables, variances, centers, weights, platform,
properties):
num_particles = len(variables) // 3 + 1
coordinates = [
f'{x}{i+1}' for i in range(num_particles) for x in 'xyz'
]
exponents = []
for v, variance, x in zip(variables, variances, coordinates):
if v.periodic: # von Mises
factor = 2 * np.pi / v._range
exponents.append(
f'{1.0/(variance*factor**2)}*(cos({factor}*({v.id}-{x}))-1)'
)
else: # Gauss
exponents.append(f'(-{0.5/variance})*({v.id}-{x})^2')
expression = f'weight*exp({"+".join(exponents)})'
force = openmm.CustomCompoundBondForce(num_particles, expression)
force.addPerBondParameter('weight')
for v in variables:
force.addGlobalParameter(v.id, 0)
force.addEnergyParameterDerivative(v.id)
system = openmm.System()
positions = []
for i, (center, weight) in enumerate(zip(centers, weights)):
for position in np.resize(center, (num_particles, 3)):
system.addParticle(1.0)
positions.append(openmm.Vec3(*position))
force.addBond(range(i * num_particles, (i + 1) * num_particles),
[weight])
system.addForce(force)
integrator = openmm.CustomIntegrator(0)
super().__init__(system, integrator, platform, properties)
self.parameters = [v.id for v in variables]
self.setPositions(positions)
def export(system):
'''
Generate OpenMM system from a system
Parameters
----------
system : System
Returns
-------
omm_system : simtk.openmm.System
'''
try:
import simtk.openmm as mm
except ImportError:
raise ImportError('Can not import OpenMM')
supported_terms = {
LJ126Term, MieTerm, HarmonicBondTerm, HarmonicAngleTerm,
SDKAngleTerm, PeriodicDihedralTerm, OplsImproperTerm,
HarmonicImproperTerm, DrudeTerm
}
unsupported = system.ff_classes - supported_terms
if unsupported != set():
raise Exception(
'Unsupported FF terms: %s' %
(', '.join(map(lambda x: x.__name__, unsupported))))
if system.vsite_types - {TIP4PSite} != set():
raise Exception(
'Virtual sites other than TIP4PSite haven\'t been implemented')
top = system.topology
ff = system.ff
omm_system = mm.System()
if system.use_pbc:
omm_system.setDefaultPeriodicBoxVectors(*top.cell.vectors)
for atom in top.atoms:
omm_system.addParticle(atom.mass)
### Set up bonds #######################################################################
for bond_class in system.bond_classes:
if bond_class == HarmonicBondTerm:
logger.info('Setting up harmonic bonds...')
bforce = mm.HarmonicBondForce()
for bond in top.bonds:
if bond.is_drude:
# DrudeForce will handle the bond between Drude pair
continue
bterm = system.bond_terms[id(bond)]
if type(bterm) != HarmonicBondTerm:
continue
bforce.addBond(bond.atom1.id, bond.atom2.id, bterm.length,
bterm.k * 2)
else:
raise Exception('Bond terms other that HarmonicBondTerm '
'haven\'t been implemented')
bforce.setUsesPeriodicBoundaryConditions(system.use_pbc)
bforce.setForceGroup(ForceGroup.BOND)
omm_system.addForce(bforce)
### Set up angles #######################################################################
for angle_class in system.angle_classes:
if angle_class == HarmonicAngleTerm:
logger.info('Setting up harmonic angles...')
aforce = mm.HarmonicAngleForce()
for angle in top.angles:
aterm = system.angle_terms[id(angle)]
if type(aterm) == HarmonicAngleTerm:
aforce.addAngle(angle.atom1.id, angle.atom2.id,
angle.atom3.id, aterm.theta * PI / 180,
aterm.k * 2)
elif angle_class == SDKAngleTerm:
logger.info('Setting up SDK angles...')
aforce = mm.CustomCompoundBondForce(
3, 'k*(theta-theta0)^2+step(rmin-r)*LJ96;'
'LJ96=6.75*epsilon*((sigma/r)^9-(sigma/r)^6)+epsilon;'
'theta=angle(p1,p2,p3);'
'r=distance(p1,p3);'
'rmin=1.144714*sigma')
aforce.addPerBondParameter('theta0')
aforce.addPerBondParameter('k')
aforce.addPerBondParameter('epsilon')
aforce.addPerBondParameter('sigma')
for angle in top.angles:
aterm = system.angle_terms[id(angle)]
if type(aterm) != SDKAngleTerm:
continue
vdw = ff.get_vdw_term(ff.atom_types[angle.atom1.type],
ff.atom_types[angle.atom2.type])
if type(
vdw
) != MieTerm or vdw.repulsion != 9 or vdw.attraction != 6:
raise Exception(
f'Corresponding 9-6 MieTerm for {aterm} not found in FF'
)
aforce.addBond(
[angle.atom1.id, angle.atom2.id, angle.atom3.id], [
aterm.theta * PI / 180, aterm.k, vdw.epsilon,
vdw.sigma
])
else:
raise Exception(
'Angle terms other that HarmonicAngleTerm and SDKAngleTerm '
'haven\'t been implemented')
aforce.setUsesPeriodicBoundaryConditions(system.use_pbc)
aforce.setForceGroup(ForceGroup.ANGLE)
omm_system.addForce(aforce)
### Set up constraints #################################################################
logger.info(
f'Setting up {len(system.constrain_bonds)} bond constraints...')
for bond in top.bonds:
if id(bond) in system.constrain_bonds:
omm_system.addConstraint(bond.atom1.id, bond.atom2.id,
system.constrain_bonds[id(bond)])
logger.info(
f'Setting up {len(system.constrain_angles)} angle constraints...')
for angle in top.angles:
if id(angle) in system.constrain_angles:
omm_system.addConstraint(angle.atom1.id, angle.atom3.id,
system.constrain_angles[id(angle)])
### Set up dihedrals ###################################################################
for dihedral_class in system.dihedral_classes:
if dihedral_class == PeriodicDihedralTerm:
logger.info('Setting up periodic dihedrals...')
dforce = mm.PeriodicTorsionForce()
for dihedral in top.dihedrals:
dterm = system.dihedral_terms[id(dihedral)]
ia1, ia2, ia3, ia4 = dihedral.atom1.id, dihedral.atom2.id, dihedral.atom3.id, dihedral.atom4.id
if type(dterm) == PeriodicDihedralTerm:
for par in dterm.parameters:
dforce.addTorsion(ia1, ia2, ia3, ia4, par.n,
par.phi * PI / 180, par.k)
else:
continue
else:
raise Exception(
'Dihedral terms other that PeriodicDihedralTerm '
'haven\'t been implemented')
dforce.setUsesPeriodicBoundaryConditions(system.use_pbc)
dforce.setForceGroup(ForceGroup.DIHEDRAL)
omm_system.addForce(dforce)
### Set up impropers ####################################################################
for improper_class in system.improper_classes:
if improper_class == OplsImproperTerm:
logger.info('Setting up periodic impropers...')
iforce = mm.CustomTorsionForce('k*(1-cos(2*theta))')
iforce.addPerTorsionParameter('k')
for improper in top.impropers:
iterm = system.improper_terms[id(improper)]
if type(iterm) == OplsImproperTerm:
# in OPLS convention, the third atom is the central atom
iforce.addTorsion(improper.atom2.id, improper.atom3.id,
improper.atom1.id, improper.atom4.id,
[iterm.k])
elif improper_class == HarmonicImproperTerm:
logger.info('Setting up harmonic impropers...')
iforce = mm.CustomTorsionForce(f'k*min(dtheta,2*pi-dtheta)^2;'
f'dtheta=abs(theta-phi0);'
f'pi={PI}')
iforce.addPerTorsionParameter('phi0')
iforce.addPerTorsionParameter('k')
for improper in top.impropers:
iterm = system.improper_terms[id(improper)]
if type(iterm) == HarmonicImproperTerm:
iforce.addTorsion(improper.atom1.id, improper.atom2.id,
improper.atom3.id, improper.atom4.id,
[iterm.phi * PI / 180, iterm.k])
else:
raise Exception(
'Improper terms other that PeriodicImproperTerm and '
'HarmonicImproperTerm haven\'t been implemented')
iforce.setUsesPeriodicBoundaryConditions(system.use_pbc)
iforce.setForceGroup(ForceGroup.IMPROPER)
omm_system.addForce(iforce)
### Set up non-bonded interactions #########################################################
# NonbonedForce is not flexible enough. Use it only for Coulomb interactions (including 1-4 Coulomb exceptions)
# CustomNonbondedForce handles vdW interactions (including 1-4 LJ exceptions)
cutoff = ff.vdw_cutoff
logger.info('Setting up Coulomb interactions...')
nbforce = mm.NonbondedForce()
if system.use_pbc:
nbforce.setNonbondedMethod(mm.NonbondedForce.PME)
nbforce.setEwaldErrorTolerance(5E-4)
nbforce.setCutoffDistance(cutoff)
# dispersion will be handled by CustomNonbondedForce
nbforce.setUseDispersionCorrection(False)
try:
nbforce.setExceptionsUsePeriodicBoundaryConditions(True)
except:
logger.warning('Cannot apply PBC for Coulomb 1-4 exceptions')
else:
nbforce.setNonbondedMethod(mm.NonbondedForce.NoCutoff)
nbforce.setForceGroup(ForceGroup.COULOMB)
omm_system.addForce(nbforce)
for atom in top.atoms:
nbforce.addParticle(atom.charge, 1.0, 0.0)
### Set up vdW interactions #########################################################
atom_types = list(ff.atom_types.values())
type_names = list(ff.atom_types.keys())
n_type = len(atom_types)
for vdw_class in system.vdw_classes:
if vdw_class == LJ126Term:
logger.info('Setting up LJ-12-6 vdW interactions...')
if system.use_pbc and ff.vdw_long_range == ForceField.VDW_LONGRANGE_SHIFT:
invRc6 = 1 / cutoff**6
cforce = mm.CustomNonbondedForce(
f'A(type1,type2)*(invR6*invR6-{invRc6 * invRc6})-'
f'B(type1,type2)*(invR6-{invRc6});'
f'invR6=1/r^6')
else:
cforce = mm.CustomNonbondedForce(
'A(type1,type2)*invR6*invR6-B(type1,type2)*invR6;'
'invR6=1/r^6')
cforce.addPerParticleParameter('type')
A_list = [0.0] * n_type * n_type
B_list = [0.0] * n_type * n_type
for i, atype1 in enumerate(atom_types):
for j, atype2 in enumerate(atom_types):
vdw = ff.get_vdw_term(atype1, atype2)
if type(vdw) == LJ126Term:
A = 4 * vdw.epsilon * vdw.sigma**12
B = 4 * vdw.epsilon * vdw.sigma**6
else:
A = B = 0
A_list[i + n_type * j] = A
B_list[i + n_type * j] = B
cforce.addTabulatedFunction(
'A', mm.Discrete2DFunction(n_type, n_type, A_list))
cforce.addTabulatedFunction(
'B', mm.Discrete2DFunction(n_type, n_type, B_list))
for atom in top.atoms:
id_type = type_names.index(atom.type)
cforce.addParticle([id_type])
elif vdw_class == MieTerm:
logger.info('Setting up Mie vdW interactions...')
if system.use_pbc and ff.vdw_long_range == ForceField.VDW_LONGRANGE_SHIFT:
cforce = mm.CustomNonbondedForce(
'A(type1,type2)/r^REP(type1,type2)-'
'B(type1,type2)/r^ATT(type1,type2)-'
'SHIFT(type1,type2)')
else:
cforce = mm.CustomNonbondedForce(
'A(type1,type2)/r^REP(type1,type2)-'
'B(type1,type2)/r^ATT(type1,type2)')
cforce.addPerParticleParameter('type')
A_list = [0.0] * n_type * n_type
B_list = [0.0] * n_type * n_type
REP_list = [0.0] * n_type * n_type
ATT_list = [0.0] * n_type * n_type
SHIFT_list = [0.0] * n_type * n_type
for i, atype1 in enumerate(atom_types):
for j, atype2 in enumerate(atom_types):
vdw = ff.get_vdw_term(atype1, atype2)
if type(vdw) == MieTerm:
A = vdw.factor_energy(
) * vdw.epsilon * vdw.sigma**vdw.repulsion
B = vdw.factor_energy(
) * vdw.epsilon * vdw.sigma**vdw.attraction
REP = vdw.repulsion
ATT = vdw.attraction
SHIFT = A / cutoff**REP - B / cutoff**ATT
else:
A = B = REP = ATT = SHIFT = 0
A_list[i + n_type * j] = A
B_list[i + n_type * j] = B
REP_list[i + n_type * j] = REP
ATT_list[i + n_type * j] = ATT
SHIFT_list[i + n_type * j] = SHIFT
cforce.addTabulatedFunction(
'A', mm.Discrete2DFunction(n_type, n_type, A_list))
cforce.addTabulatedFunction(
'B', mm.Discrete2DFunction(n_type, n_type, B_list))
cforce.addTabulatedFunction(
'REP', mm.Discrete2DFunction(n_type, n_type, REP_list))
cforce.addTabulatedFunction(
'ATT', mm.Discrete2DFunction(n_type, n_type, ATT_list))
if system.use_pbc and ff.vdw_long_range == ForceField.VDW_LONGRANGE_SHIFT:
cforce.addTabulatedFunction(
'SHIFT',
mm.Discrete2DFunction(n_type, n_type, SHIFT_list))
for atom in top.atoms:
id_type = type_names.index(atom.type)
cforce.addParticle([id_type])
else:
raise Exception('vdW terms other than LJ126Term and MieTerm '
'haven\'t been implemented')
if system.use_pbc:
cforce.setNonbondedMethod(
mm.CustomNonbondedForce.CutoffPeriodic)
cforce.setCutoffDistance(cutoff)
if ff.vdw_long_range == ForceField.VDW_LONGRANGE_CORRECT:
cforce.setUseLongRangeCorrection(True)
else:
cforce.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
cforce.setForceGroup(ForceGroup.VDW)
omm_system.addForce(cforce)
### Set up 1-2, 1-3 and 1-4 exceptions ##################################################
logger.info('Setting up 1-2, 1-3 and 1-4 exceptions...')
custom_nb_forces = [
f for f in omm_system.getForces()
if type(f) == mm.CustomNonbondedForce
]
pair12, pair13, pair14 = top.get_12_13_14_pairs()
for atom1, atom2 in pair12 + pair13:
nbforce.addException(atom1.id, atom2.id, 0.0, 1.0, 0.0)
for f in custom_nb_forces:
f.addExclusion(atom1.id, atom2.id)
# As long as 1-4 LJ OR Coulomb need to be scaled, then this pair should be excluded from ALL non-bonded forces.
# This is required by OpenMM's internal implementation.
# Even though NonbondedForce can handle 1-4 vdW, we use it only for 1-4 Coulomb.
# And use CustomBondForce to handle 1-4 vdW, which makes it more clear for energy decomposition.
if ff.scale_14_vdw != 1 or ff.scale_14_coulomb != 1:
pair14_forces = {} # {VdwTerm: mm.NbForce}
for atom1, atom2 in pair14:
charge_prod = atom1.charge * atom2.charge * ff.scale_14_coulomb
nbforce.addException(atom1.id, atom2.id, charge_prod, 1.0, 0.0)
for f in custom_nb_forces:
f.addExclusion(atom1.id, atom2.id)
if ff.scale_14_vdw == 0:
continue
vdw = ff.get_vdw_term(ff.atom_types[atom1.type],
ff.atom_types[atom2.type])
# We generalize LJ126Term and MieTerm because of minimal computational cost for 1-4 vdW
if type(vdw) in (LJ126Term, MieTerm):
cbforce = pair14_forces.get(MieTerm)
if cbforce is None:
cbforce = mm.CustomBondForce(
'C*epsilon*((sigma/r)^n-(sigma/r)^m);'
'C=n/(n-m)*(n/m)^(m/(n-m))')
cbforce.addPerBondParameter('epsilon')
cbforce.addPerBondParameter('sigma')
cbforce.addPerBondParameter('n')
cbforce.addPerBondParameter('m')
cbforce.setUsesPeriodicBoundaryConditions(
system.use_pbc)
cbforce.setForceGroup(ForceGroup.VDW)
omm_system.addForce(cbforce)
pair14_forces[MieTerm] = cbforce
epsilon = vdw.epsilon * ff.scale_14_vdw
if type(vdw) == LJ126Term:
cbforce.addBond(atom1.id, atom2.id,
[epsilon, vdw.sigma, 12, 6])
elif type(vdw) == MieTerm:
cbforce.addBond(atom1.id, atom2.id, [
epsilon, vdw.sigma, vdw.repulsion, vdw.attraction
])
else:
raise Exception(
'1-4 scaling for vdW terms other than LJ126Term and MieTerm '
'haven\'t been implemented')
### Set up Drude particles ##############################################################
for polar_class in system.polarizable_classes:
if polar_class == DrudeTerm:
logger.info('Setting up Drude polarizations...')
pforce = mm.DrudeForce()
pforce.setForceGroup(ForceGroup.DRUDE)
omm_system.addForce(pforce)
parent_idx_thole = {
} # {parent: (index in DrudeForce, thole)} for addScreenPair
for parent, drude in system.drude_pairs.items():
pterm = system.polarizable_terms[parent]
n_H = len([
atom for atom in parent.bond_partners
if atom.symbol == 'H'
])
alpha = pterm.alpha + n_H * pterm.merge_alpha_H
idx = pforce.addParticle(drude.id, parent.id, -1, -1, -1,
drude.charge, alpha, 0, 0)
parent_idx_thole[parent] = (idx, pterm.thole)
# exclude the non-boned interactions between Drude and parent
# and those concerning Drude particles in 1-2 and 1-3 pairs
# pairs formed by real atoms have already been handled above
# also apply thole screening between 1-2 and 1-3 Drude dipole pairs
drude_exclusions = list(system.drude_pairs.items())
for atom1, atom2 in pair12 + pair13:
drude1 = system.drude_pairs.get(atom1)
drude2 = system.drude_pairs.get(atom2)
if drude1 is not None:
drude_exclusions.append((drude1, atom2))
if drude2 is not None:
drude_exclusions.append((atom1, drude2))
if drude1 is not None and drude2 is not None:
drude_exclusions.append((drude1, drude2))
idx1, thole1 = parent_idx_thole[atom1]
idx2, thole2 = parent_idx_thole[atom2]
pforce.addScreenedPair(idx1, idx2,
(thole1 + thole2) / 2)
for a1, a2 in drude_exclusions:
nbforce.addException(a1.id, a2.id, 0, 1.0, 0)
for f in custom_nb_forces:
f.addExclusion(a1.id, a2.id)
# scale the non-boned interactions concerning Drude particles in 1-4 pairs
# pairs formed by real atoms have already been handled above
drude_exceptions14 = []
for atom1, atom2 in pair14:
drude1 = system.drude_pairs.get(atom1)
drude2 = system.drude_pairs.get(atom2)
if drude1 is not None:
drude_exceptions14.append((drude1, atom2))
if drude2 is not None:
drude_exceptions14.append((atom1, drude2))
if drude1 is not None and drude2 is not None:
drude_exceptions14.append((drude1, drude2))
for a1, a2 in drude_exceptions14:
charge_prod = a1.charge * a2.charge * ff.scale_14_coulomb
nbforce.addException(a1.id, a2.id, charge_prod, 1.0, 0.0)
for f in custom_nb_forces:
f.addExclusion(a1.id, a2.id)
else:
raise Exception(
'Polarizable terms other that DrudeTerm haven\'t been implemented'
)
### Set up virtual sites ################################################################
if top.has_virtual_site:
logger.info('Setting up virtual sites...')
for atom in top.atoms:
vsite = atom.virtual_site
if type(vsite) == TIP4PSite:
O, H1, H2 = vsite.parents
coeffs = system.get_TIP4P_linear_coeffs(atom)
omm_vsite = mm.ThreeParticleAverageSite(
O.id, H1.id, H2.id, *coeffs)
omm_system.setVirtualSite(atom.id, omm_vsite)
elif vsite is not None:
raise Exception(
'Virtual sites other than TIP4PSite haven\'t been implemented'
)
# exclude the non-boned interactions between virtual sites and parents
# and particles (atoms, drude particles, virtual sites) in 1-2 and 1-3 pairs
# TODO Assume no more than one virtual site is attached to each atom
vsite_exclusions = list(system.vsite_pairs.items())
for atom, vsite in system.vsite_pairs.items():
drude = system.drude_pairs.get(atom)
if drude is not None:
vsite_exclusions.append((vsite, drude))
for atom1, atom2 in pair12 + pair13:
vsite1 = system.vsite_pairs.get(atom1)
vsite2 = system.vsite_pairs.get(atom2)
drude1 = system.drude_pairs.get(atom1)
drude2 = system.drude_pairs.get(atom2)
if vsite1 is not None:
vsite_exclusions.append((vsite1, atom2))
if drude2 is not None:
vsite_exclusions.append((vsite1, drude2))
if vsite2 is not None:
vsite_exclusions.append((vsite2, atom1))
if drude1 is not None:
vsite_exclusions.append((vsite2, drude1))
if None not in [vsite1, vsite2]:
vsite_exclusions.append((vsite1, vsite2))
for a1, a2 in vsite_exclusions:
nbforce.addException(a1.id, a2.id, 0, 1.0, 0)
for f in custom_nb_forces:
f.addExclusion(a1.id, a2.id)
# scale the non-boned interactions between virtual sites and particles in 1-4 pairs
# TODO Assume no 1-4 LJ interactions on virtual sites
vsite_exceptions14 = []
for atom1, atom2 in pair14:
vsite1 = system.vsite_pairs.get(atom1)
vsite2 = system.vsite_pairs.get(atom2)
drude1 = system.drude_pairs.get(atom1)
drude2 = system.drude_pairs.get(atom2)
if vsite1 is not None:
vsite_exceptions14.append((vsite1, atom2))
if drude2 is not None:
vsite_exceptions14.append((vsite1, drude2))
if vsite2 is not None:
vsite_exceptions14.append((vsite2, atom1))
if drude1 is not None:
vsite_exceptions14.append((vsite2, drude1))
if None not in [vsite1, vsite2]:
vsite_exceptions14.append((vsite1, vsite2))
for a1, a2 in vsite_exceptions14:
charge_prod = a1.charge * a2.charge * ff.scale_14_coulomb
nbforce.addException(a1.id, a2.id, charge_prod, 1.0, 0.0)
for f in custom_nb_forces:
f.addExclusion(a1.id, a2.id)
### Remove COM motion ###################################################################
logger.info('Setting up COM motion remover...')
omm_system.addForce(mm.CMMotionRemover(10))
return omm_system
def get_Ucv_force(n_beads, Ucv_ext, cv_grid_ext, cv_coeff, cv_mean, feat_types, feat_idxs):
feature_funcs = {"dist":"distance(p{}, p{})", "invdist":"(1/distance(p{}, p{}))",
"angle":"angle(p{}, p{}, p{})", "dih":"dihedral(p{}, p{}, p{}, p{})"}
cv_expr = "Table("
pr_cv_expr = "Table("
#cv_expr = ""
#pr_cv_expr = ""
feat_idx = 0
atm_to_p_idx = []
p_idx = -np.ones(n_beads, int)
curr_p_idx = 1
for i in range(len(feat_types)):
feat_fun = feature_funcs[feat_types[i]]
for j in range(cv_coeff.shape[0]):
idxs = feat_idxs[i][j]
expr_idxs = []
for n in range(len(idxs)):
if p_idx[idxs[n]] == -1:
p_idx[idxs[n]] = curr_p_idx
atm_to_p_idx.append(int(idxs[n]))
curr_p_idx += 1
expr_idxs.append(p_idx[idxs[n]])
#print("{} -> {} {} -> {}".format(idxs[0], expr_idxs[0], idxs[1], expr_idxs[1])) # DEBUGGING
feat_explicit = feat_fun.format(*expr_idxs)
feat_coeff = cv_coeff[feat_idx,0]
feat_mean = cv_mean[feat_idx]
if feat_idx > 0:
if feat_coeff < 0:
cv_expr += " - {:.5f}*(".format(abs(feat_coeff))
pr_cv_expr += "\n - {:.5f}*(".format(abs(feat_coeff))
else:
cv_expr += " + {:.5f}*(".format(abs(feat_coeff))
pr_cv_expr += "\n + {:.5f}*(".format(abs(feat_coeff))
else:
if feat_coeff < 0:
cv_expr += "-{:.5f}*(".format(abs(feat_coeff))
pr_cv_expr += "-{:.5f}*(".format(abs(feat_coeff))
else:
cv_expr += "{:.5f}*(".format(abs(feat_coeff))
pr_cv_expr += "{:.5f}*(".format(abs(feat_coeff))
cv_expr += feat_explicit
pr_cv_expr += feat_explicit
if feat_mean < 0:
cv_expr += " + {:.5f})".format(abs(feat_mean))
pr_cv_expr += " + {:.5f})".format(abs(feat_mean))
else:
cv_expr += " - {:.5f})".format(abs(feat_mean))
pr_cv_expr += " - {:.5f})".format(abs(feat_mean))
feat_idx += 1
cv_expr += ");"
pr_cv_expr += ");"
#cv_expr += ";"
#pr_cv_expr += ";"
Ucv_force = omm.CustomCompoundBondForce(n_beads, cv_expr)
Ucv_force.addBond(atm_to_p_idx)
#Ucv_force.addFunction("Table", Ucv_ext, cv_grid_ext[0], cv_grid_ext[-1])
Table_func = omm.Continuous1DFunction(Ucv_ext, cv_grid_ext[0], cv_grid_ext[-1])
Ucv_force.addTabulatedFunction("Table", Table_func)
return Ucv_force, pr_cv_expr
quartets.append([t[0], t[1], t[2], k])
quintents = []
for q in quartets:
for k in heavy_atoms:
if (k not in q) and all(is_neighbor[atom_index[atom], atom_index[k]]
for atom in q):
quintents.append([q[0], q[1], q[2], q[3], k])
end = time.time()
print(f"Runtime is {end - start}")
height = 1
pairs_energy = 'height*p^2*(pi/(alpha1+alpha2))^(3/2)*exp(-alpha1*alpha2*distance(p1,p2)^2/(alpha1 + alpha2)); p=2.7; pi=3.14159265'
pairs_vol = openmm.CustomCompoundBondForce(2, pairs_energy)
pairs_vol.addPerBondParameter('alpha1')
pairs_vol.addPerBondParameter('alpha2')
pairs_vol.addGlobalParameter('height', height)
for pair in pairs:
pairs_vol.addBond([pair[0], pair[1]], [alpha[pair[0]], alpha[pair[1]]])
triplets_energy = '-height*p^3*(pi/(alpha1+alpha2+alpha3))^(3/2)*exp(-((alpha1*alpha2)*distance(p1,p2)^2 + (alpha1*alpha3)*distance(p1,p3)^2 + (alpha2*alpha3)*distance(p2,p3)^2)/(alpha1 + alpha2 + alpha3) ); p=2.7; pi=3.14159265'
triplets_vol = openmm.CustomCompoundBondForce(3, triplets_energy)
triplets_vol.addPerBondParameter('alpha1')
triplets_vol.addPerBondParameter('alpha2')
triplets_vol.addPerBondParameter('alpha3')
triplets_vol.addGlobalParameter('height', height)
for triplet in triplets:
triplets_vol.addBond(
[triplet[0], triplet[1], triplet[2]],