Source code for schrodinger.application.matsci.parserutils

"""
Utilities for working with argument parsers

Copyright Schrodinger, LLC. All rights reserved.
"""

import argparse
import functools
import json
import operator
import os
import random
import sys

import yaml

import pymmlibs
from schrodinger import structure
from schrodinger.application.desmond import cms
from schrodinger.application.matsci import atomicsymbols
from schrodinger.application.matsci import desmondutils
from schrodinger.application.matsci import jaguarworkflows
from schrodinger.application.matsci.nano import xtal
from schrodinger.infra import mm
from schrodinger.structutils.analyze import validate_asl
from schrodinger.test import ioredirect
from schrodinger.utils import fileutils
from schrodinger.utils import sea

OLD_TRJ_END = '-out.idx'
NEW_TRJ_END = '_trj'
RANDOM_SEED_MIN = 0
# MATSCI-4747 largest integer on Windows
RANDOM_SEED_MAX = min(2**31 - 1, sys.maxsize)
RANDOM_SEED_DEFAULT = 1234
RANDOM_SEED_RANDOM = 'random'
YAML_EXT = '.yaml'

# Desmond Trajectory
FLAG_CMS_PATH = '-cms_file'
FLAG_TRJ_PATH = '-trj'
FLAG_TRJ_MIN = '-trj_min'
FLAG_TRJ_MAX = '-trj_max'

# ASL Flags
FLAG_ASL = '-asl'
FLAG_CENTER_ASL = '-center_asl'

# Distance Flags
FLAG_RESOLUTION = '-resolution'


