files: #!/usr/local/bin/env python
  
  # =============================================================================================
  # MODULE DOCSTRING
  # =============================================================================================
  
  """
  Set up and run YANK calculation from script.
  
  """
  
  # =============================================================================================
  # GLOBAL IMPORTS
  # =============================================================================================
  
  import os
  from ..experiment import ExperimentBuilder
  
  
  # =============================================================================================
  # COMMAND-LINE INTERFACE
  # =============================================================================================
  
  usage = """
  YANK script
  
  Usage:
    yank script (-y FILEPATH | --yaml=FILEPATH) [--jobid=INTEGER] [--njobs=INTEGER] [-o OVERRIDE]...
  
  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:
    --jobid=INTEGER               You can run only a subset of the experiments by specifying jobid and njobs, where
                                  0 <= job_id <= n_jobs-1. 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.
    --njobs=INTEGER               Specify the total number of parallel executions. jobid has to be specified too.
    -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
  """
  
  # =============================================================================================
  # COMMAND DISPATCH
  # =============================================================================================
  
  
  def dispatch(args):
      """
      Set up and run YANK calculation from a script.
  
      Parameters
      ----------
      args : dict
         Command-line arguments from docopt.
  
      """
      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']
      else:
          job_id = None
      if args['--njobs']:
          n_jobs = args['--njobs']
      else:
          n_jobs = None
  
      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)
          yaml_builder.run_experiments()
          return True
  
      return False
  
  
  [... the rest is omitted ...]
files: Yank/commands/script.py
files: #!/usr/local/bin/env python
  
  # =============================================================================================
  # MODULE DOCSTRING
  # =============================================================================================
  
  """
  Set up and run YANK calculation from script.
  
  """
  
  # =============================================================================================
  # GLOBAL IMPORTS
  # =============================================================================================
  
  import os
  from ..experiment import ExperimentBuilder
  
  
  # =============================================================================================
  # COMMAND-LINE INTERFACE
  # =============================================================================================
  
  usage = """
  YANK script
  
  Usage:
    yank script (-y FILEPATH | --yaml=FILEPATH) [-o OVERRIDE]...
  
  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:
  
    -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
  """
  
  # =============================================================================================
  # COMMAND DISPATCH
  # =============================================================================================
  
  
  def dispatch(args):
      """
      Set up and run YANK calculation from a script.
  
      Parameters
      ----------
      args : dict
         Command-line arguments from docopt.
  
      """
      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['--yaml']:
          yaml_path = args['--yaml']
  
          if not os.path.isfile(yaml_path):
              raise ValueError('Cannot find YAML script "{}"'.format(yaml_path))
  
          yaml_builder = ExperimentBuilder(yaml_source=yaml_path)
          if override:  # Parse the string present.
              yaml_builder.update_yaml(override)
          yaml_builder.run_experiments()
          return True
  
      return False
  
  
  [... the rest is omitted ...]
