Shooting Point Selectors

Shooting Point Selectors#

class aimmd.distributed.spselectors.UniformSPSelector(exclude_frames: int = 0)#

Select shooting points uniformly from the given in-trajectory.

Initialize a UniformSPSelector.

Parameters:

exclude_frames (int, optional) – How many frames to exclude from the selection at the start and end of the trajectory, e.g. if exclude_frames=2 we exclude the first and last 2 frames, by default 0

async pick(outfile: str, frame_extractor: FrameExtractor, simstate_info: PathSamplingSimStateInfo, *, model: RCModelAsyncMixin, trajectory: Trajectory | None = None) → Trajectory#

Pick and return a random snapshot from trajectory.

Write out the picked snapshot using frame_extractor as outfile.

Parameters:

  • outfile (str) – Absolute or relative path to write the picked snapshot to.

  • frame_extractor (FrameExtractor) – asyncmd.FrameExtractor subclass used to extract the chosen frame from trajectory.

  • simstate_info (PathSamplingSimStateInfo) – Dataclass carrying the current state of the pathsampling simulation and information about the current step, e.g. the step_num, workdir.

  • model (RCModelAsyncMixin) – Completely ignored in uniform selection.

  • trajectory (Trajectory) – The trajectory to pick the snapshot from, required although by default None.

Returns:

The chosen snapshot.

Return type:

Trajectory

Raises:

  • ValueError – If no input trajectory is given to pick a snapshot from.

  • ValueError – If the input trajectory is to short to pick a snapshot from after discarding the first and last exclude_frames frames.

async probability(snapshot: Trajectory, trajectory: Trajectory, simstate_info: PathSamplingSimStateInfo, *, model: RCModelAsyncMixin) → float#

Return proposal probability to pick (any) snapshot from trajectory.

For the UniformSPSelector this just returns 1 / (len(trajectory) - 2 * exclude_frames), because every snapshot is equally likely.

Parameters:

  • snapshot (Trajectory) – The snapshot in question.

  • trajectory (Trajectory) – The trajectory to pick snapshot from.

  • simstate_info (PathSamplingSimStateInfo) – Dataclass carrying the current state of the pathsampling simulation and information about the current step, e.g. the step_num, workdir.

  • model (RCModelAsyncMixin) – Completely ignored in uniform selection.

Returns:

The probability to pick any snapshot from trajectory.

Return type:

float

class aimmd.distributed.dataclasses.DensityAdaptionParameters(n_bins: int = 10, scheme: Literal['lazzeri', 'p_x_tp'] | None=None, reevaluate_density_interval: int | None = None, reset_before_pick: bool = True, add_trajectories_from_sampler: bool = True, trajectories_to_flatten: list[Trajectory] = <factory>, trajectories_to_flatten_weights: list[npt.NDArray[np.floating]] | None = None)#

Dataclass to specify parameters for density adaption in RCModelSPSelector.

This dataclass includes predefined schemes that will ensure consistency of the arguments used for this scheme. Currently these are “lazzeri” and “p_x_tp”, the former flattens the committor observed on the input transition path while the latter flattens the committor on the ensemble of transition paths observed so far.

Parameters:

  • n_bins (int) – The number of bins to use (in each dimension of the model prediction) in the histogram used for density estimation.

  • scheme (Literal["lazzeri", "p_x_tp"] | None) – A string indicating one of the aforementioned predefined schemes or None. None means no parameter consistency checking will be performed.

  • reevaluate_density_interval (int | None) – The interval (in steps done by the sampler) in which to reevaluate the density estimate using the current model. None means do not reevaluate ever, it should only be used if your model prediction does not change. It can be ignored for schemes that reset before every pick, e.g. “lazzeri”, since the current model is always used to add the trajectory.

  • reset_before_pick (bool) – Whether the density estimate should be cleared before every pick and before (potentially) adding the trajectory from this pick. Setting this to True and adding only the trajectory from each pick one arrives at the “lazzeri” scheme. By default True.

  • add_trajectories_from_sampler (bool) – Whether to add the trajectories produced in the Markov Chain to the density estimate. By default True.

  • trajectories_to_flatten (list[Trajectory]) – An (initial) list of trajectories to flatten the density from. When selecting configurations from a predefined reservoir, these trajectories should be added here. Could also be used to kickstart density adaption with an initial guess. By default an empty list is used.

  • trajectories_to_flatten_weights (list[np.ndarray] | None) – List of numpy arrays with weights for the trajectories_to_flatten, the list must contain as many numpy arrays as there are trajectories and the lengths of the arrays must match the corresponding trajectories. Can also be None, in that case equal weights for each configuration will be used (in case trajectories_to_flatten is not empty). By default None.

