import os
from contextlib import ExitStack
from contextlib import contextmanager
from typing import Iterable
from typing import List
from typing import NamedTuple
from typing import Optional
import pymmlibs
from schrodinger import structure
from schrodinger.forcefield import OPLSVersion
from schrodinger.forcefield import common
from schrodinger.forcefield import constants
from schrodinger.infra import mm
from schrodinger.infra import structure as cppstructure
from schrodinger.structutils import analyze
from schrodinger.utils import log
logger = log.get_output_logger(__file__)
[docs]class Restraint:
"""
Class representing restraints.
"""
[docs] def __init__(self, force_constant, target_value, atom_index1, atom_index2,
atom_index3, atom_index4, flat_bottom_half_width):
self.force_constant = float(force_constant)
self.target_value = float(target_value)
self.atom_index1 = int(atom_index1)
self.atom_index2 = int(atom_index2)
self.atom_index3 = int(atom_index3)
self.atom_index4 = int(atom_index4)
self.flat_bottom_half_width = float(flat_bottom_half_width)
[docs]class Constraint:
"""
Class representing constraints.
"""
[docs] def __init__(self,
atom_index1,
atom_index2,
type_index1,
type_index2,
use_rigorous=False):
self.atom_index1 = int(atom_index1)
self.atom_index2 = int(atom_index2)
self.type_index1 = int(type_index1)
self.type_index2 = int(type_index2)
self.use_rigorous = int(use_rigorous)
if self.type_index1 <= 0 and self.type_index2 == 1:
self.torsion_type = mm.MMFfldTorsionCoordCons
else:
self.torsion_type = mm.MMFfldTorsionCons
[docs]class MinimizationOptions(NamedTuple):
"""
Class for assigning the mmffld minimization options.
"""
opls_version: int = OPLSVersion.F16
nonbond_cutoff: float = mm.DEFAULT_SYS_CUTOFF
max_step: int = mm.DEFAULT_MAX_STEP
energy_convergence: float = mm.DEFAULT_CONV_ENER
gradient_convergence: float = mm.DEFAULT_CONV_GRAD
line_search_method: int = mm.MMFfldMinGradLS
min_method: int = mm.MMFfldMinAUTO
no_restrain_zob: bool = False
bend_conj_amines: bool = False
perturb: bool = False
restraints: Optional[List[Restraint]] = None
constraints: Optional[List[Constraint]] = None
debug_outfile: Optional[str] = None
[docs]def minimize_structure(st: structure.Structure,
options: Optional[MinimizationOptions] = None,
archive_path: Optional[os.PathLike] = None):
"""
Minimizes a single structure.
For minimizing multiple structures, the performance improves drastically if
the forcefield handle instantiation happens outside the iteration loop.
:param st: structure
:param options: mmffld minimization options
:param archive_path: path to specific .opls file
:return: minimization results
:rtype: mm.MinimizeResults
"""
options = options or MinimizationOptions()
with common.opls_force_field(version=options.opls_version,
archive_path=archive_path) \
as mmffld_handle:
with _minimization_stack(mmffld_handle, st, options=options):
min_res = mm.mmffld_minimize(mmffld_handle)
if options.debug_outfile:
_write_debug_outfile(mmffld_handle, options.debug_outfile)
return min_res
[docs]def iterate_minimize_structure(st, iteration, options=None):
"""
Minimize a single structure multiple times.
This function instantiates the forcefield handle before the iteration
loop and hence increases the performance by eliminating instantiation
repetitions.
:param st: structure
:type st: schrodinger.structure object
:param iteration: iteration number
:type iteration: int
:param options: mmffld minimization options
:type options: MinimizationOptions namedtuple
:return: minimization results
:rtype: mm.MinimizeResults
"""
options = options or MinimizationOptions()
with common.opls_force_field():
for _ in range(iteration):
min_res = minimize_structure(st, options=options)
yield min_res
[docs]def minimize_substructure(
st: structure.Structure,
atoms_to_minimize: Iterable[structure._StructureAtom]
) -> structure.Structure:
"""
Given a structure, minimizes the specified atoms in it, and returns it
:param st: The structure containing the atoms to be minimized
:param atoms_to_minimize: The subset of atoms in the structure to be
minimize
:return: The structure with the atoms minimized
"""
atoms_to_minimize = set(atoms_to_minimize)
with common.opls_force_field() as mmffld_handle:
with common.assign_force_field(mmffld_handle, st):
for at in st.atom:
if at in atoms_to_minimize:
continue
mm.mmffld_enterRestraint(mmffld_handle, mm.MMFfldPosFrozen, st,
at.index, 0, 0, 0, 0.0, 0.0, 0.0)
min_res = mm.mmffld_minimize(mmffld_handle)
return st
[docs]def minimize_ligands(st: structure.Structure, **kwargs) -> structure.Structure:
"""
Given a structure, returns that structure with all the ligands minimized
:param st: The structure containing ligands to be minimized
:param `**kwargs`: Any options to supply to `analyze.find_ligands`
:return: The structure with the ligands minimized
"""
ligs = analyze.find_ligands(st, **kwargs)
atom_indices_per_ligand = (lig.atom_indexes for lig in ligs)
atom_indices_in_structure = set(*atom_indices_per_ligand)
ligand_atoms = {st.atom[i] for i in atom_indices_in_structure}
return minimize_substructure(st, ligand_atoms)
[docs]def write_restraints(filename, restraints):
"""
Writes restraints into filename.
:param filename: restraint filename
:type filename: string
:param restraints: a list to be filled with restraints
:type restraints: list of Restraints
"""
rest_lines = []
rst_format = '%13.8lf %13.6lf %4d %4d %4d %4d %10.4lf \n'
if not all(isinstance(rest, Restraint) for rest in restraints):
raise RuntimeError('restraint object should be a member'
' of minimizer.Restraint class. ')
for rest in restraints:
rest_line = rst_format % (rest.force_constant, rest.target_value,
rest.atom_index1, rest.atom_index2,
rest.atom_index3, rest.atom_index4,
rest.flat_bottom_half_width)
if rest_line not in rest_lines:
rest_lines.append(rest_line)
with open(filename, 'w') as fh:
fh.writelines(rest_lines)
[docs]def read_restraints(filename, restraints=None, constraints=None):
"""
Reads restraints and constraints arguments
from a file and creates separate lists for
Restraint and Constraint objects.
:param filename: restraint/constraint filename
:type filename: string
:param restraints: a list to be filled with restraints
:type restraints: list of Restraints
:param constraints: a list to be filled with constraints
:type constraints: list of Constraints
"""
if restraints is None and constraints is None:
raise RuntimeError('missing required positional arguments:'
' restraints and/or constraints')
with open(filename) as fh:
list_args = [line.split() for line in fh]
for list_arg in list_args:
if restraints is not None and len(
list_arg) == constants.RESTRAINT_LENGTH:
restraints.append(Restraint(*list_arg))
elif constraints is not None and len(
list_arg) == constants.CONSTRAINT_LENGTH:
constraints.append(Constraint(*list_arg))
else:
raise RuntimeError(
f'Unknown restraint or constraint format: {list_arg}')
def _write_debug_outfile(mmffld_handle, outfile):
"""
Writes output files if write_output in minimize_structure is set to True.
:param mmffld_handle: force field handle
:type mmffld_handle: integer
:param outfile: mae output filename
:type outfile: string
"""
pymmlibs._mmffld_openEncryptedOutputFile(mmffld_handle, outfile)
mm.mmffld_printDetailedEnergyComponents(mmffld_handle)
pymmlibs._mmffld_writeEncryptedOutputFile(mmffld_handle)
def _set_mmffld_box(mmffld_handle,
bx=mm.INVALID_SYS_BOX_DIMENSION,
by=mm.INVALID_SYS_BOX_DIMENSION,
bz=mm.INVALID_SYS_BOX_DIMENSION,
alpha=mm.DEFAULT_SYS_BOX_ANGLE,
beta=mm.DEFAULT_SYS_BOX_ANGLE,
gamma=mm.DEFAULT_SYS_BOX_ANGLE,
pbc_bonds=mm.DEFAULT_TYPER_ADD_PBC_BONDS):
"""
Sets the PBC box information.
:param mmffld_handle: force field handle
:type mmffld_handle: integer
:params bx, by, bz: box dimensions
:type bx, by, bz: float
:params alpha, beta, gamma: box angles
:type alpha, beta, gamma: float
:param pbc_bonds: setup PBC for infinite systems
:type pbc_bonds: int
"""
mm.mmffld_setOption(mmffld_handle, mm.MMFfldOption_SYS_BX, 0, bx, "")
mm.mmffld_setOption(mmffld_handle, mm.MMFfldOption_SYS_BY, 0, by, "")
mm.mmffld_setOption(mmffld_handle, mm.MMFfldOption_SYS_BZ, 0, bz, "")
mm.mmffld_setOption(mmffld_handle, mm.MMFfldOption_SYS_ALPHA, 0, alpha, "")
mm.mmffld_setOption(mmffld_handle, mm.MMFfldOption_SYS_BETA, 0, beta, "")
mm.mmffld_setOption(mmffld_handle, mm.MMFfldOption_SYS_GAMMA, 0, gamma, "")
mm.mmffld_setOption(mmffld_handle, mm.MMFfldOption_TYPER_addPBCbonds,
pbc_bonds, 0.0, "")
@contextmanager
def _apply_restraints(mmffld_handle, restraints):
"""
A context manager to apply the restraints to the mmffld_handle
:param mmffld_handle: force field handle
:type mmffld_handle: integer
:param restraints: minimization restraint list
:type restraints: list of Restraints
"""
for restraint in restraints or []:
mm.mmffld_enterRestraint(
mmffld_handle, mm.MMFfldTorsionRest, mm.MMCT_INVALID_CT,
restraint.atom_index1, restraint.atom_index2, restraint.atom_index3,
restraint.atom_index4, restraint.force_constant,
restraint.target_value, restraint.flat_bottom_half_width)
try:
yield
finally:
for restraint in restraints or []:
mm.mmffld_deleteRestraint(mmffld_handle, mm.MMFfldTorsionRest,
mm.MMCT_INVALID_CT, restraint.atom_index1,
restraint.atom_index2,
restraint.atom_index3,
restraint.atom_index4)
@contextmanager
def _apply_constraints(mmffld_handle, constraints):
"""
A context manager to apply the constraints to the mmffld_handle
:param mmffld_handle: force field handle
:type mmffld_handle: integer
:param constraints: minimization constraint list
:type constraints: list of Constraints
"""
for constraint in constraints or []:
mm.mmffld_enterConstraint(mmffld_handle, constraint.torsion_type,
mm.MMCT_INVALID_CT, constraint.atom_index1,
constraint.atom_index2,
constraint.type_index1,
constraint.type_index2,
constraint.use_rigorous)
try:
yield
finally:
for constraint in constraints or []:
mm.mmffld_deleteConstraint(mmffld_handle, constraint.torsion_type,
mm.MMCT_INVALID_CT,
constraint.atom_index1,
constraint.atom_index2,
constraint.type_index1,
constraint.type_index2)
@contextmanager
def _apply_pbc_options(mmffld_handle, st):
"""
A context manager to apply periodic boundary conditions to mmffld_handle
:param mmffld_handle: force field handle
:type mmffld_handle: int
:param st: structure
:type st: structure.Structure
"""
try:
pbc = cppstructure.PBC(st)
except IndexError as err:
yield
return
bx, by, bz = pbc.getBoxLengths()
alpha, beta, gamma = pbc.getBoxAngles()
_set_mmffld_box(mmffld_handle,
bx=bx,
by=by,
bz=bz,
alpha=alpha,
beta=beta,
gamma=gamma,
pbc_bonds=1)
max_cutoff = pbc.getHalfShortOrthogonalBoxLength()
_, nonbond_cutoff, _ = mm.mmffld_getOption(mmffld_handle,
mm.MMFfldOption_SYS_CUTOFF)
if max_cutoff < nonbond_cutoff:
raise RuntimeError(
f'The nonbonded cutoff {nonbond_cutoff} cannot be '
f'larger than half the smallest box size {max_cutoff}')
try:
yield
finally:
# restore default values:
_set_mmffld_box(mmffld_handle)
@contextmanager
def _assign_option(mmffld_handle, option, value):
"""
A context manager to assign the value of an option.
:param mmffld_handle: force field handle
:type mmffld_handle: integer
:param option: mmffld option
:type option: int
:param value: value of the option
:type value: integer, float, or string
"""
init_iopt, init_dopt, init_sopt = mm.mmffld_getOption(mmffld_handle, option)
if isinstance(value, int):
mm.mmffld_setOption(mmffld_handle, option, value, 0.0, "")
if isinstance(value, float):
mm.mmffld_setOption(mmffld_handle, option, 0, value, "")
if isinstance(value, str):
mm.mmffld_setOption(mmffld_handle, option, 0, 0.0, value)
try:
yield
finally:
mm.mmffld_setOption(mmffld_handle, option, init_iopt, init_dopt,
str(init_sopt or ''))
@contextmanager
def _minimization_stack(mmffld_handle, st, options):
"""
A context manager to combine different minimization context managers
:param mmffld_handle: force field handle
:type mmffld_handle: int
:param st: structure
:type st: structure.Structure
:param options: mmffld minimization options
:type options: MinimizationOptions namedtuple
"""
options_dict = {
mm.MMFfldOption_SYS_CUTOFF: options.nonbond_cutoff,
mm.MMFfldOption_MIN_PERTURB: int(options.perturb),
mm.MMFfldOption_MIN_MAX_STEP: options.max_step,
mm.MMFfldOption_MIN_CONV_ENER: options.energy_convergence,
mm.MMFfldOption_MIN_CONV_GRAD: options.gradient_convergence,
mm.MMFfldOption_MIN_LS_METHOD: options.line_search_method,
mm.MMFfldOption_MIN_METHOD: options.min_method,
mm.MMFfldOption_TYPER_NO_RESTRAIN_ZOB: int(options.no_restrain_zob),
mm.MMFfldOption_TYPER_BEND_CONJ_AMINES: int(options.bend_conj_amines)
}
with ExitStack() as stack:
for key, val in options_dict.items():
stack.enter_context(_assign_option(mmffld_handle, key, val))
stack.enter_context(_apply_pbc_options(mmffld_handle, st))
stack.enter_context(common.assign_force_field(mmffld_handle, st))
stack.enter_context(_apply_restraints(mmffld_handle,
options.restraints))
stack.enter_context(
_apply_constraints(mmffld_handle, options.constraints))
yield