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 Union

from schrodinger import structure
from schrodinger.infra import mm
from schrodinger.forcefield.ffld_options import ForceFieldOptions
from schrodinger.forcefield.ffld_options import yield_enum_and_typed_val

ADVANCED_FFB_ENV = "SCHRODINGER_FFBUILDER_ADVANCED"

ForceFieldSpec = Union[int, ForceFieldOptions]


# 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


@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: options that affect ffld params
    :param kwargs: Allows passing `ForceFieldOptions` keywords seaparately to
                   the function. See keywords in `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:
        ffld_options = ForceFieldOptions(**kwargs)

    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()

    for enum, typed_val in yield_enum_and_typed_val(ffld_options):
        set_mmffld_option(handle, enum, typed_val)

    yield handle


def set_mmffld_option(mmffld_handle: int, mm_option_enum: int,
                      typed_value: Union[int, float, str]):
    """
    Assign MMFfld value for a given handle and value.

    :param mmffld_handle: force field handle number
    :param mm_option_enum: enumerated name of MMFfld option for a parameter
    :param typed_value: parameter value with explicit type setting
    """
    init_iopt, init_dopt, init_sopt = mm.mmffld_getOption(
        mmffld_handle, mm_option_enum)
    if isinstance(typed_value, int):
        mm.mmffld_setOption(mmffld_handle, mm_option_enum, typed_value, 0.0, "")
    if isinstance(typed_value, float):
        mm.mmffld_setOption(mmffld_handle, mm_option_enum, 0, typed_value, "")
    if isinstance(typed_value, str):
        mm.mmffld_setOption(mmffld_handle, mm_option_enum, 0, 0.0, typed_value)

    return init_iopt, init_dopt, init_sopt


@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