Source code for schrodinger.forcefield.ffld_options

"""
Options for force field modules

Copyright Schrodinger LLC, All Rights Reserved.
"""

import os
from typing import List
from typing import NamedTuple
from typing import Optional
from typing import TYPE_CHECKING
from typing import Union
from schrodinger.forcefield import OPLSVersion
from schrodinger.infra import mm

if TYPE_CHECKING:
    from schrodinger.forcefield.minimizer import Restraint
    from schrodinger.forcefield.minimizer import Constraint


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

    All options have mappings to mmffld enum names to enforce its values except
    'version' and 'archive_path' which do not need to be mapped.
    """

    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

    _options_map = {
        'bend_conj_amines': mm.MMFfldOption_TYPER_BEND_CONJ_AMINES,
        'no_cm1a_bcc': mm.MMFfldOption_OPT_NO_CM1ABCC,
        'charges_from_ct': mm.MMFfldOption_ENE_CHARGES_FROM_CT,
        'no_restrain_zob': mm.MMFfldOption_TYPER_NO_RESTRAIN_ZOB,
        'nonbond_cutoff': mm.MMFfldOption_SYS_CUTOFF,
    }

    _unmapped_options = ('version', 'archive_path')


class MinimizationOptions(NamedTuple):
    """
    Class for assigning the mmffld minimization options.

    All options have mappings to mmffld enum names to enforce its values except
    'restraints', 'constraints' and 'debug_outfile' which do not need to be
    mapped.
    """

    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
    perturb: bool = False
    restraints: Optional[List['Restraint']] = None
    constraints: Optional[List['Constraint']] = None
    debug_outfile: Optional[str] = None

    _options_map = {
        'perturb': mm.MMFfldOption_MIN_PERTURB,
        'max_step': mm.MMFfldOption_MIN_MAX_STEP,
        'energy_convergence': mm.MMFfldOption_MIN_CONV_ENER,
        'gradient_convergence': mm.MMFfldOption_MIN_CONV_GRAD,
        'line_search_method': mm.MMFfldOption_MIN_LS_METHOD,
        'min_method': mm.MMFfldOption_MIN_METHOD
    }

    _unmapped_options = ('restraints', 'constraints', 'debug_outfile')


FfldMappedOptions = Union[ForceFieldOptions, MinimizationOptions]


def yield_enum_and_typed_val(options: FfldMappedOptions):
    """
    Returns the enum key and type-enforced value for each option in class.

    :param options: instance of FfldMappedOptions containing the user-defined
                    options and the corresponding maps to mmffld enum names
    """

    for key, val in options._asdict().items():
        if key not in options._unmapped_options:
            enum = options._options_map[key]
            type_ = options._field_types[key]
            yield enum, type_(val)