From 9c2f9a7dfbbe1c1dfa03519b09b23c42a1bdf6e9 Mon Sep 17 00:00:00 2001
From: Ana Silveira <silveira@lilac.hpc.private>
Date: Fri, 4 Oct 2019 16:24:46 -0400
Subject: [PATCH 1/8] add support for relative binding free energy

---
 Yank/experiment.py       |  4 ++--
 Yank/restraints.py       | 39 +++++++++++++++++++++-----------
 Yank/schema/validator.py |  4 +++-
 Yank/yank.py             | 48 +++++++++++++++++++++++++---------------
 4 files changed, 61 insertions(+), 34 deletions(-)

diff --git a/Yank/experiment.py b/Yank/experiment.py
index 80938058..3f50d73d 100644
--- a/Yank/experiment.py
+++ b/Yank/experiment.py
@@ -1664,7 +1664,7 @@ def _is_pipeline_solvent_with_receptor(field, solvent_id, error):
         # DSL Schema
         ligand_dsl:
             required: no
-            type: string
+            type: list
             dependencies: [phase1_path, phase2_path]
         solvent_dsl:
             required: no
@@ -3009,7 +3009,7 @@ def _build_experiment(self, experiment_path, experiment, use_dummy_protocol=Fals
             phase_protocol = protocol[phase_name]['alchemical_path']
             alchemical_region = AlchemicalPhase._build_default_alchemical_region(system, topography,
                                                                                  phase_protocol)
-            alchemical_region = alchemical_region._replace(**alchemical_region_opts)
+            alchemical_region = [region._replace(**alchemical_region_opts) for region in alchemical_region]
 
             # Apply restraint only if this is the first phase. AlchemicalPhase
             # will take care of raising an error if the phase type does not support it.
diff --git a/Yank/restraints.py b/Yank/restraints.py
index 13599f3b..be389980 100644
--- a/Yank/restraints.py
+++ b/Yank/restraints.py
@@ -434,7 +434,10 @@ def __get__(self, instance, owner_class=None):
     def __set__(self, instance, new_restrained_atoms):
         # If we set the restrained attributes to None, no reason to check things.
         if new_restrained_atoms is not None:
-            new_restrained_atoms = self._validate_atoms(new_restrained_atoms)
+            if len(new_restrained_atoms) == 2: # relative system
+                new_restrained_atoms = [self._validate_atoms(atoms) for atoms in new_restrained_atoms]
+            else:
+                new_restrained_atoms = self._validate_atoms(new_restrained_atoms)
         setattr(instance, self._attribute_name, new_restrained_atoms)
 
     @methoddispatch
@@ -572,9 +575,11 @@ def restrain_state(self, thermodynamic_state):
                                           'atoms.'.format(self.__class__.__name__))
 
         # Create restraint force.
-        restraint_force = self._get_restraint_force(self.restrained_receptor_atoms,
-                                                    self.restrained_ligand_atoms)
-
+        if len(self.restrained_ligand_atoms) == 1: # not a relative system
+            restraint_force = self._get_restraint_force(self.restrained_receptor_atoms,
+                                                        self.restrained_ligand_atoms)
+        else:
+            restraint_force = self._get_restraint_force(*self.restrained_ligand_atoms)
         # Set periodic conditions on the force if necessary.
         restraint_force.setUsesPeriodicBoundaryConditions(thermodynamic_state.is_periodic)
 
@@ -700,10 +705,15 @@ def _determine_restraint_parameters(self, thermodynamic_state, sampler_state, to
     @property
     def _are_restrained_atoms_defined(self):
         """Check if the restrained atoms are defined well enough to make a restraint"""
-        for atoms in [self.restrained_receptor_atoms, self.restrained_ligand_atoms]:
+        for atoms in [self.restrained_ligand_atoms]:
             # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
-            if atoms is None or not (isinstance(atoms, list) and len(atoms) > 0):
+            if atoms is None or any(isinstance(atom, str) for atom in atoms) or not (isinstance(atoms, list) and len(atoms) > 0):
                 return False
+        if len(self.restrained_ligand_atoms) == 1:
+            for atoms in [self.restrained_receptor_atoms]:
+                if atoms is None or not (isinstance(atoms, list) and len(atoms) > 0):
+                    return False
+
         return True
 
     def _get_restraint_force(self, particles1, particles2):
@@ -743,7 +753,8 @@ def _determine_restrained_atoms(self, sampler_state, topography):
 
         # If receptor and ligand atoms are explicitly provided, use those.
         restrained_ligand_atoms = self.restrained_ligand_atoms
-        restrained_receptor_atoms = self.restrained_receptor_atoms
+        if len(restrained_ligand_atoms) != 2:
+            restrained_receptor_atoms = self.restrained_receptor_atoms
 
         @functools.singledispatch
         def compute_atom_set(input_atoms, topography_key, mapping_function):
@@ -775,17 +786,19 @@ def compute_atom_none(_, topography_key, mapping_function):
         def compute_atom_str(input_string, topography_key, _):
             """Helper for string parsing"""
             selection = topography.select(input_string, as_set=True)
-            selection_with_top = selection & set(getattr(topography, topography_key))
+            selection_with_top = selection  & set([index for sublist in getattr(topography, topography_key) for
+                                                   index in sublist])
             # Force output to be a normal int, dont need to worry about floats at this point, there should not be any
             # If they come out as np.int64's, OpenMM complains
             return [*map(int, selection_with_top)]
 
-        self.restrained_ligand_atoms = compute_atom_set(restrained_ligand_atoms,
+        self.restrained_ligand_atoms = [compute_atom_set(list,
                                                         'ligand_atoms',
-                                                        self._closest_atom_to_centroid)
-        self.restrained_receptor_atoms = compute_atom_set(restrained_receptor_atoms,
-                                                          'receptor_atoms',
-                                                          self._closest_atom_to_centroid)
+                                                        self._closest_atom_to_centroid) for list in restrained_ligand_atoms]
+        if len(self.restrained_ligand_atoms) != 2:
+            self.restrained_receptor_atoms = compute_atom_set(restrained_receptor_atoms,
+                                                              'receptor_atoms',
+                                                              self._closest_atom_to_centroid)
 
     @staticmethod
     def _closest_atom_to_centroid(positions, indices=None, masses=None):
diff --git a/Yank/schema/validator.py b/Yank/schema/validator.py
index 83332bbd..21359f2f 100644
--- a/Yank/schema/validator.py
+++ b/Yank/schema/validator.py
@@ -188,7 +188,9 @@ def _check_with_only_with_no_cutoff(self, field, value):
     def _check_with_specify_lambda_electrostatics_and_sterics(self, field, value):
         """Check that the keys of a dictionary contain both lambda_electrostatics and lambda_sterics."""
         if ((isinstance(value, dict) or isinstance(value, collections.OrderedDict)) and
-                not ('lambda_sterics' in value and 'lambda_electrostatics' in value)):
+             not (('lambda_sterics' in value and 'lambda_electrostatics' in value) or
+             ('lambda_sterics_zero' in value and 'lambda_electrostatics_zero' in value and
+               'lambda_sterics_one' in value and 'lambda_electrostatics_one' in value))):
             self._error(field, "Missing required keys lambda_sterics and/or lambda_electrostatics")
 
     def _check_with_math_expressions_variables_are_given(self, field, value):
diff --git a/Yank/yank.py b/Yank/yank.py
index a9f8f9c3..81ddf618 100644
--- a/Yank/yank.py
+++ b/Yank/yank.py
@@ -120,8 +120,7 @@ def ligand_atoms(self):
 
     @ligand_atoms.setter
     def ligand_atoms(self, value):