[docs]def type_ranged_num(arg, top=None, bottom=0.0, top_allowed=False, bottom_allowed=False, typer=float): """ Validate that the argument is a number over a given range. By default, it is (0, +inf) :note: This function can be used to directly create argparse typing for numbers within a specified range using a lambda function such as: eqopts.add_argument( FLAG_TEMP, action='store', type=lambda x:parserutils.type_ranged_num(x, top=423, bottom=273), metavar='KELVIN', help='Simulation temperature for MD equilibration step') :type arg: str :param arg: The argument to validate :type top: float or int :param top: The upper limit of the allowed range :type bottom: float or int :param bottom: The lower limit of the allowed range :type top_allowed: bool :param top_allowed: Whether arg may take the top value or not (whether the upper limit is inclusive or exclusive) :type bottom_allowed: bool :param bottom_allowed: Whether arg may take the bottom value or not (whether the bottom limit is inclusive or exclusive) :type typer: callable :param typer: Should be one of the built-in float or int functions and defines the type of the value returned from this function. :rtype: float or int :return: The argument converted to a float or int, depending on the value of the typer keyword argument :raise `argparse.ArgumentTypeError`: If the argument cannot be converted to a ranged floating point number in the given range """ msg = '{val} is not {ntype} in the range of {bsym}{bbound}, {tbound}{tsym}' if typer == int: ntype = 'an integer' else: ntype = 'a real number' if bottom_allowed: bcomp = operator.lt bsym = '[' else: bcomp = operator.le bsym = '(' if bottom is None: bots = '-infinity' else: bots = bottom if top_allowed: tcomp = operator.gt tsym = ']' else: tcomp = operator.ge tsym = ')' if top is None: tops = 'infinity' else: tops = top error = msg.format( val=arg, ntype=ntype, bsym=bsym, bbound=bots, tbound=tops, tsym=tsym) try: val = typer(arg) except (TypeError, ValueError): raise argparse.ArgumentTypeError(error) if bottom is not None and bcomp(val, bottom): raise argparse.ArgumentTypeError(error) if top is not None and tcomp(val, top): raise argparse.ArgumentTypeError(error) return val
[docs]def type_ranged_int(*args, **kwargs): """ Validate that the argument is an int over a given range see type_ranged_num for documentation of arguments, return values and exceptions raised """ kwargs['typer'] = int return type_ranged_num(*args, **kwargs)
[docs]def get_ranged_int_typer(bottom_allowed=True, top_allowed=True, **kwargs): """ Get a validator that will validate an integer over a custom range Usage:: parser.add_argument( A_FLAG, type=get_ranged_int_typer(bottom=4, top=12, bottom_allowed=False)) See type_ranged_num for documentation of arguments, return values and exceptions raised """ return functools.partial( type_ranged_int, bottom_allowed=bottom_allowed, top_allowed=top_allowed, **kwargs)
[docs]def get_ranged_float_typer(bottom_allowed=True, top_allowed=True, **kwargs): """ Get a validator that will validate a float over a custom range Usage:: parser.add_argument( A_FLAG, type=get_ranged_float_typer(bottom=-5.0, top=5.0, top_allowed=False)) See type_ranged_num for documentation of arguments, return values and exceptions raised """ return functools.partial( type_ranged_num, bottom_allowed=bottom_allowed, top_allowed=top_allowed, **kwargs)
[docs]def type_positive_float(arg): """ Validate that the argument is a positive float see type_ranged_num for documentation of arguments, return values and exceptions raised """ return type_ranged_num(arg, top=None, bottom=0.0, bottom_allowed=False)
[docs]def type_negative_float(arg): """ Validate that the argument is a negative float see type_ranged_num for documentation of arguments, return values and exceptions raised """ return type_ranged_num(arg, top=0.0, bottom=None, top_allowed=False)
[docs]def type_positive_int(arg): """ Validate that the argument is a positive int see type_ranged_num for documentation of arguments, return values and exceptions raised """ return type_ranged_int(arg, top=None, bottom=0, bottom_allowed=False)
[docs]def type_nonnegative_int(arg): """ Validate that the argument is a nonnegative int see type_ranged_num for documentation of arguments, return values and exceptions raised """ return type_ranged_int(arg, top=None, bottom=0, bottom_allowed=True)
[docs]def type_nonpositive_float(arg): """ Validate that the argument is a nonpositive float see type_ranged_num for documentation of arguments, return values and exceptions raised """ return type_ranged_num(arg, top=0.0, bottom=None, top_allowed=True)
[docs]def type_nonzero_percent(arg): """ Validate that the argument is a percent but not zero see type_ranged_num for documentation of arguments, return values and exceptions raised """ return type_ranged_num( arg, top=100.0, bottom=0.0, bottom_allowed=False, top_allowed=True)
[docs]def type_random_seed(arg, seed_min=RANDOM_SEED_MIN, seed_max=RANDOM_SEED_MAX): """ Validate that the argument is a valid random seed value. If a random value is requested, that value is generated and returned. :param str arg: The argument to validate :param int seed_min: The minimum allowed random number seed :param int seed_max: The maximum allowed random number seed :rtype: int :return: The random seed :raise `argparse.ArgumentTypeError`: If the argument is not RANDOM_SEED_RANDOM or an integer """ if arg == RANDOM_SEED_RANDOM: return random.randint(seed_min, seed_max) return type_ranged_int( arg, top=seed_max, bottom=seed_min, bottom_allowed=True, top_allowed=True)
[docs]def type_file(arg): """ Validate that the argument is an existing filename :type arg: str or unicode :param arg: The argument to validate :rtype: str :return: The str-ed argument :raise `argparse.ArgumentTypeError`: If the given filename does not exist """ # if called directly from argparse.ArgumentParser.add_argument then # arg is a unicode which is typically typed as a str via the type=str # kwarg, do the same in this custom typer arg = str(arg) file_found = fileutils.get_existing_filepath(arg) if not file_found: raise argparse.ArgumentTypeError('File does not exist: %s' % arg) return file_found
[docs]def type_nonnegative_float(arg): """ Validate that the argument is a nonnegative float see type_ranged_num for documentation of arguments, return values and exceptions raised """ return type_ranged_num(arg, top=None, bottom=0.0, bottom_allowed=True)
[docs]def get_trj_dir_name(st): """ Get the trajectory directory name from the given structure. :type st: schrodinger.structure.Structure :param st: the structure :rtype: str or None :return: the trajectory directory name or None if there isn't one """ trj_prop = st.property.get(cms.Cms.PROP_TRJ) if not trj_prop: return elif trj_prop.endswith(OLD_TRJ_END): trj_prop = trj_prop.replace(OLD_TRJ_END, NEW_TRJ_END) if trj_prop.endswith(NEW_TRJ_END): return trj_prop
[docs]def type_cms_file(arg, ensure_trj=False): """ Validate that the argument is a cms file. :type arg: str or unicode :param arg: the argument to validate :type ensure_trj: bool :param ensure_trj: ensure that the cms file has trajectory information, for example is a Desmond output file :raise argparse.ArgumentTypeError: if the given argument is not a cms file :rtype: str :return: the str-ed argument """ arg = type_file(arg) ext = fileutils.splitext(arg)[1] if ext not in fileutils.EXTENSIONS['maestro']: msg = ('File {arg} is not a cms file.').format(arg=arg) raise argparse.ArgumentTypeError(msg) try: obj = cms.Cms(arg) except Exception: msg = ('File {arg} is not a valid cms file.').format(arg=arg) raise argparse.ArgumentTypeError(msg) if ensure_trj: traj_dir = get_trj_dir_name(obj.fsys_ct) orig_cms = obj.fsys_ct.property.get(cms.Cms.PROP_CMS) if not (traj_dir and orig_cms): msg = ('File {arg} is missing trajectory information.' ).format(arg=arg) raise argparse.ArgumentTypeError(msg) adir = os.path.join(os.path.split(arg)[0], traj_dir) if not os.path.isdir(adir): msg = ('File {arg} points to a trajectory directory that does not ' 'exist.').format(arg=arg) raise argparse.ArgumentTypeError(msg) return arg
[docs]def type_xyz_file(arg): """ Validate that the argument is an existing filename with xyz extension and can be converted into sdf format. :param arg: str :type arg: The argument to validate :return: str :rtype: The validated argument :raise `argparse.ArgumentTypeError`: If the given filename does not exist :raise `argparse.ArgumentTypeError`: If the given filename does not have xyz as file extension. :raise `argparse.ArgumentTypeError`: If the given filename cannot be converted """ file_path = type_file(arg) try: fileutils.xyz_to_sdf(file_path, save_file=False) except ValueError as err: raise argparse.ArgumentTypeError(f"{err}") except RuntimeError as err: raise argparse.ArgumentTypeError( f"Open Babel cannot convert the input {file_path} with stderr being {err}" ) return file_path
[docs]def type_cms_with_trj(arg): """ Validate that the argument is a cms file with trajectory information. :type arg: str or unicode :param arg: the argument to validate :rtype: str :return: the str-ed argument """ return type_cms_file(arg, ensure_trj=True)
[docs]def type_json_file(arg): """ Validate that the argument is a json file. :type arg: str or unicode :param arg: the argument to validate :raise argparse.ArgumentTypeError: if the given argument is not a json file :rtype: str :return: the str-ed argument """ arg = type_file(arg) ext = fileutils.splitext(arg)[1] if ext != '.json': msg = ('File {arg} is not a json file.').format(arg=arg) raise argparse.ArgumentTypeError(msg) try: with open(arg, 'r') as afile: json.load(afile) except Exception: msg = ('File {arg} is not a valid json file.').format(arg=arg) raise argparse.ArgumentTypeError(msg) return arg
[docs]def type_yaml_file(arg): """ Validate that the argument is a yaml file. :type arg: str or unicode :param arg: the argument to validate :raise argparse.ArgumentTypeError: if the given argument is not a yaml file :rtype: str :return: the str-ed argument """ arg = type_file(arg) ext = fileutils.splitext(arg)[1] if ext != YAML_EXT: msg = (f"File {arg} doesn't have {YAML_EXT} as extension." ).format(arg=arg) raise argparse.ArgumentTypeError(msg) try: with open(arg, 'r') as afile: yaml.load(afile) except Exception as err: msg = ( f'File {arg} is not a valid yaml file. {str(err)}').format(arg=arg) raise argparse.ArgumentTypeError(msg) return arg
[docs]def type_yaml_str(arg): """ Load the arg as a yaml string. :type arg: str or unicode :param arg: the argument to validate :raise argparse.ArgumentTypeError: if the given argument is not in yaml format :rtype: dict :return: loaded yaml string """ try: parsed_str = yaml.load(arg) except Exception as err: msg = ( f'String {arg} is not in yaml format. {str(err)}').format(arg=arg) raise argparse.ArgumentTypeError(msg) return parsed_str
[docs]def type_structure_file(arg): """ Validate that the argument is a structure file that can be read by StructureReader. :type arg: str or unicode :param arg: the argument to validate :raise argparse.ArgumentTypeError: if the given argument is not a readable structure file :rtype: str :return: the str-ed argument """ arg = type_file(arg) try: struct = structure.Structure.read(arg) except ValueError: msg = '%s is not a valid structure file' % arg raise argparse.ArgumentTypeError(msg) return arg
[docs]def type_periodic_structure_file(arg): """ Validate that the argument is a periodic structure file (priority is chorus followed by pdb). :type arg: str or unicode :param arg: the argument to validate :raise argparse.ArgumentTypeError: if the given argument is not a periodic structure file :rtype: str :return: the str-ed argument """ arg = type_structure_file(arg) st = structure.Structure.read(arg) if not xtal.has_pbc(st): msg = '%s is not a periodic structure file' % arg raise argparse.ArgumentTypeError(msg) return arg
[docs]def type_element(value): """ Validate that the argument is a valid atomic atymbol :type value: str :param value: The argument to validate :rtype: str :return: The validated argument :raise `argparse.ArgumentTypeError`: If the argument is not a valid atomic symbol """ if value not in atomicsymbols.ATOMIC_SYMBOLS: raise argparse.ArgumentTypeError( '%s is not found in the list of atomic symbols' % value) return value
[docs]def type_forcefield(value): """ Validate that the argument is a valid force field name This accepts both OPLS3 and OPLS3e for force field #16 :type value: str :param value: The argument to validate :rtype: str :return: The canonical force field name :raise `argparse.ArgumentTypeError`: If the argument is not a valid atomic symbol """ try: with ioredirect.IOSilence(): ff_int = mm.opls_name_to_version(value) except IndexError: msg = ('%s is an unrecognized force field name. %s' % (value, valid_forcefield_info())) raise argparse.ArgumentTypeError(msg) return mm.opls_version_to_name(ff_int)
[docs]def type_forcefield_and_get_number(value): """ Validate that the argument is a valid force field name and returns both the canonical force field name and number. This accepts both OPLS3 and OPLS3e for force field #16 :type value: str :param value: The argument to validate :rtype: str, int :return: The canonical force field name and the corresponding force field int :raise `argparse.ArgumentTypeError`: If the argument is not a valid atomic symbol """ name = type_forcefield(value) number = mm.opls_name_to_version(name) return name, number
[docs]def type_desmond_ensemble(value): """ Validate that the argument is a valid desmond ensemble :type value: str :param value: The argument to validate :rtype: str :return: The validated argument :raise `argparse.ArgumentTypeError`: If the argument is not a valid desmond ensemble """ if value not in desmondutils.ENSEMBLES: raise argparse.ArgumentTypeError( '%s is not found in the list of known desmond ensembles' % value) return value
[docs]def valid_forcefield_info(): """ Get an informative sentence that can be included in help messages that indicate the valid force field names. Note that these are the default names and type_forcefield will accept many variants of them, including OPLS3 for OPLS3E. :rtype: str :return: A full sentence describing the valid force field names """ ffnames = " and ".join(mm.opls_names()) return 'Valid force fields are %s.' % ffnames
[docs]class NoArgCustomAction(argparse.Action): """ An action that allows giving an argparse flag a custom action so that the flag does not require an argument One possible use is the implementation of a flag like '-doc' which immediately performs an action when it is found on the command line - such as printing out detailed documentation and then exiting. This avoids the user having to specify any other flags even if they are marked as required. """
[docs] def __init__(self, *args, **kwargs): kwargs['nargs'] = 0 super().__init__(*args, **kwargs)
def __call__(self, parser, namespace, values, option_string=""): # Overwrite this function to perform the custom action super().__call__(parser, namespace, values, option_string=option_string)
[docs]class KeyPairDict(dict): """ This is a subclass of dict so that the {a:b} can be displayed as "a=b" (e.g, as the help message" """ def __str__(self, *args, **kwargs): return ' '.join(['='.join([x, str(y)]) for x, y in self.items()])
[docs]class StoreDictKeyPair(argparse._StoreAction): """ This is the action used with nargs="+" and values='KEYWORD=VALUE KEYWORD=VALUE' By default nargs="+" and action='store' save the multiple values "k1=v1 k2=v2" as list. With this StoreDictKeyPair as the action, "k1=v1 k2=v2" will be stored as {k1:v1, k2:v2} in the options. """ def __call__(self, parser, namespace, values, option_string=None): """ See parent classes for documentation. """ if isinstance(values, list): try: keywords = ' '.join(values) values = type_keywords_to_dict(keywords) except argparse.ArgumentTypeError as err: parser.error(str(err) + f" in {keywords}") super().__call__(parser, namespace, values, option_string=None)
[docs]def type_keywords_to_dict(value): """ Convert a string of 'keyword=value keyword=value ...' to a dictionary keyed by keyword. Values are the keyword values. :note: When calling parser.add_argument for the flag that uses this function, do *not* use "nargs=+". That will result in the return value from this function being placed as an item in a list. Instead, do not use nargs at all and simply use "action='store'". All of the keyword=value pairs will then be passed into this function as a single string. Example: parser.add_argument( jobutils.FLAG_KEYWORDS, action='store', metavar='KEYWORD=VALUE KEYWORD=VALUE ...', type=parserutils.type_keywords_to_dict, help='Jaguar keywords given as a space-separated keyword=value ' ' pairs' ) will result in options.keyword = the dictionary returned by this function :type value: str :param value: The argument to validate :rtype: dict :return: A dictionary with keywords as key keyword values as values :raise `argparse.ArgumentTypeError`: If the argument is not a valid keyword string """ try: keydict = jaguarworkflows.keyword_string_to_dict(value) except ValueError as msg: raise argparse.ArgumentTypeError(str(msg)) return keydict
[docs]def type_keywords_to_string(value): """ Validate the format of a string of 'keyword=value keyword=value ...' :note: When calling parser.add_argument for the flag that uses this function, do *not* use "nargs=+". That will result in the return value from this function being placed as an item in a list. Instead, do not use nargs at all and simply use "action='store'". All of the keyword=value pairs will then be passed into this function as a single string. Example: parser.add_argument( jobutils.FLAG_KEYWORDS, action='store', metavar='KEYWORD=VALUE KEYWORD=VALUE ...', type=parserutils.type_keywords_to_string, help='Jaguar keywords given as a space-separated keyword=value ' ' pairs' ) will result in options.keyword = the string returned by this function :type value: str :param value: The argument to validate :rtype: string :return: The original command line string :raise `argparse.ArgumentTypeError`: If the argument is not a valid keyword string """ try: keydict = jaguarworkflows.keyword_string_to_dict(value) except ValueError as msg: raise argparse.ArgumentTypeError(str(msg)) return value
[docs]def type_num_list(value, delimiter=',', typer=float, bottom=None, bottom_allowed=False): """ Validate that the argument is a list of numbers separated by the delimiter :type value: str :param value: string containing numbers separated by delimiter :type delimiter: str :param delimiter: the char to split the input value :type typer: callable :param typer: Should be one of the built-in float or int functions and defines the type of the value returned from this function. :type bottom: float or int :param bottom: The lower limit of the allowed range :type bottom_allowed: bool :param bottom_allowed: Whether arg may take the bottom value or not (whether the bottom limit is inclusive or exclusive) :rtype: list :return: The validated argument converted to a list of numbers :raise `argparse.ArgumentTypeError`: If the argument is not a valid list of numbers """ value_list = [x.strip() for x in value.split(delimiter)] return [ type_ranged_num( x, bottom=bottom, bottom_allowed=bottom_allowed, typer=typer) for x in value_list ]
[docs]def type_num_comma_list(value, typer=float, bottom=None, bottom_allowed=False): """ Validate that the argument is a list of numbers separated by comma :type value: str :param value: string containing numbers separated by commas :type typer: callable :param typer: Should be one of the built-in float or int functions and defines the type of the value returned from this function. :type bottom: float or int :param bottom: The lower limit of the allowed range :type bottom_allowed: bool :param bottom_allowed: Whether arg may take the bottom value or not (whether the bottom limit is inclusive or exclusive) :rtype: list :return: The validated argument positive integers coverted to list :raise `argparse.ArgumentTypeError`: If the argument is not a valid list of floats """ return type_num_list( value, delimiter=',', bottom=bottom, bottom_allowed=bottom_allowed, typer=typer)
[docs]def type_num_colon_list(value, typer=float, bottom=None, bottom_allowed=False): """ Validate that the argument is a list of numbers separated by colon :type value: str :param value: string containing numbers separated by colon :type typer: callable :param typer: Should be one of the built-in float or int functions and defines the type of the value returned from this function. :type bottom: float or int :param bottom: The lower limit of the allowed range :type bottom_allowed: bool :param bottom_allowed: Whether arg may take the bottom value or not (whether the bottom limit is inclusive or exclusive) :rtype: list :return: The validated argument positive integers coverted to list :raise `argparse.ArgumentTypeError`: If the argument is not a valid list of floats """ return type_num_list( value, delimiter=':', bottom=bottom, bottom_allowed=bottom_allowed, typer=typer)
[docs]def type_float_comma_list(value): """ Validate that the argument is a float list separated by comma :type value: str :param value: string containing numbers separated by commas :rtype: list :return: The validated argument positive integers converted to list :raise `argparse.ArgumentTypeError`: If the argument is not a valid list of floats """ return type_num_comma_list(value)
[docs]def type_nonnegative_float_comma_list(value): """ Validate that the argument is a nonnegative float list separated by comma :type value: str :param value: string containing numbers separated by commas :rtype: list :return: The validated argument positive integers converted to list :raise `argparse.ArgumentTypeError`: If the argument is not a valid list of floats """ return type_num_comma_list(value, bottom=0.0, bottom_allowed=True)
[docs]def type_positive_integer_comma_list(value): """ Validate that the argument is a integer list separated by comma :type value: str :param value: string containing numbers separated by commas :rtype: list :return: The validated argument positive integers coverted to list :raise `argparse.ArgumentTypeError`: If the argument is not a valid list of positive integers """ return type_num_comma_list(value, typer=int, bottom=0.0)
[docs]def type_element_comma_list(value): """ Validate that the argument is a element separated by comma :type str value: element separated by commas :rtype: list :return: The validated argument list of element symbols :raise `argparse.ArgumentTypeError`: If the argument is not a valid list of strings """ value = [x.strip() for x in value.split(',')] return [type_element(x) for x in value]
[docs]def type_percentage(value): """ Validate that the argument is a percentage between 0 and 100 see type_ranged_num for documentation of arguments, return values and exceptions raised """ return type_ranged_num( value, top=100, bottom=0.0, top_allowed=True, bottom_allowed=True)
[docs]def type_asl(asl): """ Validate ASL. :rtype: str :return: ASL :raise `argparse.ArgumentTypeError`: If ASL is not valid """ with ioredirect.IOSilence(): if not validate_asl(asl): raise argparse.ArgumentTypeError('Invalid ASL: %s' % asl) return asl
[docs]def type_desmond_cfg(arg): """ Validate that the argument is a desmond CFG file. :type arg: str or unicode :param arg: the argument to validate :raise argparse.ArgumentTypeError: if the given argument is not a readable CFG file :rtype: str :return: the str-ed argument """ arg = type_file(arg) # Took from analyze_simulation.py :: _parse_ark try: with open(arg) as cfg_fh: sea.Map(cfg_fh.read()) except Exception as err: raise argparse.ArgumentTypeError('Invalid desmond cfg file (%s): %s' % (arg, str(err))) return arg
[docs]class DriverParser(argparse.ArgumentParser): """ Subclass that shows driver usage relative to python/scripts or python/common directories and adds a help argument by default """
[docs] def __init__(self, *args, **kwargs): """ Accepts all arguments normally given for an ArgumentParser """ # Gets the full path to the driver, and calculates the path shown as # "usage" in the driver help message by removing the extra parts. # If the driver is located in python/scripts or python/common, usage just # includes the driver file name. If the driver is in a subdirectory of # python/scripts or python/common, usage also includes the subdirectory's name. # This is so the driver can be found when using: $SCHRODINGER/run <usage> if 'prog' not in kwargs: file_path = os.path.abspath(sys.argv[0]) base_name = os.path.basename(file_path) dirname = os.path.basename(os.path.dirname(file_path)) if dirname in ('common', 'scripts'): relative_path = base_name else: relative_path = os.path.join(dirname, base_name) kwargs['prog'] = "$SCHRODINGER/run " + relative_path super().__init__(*args, **kwargs) self.add_argument( '-h', '-help', action='help', default=argparse.SUPPRESS, help='Show this help message and exit.')