#
# Copyright 2015 by Justin MacCallum, Alberto Perez, Ken Dill
# All rights reserved
#
"""
This module implements a number of interfaces that help to minimize coupling
and eliminate circular imports
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import List, NamedTuple, Optional, Sequence
import numpy as np
import openmm as mm # type: ignore
from openmm import app # type: ignore
from meld.system import (
density,
indexing,
mapping,
param_sampling,
pdb_writer,
restraints,
temperature,
)
[docs]class ICommunicator(ABC):
"""
Interface for communication between leader and workers
"""
[docs] @abstractmethod
def initialize(self) -> None:
"""
Initialize and start MPI
"""
pass
[docs] @abstractmethod
def is_leader(self) -> bool:
"""
Is this the leader node?
Returns:
:const:`True` if we are the leader, otherwise :const:`False`
"""
pass
[docs] @abstractmethod
def barrier(self) -> None:
"""
Wait until all workers reach this point
"""
pass
[docs] @abstractmethod
def distribute_alphas_to_workers(self, all_alphas: List[float]) -> List[float]:
"""
Distribute alphas to workers
Args:
all_alphas: the alpha values to be distributed
Returns:
the block of alpha values for the leader
"""
pass
[docs] @abstractmethod
def receive_alphas_from_leader(self) -> List[float]:
"""
Receive a block of alphas from leader.
Returns:
the block of alpha values for this worker
"""
pass
[docs] @abstractmethod
def distribute_states_to_workers(
self, all_states: Sequence[IState]
) -> List[IState]:
"""
Distribute a block of states to each worker.
Args:
all_states: states to be distributed
Returns:
the block of states to run on the leader node
"""
pass
[docs] @abstractmethod
def receive_states_from_leader(self) -> List[IState]:
"""
Get the block of states to run for this step
Returns:
the block of states to run for this step
"""
pass
[docs] @abstractmethod
def gather_states_from_workers(self, state_on_leader: List[IState]) -> List[IState]:
"""
Receive states from all workers
Args:
states_on_leader: the block of states on the leader after simulating
Returns:
A list of states, one from each replica.
"""
pass
[docs] @abstractmethod
def send_states_to_leader(self, block: Sequence[IState]) -> None:
"""
Send a block of states to the leader
Args:
block: block of states to send to the leader.
"""
pass
[docs] @abstractmethod
def broadcast_all_states_to_workers(self, states: Sequence[IState]) -> None:
"""
Broadcast all states to all workers.
Args:
states: a list of states
"""
pass
[docs] @abstractmethod
def receive_all_states_from_leader(self) -> Sequence[IState]:
"""
Receive all states from leader.
Returns:
a list of states to calculate the energy of
"""
pass
[docs] @abstractmethod
def gather_energies_from_workers(
self, energies_on_leader: np.ndarray
) -> np.ndarray:
"""
Receive energies from each worker.
Args:
energies_on_leader: the energies from the leader
Returns:
a square matrix of every state on every replica to be used for replica exchange
Note:
Each row of the output matrix represents a different Hamiltonian. Each column
represents a different state. Each worker will compute multiple rows of the
output matrix.
"""
pass
[docs] @abstractmethod
def send_energies_to_leader(self, energies: np.ndarray) -> None:
"""
Send a block of energies to the leader.
Args:
energies: block of energies to send to the leader
Note:
Each row represents a different Hamiltonian. Each column represents a
different state.
"""
pass
[docs] @abstractmethod
def negotiate_device_id(self) -> int:
"""
Negotiate CUDA device id
Returns:
the cuda device id to use
"""
pass
@property
@abstractmethod
def n_replicas(self) -> int:
"""number of replicas"""
pass
@property
@abstractmethod
def n_atoms(self) -> int:
"""number of atoms"""
pass
@property
@abstractmethod
def n_workers(self) -> int:
"""number of workers"""
pass
@property
@abstractmethod
def rank(self) -> int:
"""rank of this worker"""
pass
[docs]class IState(ABC):
"""
Interface for SystemState
"""
positions: np.ndarray
velocities: np.ndarray
alpha: float
energy: float
group_energies: np.ndarray
box_vector: np.ndarray
parameters: param_sampling.ParameterState
mappings: np.ndarray
rdc_alignments: np.ndarray
[docs]class IRunner(ABC):
"""
Interface for replica runners
"""
temperature_scaler: Optional[temperature.TemperatureScaler]
@abstractmethod
def minimize_then_run(self, state: IState) -> IState:
pass
@abstractmethod
def run(self, state: IState) -> IState:
pass
@abstractmethod
def get_energy(self, state: IState) -> float:
pass
@abstractmethod
def get_group_energies(self, state: IState) -> np.ndarray:
pass
@abstractmethod
def prepare_for_timestep(self, state: IState, alpha: float, timestep: int):
pass
[docs]class ISystem(ABC):
"""
An interface for MELD systems
"""
restraints: restraints.RestraintManager
index: indexing.Indexer
temperature_scaler: Optional[temperature.TemperatureScaler]
param_sampler: param_sampling.ParameterManager
mapper: mapping.PeakMapManager
density: density.DensityManager
builder_info: dict
extra_bonds: List[ExtraBondParam]
extra_restricted_angles: List[ExtraAngleParam]
extra_torsions: List[ExtraTorsParam]
@property
@abstractmethod
def num_alignments(self) -> int:
"""
Number of rdc alignments
"""
pass
@property
@abstractmethod
def omm_system(self) -> mm.System:
"""
Get the openmm system
"""
pass
@property
@abstractmethod
def topology(self) -> app.Topology:
"""
Get the openmm topology
"""
pass
@property
@abstractmethod
def integrator(self) -> mm.LangevinIntegrator:
"""
Get the integrator
"""
pass
@property
@abstractmethod
def barostat(self) -> mm.MonteCarloBarostat:
"""
Get the barostat
"""
pass
@property
@abstractmethod
def solvation(self) -> str:
"""
Get the solvation model
"""
pass
@property
@abstractmethod
def n_atoms(self) -> int:
"""
number of atoms
"""
pass
@property
@abstractmethod
def template_coordinates(self) -> np.ndarray:
"""
Get the template coordinates
"""
pass
@property
@abstractmethod
def template_velocities(self) -> np.ndarray:
"""
Get the template velocities
"""
pass
@property
@abstractmethod
def template_box_vectors(self) -> Optional[np.ndarray]:
"""
Get the template box vectors
"""
pass
@property
@abstractmethod
def atom_names(self) -> List[str]:
"""
names for each atom
"""
pass
@property
@abstractmethod
def residue_numbers(self) -> List[int]:
"""
residue numbers for each atom
"""
pass
@property
@abstractmethod
def residue_names(self) -> List[str]:
"""
residue names for each atom
"""
pass
[docs] @abstractmethod
def get_state_template(self) -> IState:
"""
Get a template state for this system
"""
pass
[docs] @abstractmethod
def get_pdb_writer(self) -> pdb_writer.PDBWriter:
"""
Get the PDBWriter
"""
pass