class aimmd.distributed.spselectors.RCModelSPSelectorFromTraj(scale: float = 1.0, distribution: str = 'lorentzian', *, density_adaptation_params: DensityAdaptionParameters | None = DensityAdaptionParameters(n_bins=10, scheme='p_x_TP', reevaluate_density_interval=10, reset_before_pick=False, add_trajectories_from_sampler=True, trajectories_to_flatten=[], trajectories_to_flatten_weights=None), exclude_frames: int = 1, f_sel: Callable | None = None)#

Select SPs from a given in-trajectory (usually a TP) using a RCModel.

SP selection is biased towards the current best guess of the transition state of the RCModel according to the models z_sel method. The SPs are selected according to bias weights according to distribution f_sel(z_sel) centered around z_sel = 0.

Initialize a RCModelSPSelectorFromTraj.

Parameters:

  • scale (float, optional) – Scale of the distribution (default=1), smaller values result in a more peaked selection of shooting points around the predicted transition state. For the Lorentzian distribution scale = gamma, while for the Gaussian distribution scale = 2 * sigma**2.

  • distribution (str) – A string indicating the distribution to use when selecting SPs around the transition state, can be “lorentzian”, “gaussian”, “uniform_phi”, “uniform”, or “custom” (see also f_sel).

  • density_adaptation_params (DensityAdaptionParameters or None) – Parameters for density adaption, which includes an additional correction factor to flatten the density of potential SP configurations along the predicted committor. See also the DensityAdaptionParameters for more. If set to None, no density adaption will be performed. By default density adaption according to the density of points on transition paths, p(x|TP), will be performed.

  • exclude_frames (int, optional) – How many frames to exclude from the selection at the start and end of the trajectory, e.g. if exclude_frames=2 we exclude the first and last 2 frames, by default 1.

  • f_sel (Callable or None, optional) – If given, sets the shooting point selection distribution function to any Callable taking a single np.array of shape (n_frames,) and returning an array of the same shape containing bias weights. Note that, in this case distribution will be ignored and set to “custom”.

async biases(trajectory: Trajectory, model: RCModelAsyncMixin, exclude_frames: int = 0) → ndarray[Any, dtype[floating]]#

Return array with bias values for each configuration in trajectory.

Note that the returned array is shorter than trajectory if exclude_frames > 0, because then the first and last exclude_frames frames will be discarded.

Parameters:

  • trajectory (Trajectory) – The trajectory to calculate bias values for.

  • model (RCModelAsyncMixin) – The reaction coordinate model from which the biases are calculated.

  • exclude_frames (int) – How many frames to exclude at the start and end of the trajectory, e.g., if exclude_frames=2 we exclude the first and last 2 frames, by default 0.

Returns:

One dimensional array (of the same length as trajectory - 2 * exclude_frames) with the bias values.

Return type:

npt.NDArray[np.floating]

async pick(outfile: str, frame_extractor: FrameExtractor, simstate_info: PathSamplingSimStateInfo, *, model: RCModelAsyncMixin, trajectory: Trajectory | None = None) → Trajectory#

Pick and return a snapshot to shot from.

