diff --git a/Yank/experiment.py b/Yank/experiment.py index 2b266403..5e4d3f73 100644 --- a/Yank/experiment.py +++ b/Yank/experiment.py @@ -3142,7 +3142,8 @@ def _build_experiment(self, experiment_path, experiment, use_dummy_protocol=Fals # Start from AlchemicalPhase default alchemical region # and modified it according to the user options. phase_protocol = protocol[phase_name]['alchemical_path'] - alchemical_region = AlchemicalPhase._build_default_alchemical_region(system, topography, + alchemical_region = AlchemicalPhase._build_default_alchemical_region(system, topography, + sampler_state, phase_protocol) alchemical_region = alchemical_region._replace(**alchemical_region_opts) diff --git a/Yank/pipeline.py b/Yank/pipeline.py index 41bbe935..5649e850 100644 --- a/Yank/pipeline.py +++ b/Yank/pipeline.py @@ -233,15 +233,55 @@ def compute_net_charge(system, atom_indices): net_charge = int(round(net_charge / unit.elementary_charge)) return net_charge +def order_counterions(system, topography, sampler_state): + """Order the counterions in decreasing distance to the ligand/solute. -def find_alchemical_counterions(system, topography, region_name): + Parameters + ---------- + system : simtk.openmm.System + The system object containing the atoms of interest. + topography : Yank topography + The topography object holding the indices of the ions and the + ligand (for binding free energy) or solute (for transfer free + energy). + sampler_state : openmmtools.mdtools.sampler_state + State with positions + + Returns + ------- + ion_index, charge : list + List of ion index and charge + + """ + pos_unit = sampler_state.positions.unit + # atoms of the ligand/solute which are part of the alchemical transformation + alchemical_atoms = topography.ligand_atoms + # if no ligand present but a charged solute + if not alchemical_atoms: + alchemical_atoms = topography.solute_atoms + alchemical_positions = sampler_state.positions[alchemical_atoms] / pos_unit + # compute minimal distance between each ion and the alchemical atoms + ions_distances = np.zeros(len(topography.ions_atoms)) + for i, ion_id in enumerate(topography.ions_atoms): + ion_position = sampler_state.positions[ion_id] / pos_unit + ions_distances[i] = compute_min_dist([ion_position], alchemical_positions) + logger.debug('Min distance of ion {} {}: {}'.format(ion_id, + topography.topology.atom(ion_id).residue.name, + ions_distances[i])) + # list with ion_id ordered in decreasing distance of the ligand + ordered_ion_id = [topography.ions_atoms[ion_id] for ion_id in np.argsort(ions_distances)[::-1]] + return [(ion_id, compute_net_charge(system, [ion_id])) + for ion_id in ordered_ion_id] + +def find_alchemical_counterions(system, topography, sampler_state, region_name): """Return the atom indices of the ligand or solute counter ions. In periodic systems, the solvation box needs to be neutral, and if the decoupled molecule is charged, it will cause trouble. This can be used to find a set of ions in the system that neutralize the molecule, so that the solvation box will remain neutral all - the time. + the time. Ions are selected starting with the ones with the largest + distance to the ligand. Parameters ---------- @@ -251,6 +291,8 @@ def find_alchemical_counterions(system, topography, region_name): The topography object holding the indices of the ions and the ligand (for binding free energy) or solute (for transfer free energy). + sampler_state : openmmtools.mdtools.sampler_state + State with positions region_name : str The region name in the topography (e.g. "ligand_atoms") for which to find counter ions. @@ -280,19 +322,20 @@ def find_alchemical_counterions(system, topography, region_name): 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] + # Find net charge of all ions in the system and order them according to the + # largest distance from the ligand/solute + ions_net_charges = order_counterions(system, topography, sampler_state) 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)) + # 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: + logger.debug('Index of alchemical counter ion {}'.format(counterions_indices)) return counterions_indices # We couldn't find any subset of counterions neutralizing the region. diff --git a/Yank/tests/test_pipeline.py b/Yank/tests/test_pipeline.py index 780b93ed..838e8547 100644 --- a/Yank/tests/test_pipeline.py +++ b/Yank/tests/test_pipeline.py @@ -14,7 +14,6 @@ # ============================================================================= from yank.pipeline import * -from yank.pipeline import _redistribute_trailblaze_states from nose.tools import assert_raises_regexp diff --git a/Yank/tests/test_yank.py b/Yank/tests/test_yank.py index 7ae42e85..c5d502d8 100644 --- a/Yank/tests/test_yank.py +++ b/Yank/tests/test_yank.py @@ -106,6 +106,7 @@ def test_topography_subset_regions(): selected = topography.select(selection, sort_by=sort_by, subset=subset, as_set=False) assert result == selected, (f"Failed to match {selection}, subset {subset}, " f"by {sort_by} to {result},\ngot {selected}") + selected) def test_topography_regions(): @@ -393,14 +394,14 @@ def test_default_alchemical_region(self): # Test case: (system, topography, topography_region) test_cases = [ - (self.host_guest_implicit[1].system, self.host_guest_implicit[3], 'ligand_atoms'), - (self.host_guest_explicit[1].system, self.host_guest_explicit[3], 'ligand_atoms'), - (self.alanine_explicit[1].system, self.alanine_explicit[3], 'solute_atoms'), - (sodium_chloride.system, negative_solute, 'solute_atoms'), - (sodium_chloride.system, positive_solute, 'solute_atoms') + (self.host_guest_implicit[1].system, self.host_guest_implicit[3], self.host_guest_implicit[2], 'ligand_atoms'), + (self.host_guest_explicit[1].system, self.host_guest_explicit[3], self.host_guest_explicit[2], 'ligand_atoms'), + (self.alanine_explicit[1].system, self.alanine_explicit[3], self.alanine_explicit[2], 'solute_atoms'), + (sodium_chloride.system, negative_solute, sodium_chloride[2], 'solute_atoms'), + (sodium_chloride.system, positive_solute, sodium_chloride[2], 'solute_atoms') ] - for system, topography, alchemical_region_name in test_cases: + for system, topography, sampler_state, alchemical_region_name in test_cases: expected_alchemical_atoms = getattr(topography, alchemical_region_name) # Compute net charge of alchemical atoms. @@ -411,7 +412,7 @@ def test_default_alchemical_region(self): for protocol in protocols: alchemical_region = AlchemicalPhase._build_default_alchemical_region( - system, topography, protocol) + system, topography, sampler_state, protocol) assert alchemical_region.alchemical_atoms == expected_alchemical_atoms # The alchemical region must be neutralized. diff --git a/Yank/yank.py b/Yank/yank.py index b117365a..40f718dd 100644 --- a/Yank/yank.py +++ b/Yank/yank.py @@ -325,11 +325,11 @@ def select(self, selection, as_set=False, sort_by='auto', subset=None) -> Union[ no effect. """ + selection = copy.deepcopy(selection) topology = self.topology - - # Make sure that the subset of atoms is a list of atom indices - # and slice the topology will be matching against accordingly. + # Handle subset. Define a subset topology to manipulate, then define a common call to convert subset atom + # into absolute atom if subset is not None: subset = self.select(subset, sort_by='index', as_set=False, subset=None) topology = self.topology.subset(subset) @@ -1007,7 +1007,7 @@ def create(self, thermodynamic_state, sampler_states, topography, protocol, # Handle default alchemical region. if alchemical_regions is None: alchemical_regions = self._build_default_alchemical_region( - reference_system, topography, protocol) + reference_system, topography, sampler_states[0], protocol) # Check that we have atoms to alchemically modify. if len(alchemical_regions.alchemical_atoms) == 0: @@ -1331,7 +1331,7 @@ 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, sampler_state, 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 = {} @@ -1348,8 +1348,8 @@ def _build_default_alchemical_region(system, topography, protocol): # 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) + system, topography, sampler_state, + alchemical_region_name, broadcast_result=True) alchemical_atoms += alchemical_counterions # Sort them by index for safety. We don't want to