[WIP] add support for relative binding free energies

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!')
@@ -1678,7 +1677,7 @@ def _is_pipeline_solvent_with_receptor(field, solvent_id, error):
         # DSL Schema
         ligand_dsl:
             required: no
-            type: string
+            type: [string, list]
             dependencies: [phase1_path, phase2_path]
         solvent_dsl:
             required: no
@@ -1746,7 +1745,7 @@ def _is_pipeline_solvent_with_receptor(field, solvent_id, error):
             check_with: REGION_CLASH_DETERMINED_AT_RUNTIME_WITH_LIGAND
         ligand:
             required: no
-            type: string
+            type: [string, list]
             dependencies: [receptor, solvent]
             allowed: MOLECULE_IDS_POPULATED_AT_RUNTIME
             excludes: [solute, phase1_path, phase2_path]
@@ -2617,16 +2616,16 @@ def _generate_auto_state_parameters(phase_factory):
             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]))
+           state_parameters.append(('lambda_electrostatics', [1.0, 1.0]))
         else:
-            state_parameters.append(('lambda_electrostatics', [1.0, 0.0]))
+           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]))
+           state_parameters.append(('lambda_sterics', [1.0, 1.0]))
         else:
-            state_parameters.append(('lambda_sterics', [1.0, 0.0]))
+           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]))
+           state_parameters.append(('lambda_restraints', [1.0, 0.0]))
 
         return state_parameters
 
@@ -3089,6 +3088,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.
@@ -3112,8 +3112,22 @@ 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,
+                                    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)
 
@@ -3134,8 +3148,8 @@ 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)
-            alchemical_region = alchemical_region._replace(**alchemical_region_opts)
+                                                                                 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
             # will take care of raising an error if the phase type does not support it.

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,32 +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