async pick_snapshot(outfile: str, frame_extractor: FrameExtractor, simstate_info: PathSamplingSimStateInfo, *, model: RCModelAsyncMixin, trajectory: Trajectory | None = None) → Trajectory#

Pick and return a shooting snapshot from trajectory.

The selected snapshot is biased towards model.z_sel = 0 using the selection distribution. Write out the picked snapshot using frame_extractor as outfile.

Parameters:

  • outfile (str) – Path to write out the selected snapshot.

  • frame_extractor (FrameExtractor) – The frame extractor to use when getting the snapshot from the trajectory, e.g. RandomVelocities for TwoWayShooting.

  • simstate_info (PathSamplingSimStateInfo) – Dataclass carrying the current state of the pathsampling simulation and information about the current step, e.g. the step_num, workdir.

  • model (RCModelAsyncMixin) – The RCModel that is used to bias the selection towards the current best guess of the transition state.

  • trajectory (Trajectory) – The input trajectory from which we select the snapshot.

Returns:

The selected shooting snapshot.

Return type:

Trajectory

Raises:

ValueError – If the input trajectory is to short to pick a snapshot from after discarding the first and last exclude_frames frames.

async probability(snapshot: Trajectory, trajectory: Trajectory, simstate_info: PathSamplingSimStateInfo, *, model: RCModelAsyncMixin) → float#

Return proposal probability to pick snapshot from trajectory.

For the RCModelSPSelectorFromTraj the proposal probability is equal to f(snap) / sum_bias(traj), where f(snap) = f_sel(z_sel_snap) * density_adapt and sum_bias(traj) is the sum over f(snap) for all snapshots in the trajectory, i.e. the normalizing sum.

Parameters:

  • snapshot (Trajectory) – The snapshot in question.

  • trajectory (Trajectory) – The trajectory to pick snapshot from.

  • simstate_info (PathSamplingSimStateInfo) – Dataclass carrying the current state of the pathsampling simulation and information about the current step, e.g. the step_num, workdir.

  • model (RCModelAsyncMixin) – The reaction coordinate model used to pick the snapshot.

Returns:

Proposal probability to pick snapshot from trajectory.

Return type:

float

async sum_bias(trajectory: Trajectory, model: RCModelAsyncMixin, exclude_frames: int = 0) → floating#

Return the sum of all selection biases for trajectory under model.

Parameters:

  • trajectory (Trajectory) – The trajectory over which the bias values are summed.

  • model (RCModelAsyncMixin) – The reaction coordinate model used to calculate the biases.

  • exclude_frames (int) – How many frames to exclude at the start and end of the trajectory, e.g., if exclude_frames=2 we exclude the first and last 2 frames, by default 0.

Returns:

The sum of all bias values on trajectory under model.

Return type:

float

property distribution: str#

Return/set the shooting point selection distribution f_sel(z).

Can be either “lorentzian”, “gaussian”, “uniform_phi”, or “uniform”. For “lorentzian” f_sel(z) = scale / (scale**2 + z**2). For “gaussian” f_sel(z) = exp(-z**2 / scale). For “uniform_phi” f_sel(z) = exp(-z) / (1 + exp(-z))**2 , resulting in a uniform selection weight along the committor phi = 1 / (1 + exp(-z)). For “uniform” f_sel(z) = 1, resulting in a selection that favours points close to the states, f_sel(phi) = 1 / (phi - phi**2).

property exclude_frames: int#

Return/set the number of excluded frames at the beginning and end.

property f_sel: Callable#

Return/set the shooting point selection distribution function directly.

The function must take a single argument z (a np.array with shape=(n_frames,)) and return a np.array of the same shape as z containing the (unnormalized) bias weights.

property scale: float#

Return/set scale of the shooting point selection distribution.

class aimmd.distributed.spselectors.RCModelSPSelectorFromEQ(scale: float = 1.0, distribution: str = 'lorentzian', *, trajectories: list[Trajectory], equilibrium_weights: list[ndarray[Any, dtype[floating]]], density_adaption_params: DensityAdaptionParameters | None = DensityAdaptionParameters(n_bins=10, scheme=None, reevaluate_density_interval=10, reset_before_pick=False, add_trajectories_from_sampler=False, trajectories_to_flatten=[], trajectories_to_flatten_weights=None), f_sel: Callable | None = None)#

