Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add feature to automatically select the alchemical ion with largest distance #1257

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion Yank/experiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
55 changes: 49 additions & 6 deletions Yank/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
1 change: 0 additions & 1 deletion Yank/tests/test_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
# =============================================================================

from yank.pipeline import *
from yank.pipeline import _redistribute_trailblaze_states

from nose.tools import assert_raises_regexp

Expand Down
15 changes: 8 additions & 7 deletions Yank/tests/test_yank.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand Down
14 changes: 7 additions & 7 deletions Yank/yank.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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 = {}
Expand All @@ -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
Expand Down