files: Yank/commands/script.py
files: #!/usr/bin/env python
  
  # =============================================================================
  # MODULE DOCSTRING
  # =============================================================================
  
  """
  Tools to build Yank experiments from a YAML configuration file.
  
  """
  
  # =============================================================================
  # GLOBAL IMPORTS
  # =============================================================================
  
  import os
  import copy
  import yaml
  import logging
  import collections
  
  import openmmtools as mmtools
  import openmoltools as moltools
  from simtk import unit, openmm
  from simtk.openmm.app import PDBFile, AmberPrmtopFile
  from schema import Schema, And, Or, Use, Optional, SchemaError
  
  from . import utils, pipeline, mpi, restraints, repex
  from .yank import AlchemicalPhase, Topography
  
  logger = logging.getLogger(__name__)
  
  # =============================================================================
  # CONSTANTS
  # =============================================================================
  
  HIGHEST_VERSION = '1.3'  # highest version of YAML syntax
  
  
  # =============================================================================
  # UTILITY FUNCTIONS
  # =============================================================================
  
  def to_openmm_app(input_string):
      """Converter function to be used with validate_parameters()."""
      return getattr(openmm.app, input_string)
  
  
  def convert_if_quantity(value):
  
      try:
          quantity = utils.quantity_from_string(value)
      except:
          return value
      return quantity
  
  
  def update_nested_dict(original, updated):
      """Return a copy of a (possibly) nested dict of arbitrary depth"""
      new = original.copy()
      for key, value in updated.items():
          if isinstance(value, collections.Mapping):
              replacement = update_nested_dict(new.get(key, {}), value)
              new[key] = replacement
          else:
              new[key] = updated[key]
      return new
  
  
  # ==============================================================================
  # UTILITY CLASSES
  # ==============================================================================
  
  class YamlParseError(Exception):
      """Represent errors occurring during parsing of Yank YAML file."""
      def __init__(self, message):
          super(YamlParseError, self).__init__(message)
          logger.error(message)
  
  
  class YankLoader(yaml.Loader):
      """PyYAML Loader that recognized !Combinatorial nodes and load OrderedDicts."""
      def __init__(self, *args, **kwargs):
          super(YankLoader, self).__init__(*args, **kwargs)
          self.add_constructor(u'!Combinatorial', self.combinatorial_constructor)
          self.add_constructor(u'!Ordered', self.ordered_constructor)
  
      @staticmethod
      def combinatorial_constructor(loader, node):
          """Constructor for YAML !Combinatorial entries."""
          return utils.CombinatorialLeaf(loader.construct_sequence(node))
  
      @staticmethod
      def ordered_constructor(loader, node):
          """Constructor for YAML !Ordered tag."""
          loader.flatten_mapping(node)
          return collections.OrderedDict(loader.construct_pairs(node))
  
  
  class YankDumper(yaml.Dumper):
      """PyYAML Dumper that always return sequences in flow style and maps in block style."""
      def __init__(self, *args, **kwargs):
          super(YankDumper, self).__init__(*args, **kwargs)
          self.add_representer(utils.CombinatorialLeaf, self.combinatorial_representer)
          self.add_representer(collections.OrderedDict, self.ordered_representer)
  
      def represent_sequence(self, tag, sequence, flow_style=None):
          return yaml.Dumper.represent_sequence(self, tag, sequence, flow_style=True)
  
      def represent_mapping(self, tag, mapping, flow_style=None):
          return yaml.Dumper.represent_mapping(self, tag, mapping, flow_style=False)
  
      @staticmethod
      def combinatorial_representer(dumper, data):
          """YAML representer CombinatorialLeaf nodes."""
          return dumper.represent_sequence(u'!Combinatorial', data)
  
      @staticmethod
      def ordered_representer(dumper, data):
          """YAML representer OrderedDict nodes."""
          return dumper.represent_mapping(u'!Ordered', data)
  
  
  # ==============================================================================
  # BUILDER CLASS
  # ==============================================================================
  
  class AlchemicalPhaseFactory(object):
  
      DEFAULT_OPTIONS = {
          'anisotropic_dispersion_correction': True,
          'anisotropic_dispersion_cutoff': 'auto',
          'minimize': True,
          'minimize_tolerance': 1.0 * unit.kilojoules_per_mole/unit.nanometers,
          'minimize_max_iterations': 0,
          'randomize_ligand': False,
          'randomize_ligand_sigma_multiplier': 2.0,
          'randomize_ligand_close_cutoff': 1.5 * unit.angstrom,
          'number_of_equilibration_iterations': 0,
          'equilibration_timestep': 1.0 * unit.femtosecond,
          'checkpoint_interval': 10,
      }
  
      def __init__(self, sampler, thermodynamic_state, sampler_states, topography,
                   protocol, storage, restraint=None, alchemical_regions=None,
                   alchemical_factory=None, metadata=None, **options):
          self.sampler = sampler
          self.thermodynamic_state = thermodynamic_state
          self.sampler_states = sampler_states
          self.topography = topography
          self.protocol = protocol
          self.storage = storage
          self.restraint = restraint
          self.alchemical_regions = alchemical_regions
          self.alchemical_factory = alchemical_factory
          self.metadata = metadata
          self.options = self.DEFAULT_OPTIONS.copy()
          self.options.update(options)
  
      def create_alchemical_phase(self):
          alchemical_phase = AlchemicalPhase(self.sampler)
          create_kwargs = self.__dict__.copy()
          create_kwargs.pop('options')
          create_kwargs.pop('sampler')
  
          # Create a reporter if this is only a path.
          if isinstance(self.storage, str):
              checkpoint_interval = self.options['checkpoint_interval']
              # We don't allow checkpoint file overwriting in YAML file
              reporter = repex.Reporter(self.storage, checkpoint_interval=checkpoint_interval)
              create_kwargs['storage'] = reporter
              self.storage = reporter
  
          if self.options['anisotropic_dispersion_correction'] is True:
              dispersion_cutoff = self.options['anisotropic_dispersion_cutoff']
          else:
              dispersion_cutoff = None
  
          alchemical_phase.create(anisotropic_dispersion_cutoff=dispersion_cutoff,
                                  **create_kwargs)
          return alchemical_phase
  
      def initialize_alchemical_phase(self):
          alchemical_phase = self.create_alchemical_phase()
  
          # Minimize if requested.
          if self.options['minimize']:
              tolerance = self.options['minimize_tolerance']
              max_iterations = self.options['minimize_max_iterations']
              alchemical_phase.minimize(tolerance=tolerance, max_iterations=max_iterations)
  
          # Randomize ligand if requested.
          if self.options['randomize_ligand']:
              sigma_multiplier = self.options['randomize_ligand_sigma_multiplier']
              close_cutoff = self.options['randomize_ligand_close_cutoff']
              alchemical_phase.randomize_ligand(sigma_multiplier=sigma_multiplier,
                                                close_cutoff=close_cutoff)
  
          # Equilibrate if requested.
          if self.options['number_of_equilibration_iterations'] > 0:
              n_iterations = self.options['number_of_equilibration_iterations']
              mcmc_move = mmtools.mcmc.LangevinDynamicsMove(timestep=self.options['equilibration_timestep'],
                                                            collision_rate=90.0/unit.picosecond,
                                                            n_steps=500, reassign_velocities=True,
                                                            n_restart_attempts=6)
              alchemical_phase.equilibrate(n_iterations, mcmc_moves=mcmc_move)
  
          return alchemical_phase
  
  
  class Experiment(object):
      """An experiment built by ExperimentBuilder."""
      def __init__(self, phases, number_of_iterations, switch_phase_interval):
          self.phases = phases
          self.number_of_iterations = number_of_iterations
          self.switch_phase_interval = switch_phase_interval
          self._phases_last_iterations = [None, None]
          self._are_phases_completed = [False, False]
  
      @property
      def iteration(self):
          if None in self._phases_last_iterations:
              return 0
          return self._phases_last_iterations
  
      @property
      def is_completed(self):
          return all(self._are_phases_completed)
  
      def run(self, n_iterations=None):
          # Handle default argument.
          if n_iterations is None:
              n_iterations = self.number_of_iterations
  
          # Handle case in which we don't alternate between phases.
          if self.switch_phase_interval <= 0:
              switch_phase_interval = self.number_of_iterations
  
          # Count down the
  
  [... the rest is omitted ...]