Select SPs from an equilibrium ensemble biased using a RCModel.

The equilibrium ensemble of shooting points is supplied as a list of trajectories with associated equilibrium weights.

SP selection is biased towards the current best guess of the transition state of the RCModel according to the models z_sel method. The SPs are selected according to their equilibrium weight and biased according to distribution f_sel(z_sel) centered around z_sel = 0.

Note that probability() does not return the probability to pick a snapshot on a trajectory, but instead the (unnormalized) ensemble weight of the given trajectory. This fact is also indicated by the attribute self.probability_is_ensemble_weight being True.

Initialize a RCModelSPSelectorFromEQ.

Parameters:

  • trajectories (list[Trajectory]) – A list with trajectories containing the potential shooting points, i.e. the ensemble of configurations. See also the notes below.

  • equilibrium_weights (list[npt.NDArray[np.floating]]) – A list of np.arrays with the equilibrium weights associated with each configuration in trajectories. See also the notes below.

  • scale (float, optional) – Scale of the SP selection distribution, by default 1.

  • distribution (str) – A string indicating the distribution to use when selecting SPs around the transition state, can be “lorentzian”, “gaussian”, “uniform_phi”, “uniform”, or “custom” (see also f_sel).

  • density_adaptation_params (DensityAdaptionParameters or None) – Parameters for density adaption, which includes an additional correction factor to flatten the density of potential SP configurations along the predicted committor. See also the DensityAdaptionParameters for more. If set to None, no density adaption will be performed. Note that, trajectories and equilibrium_weights will both be set in the density_adaption_parameters as trajectories and weights. Additionally it is ensured that no new trajectories are added during pick. I.e., density adaption will always flatten the committor distribution observed in the potential shooting points.

  • f_sel (Callable or None, optional) – If given, sets the shooting point selection distribution function to any Callable taking a single np.array of shape (n_frames,) and returning an array of the same shape containing bias weights. Note that, in this case distribution will be ignored and set to “custom”.

Raises:

ValueError – If the number of configurations in trajectories does not match the number of equilibrium weights supplied.

Notes

trajectories can either be a list of Trajectory objects or a single Trajectory object, similarly equilibrium_weights can be either a list of numpy arrays or a single numpy array. equilibrium_weights is expected to contain the multiplicative factors by which the frequency of observation for each configuration/frame in trajectories differs from the equilibrium ensemble. The factors do not need to be normalized (as long as they share a common partition function). Naturally it is also possible to pass the normalized equilibrium probability for each configuration in trajectories as equilibrium_weights. Depending on the origin of the configurations in trajectories, one could e.g. use for equilibrium_weights:

  • all ones if the trajectories are equilibrium trajectories.

  • exp(beta V_{bias}(x)) for each structure from a biased sampling run if all configurations where generated using the same V_{bias}.

  • the equilibrium weights from binless WHAM if the configurations are from multiple biased sampling runs using different biasing potentials V_{bias}.

async biases(trajectory: Trajectory, model: RCModelAsyncMixin, exclude_frames: int = 0) → ndarray[Any, dtype[floating]]#

Return array with bias values for each configuration in trajectory.

Note that the returned array is shorter than trajectory if exclude_frames > 0, because then the first and last exclude_frames frames will be discarded.

Parameters:

  • trajectory (Trajectory) – The trajectory to calculate bias values for.

  • model (RCModelAsyncMixin) – The reaction coordinate model from which the biases are calculated.

  • exclude_frames (int) – How many frames to exclude at the start and end of the trajectory, e.g., if exclude_frames=2 we exclude the first and last 2 frames, by default 0.

Returns:

One dimensional array (of the same length as trajectory - 2 * exclude_frames) with the bias values.

