Source code for schrodinger.forcefield.common

"""
OPLS context managers.

Copyright Schrodinger LLC, All Rights Reserved.
"""

import os
from contextlib import contextmanager
from typing import Optional
from typing import NamedTuple
from typing import Union

from schrodinger import structure
from schrodinger.infra import mm
from schrodinger.forcefield import OPLSVersion

ADVANCED_FFB_ENV = "SCHRODINGER_FFBUILDER_ADVANCED"


# Introduced for FFLD-1845
class AtomTypingFailure(mm.MmException):

    def __init__(self, msg=''):
        self.msg = msg

    def __str__(self):
        return self.msg


@contextmanager
def mmffld_environment():
    """
    A context manager for loading the mmffld environment.
    """
    mmffld_init = mm.MMFfldInitializer()
    yield


class ForceFieldOptions(NamedTuple):
    """
    Class for assigning the mmffld force field options.
    """

    version: int = OPLSVersion.F16
    archive_path: Optional[os.PathLike] = None
    bend_conj_amines: bool = False
    no_cm1a_bcc: bool = False
    charges_from_ct: bool = False
    no_restrain_zob: bool = False
    nonbond_cutoff: float = mm.DEFAULT_SYS_CUTOFF


# handle or options to create handle
ForceFieldSpec = Union[int, ForceFieldOptions]


@contextmanager
def opls_force_field(ffld_options: Optional[ForceFieldOptions] = None,
                     **kwargs) -> int:
    """
    A context manager for creating a force field handle, using
    default precedence for loading ffld datafiles.

    :param ffld_options: `common.ForceFieldOptions` object with options
                         that affect ffld params
    :param kwargs: Allows passing `ffld_option` keywords seaparately to
                   the function. See options for `common.ForceFieldOptions`
                   for available options.

    :return: mmffld handle
    """
    if ffld_options is not None and kwargs:
        raise ValueError(
            "Only one of `force_field_options` or `kwargs` (keyword "
            "arguments to `ForceFieldOptions`) should be specified")
    if ffld_options is None:
        try:
            ffld_options = ForceFieldOptions(**kwargs)
        except TypeError as e:
            raise TypeError("Unknown option obtained in opls_force_field(). "
                            "Please check class `common.ForceFieldOptions` "
                            f"for available keywords.\n{e}")

    if ffld_options.archive_path:
        mmffld_handle = mm.MMFfldHandle(os.fspath(ffld_options.archive_path),
                                        ffld_options.version)
    else:
        mmffld_handle = mm.MMFfldHandle(ffld_options.version)

    handle = mmffld_handle.getHandle()

    # CAUTION!
    # _Must_ enforce type (int, float or str) in value assignment to avoid
    # subtle hard-to-catch errors during runtime. Types not matched are
    # initialized to default value (zero or "" str) leading to potential errors.
    # bool is handled automatically.
    options_dict = {
        mm.MMFfldOption_TYPER_BEND_CONJ_AMINES: ffld_options.bend_conj_amines,
        mm.MMFfldOption_OPT_NO_CM1ABCC: ffld_options.no_cm1a_bcc,
        mm.MMFfldOption_ENE_CHARGES_FROM_CT: ffld_options.charges_from_ct,
        mm.MMFfldOption_TYPER_NO_RESTRAIN_ZOB: ffld_options.no_restrain_zob,
        mm.MMFfldOption_SYS_CUTOFF: float(ffld_options.nonbond_cutoff)
    }

    for option, value in options_dict.items():
        if isinstance(value, int):
            mm.mmffld_setOption(handle, option, value, 0.0, "")
        elif isinstance(value, float):
            mm.mmffld_setOption(handle, option, 0, value, "")
        elif isinstance(value, str):
            mm.mmffld_setOption(handle, option, 0, 0.0, value)

    yield handle


@contextmanager
def assign_force_field(mmffld_handle: mm.MMFfldHandle,
                       st: structure.Structure,
                       apply_mmlewis: bool = True):
    """
    A context manager to runs typing by loading the structure into mmffld.
    Note: If cm1a is turned on, enterMol calculates charges which is a fairly
    lengthy operation.

    :param handle: mmffld handle
    :param st: structure to assign atom types
    :param apply_mmlewis: whether to apply mmlewis to the incoming structure

    :raise AtomTypingFailure: if mmffld_enterMol() fails
    """

    try:
        mmffld_assign = mm.MMFfldAssign(mmffld_handle, st, apply_mmlewis)
    except RuntimeError as e:
        # FFLD-1845: Allow atom typing failure to be caught specifically
        raise AtomTypingFailure(msg=str(e))
    yield


def generate_partial_charges(st: structure.Structure,
                             version: mm.OPLSVersion = mm.OPLSVersion.F16):
    """
    Generates partial charges for a given structure. Results will be stored in
    the `PARTIAL_CHARGE_PROP` atom-level property.

    :param st: Structure for which partial charges are generated.
    """
    with opls_force_field(version=version) as handle:
        with assign_force_field(handle, st):
            pass