files: Yank/experiment.py
files: #!/usr/bin/env python
  
  # =============================================================================
  # MODULE DOCSTRING
  # =============================================================================
  
  """
  Tools to build Yank experiments from a YAML configuration file.
  
  """
  
  # =============================================================================
  # GLOBAL IMPORTS
  # =============================================================================
  
  import os
  import copy
  import yaml
  import logging
  import collections
  
  import openmmtools as mmtools
  import openmoltools as moltools
  from simtk import unit, openmm
  from simtk.openmm.app import PDBFile, AmberPrmtopFile
  from schema import Schema, And, Or, Use, Optional, SchemaError
  
  from . import utils, pipeline, mpi, restraints, repex
  from .yank import AlchemicalPhase, Topography
  
  logger = logging.getLogger(__name__)
  
  # =============================================================================
  # CONSTANTS
  # =============================================================================
  
  HIGHEST_VERSION = '1.3'  # highest version of YAML syntax
  
  
  # =============================================================================
  # UTILITY FUNCTIONS
  # =============================================================================
  
  def to_openmm_app(input_string):
      """Converter function to be used with validate_parameters()."""
      return getattr(openmm.app, input_string)
  
  
  def convert_if_quantity(value):
  
      try:
          quantity = utils.quantity_from_string(value)
      except:
          return value
      return quantity
  
  
  def update_nested_dict(original, updated):
      """Return a copy of a (possibly) nested dict of arbitrary depth"""
      new = original.copy()
      for key, value in updated.items():
          if isinstance(value, collections.Mapping):
              replacement = update_nested_dict(new.get(key, {}), value)
              new[key] = replacement
          else:
              new[key] = updated[key]
      return new
  
  
  # ==============================================================================
  # UTILITY CLASSES
  # ==============================================================================
  
  class YamlParseError(Exception):
      """Represent errors occurring during parsing of Yank YAML file."""
      def __init__(self, message):
          super(YamlParseError, self).__init__(message)
          logger.error(message)
  
  
  class YankLoader(yaml.Loader):
      """PyYAML Loader that recognized !Combinatorial nodes and load OrderedDicts."""
      def __init__(self, *args, **kwargs):
          super(YankLoader, self).__init__(*args, **kwargs)
          self.add_constructor(u'!Combinatorial', self.combinatorial_constructor)
          self.add_constructor(u'!Ordered', self.ordered_constructor)
  
      @staticmethod
      def combinatorial_constructor(loader, node):
          """Constructor for YAML !Combinatorial entries."""
          return utils.CombinatorialLeaf(loader.construct_sequence(node))
  
      @staticmethod
      def ordered_constructor(loader, node):
          """Constructor for YAML !Ordered tag."""
          loader.flatten_mapping(node)
          return collections.OrderedDict(loader.construct_pairs(node))
  
  
  class YankDumper(yaml.Dumper):
      """PyYAML Dumper that always return sequences in flow style and maps in block style."""
      def __init__(self, *args, **kwargs):
          super(YankDumper, self).__init__(*args, **kwargs)
          self.add_representer(utils.CombinatorialLeaf, self.combinatorial_representer)
          self.add_representer(collections.OrderedDict, self.ordered_representer)
  
      def represent_sequence(self, tag, sequence, flow_style=None):
          return yaml.Dumper.represent_sequence(self, tag, sequence, flow_style=True)
  
      def represent_mapping(self, tag, mapping, flow_style=None):
          return yaml.Dumper.represent_mapping(self, tag, mapping, flow_style=False)
  
      @staticmethod
      def combinatorial_representer(dumper, data):
          """YAML representer CombinatorialLeaf nodes."""
          return dumper.represent_sequence(u'!Combinatorial', data)
  
      @staticmethod
      def ordered_representer(dumper, data):
          """YAML representer OrderedDict nodes."""
          return dumper.represent_mapping(u'!Ordered', data)
  
  
  # ==============================================================================
  # BUILDER CLASS
  # ==============================================================================
  
  class AlchemicalPhaseFactory(object):
  
      DEFAULT_OPTIONS = {
          'anisotropic_dispersion_correction': True,
          'anisotropic_dispersion_cutoff': 'auto',
          'minimize': True,
          'minimize_tolerance': 1.0 * unit.kilojoules_per_mole/unit.nanometers,
          'minimize_max_iterations': 0,
          'randomize_ligand': False,
          'randomize_ligand_sigma_multiplier': 2.0,
          'randomize_ligand_close_cutoff': 1.5 * unit.angstrom,
          'number_of_equilibration_iterations': 0,
          'equilibration_timestep': 1.0 * unit.femtosecond,
          'checkpoint_interval': 10,
      }
  
      def __init__(self, sampler, thermodynamic_state, sampler_states, topography,
                   protocol, storage, restraint=None, alchemical_regions=None,
                   alchemical_factory=None, metadata=None, **options):
          self.sampler = sampler
          self.thermodynamic_state = thermodynamic_state
          self.sampler_states = sampler_states
          self.topography = topography
          self.protocol = protocol
          self.storage = storage
          self.restraint = restraint
          self.alchemical_regions = alchemical_regions
          self.alchemical_factory = alchemical_factory
          self.metadata = metadata
          self.options = self.DEFAULT_OPTIONS.copy()
          self.options.update(options)
  
      def create_alchemical_phase(self):
          alchemical_phase = AlchemicalPhase(self.sampler)
          create_kwargs = self.__dict__.copy()
          create_kwargs.pop('options')
          create_kwargs.pop('sampler')
  
          # Create a reporter if this is only a path.
          if isinstance(self.storage, str):
              checkpoint_interval = self.options['checkpoint_interval']
              # We don't allow checkpoint file overwriting in YAML file
              reporter = repex.Reporter(self.storage, checkpoint_interval=checkpoint_interval)
              create_kwargs['storage'] = reporter
              self.storage = reporter
  
          if self.options['anisotropic_dispersion_correction'] is True:
              dispersion_cutoff = self.options['anisotropic_dispersion_cutoff']
          else:
              dispersion_cutoff = None
  
          alchemical_phase.create(anisotropic_dispersion_cutoff=dispersion_cutoff,
                                  **create_kwargs)
          return alchemical_phase
  
      def initialize_alchemical_phase(self):
          alchemical_phase = self.create_alchemical_phase()
  
          # Minimize if requested.
          if self.options['minimize']:
              tolerance = self.options['minimize_tolerance']
              max_iterations = self.options['minimize_max_iterations']
              alchemical_phase.minimize(tolerance=tolerance, max_iterations=max_iterations)
  
          # Randomize ligand if requested.
          if self.options['randomize_ligand']:
              sigma_multiplier = self.options['randomize_ligand_sigma_multiplier']
              close_cutoff = self.options['randomize_ligand_close_cutoff']
              alchemical_phase.randomize_ligand(sigma_multiplier=sigma_multiplier,
                                                close_cutoff=close_cutoff)
  
          # Equilibrate if requested.
          if self.options['number_of_equilibration_iterations'] > 0:
              n_iterations = self.options['number_of_equilibration_iterations']
              mcmc_move = mmtools.mcmc.LangevinDynamicsMove(timestep=self.options['equilibration_timestep'],
                                                            collision_rate=90.0/unit.picosecond,
                                                            n_steps=500, reassign_velocities=True,
                                                            n_restart_attempts=6)
              alchemical_phase.equilibrate(n_iterations, mcmc_moves=mcmc_move)
  
          return alchemical_phase
  
  
  class Experiment(object):
      """An experiment built by ExperimentBuilder."""
      def __init__(self, phases, number_of_iterations, switch_phase_interval):
          self.phases = phases
          self.number_of_iterations = number_of_iterations
          self.switch_phase_interval = switch_phase_interval
          self._phases_last_iterations = [None, None]
          self._are_phases_completed = [False, False]
  
      @property
      def iteration(self):
          if None in self._phases_last_iterations:
              return 0
          return self._phases_last_iterations
  
      @property
      def is_completed(self):
          return all(self._are_phases_completed)
  
      def run(self, n_iterations=None):
          # Handle default argument.
          if n_iterations is None:
              n_iterations = self.number_of_iterations
  
          # Handle case in which we don't alternate between phases.
          if self.switch_phase_interval <= 0:
              switch_phase_interval = self.number_of_iterations
  
          # Count down the
  
  [... the rest is omitted ...]
files: Yank/experiment.py