Return type:

npt.NDArray[np.floating]

async pick(outfile: str, frame_extractor: FrameExtractor, simstate_info: PathSamplingSimStateInfo, *, model: RCModelAsyncMixin, trajectory: Trajectory | None = None) → Trajectory#

Pick and return a snapshot to shot from.

async pick_snapshot(outfile: str, frame_extractor: FrameExtractor, simstate_info: PathSamplingSimStateInfo, *, model: RCModelAsyncMixin, trajectory: Trajectory | None = None) → Trajectory#

Pick and return a snapshot from the ensemble of potential snapshots.

The selected snapshot is biased towards model.z_sel = 0 using the selection distribution. Write out the picked snapshot using frame_extractor as outfile.

Parameters:

  • outfile (str) – Path to write out the selected snapshot.

  • frame_extractor (FrameExtractor) – The frame extractor to use when getting the snapshot from the trajectory, e.g. RandomVelocities for TwoWayShooting.

  • simstate_info (PathSamplingSimStateInfo) – Dataclass carrying the current state of the pathsampling simulation and information about the current step, e.g. the step_num, workdir.

  • model (RCModelAsyncMixin) – The RCModel that is used to bias the selection towards the current best guess of the transition state.

  • trajectory (Optional[Trajectory], optional) – Ignored, by default None. No trajectory is needed to pick a snapshot from. The argument is only retained to keep API compatibility with the other SPSelectors.

Returns:

The selected shooting snapshot.

Return type:

Trajectory

async probability(snapshot: Trajectory, trajectory: Trajectory, simstate_info: PathSamplingSimStateInfo, *, model: RCModelAsyncMixin) → float#

Return the equilibrium ensemble weight for trajectory.

Note that, this method does not return the probability to pick snapshot from trajectory (as for SPSelectors that pick the shooting snapshot from an input trajectory). It still needs to be called with a value for snapshot due to API consistency reasons.

Parameters:

  • snapshot (Trajectory) – Ignored, not needed to calculate ensemble weight of trajectory.

  • trajectory (Trajectory) – The trajectory to calculate the ensemble weight for.

  • simstate_info (PathSamplingSimStateInfo) – Dataclass carrying the current state of the pathsampling simulation and information about the current step, e.g. the step_num, workdir.

  • model (RCModelAsyncMixin) – The reaction coordinate model used to pick the snapshot.

Returns:

Ensemble weight for trajectory.

Return type:

float

async sum_bias(trajectory: Trajectory, model: RCModelAsyncMixin, exclude_frames: int = 0) → floating#

Return the sum of all selection biases for trajectory under model.

Parameters:

  • trajectory (Trajectory) – The trajectory over which the bias values are summed.

  • model (RCModelAsyncMixin) – The reaction coordinate model used to calculate the biases.

  • exclude_frames (int) – How many frames to exclude at the start and end of the trajectory, e.g., if exclude_frames=2 we exclude the first and last 2 frames, by default 0.

Returns:

The sum of all bias values on trajectory under model.

Return type:

float

property distribution: str#

Return/set the shooting point selection distribution f_sel(z).

Can be either “lorentzian”, “gaussian”, “uniform_phi”, or “uniform”. For “lorentzian” f_sel(z) = scale / (scale**2 + z**2). For “gaussian” f_sel(z) = exp(-z**2 / scale). For “uniform_phi” f_sel(z) = exp(-z) / (1 + exp(-z))**2 , resulting in a uniform selection weight along the committor phi = 1 / (1 + exp(-z)). For “uniform” f_sel(z) = 1, resulting in a selection that favours points close to the states, f_sel(phi) = 1 / (phi - phi**2).

property f_sel: Callable#

Return/set the shooting point selection distribution function directly.

The function must take a single argument z (a np.array with shape=(n_frames,)) and return a np.array of the same shape as z containing the (unnormalized) bias weights.

property scale: float#

Return/set scale of the shooting point selection distribution.