#
# Copyright 2015 by Justin MacCallum, Alberto Perez, Ken Dill
# All rights reserved
#
"""
Module to handle replica exchange swaps
"""
import logging
import math
import random
from typing import List
import numpy as np # type: ignore
from meld import util
from meld.remd import adaptor
logger = logging.getLogger(__name__)
[docs]class NearestNeighborLadder:
"""
Class to compute replica exchange swaps between neighboring replicas.
"""
[docs] def __init__(self, n_trials: int) -> None:
"""
Initialize a NearestNeighborLadder
Args:
n_trials: number of swaps to attempt
"""
self.n_trials = n_trials
[docs] @util.log_timing(logger)
def compute_exchanges(
self, energies: np.ndarray, adaptor: adaptor.Adaptor
) -> List[int]:
"""
Compute the exchanges from a given energy matrix.
Args:
energies: numpy array of energies (see below for details)
adaptor: replica exchange adaptor that is updated every attempted swap
Returns:
a permutation vector (see below for details)
The energy matrix should be an n_rep x n_rep numpy array. All energies
are in dimensionless units (unit of kT). Each column represents a
particular structure, while each row represents a particular
combination of temperature and Hamiltonian. So, ``energies[i,j]`` is
the energy of structure j with temperature and hamiltonian i. The
diagonal ``energies[i,i]`` are the energies that were actually
simulated.
This method will attempt ``self.n_trials`` swaps between randomly
chosen pairs of adjacent replicas. It will return a permutation vector
that describes which index each structure should be at after swapping.
So, if the output was ``[2, 0, 1]``, it would mean that replica 0
should now be at index 2, replica 1 should now be at index 0, and
replica 2 should now be at index 1. Output of ``[0, 1, 2]`` would
mean no change.
The adaptor object is called once for each attempted exchange with the
indices of the attempted swap and the success or failure of the swap.
"""
assert len(energies.shape) == 2
assert energies.shape[0] == energies.shape[1]
n_replicas = energies.shape[0]
permutation_vector = list(range(n_replicas))
choices = range(n_replicas - 1)
for iteration in range(self.n_trials):
i = random.choice(choices)
j = i + 1
self._do_trial(i, j, permutation_vector, energies, adaptor)
return permutation_vector
def _do_trial(
self,
i: int,
j: int,
permutation_vector: List[int],
energies: np.ndarray,
adaptor: adaptor.Adaptor,
) -> None:
"""Perform a replica exchange trial"""
delta = energies[i, i] - energies[j, i] + energies[j, j] - energies[i, j]
accepted = False
if delta >= 0:
accepted = True
else:
metrop = math.exp(delta)
rand = random.random()
if rand < metrop:
accepted = True
if accepted:
self._swap_permutation(i, j, permutation_vector)
self._swap_energies(i, j, energies)
adaptor.update(i, True)
else:
adaptor.update(i, False)
@staticmethod
def _swap_permutation(i: int, j: int, permutation_vector: List[int]) -> None:
"""Swap two elements of the permutation matrix"""
permutation_vector[i], permutation_vector[j] = (
permutation_vector[j],
permutation_vector[i],
)
@staticmethod
def _swap_energies(i: int, j: int, energies: np.ndarray) -> None:
"""Swap two columns of the energy matrix"""
energies[:, [i, j]] = energies[:, [j, i]]