import numpy as np
from ensembler.util.ensemblerTypes import NoReturn, Dict, List, Tuple
[docs]class Exchange_pattern:
"""
general Scaffold for an exchange pattern
"""
replica_graph = None
[docs] def __init__(self, replica_graph:dict):
"""
build the exchange scaffold for replica_graph
Parameters
----------
replica_graph: dict
representation of the replica graph
"""
self.replica_graph = replica_graph
[docs] def exchange(self, verbose: bool = False):
"""
@ interface needs to be implemented in subclass
Parameters
----------
verbose
Returns
-------
"""
raise NotImplementedError("Not Implemented")
[docs] def _do_exchange(self, exchanges_to_make: Dict[Tuple[int, int], bool], verbose: bool = False)->NoReturn:
"""
post exchange decision step, executing the exchange
Parameters
----------
exchanges_to_make: Dict[Tuple[int, int], bool]
the key of the dictionary is a tuple containing the replica identifiers, which exchange parameters. The value is stating if they should exchang.
verbose: bool, optional
DADADA (default: False)
"""
for (I, J), exchange in exchanges_to_make.items():
if (exchange):
if (verbose): print("Exchanging: " + str(I) + "\t" + str(J) + "\t" + str(exchange) + "\n")
partnerI = self.replica_graph.replicas[I]
partnerJ = self.replica_graph.replicas[J]
# here the real exchange happens!
tmp_exchange_parameter = getattr(partnerI, self.replica_graph.exchange_param)
setattr(partnerI, self.replica_graph.exchange_param,
getattr(partnerJ, self.replica_graph.exchange_param))
setattr(partnerJ, self.replica_graph.exchange_param, tmp_exchange_parameter)
# here we follow the exchanges with the replicaID
tmp_id = getattr(partnerI, "replicaID")
setattr(partnerI, "replicaID", getattr(partnerJ, "replicaID"))
setattr(partnerJ, "replicaID", tmp_id)
# This is traj specific.... - might overwrite other params
if (self.replica_graph.exchange_param == "trajectory"):
partnerI._update_state_from_traj()
else:
if (verbose): print("not Exchanging: " + str(I) + " / " + str(J) + " \n")
[docs]class localExchangeScheme(Exchange_pattern):
"""
This scheme is exchanging a replica with its neighbour if possible
"""
exchange_offset: int = 0
[docs] def exchange(self, verbose: bool = False):
"""
Exchange the Trajectory of T-replicas in pairwise fashion
Parameters
----------
verbose: bool, optional
Make some music (Default: False)
"""
self.replica_graph._currentTrial += 1
# Get Potential Energies
original_exchange_coordinates, original_totPots, swapped_exchange_coordinates, swapped_totPots = self._collect_replica_energies(
verbose)
if (verbose):
print("origTotE ", [original_totPots[key] for key in sorted(original_exchange_coordinates)])
print("SWPTotE ", [swapped_totPots[key] for key in sorted(swapped_exchange_coordinates)])
# decide exchange
exchanges_to_make = {}
for partner1, partner2 in zip(original_exchange_coordinates[self.exchange_offset::2],
original_exchange_coordinates[1 + self.exchange_offset::2]):
originalEnergies = np.add(original_totPots.get(partner1), original_totPots.get(partner2))
swapEnergies = np.add(swapped_totPots.get(partner1), swapped_totPots.get(partner2))
exchanges_to_make.update(
{(partner1, partner2): self.replica_graph.exchange_criterium(originalEnergies, swapEnergies)})
# Acutal Exchange of params (actually trajs here
if (verbose):
print("Exchange: ", exchanges_to_make)
print("exchaning param: ", self.replica_graph.exchange_param)
# execute exchange
self._do_exchange(exchanges_to_make=exchanges_to_make)
# update exchange info
self.replica_graph._current_exchanges = exchanges_to_make
self.update_exchange_information(original_exchange_coordinates=original_exchange_coordinates,
original_totPots=original_totPots, swapped_totPots=swapped_totPots)
# update the offset
self.exchange_offset = (self.exchange_offset + 1) % 2
[docs] def _swap_coordinates(self, original_exCoord:List[int])->List[int]:
"""
swap the coordinates pairwise of the original exchange coord list
Parameters
----------
original_exCoord: List
Of replicas in order
Returns
-------
List
return list with swapped replica IDs
"""
##take care of offset situations and border replicas
swapped_exCoord = [] if self.exchange_offset == 0 else [original_exCoord[0]]
##generate sequence with swapped params
for partner1, partner2 in zip(original_exCoord[self.exchange_offset::2],
original_exCoord[1 + self.exchange_offset::2]):
swapped_exCoord.extend([partner2, partner1])
##last replica on the border?
if (self.exchange_offset == 0):
swapped_exCoord.append(original_exCoord[-1])
return swapped_exCoord
##Exchange functions
[docs] def _collect_replica_energies(self, verbose: bool=False)->Tuple[List[int], List[float], List[int], List[float]]:
"""
_collect_replica_energies
collect all rep
Parameters
----------
verbose: bool, optional
More noise! More noise!
Returns
-------
Tuple[List[int], List[float], List[int], List[float]]
original_exCoord, original_totPots, swapped_exCoord, swapped_totPots
"""
original_totPots = self.replica_graph.get_replica_total_energies()
original_exCoord = list(sorted(original_totPots.keys()))
#replica_values = list([self.replica_graph.replicas[key] for key in original_exCoord])
# SWAP exchange_coordinates params pairwise
swapped_exCoord = self._swap_coordinates(original_exCoord)
if (verbose):
print("original Coords ", original_exCoord)
print("swapped Coords ", swapped_exCoord)
# SWAP Params to calculate energies in swapped case
self.replica_graph.set_parameter_set(coordinates=swapped_exCoord, replica_indices=original_exCoord) # swap parameters
##scaleVel - adapt to new parameter set
self.replica_graph._adapt_system_to_exchange_coordinate()
##get_swapped energies
swapped_totPots = self.replica_graph.get_replica_total_energies() # calc swapped parameter Energies
#Swap back
self.replica_graph.set_parameter_set(coordinates=original_exCoord,
replica_indices=original_exCoord) # swap back parameters
##scaleVel - adapt to old parameter set
self.replica_graph._adapt_system_to_exchange_coordinate()
return original_exCoord, original_totPots, swapped_exCoord, swapped_totPots
[docs]class globalExchangeScheme(Exchange_pattern):
"""
Implement a grlobal exchange scheme here :)
"""
[docs] def exchange(self, verbose: bool = False):
pass