Source code for schrodinger.forcefield.custom_params

"""
Utility script used to upgrade and merge custom OPLS parameters.

Copyright Schrodinger LLC, All Rights Reserved.
"""

import argparse
import os
import shutil
import sys
from argparse import SUPPRESS
from pathlib import Path

from schrodinger.infra import mm
from schrodinger.utils import cmdline
from schrodinger.utils import env
from schrodinger.utils import log

logger = log.get_output_logger(__file__)

CREATE_TASK = "create"
UPGRADE_TASK = "upgrade"
MERGE_TASK = "merge"


def create_archive_from_oplsdir(oplsdir: str,
                                output_path: str,
                                force: bool = False,
                                test: bool = False):
    """
    Creates a new custom parameter file at the specified path containing
    datafiles found in the given OPLSDIR.
    """

    datafiles = []
    for search_dir in (
            Path(oplsdir),  # toplevel includes nonversioned datafiles
            Path(mm.get_OPLS_data_subdirectory(oplsdir, mm.OPLSVersion.F14)),
            Path(mm.get_OPLS_data_subdirectory(oplsdir, mm.OPLSVersion.F16)),
    ):
        if search_dir.is_dir():
            datafiles.extend(f for f in search_dir.iterdir() if f.is_file())

    if not datafiles:
        raise RuntimeError("No valid parameter files found.")

    # If a directory was specified, create a release-named archive within it
    if os.path.isdir(output_path):
        output_path = mm.get_archive_path(output_path)

    if force and os.path.exists(output_path):
        os.remove(output_path)

    env_value = "true" if test else None  # else unset the env var
    with env.EnvironmentContext(mm.CUSTOM_OPLS_CREATE_FOR_TEST, env_value):
        archive = mm.OPLSArchive.create(output_path)

    for datafile in datafiles:
        try:
            archive.add(datafile)
        except IndexError:
            logger.error(f"Unable to add unknown datafile: {datafile}")


def upgrade_oplsdir(oplsdir: str):
    """
    If not available, creates a new custom parameter file carrying over any
    previous compatible parameters with the current version.
    """

    archive_path = mm.get_archive_path(oplsdir)

    # If current parameters exist, validate it
    if os.path.exists(archive_path):
        mm.OPLSArchive.validate(archive_path)
        logger.info("Custom parameters are up to date.")
        return

    # Otherwise copy latest compatible parameters (and update metadata)
    compatible_archive_path = mm.get_compatible_archive_path(oplsdir)
    shutil.copyfile(compatible_archive_path, archive_path)
    mm.OPLSArchive.upgrade(archive_path)


def merge_oplsdir(in_oplsdir: str, out_oplsdir: str):
    """
    Merge the current custom parameters from one OPLSDIR into another
    """

    mm.validate_opls_directory(in_oplsdir)
    mm.validate_opls_directory(out_oplsdir)
    src_archive_path = mm.get_archive_path(in_oplsdir)
    dest_archive_path = mm.get_archive_path(out_oplsdir)
    dest_archive = mm.OPLSArchive(dest_archive_path)
    dest_archive.merge_archive(src_archive_path)


def parse_args(argv=None):
    """
    Parses command-line arguments.

    :return: argument namespace with command line options
    :rtype: argparse.Namespace
    """

    parser = argparse.ArgumentParser(
        description=__doc__,
        prog="$SCHRODINGER/utilities/custom_params",
        formatter_class=argparse.RawDescriptionHelpFormatter)
    subparsers = parser.add_subparsers(dest="task", metavar="<task>")
    subparsers.required = True

    # Hidden tasks to create archive from OPLSDIR with individual datafiles
    create_parser = subparsers.add_parser(CREATE_TASK, description=SUPPRESS)
    create_parser.add_argument("in_oplsdir", metavar="<OPLSDIR>")
    create_parser.add_argument("out_oplsfile", metavar="<name>.opls")
    create_parser.add_argument("-force", action="store_true")
    create_parser.add_argument("-test", action="store_true")

    msg = "Upgrade OPLSDIR parameters for use in the current release."
    upgrade_parser = subparsers.add_parser(
        UPGRADE_TASK,
        help=msg,
        description=msg,
        formatter_class=argparse.RawDescriptionHelpFormatter)

    msg = "OPLSDIR containing custom parameters to be upgraded."
    upgrade_parser.add_argument("oplsdir", metavar="<OPLSDIR>", help=msg)

    msg = "Merge parameters from an alternate OPLSDIR into a primary OPLSDIR."
    merge_parser = subparsers.add_parser(
        MERGE_TASK,
        help=msg,
        description=msg,
        formatter_class=argparse.RawDescriptionHelpFormatter)

    msg = "OPLSDIR containing custom parameters to be added."
    merge_parser.add_argument("in_oplsdir", metavar="<a_OPLSDIR>", help=msg)

    msg = "Primary OPLSDIR that will receive new parameters."
    merge_parser.add_argument("out_oplsdir", metavar="<p_OPLSDIR>", help=msg)

    return parser.parse_args(argv)


def main(argv=None):

    args = parse_args(argv)

    try:
        if args.task == CREATE_TASK:
            logger.info(f"Collecting parameter files from {args.in_oplsdir}")
            logger.info(f"Creating {args.out_oplsfile} ...")
            create_archive_from_oplsdir(args.in_oplsdir, args.out_oplsfile,
                                        args.force, args.test)

        elif args.task == UPGRADE_TASK:
            logger.info(f"Upgrading {args.oplsdir} ...")
            upgrade_oplsdir(args.oplsdir)

        elif args.task == MERGE_TASK:
            logger.info(f"Merging {args.in_oplsdir} -> {args.out_oplsdir} ...")
            merge_oplsdir(args.in_oplsdir, args.out_oplsdir)

        else:
            raise RuntimeError(f"Unknown task: {args.task}")

    except RuntimeError as exc:
        logger.error(f"ERROR: {exc}")
        sys.exit(1)

    logger.info(f"OPLS custom parameter {args.task} completed.")


if __name__ == "__main__":
    cmdline.main_wrapper(main)