Yank/commands/script.py
#!/usr/local/bin/env python
# =============================================================================================
# MODULE DOCSTRING
# =============================================================================================
"""
Set up and run YANK calculation from script.
"""
# =============================================================================================
# GLOBAL IMPORTS
# =============================================================================================
# Module imports handled in individual functions since CLI should be faster to boot up
# =============================================================================================
# COMMAND-LINE INTERFACE
# =============================================================================================
usage = """
YANK script
Usage:
yank script (-y FILEPATH | --yaml=FILEPATH) [--setup-only] [--jobid=INTEGER] [--njobs=INTEGER] [-o OVERRIDE] [--status] ...
Description:
Set up and run free energy calculations from a YAML script. All options can be specified in the YAML script.
Required Arguments:
-y, --yaml=FILEPATH Path to the YAML script specifying options and/or how to set up and run the experiment.
Optional Arguments:
--setup-only If given, only the automatic setup pipeline is run without running the free energy
calculation.
--jobid=INTEGER You can run only a subset of the experiments by specifying jobid and njobs, where
1 <= job_id <= n_jobs. In this case, njobs must be specified as well and YANK will
run only 1/n_jobs of the experiments. This can be used to run several separate YANK
executions in parallel starting from the same script.
Alternatively, a zero-indexed environment variable name can be specified, in which
case the jobid will be taken from this variable (0 <= environment variable < n_jobs).
--njobs=INTEGER Specify the total number of parallel executions. jobid has to be specified too.
Alternatively, an environment variable name can be specified, in which
case the number of jobs jobid will be taken from this variable.
-o, --override=OVERRIDE Override a single option in the script file. May be specified multiple times.
Specified as a nested dictionary of the form:
top_option:sub_option:value
of any depth level where the last entry is always the value, do not use quotes.
Please see script file documentation for valid options.
Specifying the same option multiple times results in an error.
This method is not recommended for complex options such as lists or combinations
--status Generate status.pkl pickle file with current estimates of progress,
updated each time experiment is switched.
This feature is EXPERIMENTAL and its invocation scheme is subject to change in
future releases.
"""
# =============================================================================================
# COMMAND DISPATCH
# =============================================================================================
def dispatch(args):
"""
Set up and run YANK calculation from a script.
Parameters
----------
args : dict
Command-line arguments from docopt.
"""
import os
from ..experiment import ExperimentBuilder
override = None
if args['--override']: # Is False for None and [] (empty list)
over_opts = args['--override']
# Check for duplicates
if len(over_opts) != len(set(over_opts)):
raise ValueError("There were duplicate override options, result will be ambiguous!")
# Create a dict of strings
override_dict = {}
for opt in over_opts:
split_opt = opt.split(':')
try:
key, value = split_opt[-2:]
except IndexError:
raise IndexError("Override option {} is not of the form top_dict:key:value. Minimal accepted case is "
"key:value".format(opt))
top_dict = override_dict
for opt_level in split_opt[:-2]:
try:
# Dict already exists
top_dict[opt_level]
except KeyError:
top_dict[opt_level] = {}
finally:
top_dict = top_dict[opt_level]
top_dict[key] = value
# Create a string that looks like a dictionary, removing quotes
# This is done to avoid input type ambiguity and instead let the parser handle it as though it were a file
override = str(override_dict).replace("'", "").replace('"', '')
if args['--jobid']:
job_id = args['--jobid']
try:
# Try processing it as an integer
job_id = int(job_id)
except ValueError:
if job_id in os.environ:
# Take the value of the environment variable
job_id = int(os.environ[job_id]) + 1
else:
raise ValueError("--jobid '%s' unknown" % job_id)
else:
job_id = None
if args['--njobs']:
n_jobs = args['--njobs']
try:
# Try processing it as an integer
n_jobs = int(n_jobs)
except ValueError:
if n_jobs in os.environ:
# Take the value of the environment variable
n_jobs = int(os.environ[n_jobs])
else:
raise ValueError("--njobs '%s' unknown" % n_jobs)
else:
n_jobs = None
# Handle specification of status file
write_status = False
if args['--status']:
write_status = args['--status']
if args['--yaml']:
yaml_path = args['--yaml']
if not os.path.isfile(yaml_path):
raise ValueError('Cannot find YAML script "{}"'.format(yaml_path))
yaml_builder = ExperimentBuilder(script=yaml_path, job_id=job_id, n_jobs=n_jobs)
if override: # Parse the string present.
yaml_builder.update_yaml(override)
# Check if we need to run only the automatic setup pipeline or also the experiments.
if args['--setup-only']:
yaml_builder.setup_experiments()
else:
yaml_builder.run_experiments(write_status=write_status)
return True
return False