from contextlib import ExitStack
from contextlib import contextmanager
from typing import Iterable
from typing import List
from typing import Tuple
from typing import Optional
from typing import Union
import pymmlibs
from schrodinger import structure
from schrodinger.forcefield import common
from schrodinger.forcefield.ffld_options import ForceFieldOptions
from schrodinger.forcefield.ffld_options import MinimizationOptions
from schrodinger.forcefield.ffld_options import yield_enum_and_typed_val
from schrodinger.forcefield import constants
from schrodinger.infra import mm
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] @classmethod
def from_atom_indices(
cls,
st: structure.Structure,
atom_indices: Tuple[int, int, int, int],
force_constant: float = 5.0,
flat_bottom_half_width: float = 0.0,
target_torsion_angle: Optional[float] = None,
):
"""
Given structure and atom indices, return a restraint for
mmffld minimization with default force constant, flat
bottom half width, and target torsion angle.
:param st: structure
:param atom_indices: atom indices for a torsion
:param force_constant: restraint force constant
:param flat_bottom_half_width: flat bottom half width for restraints
:param target_torsion_angle: target torsion angle in degrees
if None is given, the current torsion angle will be used
:return: torsion restraint for minimization
"""
if target_torsion_angle is None:
target_torsion_angle = st.measure(*atom_indices)
restraint = (force_constant, target_torsion_angle, *atom_indices,
flat_bottom_half_width)
return cls(*restraint)
[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]def minimize_structure(st: structure.Structure,
options: Optional[MinimizationOptions] = None,
force_field: Optional[common.ForceFieldSpec] = None):
"""
Minimizes a single structure.
For minimizing multiple structures, the performance improves drastically if
the forcefield environment instantiation happens outside the iteration loop.
:param st: structure
:param options: mmffld minimization options
:param force_field: either an existing force field handle or ForceFieldOptions,
from which a handle will be created
:return: minimization results
:rtype: mm.MinimizeResults
"""
options = options or MinimizationOptions()
if force_field is None:
force_field = ForceFieldOptions()
with ExitStack() as stack:
if isinstance(force_field, ForceFieldOptions):
mmffld_handle = stack.enter_context(
common.opls_force_field(force_field))
else:
assert isinstance(force_field, int)
mmffld_handle = force_field
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: structure.Structure,
iteration: int,
options: Optional[MinimizationOptions] = None) -> mm.MinimizeResults:
"""
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
:param iteration: iteration number
:param options: mmffld minimization options
:return: minimization results
"""
options = options or MinimizationOptions()
with common.mmffld_environment():
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: str, restraints: List[Restraint]):
"""
Writes restraints into filename.
:param filename: restraint filename
:param restraints: a list to be filled with 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: str,
restraints: Optional[List[Restraint]] = None,
constraints: Optional[List[Constraint]] = None):
"""
Reads restraints and constraints arguments
from a file and creates separate lists for
Restraint and Constraint objects.
:param filename: restraint/constraint filename
:param restraints: a list to be filled with restraints
:param constraints: a list to be filled with 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: int, outfile: str):
"""
Writes output files if write_output in minimize_structure is set to True.
:param mmffld_handle: force field handle
:param outfile: mae output filename
"""
pymmlibs._mmffld_openEncryptedOutputFile(mmffld_handle, outfile)
mm.mmffld_printDetailedEnergyComponents(mmffld_handle)
pymmlibs._mmffld_writeEncryptedOutputFile(mmffld_handle)
def _set_mmffld_box(mmffld_handle: int,
bx: float = mm.INVALID_SYS_BOX_DIMENSION,
by: float = mm.INVALID_SYS_BOX_DIMENSION,
bz: float = mm.INVALID_SYS_BOX_DIMENSION,
alpha: float = mm.DEFAULT_SYS_BOX_ANGLE,
beta: float = mm.DEFAULT_SYS_BOX_ANGLE,
gamma: float = mm.DEFAULT_SYS_BOX_ANGLE,
pbc_bonds: bool = mm.DEFAULT_TYPER_ADD_PBC_BONDS):
"""
Sets the PBC box information.
:param mmffld_handle: force field handle
:params bx, by, bz: box dimensions
:params alpha, beta, gamma: box angles
:param pbc_bonds: setup PBC for infinite systems
"""
for enum, val in {
mm.MMFfldOption_SYS_BX: bx,
mm.MMFfldOption_SYS_BY: by,
mm.MMFfldOption_SYS_BZ: bz,
mm.MMFfldOption_SYS_ALPHA: alpha,
mm.MMFfldOption_SYS_BETA: beta,
mm.MMFfldOption_SYS_GAMMA: gamma,
}.items():
mm.mmffld_setOption(mmffld_handle, enum, 0, float(val), "")
mm.mmffld_setOption(mmffld_handle, mm.MMFfldOption_TYPER_addPBCbonds,
pbc_bonds, 0.0, "")
@contextmanager
def _apply_restraints(mmffld_handle: int,
restraints: Optional[List[Restraint]]):
"""
A context manager to apply the restraints to the mmffld_handle.
:param mmffld_handle: force field handle
:param restraints: minimization restraint list
"""
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: int,
constraints: Optional[List[Constraint]]):
"""
A context manager to apply the constraints to the mmffld_handle.
:param mmffld_handle: force field handle
:param constraints: minimization constraint list
"""
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: int, st: structure.Structure):
"""
A context manager to apply periodic boundary conditions to mmffld_handle.
:param mmffld_handle: force field handle
:param st: structure
"""
pbc = st.pbc
if pbc is None:
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: int, mm_option_enum: int,
typed_value: Union[int, float, str]):
"""
A context manager to assign the value of an option.
:param mmffld_handle: force field handle
:param enum_option: mmffld option
:param typed_value: typed value of the option
"""
init_iopt, init_dopt, init_sopt = common.set_mmffld_option(
mmffld_handle, mm_option_enum, typed_value)
try:
yield
finally:
mm.mmffld_setOption(mmffld_handle, mm_option_enum, init_iopt, init_dopt,
str(init_sopt or ''))
@contextmanager
def _minimization_stack(mmffld_handle: int, st: structure.Structure,
options: MinimizationOptions):
"""
A context manager to combine different minimization context managers.
:param mmffld_handle: force field handle
:param st: structure
:param options: mmffld minimization options
"""
with ExitStack() as stack:
for enum, typed_val in yield_enum_and_typed_val(options):
stack.enter_context(_assign_option(mmffld_handle, enum, typed_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