# Source code for propertyestimator.protocols.yank

    """
    A collection of protocols for performing free energy calculations
    using the YANK package.
    """
    import logging
    import os
    import shutil
    import threading
    import traceback
    from enum import Enum

    import yaml

    from propertyestimator import unit
    from propertyestimator.forcefield import SmirnoffForceFieldSource
    from propertyestimator.substances import Substance
    from propertyestimator.thermodynamics import ThermodynamicState
    from propertyestimator.utils.exceptions import PropertyEstimatorException
    from propertyestimator.utils.openmm import setup_platform_with_resources, openmm_quantity_to_pint, \
        pint_quantity_to_openmm, disable_pbc
    from propertyestimator.utils.quantities import EstimatedQuantity
    from propertyestimator.utils.utils import temporarily_change_directory
    from propertyestimator.workflow.decorators import protocol_input, protocol_output, InequalityMergeBehaviour, UNDEFINED
    from propertyestimator.workflow.plugins import register_calculation_protocol
    from propertyestimator.workflow.protocols import BaseProtocol

    @register_calculation_protocol()
    class BaseYankProtocol(BaseProtocol):
        """An abstract base class for protocols which will performs a set of alchemical
        free energy simulations using the YANK framework.

        Protocols which inherit from this base must implement the abstract `_get_yank_options`
        methods.
        """

        thermodynamic_state = protocol_input(
            docstring='The state at which to run the calculations.',
            type_hint=ThermodynamicState,
            default_value=UNDEFINED
        )

        number_of_equilibration_iterations = protocol_input(
            docstring='The number of iterations used for equilibration before production run. '
                      'Only post-equilibration iterations are written to file.',
            type_hint=int, merge_behavior=InequalityMergeBehaviour.LargestValue,
            default_value=1
        )
        number_of_iterations = protocol_input(
            docstring='The number of YANK iterations to perform.',
            type_hint=int, merge_behavior=InequalityMergeBehaviour.LargestValue,
            default_value=5000
        )
        steps_per_iteration = protocol_input(
            docstring='The number of steps per YANK iteration to perform.',
            type_hint=int, merge_behavior=InequalityMergeBehaviour.LargestValue,
            default_value=500
        )
        checkpoint_interval = protocol_input(
            docstring='The number of iterations between saving YANK checkpoint files.',
            type_hint=int, merge_behavior=InequalityMergeBehaviour.SmallestValue,
            default_value=50
        )

        timestep = protocol_input(
            docstring='The length of the timestep to take.',
            type_hint=unit.Quantity, merge_behavior=InequalityMergeBehaviour.SmallestValue,
            default_value=2 * unit.femtosecond
        )

        verbose = protocol_input(
            docstring='Controls whether or not to run YANK at high verbosity.',
            type_hint=bool,
            default_value=False
        )
        setup_only = protocol_input(
            docstring='If true, YANK will only create and validate the setup files, '
                      'but not actually run any simulations. This argument is mainly '
                      'only to be used for testing purposes.',
            type_hint=bool,
            default_value=False
        )

        estimated_free_energy = protocol_output(
            docstring='The estimated free energy value and its uncertainty '
                      'returned by YANK.',
            type_hint=EstimatedQuantity
        )

        @staticmethod
        def _get_residue_names_from_role(substances, coordinate_path, role):
            """Returns a list of all of the residue names of
            components which have been assigned a given role.

            Parameters
            ----------
            substances: list of Substance
                The substances which contains the components.
            coordinate_path: str
                The path to the coordinates which describe the systems
                topology.
            role: Substance.ComponentRole
                The role of the component to identify.

            Returns
            -------
            set of str
                The identified residue names.
            """

            from simtk.openmm import app
            from openforcefield.topology import Molecule, Topology

            if role == Substance.ComponentRole.Undefined:
                return 'all'

            unique_molecules = [Molecule.from_smiles(component.smiles) for
                                substance in substances for
                                component in substance.components]

            openmm_topology = app.PDBFile(coordinate_path).topology
            topology = Topology.from_openmm(openmm_topology, unique_molecules)

            # Determine the smiles of all molecules in the system. We need to use
            # the toolkit to re-generate the smiles as later we will compare these
            # against more toolkit generated smiles.
            components = [component for substance in substances for
                          component in substance.components if
                          component.role == role]

            component_smiles = [Molecule.from_smiles(component.smiles).to_smiles() for
                                component in components]

            residue_names = set()

            all_openmm_atoms = list(openmm_topology.atoms())

            # Find the resiude names of the molecules which have the correct
            # role.
            for topology_molecule in topology.topology_molecules:

                molecule_smiles = topology_molecule.reference_molecule.to_smiles()

                if molecule_smiles not in component_smiles:
                    continue

                molecule_residue_names = set([all_openmm_atoms[topology_atom.topology_atom_index].residue.name for
                                              topology_atom in topology_molecule.atoms])

                assert len(molecule_residue_names) == 1
                residue_names.update(molecule_residue_names)

            return residue_names

        @staticmethod
        def _get_dsl_from_role(substances, coordinate_path, role):
            """Returns an MDTraj DSL string which identifies those
            atoms which belong to components flagged with a specific
            role.

            Parameters
            ----------
            substances: list of Substance
                The substances which contains the components.
            coordinate_path: str
                The path to the coordinates which describe the systems
                topology.
            role: Substance.ComponentRole
                The role of the component to identify.

            Returns
            -------
            str
                The DSL string.
            """

            residue_names = BaseYankProtocol._get_residue_names_from_role(substances, coordinate_path, role)

            dsl_string = ' or '.join([f'resname {residue_name}' for residue_name in residue_names])
            return dsl_string

        def _get_options_dictionary(self, available_resources):
            """Returns a dictionary of options which will be serialized
            to a yaml file and passed to YANK.

            Parameters
            ----------
            available_resources: ComputeResources
                The resources available to execute on.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of YANK options.
            """

            from openforcefield.utils import quantity_to_string

            platform_name = 'CPU'

            if available_resources.number_of_gpus > 0:

                # A platform which runs on GPUs has been requested.
                from propertyestimator.backends import ComputeResources
                toolkit_enum = ComputeResources.GPUToolkit(available_resources.preferred_gpu_toolkit)

                # A platform which runs on GPUs has been requested.
                platform_name = 'CUDA' if toolkit_enum == ComputeResources.GPUToolkit.CUDA else \
                                                          ComputeResources.GPUToolkit.OpenCL

            return {
                'verbose': self.verbose,
                'output_dir': '.',

                'temperature': quantity_to_string(pint_quantity_to_openmm(self.thermodynamic_state.temperature)),
                'pressure': quantity_to_string(pint_quantity_to_openmm(self.thermodynamic_state.pressure)),

                'minimize': True,

                'number_of_equilibration_iterations': self.number_of_equilibration_iterations,
                'default_number_of_iterations': self.number_of_iterations,
                'default_nsteps_per_iteration': self.steps_per_iteration,
                'checkpoint_interval': self.checkpoint_interval,

                'default_timestep': quantity_to_string(pint_quantity_to_openmm(self.timestep)),

                'annihilate_electrostatics': True,
                'annihilate_sterics': False,

                'platform': platform_name
            }

        def _get_system_dictionary(self):
            """Returns a dictionary of the system which will be serialized
            to a yaml file and passed to YANK. Only a single system may be
            specified.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of YANK systems.
            """
            raise NotImplementedError()

        def _get_protocol_dictionary(self):
            """Returns a dictionary of the protocol which will be serialized
            to a yaml file and passed to YANK. Only a single protocol may be
            specified.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of a YANK protocol.
            """
            raise NotImplementedError()

        def _get_experiments_dictionary(self):
            """Returns a dictionary of the experiments which will be serialized
            to a yaml file and passed to YANK. Only a single experiment may be
            specified.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of a YANK experiment.
            """

            system_dictionary = self._get_system_dictionary()
            system_key = next(iter(system_dictionary))

            protocol_dictionary = self._get_protocol_dictionary()
            protocol_key = next(iter(protocol_dictionary))

            return {
                'system': system_key,
                'protocol': protocol_key
            }

        def _get_full_input_dictionary(self, available_resources):
            """Returns a dictionary of the full YANK inputs which will be serialized
            to a yaml file and passed to YANK

            Parameters
            ----------
            available_resources: ComputeResources
                The resources available to execute on.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of a YANK input file.
            """

            return {
                'options': self._get_options_dictionary(available_resources),

                'systems': self._get_system_dictionary(),
                'protocols': self._get_protocol_dictionary(),

                'experiments': self._get_experiments_dictionary()
            }

        @staticmethod
        def _extract_trajectory(checkpoint_path, output_trajectory_path):
            """Extracts the stored trajectory of the 'initial' state from a
            yank `.nc` checkpoint file and stores it to disk as a `.dcd` file.

            Parameters
            ----------
            checkpoint_path: str
                The path to the yank `.nc` file
            output_trajectory_path: str
                The path to store the extracted trajectory at.
            """

            from yank.analyze import extract_trajectory

            mdtraj_trajectory = extract_trajectory(checkpoint_path, state_index=0, image_molecules=True)
            mdtraj_trajectory.save_dcd(output_trajectory_path)

        @staticmethod
        def _run_yank(directory, available_resources, setup_only):
            """Runs YANK within the specified directory which contains a `yank.yaml`
            input file.

            Parameters
            ----------
            directory: str
                The directory within which to run yank.
            available_resources: ComputeResources
                The compute resources available to yank.
            setup_only: bool
                If true, YANK will only create and validate the setup files,
                but not actually run any simulations. This argument is mainly
                only to be used for testing purposes.

            Returns
            -------
            simtk.unit.Quantity
                The free energy returned by yank.
            simtk.unit.Quantity
                The uncertainty in the free energy returned by yank.
            """

            from yank.experiment import ExperimentBuilder
            from yank.analyze import ExperimentAnalyzer

            from simtk import unit as simtk_unit

            with temporarily_change_directory(directory):

                # Set the default properties on the desired platform
                # before calling into yank.
                setup_platform_with_resources(available_resources)

                exp_builder = ExperimentBuilder('yank.yaml')

                if setup_only is True:
                    return 0.0 * simtk_unit.kilojoule_per_mole, 0.0 * simtk_unit.kilojoule_per_mole

                exp_builder.run_experiments()

                analyzer = ExperimentAnalyzer('experiments')
                output = analyzer.auto_analyze()

                free_energy = output['free_energy']['free_energy_diff_unit']
                free_energy_uncertainty = output['free_energy']['free_energy_diff_error_unit']

            return free_energy, free_energy_uncertainty

        @staticmethod
        def _run_yank_as_process(queue, directory, available_resources, setup_only):
            """A wrapper around the `_run_yank` method which takes
            a `multiprocessing.Queue` as input, thereby allowing it
            to be launched from a separate process and still return
            it's output back to the main process.

            Parameters
            ----------
            queue: multiprocessing.Queue
                The queue object which will communicate with the
                launched process.
            directory: str
                The directory within which to run yank.
            available_resources: ComputeResources
                The compute resources available to yank.
            setup_only: bool
                If true, YANK will only create and validate the setup files,
                but not actually run any simulations. This argument is mainly
                only to be used for testing purposes.

            Returns
            -------
            simtk.unit.Quantity
                The free energy returned by yank.
            simtk.unit.Quantity
                The uncertainty in the free energy returned by yank.
            str, optional
                The stringified errors which occurred on the other process,
                or `None` if no exceptions were raised.
            """

            free_energy = None
            free_energy_uncertainty = None

            error = None

            try:
                free_energy, free_energy_uncertainty = BaseYankProtocol._run_yank(directory,
                                                                                  available_resources,
                                                                                  setup_only)
            except Exception as e:
                error = traceback.format_exception(None, e, e.__traceback__)

            queue.put((free_energy, free_energy_uncertainty, error))

        def execute(self, directory, available_resources):

            yaml_filename = os.path.join(directory, 'yank.yaml')

            # Create the yank yaml input file from a dictionary of options.
            with open(yaml_filename, 'w') as file:
                yaml.dump(self._get_full_input_dictionary(available_resources), file, sort_keys=False)

            setup_only = self.setup_only

            # Yank is not safe to be called from anything other than the main thread.
            # If the current thread is not detected as the main one, then yank should
            # be spun up in a new process which should itself be safe to run yank in.
            if threading.current_thread() is threading.main_thread():
                logging.info('Launching YANK in the main thread.')
                free_energy, free_energy_uncertainty = self._run_yank(directory, available_resources, setup_only)
            else:

                from multiprocessing import Process, Queue

                logging.info('Launching YANK in a new process.')

                # Create a queue to pass the results back to the main process.
                queue = Queue()
                # Create the process within which yank will run.
                process = Process(target=BaseYankProtocol._run_yank_as_process, args=[queue, directory,
                                                                                      available_resources, setup_only])

                # Start the process and gather back the output.
                process.start()
                free_energy, free_energy_uncertainty, error = queue.get()
                process.join()

                if error is not None:
                    return PropertyEstimatorException(directory, error)

            self.estimated_free_energy = EstimatedQuantity(openmm_quantity_to_pint(free_energy),
                                                           openmm_quantity_to_pint(free_energy_uncertainty),
                                                           self._id)

            return self._get_output_dictionary()

    @register_calculation_protocol()
    class LigandReceptorYankProtocol(BaseYankProtocol):
        """A protocol for performing ligand-receptor alchemical free energy
        calculations using the YANK framework.
        """

        class RestraintType(Enum):
            """The types of ligand restraints available within yank.
            """
            Harmonic = 'Harmonic'
            FlatBottom = 'FlatBottom'

        ligand_residue_name = protocol_input(
            docstring='The residue name of the ligand.',
            type_hint=str,
            default_value=UNDEFINED
        )
        receptor_residue_name = protocol_input(
            docstring='The residue name of the receptor.',
            type_hint=str,
            default_value=UNDEFINED
        )

        solvated_ligand_coordinates = protocol_input(
            docstring='The file path to the solvated ligand coordinates.',
            type_hint=str,
            default_value=UNDEFINED
        )
        solvated_ligand_system = protocol_input(
            docstring='The file path to the solvated ligand system object.',
            type_hint=str,
            default_value=UNDEFINED
        )

        solvated_complex_coordinates = protocol_input(
            docstring='The file path to the solvated complex coordinates.',
            type_hint=str,
            default_value=UNDEFINED
        )
        solvated_complex_system = protocol_input(
            docstring='The file path to the solvated complex system object.',
            type_hint=str,
            default_value=UNDEFINED
        )

        force_field_path = protocol_input(
            docstring='The path to the force field which defines the charge method '
                      'to use for the calculation.',
            type_hint=str,
            default_value=UNDEFINED
        )

        apply_restraints = protocol_input(
            docstring='Determines whether the ligand should be explicitly restrained to the '
                      'receptor in order to stop the ligand from temporarily unbinding.',
            type_hint=bool,
            default_value=True
        )
        restraint_type = protocol_input(
            docstring='The type of ligand restraint applied, provided that `apply_restraints` '
                      'is `True`',
            type_hint=RestraintType,
            default_value=RestraintType.Harmonic
        )

        solvated_ligand_trajectory_path = protocol_output(
            docstring='The file path to the generated ligand trajectory.',
            type_hint=str
        )
        solvated_complex_trajectory_path = protocol_output(
            docstring='The file path to the generated ligand trajectory.',
            type_hint=str
        )

        def __init__(self, protocol_id):
            """Constructs a new LigandReceptorYankProtocol object."""

            super().__init__(protocol_id)

            self._local_ligand_coordinates = 'ligand.pdb'
            self._local_ligand_system = 'ligand.xml'

            self._local_complex_coordinates = 'complex.pdb'
            self._local_complex_system = 'complex.xml'

        def _get_solvent_dictionary(self):
            """Returns a dictionary of the solvent which will be serialized
            to a yaml file and passed to YANK. In most cases, this should
            just be passing force field settings over, such as PME settings.

            Returns
            -------
            dict of str and Any
                A yaml compatible dictionary of YANK solvents.
            """

            with open(self.force_field_path, 'r') as file:
                force_field_source = SmirnoffForceFieldSource.parse_json(file.read())

            force_field = force_field_source.to_force_field()
            charge_method = force_field.get_parameter_handler('Electrostatics').method

            if charge_method.lower() != 'pme':
                raise ValueError('Currently only PME electrostatics are supported.')

            return {'default': {
                'nonbonded_method': charge_method,
            }}

        def _get_system_dictionary(self):

            solvent_dictionary = self._get_solvent_dictionary()
            solvent_key = next(iter(solvent_dictionary))

            host_guest_dictionary = {
                'phase1_path': [self._local_complex_system, self._local_complex_coordinates],
                'phase2_path': [self._local_ligand_system, self._local_ligand_coordinates],

                'ligand_dsl': f'resname {self.ligand_residue_name}',
                'solvent': solvent_key
            }

            return {'host-guest': host_guest_dictionary}

        def _get_protocol_dictionary(self):

            absolute_binding_dictionary = {
                'complex': {'alchemical_path': 'auto'},
                'solvent': {'alchemical_path': 'auto'}
            }

            return {'absolute_binding_dictionary': absolute_binding_dictionary}

        def _get_experiments_dictionary(self):

            experiments_dictionary = super(LigandReceptorYankProtocol, self)._get_experiments_dictionary()

            if self.apply_restraints:

                experiments_dictionary['restraint'] = {
                    'restrained_ligand_atoms': f'(resname {self.ligand_residue_name}) and (mass > 1.5)',
                    'restrained_receptor_atoms': f'(resname {self.receptor_residue_name}) and (mass > 1.5)',

                    'type': self.restraint_type.value
                }

            return experiments_dictionary

        def _get_full_input_dictionary(self, available_resources):

            full_dictionary = super(LigandReceptorYankProtocol, self)._get_full_input_dictionary(available_resources)
            full_dictionary['solvents'] = self._get_solvent_dictionary()

            return full_dictionary

        def execute(self, directory, available_resources):

            # Because of quirks in where Yank looks files while doing temporary
            # directory changes, we need to copy the coordinate files locally so
            # they are correctly found.
            shutil.copyfile(self.solvated_ligand_coordinates, os.path.join(directory, self._local_ligand_coordinates))
            shutil.copyfile(self.solvated_ligand_system, os.path.join(directory, self._local_ligand_system))

            shutil.copyfile(self.solvated_complex_coordinates, os.path.join(directory, self._local_complex_coordinates))
            shutil.copyfile(self.solvated_complex_system, os.path.join(directory, self._local_complex_system))

            result = super(LigandReceptorYankProtocol, self).execute(directory, available_resources)

            if isinstance(result, PropertyEstimatorException):
                return result

            if self.setup_only:
                return self._get_output_dictionary()

            ligand_yank_path = os.path.join(directory, 'experiments', 'solvent.nc')
            complex_yank_path = os.path.join(directory, 'experiments', 'complex.nc')

            self.solvated_ligand_trajectory_path = os.path.join(directory, 'ligand.dcd')
            self.solvated_complex_trajectory_path = os.path.join(directory, 'complex.dcd')

            self._extract_trajectory(ligand_yank_path, self.solvated_ligand_trajectory_path)
            self._extract_trajectory(complex_yank_path, self.solvated_complex_trajectory_path)

            return self._get_output_dictionary()

    @register_calculation_protocol()
    class SolvationYankProtocol(BaseYankProtocol):
        """A protocol for performing solvation alchemical free energy
        calculations using the YANK framework.

        This protocol can be used for box solvation free energies (setting
        the `solvent_1` input to the solvent of interest and setting
        `solvent_2` as an empty `Substance`) or transfer free energies (setting
        both the `solvent_1` and `solvent_2` inputs to different solvents).
        """

        solute = protocol_input(
            docstring='The substance describing the composition of '
                      'the solute. This should include the solute '
                      'molecule as well as any counter ions.',
            type_hint=Substance,
            default_value=UNDEFINED
        )

        solvent_1 = protocol_input(
            docstring='The substance describing the composition of '
                      'the first solvent.',
            type_hint=Substance,
            default_value=UNDEFINED
        )
        solvent_2 = protocol_input(
            docstring='The substance describing the composition of '
                      'the second solvent.',
            type_hint=Substance,
            default_value=UNDEFINED
        )

        solvent_1_coordinates = protocol_input(
            docstring='The file path to the coordinates of the solute embedded in the '
                      'first solvent.',
            type_hint=str,
            default_value=UNDEFINED
        )
        solvent_1_system = protocol_input(
            docstring='The file path to the system object of the solute embedded in the '
                      'first solvent.',
            type_hint=str,
            default_value=UNDEFINED
        )

        solvent_2_coordinates = protocol_input(
            docstring='The file path to the coordinates of the solute embedded in the '
                      'second solvent.',
            type_hint=str,
            default_value=UNDEFINED
        )
        solvent_2_system = protocol_input(
            docstring='The file path to the system object of the solute embedded in the '
                      'second solvent.',
            type_hint=str,
            default_value=UNDEFINED
        )

        electrostatic_lambdas_1 = protocol_input(
            docstring='The list of electrostatic alchemical states that YANK should sample at. '
                      'These values will be passed to the YANK `lambda_electrostatics` option. '
                      'If no option is set, YANK will use `trailblaze` algorithm to determine '
                      'this option automatically.',
            type_hint=list,
            optional=True,
            default_value=UNDEFINED
        )
        steric_lambdas_1 = protocol_input(
            docstring='The list of steric alchemical states that YANK should sample at. '
                      'These values will be passed to the YANK `lambda_sterics` option. '
                      'If no option is set, YANK will use `trailblaze` algorithm to determine '
                      'this option automatically.',
            type_hint=list,
            optional=True,
            default_value=UNDEFINED
        )
        electrostatic_lambdas_2 = protocol_input(
            docstring='The list of electrostatic alchemical states that YANK should sample at. '
                      'These values will be passed to the YANK `lambda_electrostatics` option. '
                      'If no option is set, YANK will use `trailblaze` algorithm to determine '
                      'this option automatically.',
            type_hint=list,
            optional=True,
            default_value=UNDEFINED
        )
        steric_lambdas_2 = protocol_input(
            docstring='The list of steric alchemical states that YANK should sample at. '
                      'These values will be passed to the YANK `lambda_sterics` option. '
                      'If no option is set, YANK will use `trailblaze` algorithm to determine '
                      'this option automatically.',
            type_hint=list,
            optional=True,
            default_value=UNDEFINED
        )

        solvent_1_trajectory_path = protocol_output(
            docstring='The file path to the trajectory of the solute in the '
                      'first solvent.',
            type_hint=str
        )
        solvent_2_trajectory_path = protocol_output(
            docstring='The file path to the trajectory of the solute in the '
                      'second solvent.',
            type_hint=str
        )

        def __init__(self, protocol_id):
            super().__init__(protocol_id)

            self._local_solvent_1_coordinates = 'solvent_1.pdb'
            self._local_solvent_1_system = 'solvent_1.xml'

            self._local_solvent_2_coordinates = 'solvent_2.pdb'
            self._local_solvent_2_system = 'solvent_2.xml'

        def _get_system_dictionary(self):

            solvent_1_dsl = self._get_dsl_from_role([self.solute, self.solvent_1],
                                                    self.solvent_1_coordinates,
                                                    Substance.ComponentRole.Solvent)

            solvent_2_dsl = self._get_dsl_from_role([self.solute, self.solvent_2],
                                                    self.solvent_2_coordinates,
                                                    Substance.ComponentRole.Solvent)

            full_solvent_dsl_components = []

            if len(solvent_1_dsl) > 0:
                full_solvent_dsl_components.append(solvent_1_dsl)
            if len(solvent_2_dsl) > 0:
                full_solvent_dsl_components.append(solvent_2_dsl)

            solvation_system_dictionary = {
                'phase1_path': [self._local_solvent_1_system, self._local_solvent_1_coordinates],
                'phase2_path': [self._local_solvent_2_system, self._local_solvent_2_coordinates],

                'solvent_dsl': ' or '.join(full_solvent_dsl_components)
            }

            return {'solvation-system': solvation_system_dictionary}

        def _get_protocol_dictionary(self):

            solvent_1_protocol_dictionary = {
                'lambda_electrostatics': self.electrostatic_lambdas_1,
                'lambda_sterics': self.steric_lambdas_1
            }

            if self.electrostatic_lambdas_1 == UNDEFINED and self.steric_lambdas_1 == UNDEFINED:

                solvent_1_protocol_dictionary = 'auto'

            elif ((self.electrostatic_lambdas_1 != UNDEFINED and self.steric_lambdas_1 == UNDEFINED) or
                  (self.electrostatic_lambdas_1 == UNDEFINED and self.steric_lambdas_1 != UNDEFINED)):

                raise ValueError('Either both of `electrostatic_lambdas_1` and '
                                 '`steric_lambdas_1` must be set, or neither '
                                 'must be set.')

            solvent_2_protocol_dictionary = {
                'lambda_electrostatics': self.electrostatic_lambdas_2,
                'lambda_sterics': self.steric_lambdas_2
            }

            if self.electrostatic_lambdas_2 == UNDEFINED and self.steric_lambdas_2 == UNDEFINED:

                solvent_2_protocol_dictionary = 'auto'

            elif ((self.electrostatic_lambdas_2 != UNDEFINED and self.steric_lambdas_2 == UNDEFINED) or
                  (self.electrostatic_lambdas_2 == UNDEFINED and self.steric_lambdas_2 != UNDEFINED)):

                raise ValueError('Either both of `electrostatic_lambdas_2` and '
                                 '`steric_lambdas_2` must be set, or neither '
                                 'must be set.')

            protocol_dictionary = {
                'solvent1': {'alchemical_path': solvent_1_protocol_dictionary},
                'solvent2': {'alchemical_path': solvent_2_protocol_dictionary}
            }

            return {'solvation-protocol': protocol_dictionary}

        def execute(self, directory, available_resources):

            from simtk.openmm import XmlSerializer

            solute_components = [component for component in self.solute.components if
                                 component.role == Substance.ComponentRole.Solute]

            solvent_1_components = [component for component in self.solvent_1.components if
                                    component.role == Substance.ComponentRole.Solvent]

            solvent_2_components = [component for component in self.solvent_2.components if
                                    component.role == Substance.ComponentRole.Solvent]

            if len(solute_components) != 1:
                return PropertyEstimatorException(directory, 'There must only be a single component marked as a solute.')
            if len(solvent_1_components) == 0 and len(solvent_2_components) == 0:
                return PropertyEstimatorException(directory, 'At least one of the solvents must not be vacuum.')

            # Because of quirks in where Yank looks files while doing temporary
            # directory changes, we need to copy the coordinate files locally so
            # they are correctly found.
            shutil.copyfile(self.solvent_1_coordinates, os.path.join(directory, self._local_solvent_1_coordinates))
            shutil.copyfile(self.solvent_1_system, os.path.join(directory, self._local_solvent_1_system))

            shutil.copyfile(self.solvent_2_coordinates, os.path.join(directory, self._local_solvent_2_coordinates))
            shutil.copyfile(self.solvent_2_system, os.path.join(directory, self._local_solvent_2_system))

            # Disable the pbc of the any solvents which should be treated
            # as vacuum.
            vacuum_system_path = None

            if len(solvent_1_components) == 0:
                vacuum_system_path = self._local_solvent_1_system
            elif len(solvent_2_components) == 0:
                vacuum_system_path = self._local_solvent_2_system

            if vacuum_system_path is not None:

                logging.info(f'Disabling the periodic boundary conditions in {vacuum_system_path} '
                             f'by setting the cutoff type to NoCutoff')

                with open(os.path.join(directory, vacuum_system_path), 'r') as file:
                    vacuum_system = XmlSerializer.deserialize(file.read())

                disable_pbc(vacuum_system)

                with open(os.path.join(directory, vacuum_system_path), 'w') as file:
                    file.write(XmlSerializer.serialize(vacuum_system))

            # Set up the yank input file.
            result = super(SolvationYankProtocol, self).execute(directory, available_resources)

            if isinstance(result, PropertyEstimatorException):
                return result

            if self.setup_only:
                return self._get_output_dictionary()

            solvent_1_yank_path = os.path.join(directory, 'experiments', 'solvent1.nc')
            solvent_2_yank_path = os.path.join(directory, 'experiments', 'solvent2.nc')

            self.solvent_1_trajectory_path = os.path.join(directory, 'solvent1.dcd')
            self.solvent_2_trajectory_path = os.path.join(directory, 'solvent2.dcd')

            self._extract_trajectory(solvent_1_yank_path, self.solvent_1_trajectory_path)
            self._extract_trajectory(solvent_2_yank_path, self.solvent_2_trajectory_path)

            return self._get_output_dictionary()