-        self._ligand_atoms = self.select(value)
-
+        self._ligand_atoms = tuple([self.select(v) for v in value])
         # Safety check: with a ligand there should always be a receptor.
         if len(self._ligand_atoms) > 0 and len(self.receptor_atoms) == 0:
             raise ValueError('Specified ligand but cannot find '
@@ -141,7 +140,7 @@ def receptor_atoms(self):
             return []
 
         # Create a set for fast searching.
-        ligand_atomset = frozenset(self._ligand_atoms)
+        ligand_atomset = tuple(frozenset([index for ligand in self._ligand_atoms for index in ligand]))
         # Receptor atoms are all solute atoms that are not ligand.
         return [i for i in self.solute_atoms if i not in ligand_atomset]
 
@@ -1012,8 +1011,9 @@ def create(self, thermodynamic_state, sampler_states, topography, protocol,
                 reference_system, topography, protocol)
 
         # Check that we have atoms to alchemically modify.
-        if len(alchemical_regions.alchemical_atoms) == 0:
-            raise ValueError("Couldn't find atoms to alchemically modify.")
+        for region in alchemical_regions:
+            if len(region.alchemical_atoms) == 0:
+                raise ValueError("Couldn't find atoms to alchemically modify.")
 
         # Create alchemically-modified system using alchemical factory.
         logger.debug("Creating alchemically-modified states...")
@@ -1024,11 +1024,16 @@ def create(self, thermodynamic_state, sampler_states, topography, protocol,
 
         # Create compound alchemically modified state to pass to sampler.
         thermodynamic_state.system = alchemical_system
-        alchemical_state = mmtools.alchemy.AlchemicalState.from_system(alchemical_system)
-        if restraint_state is not None:
-            composable_states = [alchemical_state, restraint_state]
+        if len(alchemical_regions) > 1:
+            alchemical_state = [mmtools.alchemy.AlchemicalState.from_system(alchemical_system, parameters_name_suffix = 'zero'),
+                                mmtools.alchemy.AlchemicalState.from_system(alchemical_system, parameters_name_suffix = 'one')]
         else:
-            composable_states = [alchemical_state]
+            alchemical_state = [mmtools.alchemy.AlchemicalState.from_system(alchemical_system)]
+
+        composable_states = [state for state in alchemical_state]
+        if restraint_state is not None:
+            composable_states += [restraint_state]
+
         compound_state = mmtools.states.CompoundThermodynamicState(
             thermodynamic_state=thermodynamic_state, composable_states=composable_states)
 
@@ -1340,25 +1345,28 @@ def _build_default_alchemical_region(system, topography, protocol):
 
         # Modify ligand if this is a receptor-ligand phase, or
         # solute if this is a transfer free energy calculation.
-        if len(topography.ligand_atoms) > 0:
-            alchemical_region_name = 'ligand_atoms'
-        else:
+        alchemical_region_name = 'ligand_atoms'
+        if len(topography.ligand_atoms) == 0:
             alchemical_region_name = 'solute_atoms'
         alchemical_atoms = getattr(topography, alchemical_region_name)
+        if len(topography.ligand_atoms) > 1:
+            alchemical_region_name = ['ligand_atoms_zero', 'ligand_atoms_one']
 
         # In periodic systems, we alchemically modify the ligand/solute
         # counterions to make sure that the solvation box is always neutral.
         if system.usesPeriodicBoundaryConditions():
-            alchemical_counterions = mpiplus.run_single_node(0, pipeline.find_alchemical_counterions,
-                                                             system, topography, alchemical_region_name,
-                                                             broadcast_result=True)
-            alchemical_atoms += alchemical_counterions
+            if len(topography.ligand_atoms) < 2: # not a relative system
+                alchemical_counterions = mpiplus.run_single_node(0, pipeline.find_alchemical_counterions,
+                                                                 system, topography, alchemical_region_name,
+                                                                 broadcast_result=True)
+                alchemical_atoms += alchemical_counterions
 
             # Sort them by index for safety. We don't want to
             # accidentally exchange two atoms' positions.
             alchemical_atoms = sorted(alchemical_atoms)
 
-        alchemical_region_kwargs['alchemical_atoms'] = alchemical_atoms
+        if len(topography.ligand_atoms) < 2: # not a relative system
+            alchemical_region_kwargs['alchemical_atoms'] = alchemical_atoms
 
         # Check if we need to modify bonds/angles/torsions.
         for element_type in ['bonds', 'angles', 'torsions']:
@@ -1369,8 +1377,12 @@ def _build_default_alchemical_region(system, topography, protocol):
             alchemical_region_kwargs['alchemical_' + element_type] = modify_it
 
         # Create alchemical region.
-        alchemical_region = mmtools.alchemy.AlchemicalRegion(**alchemical_region_kwargs)
+        if len(topography.ligand_atoms) > 1: # relative system
+            alchemical_region = [mmtools.alchemy.AlchemicalRegion(name='zero', alchemical_atoms=alchemical_atoms[0], **alchemical_region_kwargs),
+                                 mmtools.alchemy.AlchemicalRegion(name='one', alchemical_atoms=alchemical_atoms[1], **alchemical_region_kwargs)]
 
+        else:
+            alchemical_region = [mmtools.alchemy.AlchemicalRegion(**alchemical_region_kwargs)]
         return alchemical_region
 
     @staticmethod

From 54c879b897fe502656cd163c850dd3351fceb68a Mon Sep 17 00:00:00 2001
From: Ana Silveira <anita_silveira@hotmail.com>
Date: Tue, 8 Oct 2019 16:07:29 -0400
Subject: [PATCH 2/8] change logic for determining if it is a relative system

---
 Yank/experiment.py       |  8 +++--
 Yank/pipeline.py         |  1 -
 Yank/restraints.py       | 72 ++++++++++++++++++++++++----------------
 Yank/schema/validator.py |  4 +--
 Yank/yank.py             | 53 ++++++++++++++++++-----------
 5 files changed, 85 insertions(+), 53 deletions(-)

diff --git a/Yank/experiment.py b/Yank/experiment.py
index 3f50d73d..51642d3f 100644
--- a/Yank/experiment.py
+++ b/Yank/experiment.py
@@ -651,7 +651,6 @@ def parse(self, script):
             yaml_content = script.copy()
 
         self._raw_yaml = yaml_content.copy()
-
         # Check that YAML loading was successful
         if yaml_content is None:
             raise YamlParseError('The YAML file is empty!')
@@ -1664,7 +1663,7 @@ def _is_pipeline_solvent_with_receptor(field, solvent_id, error):
         # DSL Schema
         ligand_dsl:
             required: no
-            type: list
+            type: [string, list]
             dependencies: [phase1_path, phase2_path]
         solvent_dsl:
             required: no
@@ -1732,7 +1731,10 @@ def _is_pipeline_solvent_with_receptor(field, solvent_id, error):
             check_with: REGION_CLASH_DETERMINED_AT_RUNTIME_WITH_LIGAND
         ligand:
             required: no
-            type: string
+            valueschema:
+                anyof:
+                    - type: list
+                    - type: string
             dependencies: [receptor, solvent]
             allowed: MOLECULE_IDS_POPULATED_AT_RUNTIME
             excludes: [solute, phase1_path, phase2_path]
diff --git a/Yank/pipeline.py b/Yank/pipeline.py
index 92cf14b7..59ae3166 100644
--- a/Yank/pipeline.py
+++ b/Yank/pipeline.py
@@ -272,7 +272,6 @@ def find_alchemical_counterions(system, topography, region_name):
     if len(atom_indices) == 0:
         raise ValueError("Cannot find counterions for region {}. "
                          "The region has no atoms.")
-
     # If the net charge of alchemical atoms is 0, we don't need counterions.
     mol_net_charge = compute_net_charge(system, atom_indices)
     logger.debug('{} net charge: {}'.format(region_name, mol_net_charge))
diff --git a/Yank/restraints.py b/Yank/restraints.py
index be389980..4ceff5ad 100644
--- a/Yank/restraints.py
+++ b/Yank/restraints.py
@@ -259,12 +259,11 @@ def compute_atom_intersect(self, input_atoms, topography_key: str, *additional_s
             additional_intersect = topography_set
         else:
             additional_intersect = set.intersection(*additional_sets)
-
         @functools.singledispatch
         def compute_atom_set(passed_atoms):
             """Helper function for doing set operations on heavy ligand atoms of all other types"""
             input_set = set(passed_atoms)
-            intersect_set = input_set & additional_intersect & topography_set
+            intersect_set = input_set & additional_intersect & topography_se
             if intersect_set != input_set:
                 return intersect_set
             else:
@@ -284,7 +283,7 @@ def compute_atom_str(input_string):
             # Ensure the selection is in the correct set
             set_combined = set_output & topography_set & additional_intersect
             final_output = [particle for particle in output if particle in set_combined]
-            # Force output to be a normal int, don't need to worry about floats at this point, there should not be any
+            # Force output to be a normal int, don't need to worry about floats at this point, there 	should not be any
             # If they come out as np.int64's, OpenMM complains
             return [*map(int, final_output)]
 
@@ -434,10 +433,13 @@ def __get__(self, instance, owner_class=None):
     def __set__(self, instance, new_restrained_atoms):
         # If we set the restrained attributes to None, no reason to check things.
         if new_restrained_atoms is not None:
-            if len(new_restrained_atoms) == 2: # relative system
-                new_restrained_atoms = [self._validate_atoms(atoms) for atoms in new_restrained_atoms]
-            else:
-                new_restrained_atoms = self._validate_atoms(new_restrained_atoms)
+            for element in new_restrained_atoms:
+                if isinstance(element, list): # relative system
+                    new_restrained_atoms = [self._validate_atoms(atoms) for atoms in new_restrained_atoms]
+                    break
+                else:
+                    new_restrained_atoms = self._validate_atoms(new_restrained_atoms)
+
         setattr(instance, self._attribute_name, new_restrained_atoms)
 
     @methoddispatch
@@ -574,12 +576,13 @@ def restrain_state(self, thermodynamic_state):
             raise RestraintParameterError('Restraint {}: Undefined restrained '
                                           'atoms.'.format(self.__class__.__name__))
 
-        # Create restraint force.
-        if len(self.restrained_ligand_atoms) == 1: # not a relative system
+        # Create restraint force
+        if isinstance(self.restrained_ligand_atoms[0], list) and isinstance(self.restrained_ligand_atoms[1], list): # relative system
+            restraint_force = self._get_restraint_force(*self.restrained_ligand_atoms)
+        else:
             restraint_force = self._get_restraint_force(self.restrained_receptor_atoms,
                                                         self.restrained_ligand_atoms)
-        else:
-            restraint_force = self._get_restraint_force(*self.restrained_ligand_atoms)
+
         # Set periodic conditions on the force if necessary.
         restraint_force.setUsesPeriodicBoundaryConditions(thermodynamic_state.is_periodic)
 
@@ -705,14 +708,16 @@ def _determine_restraint_parameters(self, thermodynamic_state, sampler_state, to
     @property
     def _are_restrained_atoms_defined(self):
         """Check if the restrained atoms are defined well enough to make a restraint"""
-        for atoms in [self.restrained_ligand_atoms]:
-            # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
-            if atoms is None or any(isinstance(atom, str) for atom in atoms) or not (isinstance(atoms, list) and len(atoms) > 0):
-                return False
-        if len(self.restrained_ligand_atoms) == 1:
-            for atoms in [self.restrained_receptor_atoms]:
-                if atoms is None or not (isinstance(atoms, list) and len(atoms) > 0):
-                    return False
+        if self.restrained_ligand_atoms:
+            if isinstance(self.restrained_ligand_atoms[0], list) and isinstance(self.restrained_ligand_atoms[1], list): # relative system
+                for atoms in self.restrained_ligand_atoms:
+                # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
+                    if atoms is None or any(isinstance(atom, str) for atom in atoms) or not (isinstance(atoms, list) and len(atoms) > 0):
+                        return False            
+            else:
+                for atoms in [self.restrained_receptor_atoms, self.restrained_ligand_atoms]:
+                    if atoms is None or not (isinstance(atoms, list) and len(atoms) > 0):
+                        return False
 
         return True
 
@@ -753,8 +758,9 @@ def _determine_restrained_atoms(self, sampler_state, topography):
 
         # If receptor and ligand atoms are explicitly provided, use those.
         restrained_ligand_atoms = self.restrained_ligand_atoms
-        if len(restrained_ligand_atoms) != 2:
+        if not isinstance(self.restrained_ligand_atoms[0], list) and not isinstance(self.restrained_ligand_atoms[1], list):
             restrained_receptor_atoms = self.restrained_receptor_atoms
+                
 
         @functools.singledispatch
         def compute_atom_set(input_atoms, topography_key, mapping_function):
@@ -771,7 +777,7 @@ def compute_atom_set(input_atoms, topography_key, mapping_function):
                                "Atoms not part of {0} will be ignored.".format(topography_key))
                 final_atoms = list(intersect_set)
             else:
-                final_atoms = list(input_atoms)
+                final_atoms = list(input_atoms_set)
             return final_atoms
 
         @compute_atom_set.register(type(None))
@@ -786,20 +792,30 @@ def compute_atom_none(_, topography_key, mapping_function):
         def compute_atom_str(input_string, topography_key, _):
             """Helper for string parsing"""
             selection = topography.select(input_string, as_set=True)
-            selection_with_top = selection  & set([index for sublist in getattr(topography, topography_key) for
+            if isinstance(getattr(topography, topography_key)[0], list) and isinstance(getattr(topography, topography_key)[1], list):
+                selection_with_top = selection & set([index for sublist in getattr(topography, topography_key) for
                                                    index in sublist])
+            else:
+                selection_with_top = selection & set(getattr(topography, topography_key))
             # Force output to be a normal int, dont need to worry about floats at this point, there should not be any
             # If they come out as np.int64's, OpenMM complains
             return [*map(int, selection_with_top)]
-
-        self.restrained_ligand_atoms = [compute_atom_set(list,
+         
+        if isinstance(restrained_ligand_atoms[0], list) and isinstance(restrained_ligand_atoms[1], list):
+            self.restrained_ligand_atoms = []
+            for element in restrained_ligand_atoms:
+                for value in element:
+                    self.restrained_ligand_atoms.append(compute_atom_set(value,
                                                         'ligand_atoms',
-                                                        self._closest_atom_to_centroid) for list in restrained_ligand_atoms]
-        if len(self.restrained_ligand_atoms) != 2:
+                                                         self._closest_atom_to_centroid))
+        else:
+            self.restrained_ligand_atoms = compute_atom_set(restrained_ligand_atoms,
+                                                            'ligand_atoms',
+                                                             self._closest_atom_to_centroid)
+
             self.restrained_receptor_atoms = compute_atom_set(restrained_receptor_atoms,
                                                               'receptor_atoms',
-                                                              self._closest_atom_to_centroid)
-
+                                                             self._closest_atom_to_centroid)
     @staticmethod
     def _closest_atom_to_centroid(positions, indices=None, masses=None):
         """
diff --git a/Yank/schema/validator.py b/Yank/schema/validator.py
index 21359f2f..0963a8f9 100644
--- a/Yank/schema/validator.py
+++ b/Yank/schema/validator.py
@@ -189,8 +189,8 @@ def _check_with_specify_lambda_electrostatics_and_sterics(self, field, value):
         """Check that the keys of a dictionary contain both lambda_electrostatics and lambda_sterics."""
         if ((isinstance(value, dict) or isinstance(value, collections.OrderedDict)) and
              not (('lambda_sterics' in value and 'lambda_electrostatics' in value) or
-             ('lambda_sterics_zero' in value and 'lambda_electrostatics_zero' in value and
-               'lambda_sterics_one' in value and 'lambda_electrostatics_one' in value))):
+             ('lambda_sterics_0' in value and 'lambda_electrostatics_0' in value and
+               'lambda_sterics_1' in value and 'lambda_electrostatics_1' in value))):
             self._error(field, "Missing required keys lambda_sterics and/or lambda_electrostatics")
 
     def _check_with_math_expressions_variables_are_given(self, field, value):
diff --git a/Yank/yank.py b/Yank/yank.py
index 81ddf618..1eab6167 100644
--- a/Yank/yank.py
+++ b/Yank/yank.py
@@ -120,7 +120,15 @@ def ligand_atoms(self):
 
     @ligand_atoms.setter
     def ligand_atoms(self, value):
-        self._ligand_atoms = tuple([self.select(v) for v in value])
+        relative_system = False
+        if len(value) == 2:
+            for v in value:
+                if isinstance(v, list):
+                    relative_system = True
+        if (relative_system):         
+            self._ligand_atoms = tuple([self.select(v[0]) for v in value])
+        else:
+            self._ligand_atoms = self.select(value)
         # Safety check: with a ligand there should always be a receptor.
         if len(self._ligand_atoms) > 0 and len(self.receptor_atoms) == 0:
             raise ValueError('Specified ligand but cannot find '
@@ -139,8 +147,11 @@ def receptor_atoms(self):
         if len(self._ligand_atoms) == 0:
             return []
 
-        # Create a set for fast searching.
-        ligand_atomset = tuple(frozenset([index for ligand in self._ligand_atoms for index in ligand]))
+        #  alchemical_regionsCreate a set for fast searching.
+        if isinstance(self._ligand_atoms, tuple):
+            ligand_atomset = frozenset(self._ligand_atoms[0] + self._ligand_atoms[1])
+        else:
+            ligand_atomset = frozenset(self._ligand_atoms)
         # Receptor atoms are all solute atoms that are not ligand.
         return [i for i in self.solute_atoms if i not in ligand_atomset]
 
@@ -1012,7 +1023,7 @@ def create(self, thermodynamic_state, sampler_states, topography, protocol,
 
         # Check that we have atoms to alchemically modify.
         for region in alchemical_regions:
-            if len(region.alchemical_atoms) == 0:
+            if region.alchemical_atoms == 0:
                 raise ValueError("Couldn't find atoms to alchemically modify.")
 
         # Create alchemically-modified system using alchemical factory.
@@ -1025,8 +1036,8 @@ def create(self, thermodynamic_state, sampler_states, topography, protocol,
         # Create compound alchemically modified state to pass to sampler.
         thermodynamic_state.system = alchemical_system
         if len(alchemical_regions) > 1:
-            alchemical_state = [mmtools.alchemy.AlchemicalState.from_system(alchemical_system, parameters_name_suffix = 'zero'),
-                                mmtools.alchemy.AlchemicalState.from_system(alchemical_system, parameters_name_suffix = 'one')]
+            alchemical_state = [mmtools.alchemy.AlchemicalState.from_system(alchemical_system, parameters_name_suffix = '0'),
+                                mmtools.alchemy.AlchemicalState.from_system(alchemical_system, parameters_name_suffix = '1')]
         else:
             alchemical_state = [mmtools.alchemy.AlchemicalState.from_system(alchemical_system)]
 
@@ -1342,31 +1353,35 @@ def _build_default_alchemical_region(system, topography, protocol):
         """Create a default AlchemicalRegion if the user hasn't provided one."""
         # TODO: we should probably have a second region that annihilate sterics of counterions.
         alchemical_region_kwargs = {}
-
+        
         # Modify ligand if this is a receptor-ligand phase, or
         # solute if this is a transfer free energy calculation.
-        alchemical_region_name = 'ligand_atoms'
-        if len(topography.ligand_atoms) == 0:
+        if len(topography.ligand_atoms) > 0:
+            alchemical_region_name = 'ligand_atoms'
+        else:
             alchemical_region_name = 'solute_atoms'
         alchemical_atoms = getattr(topography, alchemical_region_name)
-        if len(topography.ligand_atoms) > 1:
+
+        if isinstance(topography.ligand_atoms, tuple):
             alchemical_region_name = ['ligand_atoms_zero', 'ligand_atoms_one']
 
         # In periodic systems, we alchemically modify the ligand/solute
         # counterions to make sure that the solvation box is always neutral.
         if system.usesPeriodicBoundaryConditions():
-            if len(topography.ligand_atoms) < 2: # not a relative system
+            if not isinstance(topography.ligand_atoms, tuple): # not a relative system
                 alchemical_counterions = mpiplus.run_single_node(0, pipeline.find_alchemical_counterions,
                                                                  system, topography, alchemical_region_name,
                                                                  broadcast_result=True)
                 alchemical_atoms += alchemical_counterions
-
-            # Sort them by index for safety. We don't want to
-            # accidentally exchange two atoms' positions.
-            alchemical_atoms = sorted(alchemical_atoms)
-
-        if len(topography.ligand_atoms) < 2: # not a relative system
-            alchemical_region_kwargs['alchemical_atoms'] = alchemical_atoms
+              
+                # Sort them by index for safety. We don't want to
+                # accidentally exchange two atoms' positions.
+                alchemical_atoms = sorted(alchemical_atoms)
+                alchemical_region_kwargs['alchemical_atoms'] = alchemical_atoms
+ 
+            else:
+                alchemical_atoms = [sorted(topography.ligand_atoms[0])]
+                alchemical_atoms.append(sorted(topography.ligand_atoms[1]))
 
         # Check if we need to modify bonds/angles/torsions.
         for element_type in ['bonds', 'angles', 'torsions']:
@@ -1377,7 +1392,7 @@ def _build_default_alchemical_region(system, topography, protocol):
             alchemical_region_kwargs['alchemical_' + element_type] = modify_it
 
         # Create alchemical region.
-        if len(topography.ligand_atoms) > 1: # relative system
+        if isinstance(topography.ligand_atoms, tuple): # relative system
             alchemical_region = [mmtools.alchemy.AlchemicalRegion(name='zero', alchemical_atoms=alchemical_atoms[0], **alchemical_region_kwargs),
                                  mmtools.alchemy.AlchemicalRegion(name='one', alchemical_atoms=alchemical_atoms[1], **alchemical_region_kwargs)]
 

From 4929c599841d795194fb02f4058d859c52f49f9e Mon Sep 17 00:00:00 2001
From: Ana Silveira <anita_silveira@hotmail.com>
Date: Tue, 8 Oct 2019 18:05:58 -0400
Subject: [PATCH 3/8] some fixes

---
 Yank/experiment.py | 4 +---
 Yank/restraints.py | 4 ++--
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/Yank/experiment.py b/Yank/experiment.py
index 51642d3f..79a5d9ae 100644
--- a/Yank/experiment.py
+++ b/Yank/experiment.py
@@ -1732,9 +1732,7 @@ def _is_pipeline_solvent_with_receptor(field, solvent_id, error):
         ligand:
             required: no
             valueschema:
-                anyof:
-                    - type: list
-                    - type: string
+            type: [string, list]
             dependencies: [receptor, solvent]
             allowed: MOLECULE_IDS_POPULATED_AT_RUNTIME
             excludes: [solute, phase1_path, phase2_path]
diff --git a/Yank/restraints.py b/Yank/restraints.py
index 4ceff5ad..2c62a9c1 100644
--- a/Yank/restraints.py
+++ b/Yank/restraints.py
@@ -263,7 +263,7 @@ def compute_atom_intersect(self, input_atoms, topography_key: str, *additional_s
         def compute_atom_set(passed_atoms):
             """Helper function for doing set operations on heavy ligand atoms of all other types"""
             input_set = set(passed_atoms)
-            intersect_set = input_set & additional_intersect & topography_se
+            intersect_set = input_set & additional_intersect & topography_set
             if intersect_set != input_set:
                 return intersect_set
             else:
@@ -283,7 +283,7 @@ def compute_atom_str(input_string):
             # Ensure the selection is in the correct set
             set_combined = set_output & topography_set & additional_intersect
             final_output = [particle for particle in output if particle in set_combined]
-            # Force output to be a normal int, don't need to worry about floats at this point, there 	should not be any
+            # Force output to be a normal int, don't need to worry about floats at this point, there should not be any
             # If they come out as np.int64's, OpenMM complains
             return [*map(int, final_output)]
 

From 612cdaca5ef24c95308e1069a077447403099a2c Mon Sep 17 00:00:00 2001
From: Ana Silveira <anita_silveira@hotmail.com>
Date: Wed, 9 Oct 2019 14:55:53 -0400
Subject: [PATCH 4/8] use a function for determining if it is ligand-ligand
 restraint

---
 Yank/experiment.py |  1 -
 Yank/restraints.py | 25 ++++++++++++-------------
 Yank/yank.py       |  6 +++---
 3 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/Yank/experiment.py b/Yank/experiment.py
index 79a5d9ae..017fb3c2 100644
--- a/Yank/experiment.py
+++ b/Yank/experiment.py
@@ -1731,7 +1731,6 @@ def _is_pipeline_solvent_with_receptor(field, solvent_id, error):
             check_with: REGION_CLASH_DETERMINED_AT_RUNTIME_WITH_LIGAND
         ligand:
             required: no
-            valueschema:
             type: [string, list]
             dependencies: [receptor, solvent]
             allowed: MOLECULE_IDS_POPULATED_AT_RUNTIME
diff --git a/Yank/restraints.py b/Yank/restraints.py
index 2c62a9c1..373e473d 100644
--- a/Yank/restraints.py
+++ b/Yank/restraints.py
@@ -411,6 +411,8 @@ def _parameters(self):
                       for parameter_name in parameter_names}
         return parameters
 
+    def _is_ligand_ligand_restraint(self, restrained_ligand_atoms):
+        return len(restrained_ligand_atoms) == 2 and isinstance(restrained_ligand_atoms[0], list) and isinstance(restrained_ligand_atoms[1], list)
 
 class _RestrainedAtomsProperty(object):
     """
@@ -432,13 +434,11 @@ def __get__(self, instance, owner_class=None):
 
     def __set__(self, instance, new_restrained_atoms):
         # If we set the restrained attributes to None, no reason to check things.
-        if new_restrained_atoms is not None:
-            for element in new_restrained_atoms:
-                if isinstance(element, list): # relative system
-                    new_restrained_atoms = [self._validate_atoms(atoms) for atoms in new_restrained_atoms]
-                    break
-                else:
-                    new_restrained_atoms = self._validate_atoms(new_restrained_atoms)
+        if (new_restrained_atoms):
+            if instance._is_ligand_ligand_restraint(new_restrained_atoms):
+                new_restrained_atoms = [self._validate_atoms(atoms) for atoms in new_restrained_atoms]
+            else:
+                new_restrained_atoms = self._validate_atoms(new_restrained_atoms)
 
         setattr(instance, self._attribute_name, new_restrained_atoms)
 
@@ -451,7 +451,6 @@ def _validate_atoms(self, restrained_atoms):
             restrained_atoms = list(restrained_atoms)
         return restrained_atoms
 
-
 # ==============================================================================
 # Base class for radially-symmetric receptor-ligand restraints.
 # ==============================================================================
@@ -577,7 +576,7 @@ def restrain_state(self, thermodynamic_state):
                                           'atoms.'.format(self.__class__.__name__))
 
         # Create restraint force
-        if isinstance(self.restrained_ligand_atoms[0], list) and isinstance(self.restrained_ligand_atoms[1], list): # relative system
+        if self._is_ligand_ligand_restraint(self.restrained_ligand_atoms):
             restraint_force = self._get_restraint_force(*self.restrained_ligand_atoms)
         else:
             restraint_force = self._get_restraint_force(self.restrained_receptor_atoms,
@@ -709,7 +708,7 @@ def _determine_restraint_parameters(self, thermodynamic_state, sampler_state, to
     def _are_restrained_atoms_defined(self):
         """Check if the restrained atoms are defined well enough to make a restraint"""
         if self.restrained_ligand_atoms:
-            if isinstance(self.restrained_ligand_atoms[0], list) and isinstance(self.restrained_ligand_atoms[1], list): # relative system
+            if self._is_ligand_ligand_restraint(self.restrained_ligand_atoms):
                 for atoms in self.restrained_ligand_atoms:
                 # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
                     if atoms is None or any(isinstance(atom, str) for atom in atoms) or not (isinstance(atoms, list) and len(atoms) > 0):
@@ -758,7 +757,7 @@ def _determine_restrained_atoms(self, sampler_state, topography):
 
         # If receptor and ligand atoms are explicitly provided, use those.
         restrained_ligand_atoms = self.restrained_ligand_atoms
-        if not isinstance(self.restrained_ligand_atoms[0], list) and not isinstance(self.restrained_ligand_atoms[1], list):
+        if not self._is_ligand_ligand_restraint(restrained_ligand_atoms):
             restrained_receptor_atoms = self.restrained_receptor_atoms
                 
 
@@ -800,8 +799,8 @@ def compute_atom_str(input_string, topography_key, _):
             # Force output to be a normal int, dont need to worry about floats at this point, there should not be any
             # If they come out as np.int64's, OpenMM complains
             return [*map(int, selection_with_top)]
-         
-        if isinstance(restrained_ligand_atoms[0], list) and isinstance(restrained_ligand_atoms[1], list):
+        
+        if self._is_ligand_ligand_restraint(restrained_ligand_atoms):
             self.restrained_ligand_atoms = []
             for element in restrained_ligand_atoms:
                 for value in element:
diff --git a/Yank/yank.py b/Yank/yank.py
index 1eab6167..437ed737 100644
--- a/Yank/yank.py
+++ b/Yank/yank.py
@@ -1363,7 +1363,7 @@ def _build_default_alchemical_region(system, topography, protocol):
         alchemical_atoms = getattr(topography, alchemical_region_name)
 
         if isinstance(topography.ligand_atoms, tuple):
-            alchemical_region_name = ['ligand_atoms_zero', 'ligand_atoms_one']
+            alchemical_region_name = ['ligand_atoms_0', 'ligand_atoms_1']
 
         # In periodic systems, we alchemically modify the ligand/solute
         # counterions to make sure that the solvation box is always neutral.
@@ -1393,8 +1393,8 @@ def _build_default_alchemical_region(system, topography, protocol):
 
         # Create alchemical region.
         if isinstance(topography.ligand_atoms, tuple): # relative system
-            alchemical_region = [mmtools.alchemy.AlchemicalRegion(name='zero', alchemical_atoms=alchemical_atoms[0], **alchemical_region_kwargs),
-                                 mmtools.alchemy.AlchemicalRegion(name='one', alchemical_atoms=alchemical_atoms[1], **alchemical_region_kwargs)]
+            alchemical_region = [mmtools.alchemy.AlchemicalRegion(name='0', alchemical_atoms=alchemical_atoms[0], **alchemical_region_kwargs),
+                                 mmtools.alchemy.AlchemicalRegion(name='1', alchemical_atoms=alchemical_atoms[1], **alchemical_region_kwargs)]
 
         else:
             alchemical_region = [mmtools.alchemy.AlchemicalRegion(**alchemical_region_kwargs)]

From d404bd3c723f9b6666e2656340dd90ea6d739361 Mon Sep 17 00:00:00 2001
From: Ana Silveira <anita_silveira@hotmail.com>
Date: Thu, 2 Apr 2020 21:56:42 -0400
Subject: [PATCH 5/8] Preparing RMSD and Boresch restraints for relative free
 energies

---
 Yank/restraints.py | 68 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 47 insertions(+), 21 deletions(-)

diff --git a/Yank/restraints.py b/Yank/restraints.py
index 373e473d..633ae76c 100644
--- a/Yank/restraints.py
+++ b/Yank/restraints.py
@@ -251,7 +251,11 @@ def compute_atom_intersect(self, input_atoms, topography_key: str, *additional_s
         """
 
         topography = self.topography
-        topography_set = set(getattr(topography, topography_key))
+        if isinstance(getattr(topography, topography_key)[0], list) and isinstance(getattr(topography, topography_key)[1], list):
+            topography_set = set([index for sublist in getattr(topography, topography_key) for
+                                                   index in sublist])
+        else:
+            topography_set = set(getattr(topography, topography_key))
         # Ensure additions are sets
         additional_sets = [set(additional_set) for additional_set in additional_sets]
         if len(additional_sets) == 0:
@@ -439,7 +443,6 @@ def __set__(self, instance, new_restrained_atoms):
                 new_restrained_atoms = [self._validate_atoms(atoms) for atoms in new_restrained_atoms]
             else:
                 new_restrained_atoms = self._validate_atoms(new_restrained_atoms)
-
         setattr(instance, self._attribute_name, new_restrained_atoms)
 
     @methoddispatch
@@ -532,7 +535,7 @@ class _RadiallySymmetricRestrainedAtomsProperty(_RestrainedAtomsProperty):
         def _validate_atoms(self, restrained_atoms):
             restrained_atoms = super()._validate_atoms(restrained_atoms)
             if len(restrained_atoms) > 1:
-                logger.debug(self._CENTROID_COMPUTE_STRING.format("more than one", self._atoms_type))
+                logger.debug(self._CENTROID_COMPUTE_STRING.format("more than one", self._atoms_type))     
             return restrained_atoms
 
         @_validate_atoms.register(str)
@@ -990,6 +993,7 @@ def _create_restraint_force(self, particles1, particles2):
            The created restraint force.
 
         """
+
         # Create bond force and lambda_restraints parameter to control it.
         if len(particles1) == 1 and len(particles2) == 1:
             # CustomCentroidBondForce works only on 64bit platforms. When the
@@ -1499,7 +1503,11 @@ def restrain_state(self, thermodynamic_state):
         n_particles = 6  # number of particles involved in restraint: p1 ... p6
         restraint_force = openmm.CustomCompoundBondForce(n_particles, energy_function)
         restraint_force.addGlobalParameter('lambda_restraints', 1.0)
-        restraint_force.addBond(self.restrained_receptor_atoms + self.restrained_ligand_atoms, [])
+
+        if self._is_ligand_ligand_restraint(self.restrained_ligand_atoms):
+            restraint_force.addBond(self.restrained_ligand_atoms[0] + self.restrained_ligand_atoms[1], [])
+        else:
+            restraint_force.addBond(self.restrained_receptor_atoms + self.restrained_ligand_atoms, [])
         restraint_force.setUsesPeriodicBoundaryConditions(thermodynamic_state.is_periodic)
 
         # Get a copy of the system of the ThermodynamicState, modify it and set it back.
@@ -1733,10 +1741,17 @@ def _check_parameters_defined(self):
     @property
     def _are_restrained_atoms_defined(self):
         """Check if the restrained atoms are defined well enough to make a restraint"""
-        for atoms in [self.restrained_receptor_atoms, self.restrained_ligand_atoms]:
-            # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
-            if atoms is None or not (isinstance(atoms, list) and len(atoms) == 3):
-                return False
+        if self.restrained_ligand_atoms:
+            if self._is_ligand_ligand_restraint(self.restrained_ligand_atoms):
+                for atoms in self.restrained_ligand_atoms:
+                # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
+                    if atoms is None or any(isinstance(atom, str) for atom in atoms) or not (isinstance(atoms, list) and len(atoms) > 0):
+                        return False
+        else:
+            for atoms in [self.restrained_receptor_atoms, self.restrained_ligand_atoms]:
+                # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
+                if atoms is None or not (isinstance(atoms, list) and len(atoms) == 3):
+                    return False
         return True
 
     @staticmethod
@@ -1940,7 +1955,10 @@ def _assign_if_undefined(attr_name, attr_value):
                 setattr(self, attr_name, attr_value)
 
         # Merge receptor and ligand atoms in a single array for easy manipulation.
-        restrained_atoms = self.restrained_receptor_atoms + self.restrained_ligand_atoms
+        if self._is_ligand_ligand_restraint(self.restrained_ligand_atoms):
+            restrained_atoms = self.restrained_ligand_atoms[0] + self.restrained_ligand_atoms[1]
+        else:
+            restrained_atoms = self.restrained_receptor_atoms + self.restrained_ligand_atoms
 
         # Set spring constants uniformly, as in Ref [1] Table 1 caption.
         _assign_if_undefined('K_r', 20.0 * unit.kilocalories_per_mole / unit.angstrom**2)
@@ -2524,7 +2542,7 @@ def __init__(self, atoms_type, allowed_empty=False):
         def _validate_atoms(self, restrained_atoms):
             restrained_atoms = super()._validate_atoms(restrained_atoms)
             # TODO: Determine the minimum number of atoms needed for this restraint (can it be 0?)
-            if len(restrained_atoms) < 3 and not (len(restrained_atoms) == 0 and self._allowed_empty):
+            if len(restrained_atoms) < 3 and not (len(restrained_atoms) == 0 and self._allowed_empty) and not (isinstance(element, str) for element in restrained_atoms):
                 raise ValueError('At least three {} atoms are required to impose an '
                                  'RMSD restraint.'.format(self._atoms_type))
             return restrained_atoms
@@ -2553,10 +2571,9 @@ def restrain_state(self, thermodynamic_state):
 
         # Check if all parameters are defined.
         self._check_parameters_defined()
-
         # Merge receptor and ligand atoms in a single array for easy manipulation.
         restrained_atoms = self.restrained_receptor_atoms + self.restrained_ligand_atoms
-
+       
         # Create RMSDForce CV for all restrained atoms
         rmsd_cv = openmm.RMSDForce(self.reference_sampler_state.positions, restrained_atoms)
 
@@ -2648,7 +2665,7 @@ def _check_parameters_defined(self):
     def _are_restrained_atoms_defined(self):
         """Check if the restrained atoms are defined well enough to make a restraint"""
         for atoms in [self.restrained_receptor_atoms, self.restrained_ligand_atoms]:
-            # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
+        # Atoms should be a list or None at this point due to the _RestrainedAtomsProperty class
             if not self._are_single_atoms_defined(atoms):
                 return False
         return True
@@ -2663,14 +2680,23 @@ def _are_single_atoms_defined(atom_list):
     def _pick_restrained_atoms(self, topography):
         """Select the restrained atoms to use for this system"""
         atom_selector = _AtomSelector(topography)
-        for atom_word, top_key in zip(["restrained_ligand_atoms", "restrained_receptor_atoms"],
-                                      ["ligand_atoms",            "receptor_atoms"]):
-            atoms = getattr(self, atom_word)
-            if self._are_single_atoms_defined(atoms):
-                continue
-            defined_atoms = atom_selector.compute_atom_intersect(atoms, top_key)
-            setattr(self, atom_word, defined_atoms)
-
+        atoms = getattr(self, "restrained_receptor_atoms")
+        if not self._are_single_atoms_defined(atoms):
+       	    defined_atoms = atom_selector.compute_atom_intersect(atoms, "receptor_atoms")
+            setattr(self, "restrained_receptor_atoms", defined_atoms)
+        atoms = getattr(self, "restrained_ligand_atoms")
+        if self._is_ligand_ligand_restraint(self.restrained_ligand_atoms):
+            defined_atoms = []
+            for element in self.restrained_ligand_atoms:
+                if isinstance(element[0], str):
+                    defined_atoms += atom_selector.compute_atom_intersect(element[0], "ligand_atoms")
+                elif isinstance(element, list) and len(element) > 0:
+                    defined_atoms += element
+            setattr(self, "restrained_ligand_atoms", defined_atoms)
+        else:
+            if not self._are_single_atoms_defined(atoms):
+                defined_atoms = atom_selector.compute_atom_intersect(atoms, "ligand_atoms")
+                setattr(self, "restrained_ligand_atoms", defined_atoms)
 
 if __name__ == '__main__':
     import doctest

From 47da0c8190c03b1d5a25333d19ae3168f9d8e2dd Mon Sep 17 00:00:00 2001
From: Ana Silveira <anita_silveira@hotmail.com>
Date: Sat, 4 Apr 2020 12:23:19 -0400
Subject: [PATCH 6/8] multiple alchemical regions in
 ExperimentBuilder._generate_auto_state_parameters

---
 Yank/experiment.py | 42 ++++++++++++++++++++++++++++--------------
 1 file changed, 28 insertions(+), 14 deletions(-)

diff --git a/Yank/experiment.py b/Yank/experiment.py
index db8cfac7..9a75852e 100644
--- a/Yank/experiment.py
+++ b/Yank/experiment.py
@@ -2611,22 +2611,36 @@ def _generate_auto_state_parameters(phase_factory):
                          len(phase_factory.topography.solvent_atoms) == 0)
         state_parameters = []
 
+        if not isinstance(phase_factory.topography.ligand_atoms, tuple): # not a relative system
         # First, turn on the restraint if there are any.
-        if phase_factory.restraint is not None:
-            state_parameters.append(('lambda_restraints', [0.0, 1.0]))
-        # We support only lambda sterics and electrostatics for now.
-        if is_vacuum and not phase_factory.alchemical_regions.annihilate_electrostatics:
-            state_parameters.append(('lambda_electrostatics', [1.0, 1.0]))
-        else:
-            state_parameters.append(('lambda_electrostatics', [1.0, 0.0]))
-        if is_vacuum and not phase_factory.alchemical_regions.annihilate_sterics:
-            state_parameters.append(('lambda_sterics', [1.0, 1.0]))
+            if phase_factory.restraint is not None:
+                state_parameters.append(('lambda_restraints', [0.0, 1.0]))
+            # We support only lambda sterics and electrostatics for now.
+            if is_vacuum and not phase_factory.alchemical_regions.annihilate_electrostatics:
+                state_parameters.append(('lambda_electrostatics', [1.0, 1.0]))
+            else:
+                state_parameters.append(('lambda_electrostatics', [1.0, 0.0]))
+            if is_vacuum and not phase_factory.alchemical_regions.annihilate_sterics:
+                state_parameters.append(('lambda_sterics', [1.0, 1.0]))
+            else:
+                state_parameters.append(('lambda_sterics', [1.0, 0.0]))
+            # Turn the RMSD restraints off slowly at the end
+            if isinstance(phase_factory.restraint, restraints.RMSD):
+                state_parameters.append(('lambda_restraints', [1.0, 0.0]))
         else:
-            state_parameters.append(('lambda_sterics', [1.0, 0.0]))
-        # Turn the RMSD restraints off slowly at the end
-        if isinstance(phase_factory.restraint, restraints.RMSD):
-            state_parameters.append(('lambda_restraints', [1.0, 0.0]))
-
+            if is_vacuum and not phase_factory.alchemical_regions.annihilate_electrostatics:
+                state_parameters.append(('lambda_electrostatics_0', [1.0, 1.0]))
+                state_parameters.append(('lambda_electrostatics_1', [1.0, 1.0]))
+            else:
+                state_parameters.append(('lambda_electrostatics_0', [1.0, 0.0]))
+                state_parameters.append(('lambda_electrostatics_1', [0.0, 1.0]))
+            if is_vacuum and not phase_factory.alchemical_regions.annihilate_sterics:
+                state_parameters.append(('lambda_sterics_0', [1.0, 1.0]))
+                state_parameters.append(('lambda_sterics_1', [1.0, 1.0]))
+            else:
+                state_parameters.append(('lambda_sterics_0', [1.0, 0.0]))
+                state_parameters.append(('lambda_sterics_1', [0.0, 1.0]))
+            # TODO: Handle restraints in relative free energy calculations 
         return state_parameters
 
     @classmethod

From b66b14ed1aba65060cd285a3c4f7422d362e2e3e Mon Sep 17 00:00:00 2001
From: Ana Silveira <anita_silveira@hotmail.com>
Date: Tue, 7 Apr 2020 19:46:25 -0400
Subject: [PATCH 7/8] determine alchemical counterions for relative free energy
 calculations determine alchemical atoms for the solvent phase

---
 Yank/experiment.py |  5 ++-
 Yank/pipeline.py   | 99 +++++++++++++++++++++++++++++++++++++---------
 Yank/restraints.py | 14 +++----
 Yank/yank.py       | 51 +++++++++++++++---------
 4 files changed, 124 insertions(+), 45 deletions(-)

diff --git a/Yank/experiment.py b/Yank/experiment.py
index 9a75852e..4bd41cbc 100644
--- a/Yank/experiment.py
+++ b/Yank/experiment.py
@@ -3102,6 +3102,7 @@ def _build_experiment(self, experiment_path, experiment, use_dummy_protocol=Fals
         assert isinstance(protocol, collections.OrderedDict)
         phase_names = list(protocol.keys())
         phase_paths = self._get_nc_file_paths(experiment_path, experiment)
+        is_relative = False
         for phase_idx, (phase_name, phase_path) in enumerate(zip(phase_names, phase_paths)):
             # Check if we need to resume a phase. If the phase has been
             # already created, Experiment will resume from the storage.
@@ -3125,6 +3126,8 @@ def _build_experiment(self, experiment_path, experiment, use_dummy_protocol=Fals
             # Identify system components. There is a ligand only in the complex phase.
             if phase_idx == 0:
                 ligand_atoms = ligand_dsl
+                if len(ligand_atoms) == 2 and isinstance(ligand_atoms[0], list) and isinstance(ligand_atoms[1], list):
+                    is_relative = True
             else:
                 ligand_atoms = None
             topography = Topography(topology, ligand_atoms=ligand_atoms,
@@ -3147,7 +3150,7 @@ def _build_experiment(self, experiment_path, experiment, use_dummy_protocol=Fals
             # and modified it according to the user options.
             phase_protocol = protocol[phase_name]['alchemical_path']
             alchemical_region = AlchemicalPhase._build_default_alchemical_region(system, topography,
-                                                                                 phase_protocol)
+                                                                                 phase_protocol, is_relative)
             alchemical_region = [region._replace(**alchemical_region_opts) for region in alchemical_region]
 
             # Apply restraint only if this is the first phase. AlchemicalPhase
diff --git a/Yank/pipeline.py b/Yank/pipeline.py
index 15508ac9..bfe91dd8 100644
--- a/Yank/pipeline.py
+++ b/Yank/pipeline.py
@@ -234,7 +234,7 @@ def compute_net_charge(system, atom_indices):
     return net_charge
 
 
-def find_alchemical_counterions(system, topography, region_name):
+def find_alchemical_counterions(system, topography, region_name, is_relative):
     """Return the atom indices of the ligand or solute counter ions.
 
     In periodic systems, the solvation box needs to be neutral, and
@@ -273,31 +273,92 @@ def find_alchemical_counterions(system, topography, region_name):
     if len(atom_indices) == 0:
         raise ValueError("Cannot find counterions for region {}. "
                          "The region has no atoms.")
-    # If the net charge of alchemical atoms is 0, we don't need counterions.
-    mol_net_charge = compute_net_charge(system, atom_indices)
-    logger.debug('{} net charge: {}'.format(region_name, mol_net_charge))
-    if mol_net_charge == 0:
-        return []
-
     # Find net charge of all ions in the system.
     ions_net_charges = [(ion_id, compute_net_charge(system, [ion_id]))
-                        for ion_id in topography.ions_atoms]
+                         for ion_id in topography.ions_atoms]
     topology = topography.topology
-    ions_names_charges = [(topology.atom(ion_id).residue.name, ion_net_charge)
-                          for ion_id, ion_net_charge in ions_net_charges]
-    logger.debug('Ions net charges: {}'.format(ions_names_charges))
+    if not is_relative:
+        # If the net charge of alchemical atoms is 0, we don't need counterions.
+        mol_net_charge = compute_net_charge(system, atom_indices)
+        logger.debug('{} net charge: {}'.format(region_name, mol_net_charge))
+        if mol_net_charge == 0:
+            return []
+
+        ions_names_charges = [(topology.atom(ion_id).residue.name, ion_net_charge)
+                              for ion_id, ion_net_charge in ions_net_charges]
+        logger.debug('Ions net charges: {}'.format(ions_names_charges))
+
+        # Find minimal subset of counterions whose charges sums to -mol_net_charge.
+        return ions_subset(ions_net_charges, -mol_net_charge)
+
+        # We couldn't find any subset of counterions neutralizing the region.
+        raise ValueError('Impossible to find a solution for region {}. '
+                         'Net charge: {}, system ions: {}.'.format(
+            region_name, mol_net_charge, ions_names_charges))
+
+    else:
+        if region_name == 'solute_atoms':
+            topo = topology.subset(atom_indices)
+            resnames = []
+            for chain in topo._chains:
+                for residue in chain._residues:
+                    resnames.append(f'resname {residue.name}')
+            atom_indices = tuple([topography.select(r) for r in resnames])
+
+        ions_net_charges = [(ion_id, compute_net_charge(system, [ion_id]))
+                            for ion_id in topography.ions_atoms]
+        ignore = [ions[0] for ions in ions_net_charges] + atom_indices[0] + atom_indices[1]
+        indices = [atom.index for atom in topology.atoms if atom.index not in ignore]
+        net_charge = compute_net_charge(system, indices)
+        charge_init = net_charge + compute_net_charge(system, atom_indices[0])
+        charge_end = net_charge + compute_net_charge(system, atom_indices[1])
+        if (charge_init*charge_end) > 0:
+            if abs(charge_init) > abs(charge_end):
+                ions_to_dummies = int(charge_end - charge_init)
+                dummies_to_ions = None
+            elif abs(charge_init) < abs(charge_end):
+                ions_to_dummies = None
+                dummies_to_ions = int(-(charge_end - charge_init))
+            else:
+                ions_to_dummies = None
+                dummies_to_ions = None
+        else:
+            ions_to_dummies = int(-charge_init)
+            dummies_to_ions = int(-charge_end)
+
+        counterions = ions_subset(ions_net_charges, -charge_init)
+        avail_ions = [(id, c) for (id, c) in ions_net_charges if id not in counterions]
+        ions_to_dummies_idx = None
+        dummies_to_ions_idx = None
+        if ions_to_dummies:
+            ions_to_dummies_idx = counterions[:abs(ions_to_dummies)]
+        if dummies_to_ions:
+            dummies_to_ions_idx = ions_subset(avail_ions, dummies_to_ions)
+    
+        return ions_to_dummies_idx, dummies_to_ions_idx
+
+def ions_subset(ions_net_charges, counterions):
+    """
+    Finds minimal subset of ion indexes whose charge sums to counterions
+    Parameters
+    ----------
+        ions_net_charges : list of tuples
+            (index, charge) of ions in system.
+        counterions : int
+             Total charge.
+
+        Returns
+        -------
+            counterions_indices : list of int
+            Indices of ions.
+
+     """
 
-    # Find minimal subset of counterions whose charges sums to -mol_net_charge.
     for n_ions in range(1, len(ions_net_charges) + 1):
         for ion_subset in itertools.combinations(ions_net_charges, n_ions):
             counterions_indices, counterions_charges = zip(*ion_subset)
-            if sum(counterions_charges) == -mol_net_charge:
-                return counterions_indices
-
-    # We couldn't find any subset of counterions neutralizing the region.
-    raise ValueError('Impossible to find a solution for region {}. '
-                     'Net charge: {}, system ions: {}.'.format(
-        region_name, mol_net_charge, ions_names_charges))
+            if sum(counterions_charges) == counterions:
+                return(counterions_indices)
 
 
 # See Amber manual Table 4.1 http://ambermd.org/doc12/Amber15.pdf
diff --git a/Yank/restraints.py b/Yank/restraints.py
index 633ae76c..16fab11d 100644
--- a/Yank/restraints.py
+++ b/Yank/restraints.py
@@ -1204,13 +1204,13 @@ def _determine_restraint_parameters(self, thermodynamic_state, sampler_state, to
             The topography with labeled receptor and ligand atoms.
 
         """
-        # Determine number of atoms.
-        n_atoms = len(topography.receptor_atoms)
-
-        # Check that restrained receptor atoms are in expected range.
-        if any(atom_id >= n_atoms for atom_id in self.restrained_receptor_atoms):
-            raise ValueError('Receptor atoms {} were selected for restraint, but system '
-                             'only has {} atoms.'.format(self.restrained_receptor_atoms, n_atoms))
+        if not self._is_ligand_ligand_restraint(self.restrained_ligand_atoms):
+            # Determine number of atoms.
+            n_atoms = len(topography.receptor_atoms)
+            # Check that restrained receptor atoms are in expected range.
+            if any(atom_id >= n_atoms for atom_id in self.restrained_receptor_atoms):
+                raise ValueError('Receptor atoms {} were selected for restraint, but system '
+                                 'only has {} atoms.'.format(self.restrained_receptor_atoms, n_atoms))
 
         # Compute well radius if the user hasn't specified it in the constructor.
         if self.well_radius is None:
diff --git a/Yank/yank.py b/Yank/yank.py
index 21e7796b..69ed08a7 100644
--- a/Yank/yank.py
+++ b/Yank/yank.py
@@ -147,7 +147,7 @@ def receptor_atoms(self):
         if len(self._ligand_atoms) == 0:
             return []
 
-        #  alchemical_regionsCreate a set for fast searching.
+        #  Create a set for fast searching.
         if isinstance(self._ligand_atoms, tuple):
             ligand_atomset = frozenset(self._ligand_atoms[0] + self._ligand_atoms[1])
         else:
@@ -1347,29 +1347,34 @@ def _expand_state_cutoff(thermodynamic_state, expanded_cutoff_distance,
         return thermodynamic_state
 
     @staticmethod
-    def _build_default_alchemical_region(system, topography, protocol):
+    def _build_default_alchemical_region(system, topography, protocol, is_relative):
         """Create a default AlchemicalRegion if the user hasn't provided one."""
         # TODO: we should probably have a second region that annihilate sterics of counterions.
         alchemical_region_kwargs = {}
-        
+        alchemical_region_0_kwargs = {}
+        alchemical_region_1_kwargs = {} 
         # Modify ligand if this is a receptor-ligand phase, or
         # solute if this is a transfer free energy calculation.
         if len(topography.ligand_atoms) > 0:
             alchemical_region_name = 'ligand_atoms'
+            alchemical_atoms = getattr(topography, alchemical_region_name)
         else:
             alchemical_region_name = 'solute_atoms'
-        alchemical_atoms = getattr(topography, alchemical_region_name)
-
-        if isinstance(topography.ligand_atoms, tuple):
-            alchemical_region_name = ['ligand_atoms_0', 'ligand_atoms_1']
+            solute_atoms = getattr(topography, alchemical_region_name)
+            topo = topography.topology.subset(solute_atoms)
+            resnames = []
+            for chain in topo._chains:
+                for residue in chain._residues:
+                    resnames.append(f'resname {residue.name}')
+            alchemical_atoms = tuple([topography.select(r) for r in resnames])
 
         # In periodic systems, we alchemically modify the ligand/solute
         # counterions to make sure that the solvation box is always neutral.
-        if system.usesPeriodicBoundaryConditions():
-            if not isinstance(topography.ligand_atoms, tuple): # not a relative system
+        if system.usesPeriodicBoundaryConditions():           
+            if not is_relative:
                 alchemical_counterions = mpiplus.run_single_node(0, pipeline.find_alchemical_counterions,
                                                                  system, topography, alchemical_region_name,
-                                                                 broadcast_result=True)
+                                                                 is_relative, broadcast_result=True)
                 alchemical_atoms += alchemical_counterions
               
                 # Sort them by index for safety. We don't want to
@@ -1378,21 +1383,31 @@ def _build_default_alchemical_region(system, topography, protocol):
                 alchemical_region_kwargs['alchemical_atoms'] = alchemical_atoms
  
             else:
-                alchemical_atoms = [sorted(topography.ligand_atoms[0])]
-                alchemical_atoms.append(sorted(topography.ligand_atoms[1]))
-
+                ions_to_dummies, dummies_to_ions = mpiplus.run_single_node(0, pipeline.find_alchemical_counterions,
+                                                                           system, topography, alchemical_region_name,
+                                                                           is_relative, broadcast_result=True)                
+                if ions_to_dummies:
+                    alchemical_atoms[0] += ions_to_dummies
+                if dummies_to_ions:
+                    alchemical_atoms[1] += dummies_to_ions
+                alchemical_atoms = [sorted(atoms) for atoms in alchemical_atoms]
+                alchemical_region_0_kwargs['alchemical_atoms'] = alchemical_atoms[0]
+                alchemical_region_1_kwargs['alchemical_atoms'] = alchemical_atoms[1]
         # Check if we need to modify bonds/angles/torsions.
         for element_type in ['bonds', 'angles', 'torsions']:
             if 'lambda_' + element_type in protocol:
                 modify_it = True
             else:
                 modify_it = None
-            alchemical_region_kwargs['alchemical_' + element_type] = modify_it
-
+            if not is_relative:
+                alchemical_region_kwargs['alchemical_' + element_type] = modify_it
+            else:
+                alchemical_region_0_kwargs['alchemical_' + element_type] = modify_it
+                alchemical_region_1_kwargs['alchemical_' + element_type] = modify_it
         # Create alchemical region.
-        if isinstance(topography.ligand_atoms, tuple): # relative system
-            alchemical_region = [mmtools.alchemy.AlchemicalRegion(name='0', alchemical_atoms=alchemical_atoms[0], **alchemical_region_kwargs),
-                                 mmtools.alchemy.AlchemicalRegion(name='1', alchemical_atoms=alchemical_atoms[1], **alchemical_region_kwargs)]
+        if is_relative:
+            alchemical_region = [mmtools.alchemy.AlchemicalRegion(name='0', **alchemical_region_0_kwargs),
+                                 mmtools.alchemy.AlchemicalRegion(name='1', **alchemical_region_1_kwargs)]
 
         else:
             alchemical_region = [mmtools.alchemy.AlchemicalRegion(**alchemical_region_kwargs)]

From 7e70ca80a31a62a6b197a9f8491632ab4ec483c5 Mon Sep 17 00:00:00 2001
From: Ana Silveira <anita_silveira@hotmail.com>
Date: Mon, 27 Apr 2020 15:00:24 -0400
Subject: [PATCH 8/8] revert changes in
 ExperimentBuilder._generate_auto_state_parameters change alchemical region
 setup

---
 Yank/experiment.py       | 54 +++++++++++++++++++---------------------
 Yank/restraints.py       |  3 ++-
 Yank/schema/validator.py |  5 ++--
 Yank/yank.py             | 45 +++++++++++++++++----------------
 4 files changed, 54 insertions(+), 53 deletions(-)

diff --git a/Yank/experiment.py b/Yank/experiment.py
index 4bd41cbc..f87cfbf2 100644
--- a/Yank/experiment.py
+++ b/Yank/experiment.py
@@ -2611,36 +2611,22 @@ def _generate_auto_state_parameters(phase_factory):
                          len(phase_factory.topography.solvent_atoms) == 0)
         state_parameters = []
 
-        if not isinstance(phase_factory.topography.ligand_atoms, tuple): # not a relative system
         # First, turn on the restraint if there are any.
-            if phase_factory.restraint is not None:
-                state_parameters.append(('lambda_restraints', [0.0, 1.0]))
-            # We support only lambda sterics and electrostatics for now.
-            if is_vacuum and not phase_factory.alchemical_regions.annihilate_electrostatics:
-                state_parameters.append(('lambda_electrostatics', [1.0, 1.0]))
-            else:
-                state_parameters.append(('lambda_electrostatics', [1.0, 0.0]))
-            if is_vacuum and not phase_factory.alchemical_regions.annihilate_sterics:
-                state_parameters.append(('lambda_sterics', [1.0, 1.0]))
-            else:
-                state_parameters.append(('lambda_sterics', [1.0, 0.0]))
-            # Turn the RMSD restraints off slowly at the end
-            if isinstance(phase_factory.restraint, restraints.RMSD):
-                state_parameters.append(('lambda_restraints', [1.0, 0.0]))
+        if phase_factory.restraint is not None:
+            state_parameters.append(('lambda_restraints', [0.0, 1.0]))
+        # We support only lambda sterics and electrostatics for now.
+        if is_vacuum and not phase_factory.alchemical_regions.annihilate_electrostatics:
+           state_parameters.append(('lambda_electrostatics', [1.0, 1.0]))
         else:
-            if is_vacuum and not phase_factory.alchemical_regions.annihilate_electrostatics:
-                state_parameters.append(('lambda_electrostatics_0', [1.0, 1.0]))
-                state_parameters.append(('lambda_electrostatics_1', [1.0, 1.0]))
-            else:
-                state_parameters.append(('lambda_electrostatics_0', [1.0, 0.0]))
-                state_parameters.append(('lambda_electrostatics_1', [0.0, 1.0]))
-            if is_vacuum and not phase_factory.alchemical_regions.annihilate_sterics:
-                state_parameters.append(('lambda_sterics_0', [1.0, 1.0]))
-                state_parameters.append(('lambda_sterics_1', [1.0, 1.0]))
-            else:
-                state_parameters.append(('lambda_sterics_0', [1.0, 0.0]))
-                state_parameters.append(('lambda_sterics_1', [0.0, 1.0]))
-            # TODO: Handle restraints in relative free energy calculations 
+           state_parameters.append(('lambda_electrostatics', [1.0, 0.0]))
+        if is_vacuum and not phase_factory.alchemical_regions.annihilate_sterics:
+           state_parameters.append(('lambda_sterics', [1.0, 1.0]))
+        else:
+           state_parameters.append(('lambda_sterics', [1.0, 0.0]))
+        # Turn the RMSD restraints off slowly at the end
+        if isinstance(phase_factory.restraint, restraints.RMSD):
+           state_parameters.append(('lambda_restraints', [1.0, 0.0]))
+
         return state_parameters
 
     @classmethod
@@ -3130,6 +3116,18 @@ def _build_experiment(self, experiment_path, experiment, use_dummy_protocol=Fals
                     is_relative = True
             else:
                 ligand_atoms = None
+                topography = Topography(topology, ligand_atoms=ligand_atoms,
+                                    solvent_atoms=solvent_dsl)
+
+                solute_atoms = getattr(topography, 'solute_atoms')
+                topo = topography.topology.subset(solute_atoms)
+                resnames = []
+                for chain in topo._chains:
+                    for residue in chain._residues:
+                        resnames.append(f'resname {residue.name}')
+                if len(set(resnames)) == 2:
+                    is_relative = True
+ 
             topography = Topography(topology, ligand_atoms=ligand_atoms,
                                     solvent_atoms=solvent_dsl)
 
diff --git a/Yank/restraints.py b/Yank/restraints.py
index 16fab11d..bc7398c6 100644
--- a/Yank/restraints.py
+++ b/Yank/restraints.py
@@ -2571,9 +2571,10 @@ def restrain_state(self, thermodynamic_state):
 
         # Check if all parameters are defined.
         self._check_parameters_defined()
+
         # Merge receptor and ligand atoms in a single array for easy manipulation.
         restrained_atoms = self.restrained_receptor_atoms + self.restrained_ligand_atoms
-       
+
         # Create RMSDForce CV for all restrained atoms
         rmsd_cv = openmm.RMSDForce(self.reference_sampler_state.positions, restrained_atoms)
 
diff --git a/Yank/schema/validator.py b/Yank/schema/validator.py
index 0963a8f9..72c5dac2 100644
--- a/Yank/schema/validator.py
+++ b/Yank/schema/validator.py
@@ -187,10 +187,11 @@ def _check_with_only_with_no_cutoff(self, field, value):
 
     def _check_with_specify_lambda_electrostatics_and_sterics(self, field, value):
         """Check that the keys of a dictionary contain both lambda_electrostatics and lambda_sterics."""
+
         if ((isinstance(value, dict) or isinstance(value, collections.OrderedDict)) and
              not (('lambda_sterics' in value and 'lambda_electrostatics' in value) or
-             ('lambda_sterics_0' in value and 'lambda_electrostatics_0' in value and
-               'lambda_sterics_1' in value and 'lambda_electrostatics_1' in value))):
+             ('lambda_sterics_0' in value and 'lambda_electrostatics_0' in value and 
+              'lambda_sterics_1' in value and 'lambda_electrostatics_1' in value))):
             self._error(field, "Missing required keys lambda_sterics and/or lambda_electrostatics")
 
     def _check_with_math_expressions_variables_are_given(self, field, value):
diff --git a/Yank/yank.py b/Yank/yank.py
index 69ed08a7..9275e22e 100644
--- a/Yank/yank.py
+++ b/Yank/yank.py
@@ -1350,9 +1350,9 @@ def _expand_state_cutoff(thermodynamic_state, expanded_cutoff_distance,
     def _build_default_alchemical_region(system, topography, protocol, is_relative):
         """Create a default AlchemicalRegion if the user hasn't provided one."""
         # TODO: we should probably have a second region that annihilate sterics of counterions.
-        alchemical_region_kwargs = {}
         alchemical_region_0_kwargs = {}
-        alchemical_region_1_kwargs = {} 
+        if is_relative:
+            alchemical_region_1_kwargs = {}
         # Modify ligand if this is a receptor-ligand phase, or
         # solute if this is a transfer free energy calculation.
         if len(topography.ligand_atoms) > 0:
@@ -1367,50 +1367,51 @@ def _build_default_alchemical_region(system, topography, protocol, is_relative):
                 for residue in chain._residues:
                     resnames.append(f'resname {residue.name}')
             alchemical_atoms = tuple([topography.select(r) for r in resnames])
-
         # In periodic systems, we alchemically modify the ligand/solute
         # counterions to make sure that the solvation box is always neutral.
         if system.usesPeriodicBoundaryConditions():           
-            if not is_relative:
-                alchemical_counterions = mpiplus.run_single_node(0, pipeline.find_alchemical_counterions,
-                                                                 system, topography, alchemical_region_name,
-                                                                 is_relative, broadcast_result=True)
-                alchemical_atoms += alchemical_counterions
-              
-                # Sort them by index for safety. We don't want to
-                # accidentally exchange two atoms' positions.
-                alchemical_atoms = sorted(alchemical_atoms)
-                alchemical_region_kwargs['alchemical_atoms'] = alchemical_atoms
- 
-            else:
+            alchemical_counterions = mpiplus.run_single_node(0, pipeline.find_alchemical_counterions,
+                                                             system, topography, alchemical_region_name,
+                                                             is_relative, broadcast_result=True)  
+
+            if is_relative:
                 ions_to_dummies, dummies_to_ions = mpiplus.run_single_node(0, pipeline.find_alchemical_counterions,
                                                                            system, topography, alchemical_region_name,
-                                                                           is_relative, broadcast_result=True)                
+                                                                           is_relative, broadcast_result=True)
                 if ions_to_dummies:
                     alchemical_atoms[0] += ions_to_dummies
                 if dummies_to_ions:
                     alchemical_atoms[1] += dummies_to_ions
+                # Sort them by index for safety. We don't want to
+                # accidentally exchange two atoms' positions.
                 alchemical_atoms = [sorted(atoms) for atoms in alchemical_atoms]
                 alchemical_region_0_kwargs['alchemical_atoms'] = alchemical_atoms[0]
                 alchemical_region_1_kwargs['alchemical_atoms'] = alchemical_atoms[1]
+
+            else:
+                alchemical_atoms += alchemical_counterions            
+                alchemical_atoms = sorted(alchemical_atoms)
+                alchemical_region_0_kwargs[0]['alchemical_atoms'] = alchemical_atoms
+ 
         # Check if we need to modify bonds/angles/torsions.
         for element_type in ['bonds', 'angles', 'torsions']:
             if 'lambda_' + element_type in protocol:
                 modify_it = True
             else:
                 modify_it = None
-            if not is_relative:
-                alchemical_region_kwargs['alchemical_' + element_type] = modify_it
-            else:
+            if is_relative:
                 alchemical_region_0_kwargs['alchemical_' + element_type] = modify_it
                 alchemical_region_1_kwargs['alchemical_' + element_type] = modify_it
+            else:
+                alchemical_region_0_kwargs['alchemical_' + element_type] = modify_it
         # Create alchemical region.
-        if is_relative:
+        if is_relative:            
             alchemical_region = [mmtools.alchemy.AlchemicalRegion(name='0', **alchemical_region_0_kwargs),
                                  mmtools.alchemy.AlchemicalRegion(name='1', **alchemical_region_1_kwargs)]
-
+ 
         else:
-            alchemical_region = [mmtools.alchemy.AlchemicalRegion(**alchemical_region_kwargs)]
+            alchemical_region = [mmtools.alchemy.AlchemicalRegion(**alchemical_region_0_kwargs)]
+
         return alchemical_region
 
     @staticmethod