Network

The aimmd.network subpackage trains and applies the committor model and provides output rescaling. The committor network itself is a user-supplied torch.nn.Module passed via Params.network; AIMMD supplies the training loop, a placeholder network, and rescaling utilities.

Training

Training utilities for AIMMD committor networks.

This module provides fit(), the routine that trains/updates the neural-network model stored in params.network (a torch.nn.Module) to predict the logit committor from AIMMD path-sampling data stored in a PathEnsemble.

In an AIMMD run, workers generate trajectories and accumulate them in a PathEnsemble. Periodically, AIMMD calls a training hook to improve the model that guides subsequent sampling. By default (when Params.fit is not set), AIMMD uses aimmd.network.fit.default(), which forwards directly to fit().

Core idea Training examples are assembled from multiple categories of paths/frames (internal state frames, free trajectories, and shooting trajectories). Each frame is assigned:

• a 2-component outcome vector r = (r_to_state1, r_to_state2) encoding fractional contributions towards reaching each end state, and

• a selection probability that controls how batches are drawn.

The training objective is a log-binomial loss (optionally modified near the end states by a Bayesian-like quadratic penalty), with optional smoothness and L1 regularization terms.

Side effects fit() modifies the network in-place. Training begins after calling params.network.reset_parameters() and proceeds with Adam. Depending on stop criteria and early stopping settings, the function may restore a previously saved state dict.

aimmd.network.fit.fit(params, pathensemble, nbins=0, cutoff_min=0.5, cutoff_max=20.0, state_bins='', max_adjustment_in_bin=10.0, transition_path_upweighting=1.0, end_state_factor=1.0, augment='no', lr=0.0002, loss_bayesian_factor=0, loss_smoothening_weight=0, loss_regularization_weight=0, loss_regularization_exponent=1, epochs=500, batch_size=4096, batching_strategy='draw-replace', stop=80.0, train_validation_early_stopping=False, early_stopping_patience=10, early_stopping_min_samples=1000, early_stopping_split=0.1, in_memory=True, graphs=False, lsr_weight=0.0, lsr_lagtime=1, lsr_batch_size=None, lsr_epsilon=1e-06, mar_weight=0.0, mar_lagtime=1, mar_batch_size=None, mar_epsilon=1e-08, verbose=False, worker=None, loss_log_path=None)

Train params.network to predict the logit committor from AIMMD data.

The function builds a supervised dataset from pathensemble and performs an iterative optimization of the network parameters using Adam. The target labels are fractional shooting outcomes (not just hard 0/1), constructed from shooting results and (optionally) augmented with additional trajectory data.

Important

The network is modified in-place. Training begins after calling network.reset_parameters(); depending on termination criteria, the function may restore earlier saved weights.

Parameters:

params – AIMMD parameter container. This function uses at least the following

attributes:

• params.network : torch.nn.Module

The network to train. Must have parameters so that next(network.parameters()) is valid.

• params.sorted_states : tuple[str, str, str]

State labels in the order (state1, reactant, state2) assigned to variables a, r, b in the implementation.

• params.descriptors_function : bool or callable-like

Used to decide whether descriptors are read from 'descriptors' or 'positions' in the path storage model.

• params.descriptor_transform : callable or None

Optional transformation applied to raw descriptors to obtain network inputs. If None, an identity transform is used. :type params: aimmd.Params

pathensemble : PathEnsemble Collection of sampled paths providing the following interface:

• pathensemble.types() returning an array-like where each entry

encodes the path type as 4 characters (initial, internal, final, shoot).

• Path access compatible with aimmd.network.utils.extract_indices_and_series(),

i.e. supports indexing and provides the necessary per-path interface (type, internal(‘indices’), indices, shooting_index, get(…)).

nbins : int, default=0 If > 0, define committor-space bins (via compute_bins()) and build selection probabilities such that batches are more uniformly distributed across bins. Also regularizes rare outcomes within bins.

cutoff_min : float, default=0.5 Passed to compute_bins() as cutoff_min (lower cutoff in committor space).

cutoff_max : float, default=20.0 Passed to compute_bins() as cutoff_max (upper cutoff in logit scale).

state_bins : str, default=’’ Controls inclusion of end-state bins. Examples:

• 'AB' includes both end-state bins (where A and B are the end

state labels in params.sorted_states).

• 'all' includes all state bins.

When state bins are excluded, the corresponding end-state frames may be set to zero selection probability depending on data presence.

max_adjustment_in_bin : float, default=10.0 Caps how strongly rare outcomes inside each bin are upweighted in selection probability (and correspondingly downweighted in results).

transition_path_upweighting : float, default=1.0 Multiplier applied to selection probabilities of frames belonging to transition paths (free and shooting transition segments).

end_state_factor : float, default=1.0 Base factor used for end-state (in-a, in-b) selection probability before bin-based adjustments.

augment : {‘no’, ‘yes’, ‘experimental’}, default=’no’ Data augmentation mode.

• 'no': train only on shooting-point outcomes (using the intersection

of backward and forward masks around the shooting point).

• 'yes': include free-trajectory frames and use backward/forward

segments to assign fractional outcomes.

• 'experimental': additional heuristic augmentation with fractional

transition contributions derived from free trajectories.

lr : float, default=1e-3 Base learning rate for Adam. The effective LR is ramped up over the first ~5% of epochs (see implementation).

loss_bayesian_factor : float, default=0 If non-zero, use a modified squared-deviation loss (Bayesian-like) in logit space to better handle discrete/imbalanced data near the states. If zero, use a standard binomial (cross-entropy-like) loss.

loss_smoothening_weight : float, default=0 If non-zero, add a smoothness penalty based on the gradient of the network output w.r.t. the input descriptors (requires d.requires_grad).

loss_regularization_weight : float, default=0 If non-zero, add Ln regularization over network parameters, where n is given by loss_regularization_exponent (default: L1).

loss_regularization_exponent : float, default=1 Exponent n for the Ln regularization term (see loss_regularization_weight).

epochs : int, default=500 Target number of epochs. The loop may extend up to 1.5× epochs while tracking best losses and may stop early due to stop scale or early stopping criteria.

batch_size : int, default=4096 Number of samples per batch. With draw-with-replacement batching, the batch size is fixed even for small datasets.

batching_strategy : {‘draw-replace’, ‘loop-all’}, default=’draw-replace’ Strategy for building batches.

• 'draw-replace' draws with replacement according to selection

probabilities.

• 'loop-all' is declared but not implemented (raises NotImplementedError).

stop : float, default=50.0 Stop training if the output scale (max absolute logit) reaches or exceeds this value, or if NaNs appear.

train_validation_early_stopping : bool, default=False If True, create a validation split and stop when validation loss does not improve for early_stopping_patience epochs (after scale/range conditions are met).

early_stopping_patience : int, default=10 Number of consecutive no-improvement validation checks before stopping.

early_stopping_min_samples : int, default=1000 Minimum training set size required to enable early stopping.

early_stopping_split : float, default=0.1 Fraction of the dataset used as validation set when early stopping is on.

in_memory : bool, default=True Descriptor loading strategy:

• If True: all raw descriptors are loaded into one concatenated

numpy array at the start of training, then descriptor_transform is applied to the whole set at once and the result is kept in memory. Fast per-epoch access but requires peak RAM proportional to the full training-set size (raw array + transformed representation).

• If False: no large descriptor array is kept in memory.

Instead, only per-frame file references (npy_path, loc) are stored. Each training batch (and each LSR batch) loads its raw frames on-the-fly from the NPY cache via _load_batch_descriptors(), then applies descriptor_transform immediately. Peak RAM is proportional to one batch rather than the full dataset, at the cost of repeated disk I/O and transform overhead each epoch.

graphs : bool, default=False If True, descriptors are assumed to be graph objects in an mlcolvar-like DataDict format, requiring torch_geometric for batching.

lsr_weight : float, default=0.0 If non-zero, add a latent space regularization (LSR) term to the loss, weighted by this factor. The LSR term is computed from time-lagged pairs drawn from all continuous (non-internal) trajectories in the path ensemble and acts on the network’s latent representation (via network.forward_latent if available, otherwise the scalar network output). The term encourages the latent features to capture slow dynamical modes of the system. A value of 0.0 (default) disables LSR entirely and incurs no overhead.

lsr_lagtime : int, default=1 Lag in frames used to form time-lagged pairs (x_t, x_{t+lagtime}) within each continuous trajectory for the LSR term. Larger values target slower dynamics but reduce the number of available pairs.

lsr_batch_size : int or None, default=None Number of time-lagged pairs sampled per LSR gradient step. If None, defaults to batch_size. Should be at least 2× the latent dimension for well-conditioned cross-covariance estimates.

lsr_epsilon : float, default=1e-6 Floor value for the per-sample L2 norm during row-normalisation of latent features inside the LSR loss. Prevents division by zero for frames near the network’s zero surface.

verbose : bool, default=False If True, show progress via tqdm and print more frequent diagnostics.

loss_log_path : str or None, default=None If not None, save a CSV with one row per epoch containing the columns epoch, total_loss, committor_loss, vamp_loss, and scale. The committor and LSR terms are each already multiplied by their respective weights so that total_loss ≈ committor_loss + vamp_loss (plus any regularisation terms). The file is written only after training completes.

worker : aimmd.Worker or None, optional If provided, training periodically checks for a termination signal via getattr(worker, 'termination_signal', False) and returns empty outputs if termination is requested.

Returns:

• losses (list[float]) – Training loss per optimizer step (epoch).

• scales (list[float]) – Output scale per step, defined as max(max(q), -min(q)) over the last

computed batch output q.

• values (numpy.ndarray) – 1D array of “values” for the training points (logit committor-like

coordinate assembled from trajectory values sources). Only includes frames kept after selection probability masking.

• selection_probabilities (numpy.ndarray) – Per-sample selection probabilities (normalized to sum to 1) used for

drawing training batches.

• results (numpy.ndarray, shape (N, 2)) – Per-sample fractional outcomes to each end state. These are adjusted

in bins for imbalance and may include augmented contributions.

aimmd.network.fit.default(params, pathensemble, verbose=True, worker=None)

Default training hook used when Params.fit is not set.

AIMMD keeps the full run configuration in Params. This includes the neural-network model to be trained and all settings needed to run the simulation (system, engines, sampling, I/O).

If Params does not specify a custom training callable, AIMMD calls this function. It is a thin wrapper that forwards its arguments to fit() unchanged, so the training uses all default hyperparameters of fit() (e.g. learning rate, batch size, number of epochs, binning and augmentation settings) unless a custom training function is provided.

Parameters:

• params (aimmd.Params) – Complete AIMMD configuration, including the network model.

• pathensemble – The trajectories generated so far; used by fit() to assemble the

training set. :type pathensemble: PathEnsemble :param verbose: Controls training log/progress output. If True (default), fit() may print progress information; if False, it should run quietly. :type verbose: bool, optional :param worker: Worker context for this run. If provided, fit() may use it for coordinating with other workers. If None (default), fit() runs without a worker context. :type worker: aimmd.Worker or None, optional

Returns:

Exactly the return value of fit().

Return type:

tuple

Notes

This is a thin wrapper around fit(), which trains or updates the AIMMD model from the current PathEnsemble.

Output rescaling

Rescaling mixin for network modules.

This module defines Rescalable, a mixin that augments a torch.nn.Module by applying an optional output rescaling step after the normal forward call.

The rescaling is represented by two 1D buffers:

• rescale_knots: knot locations in the original output coordinate,

• rescale_values: corresponding target values in the rescaled coordinate.

The actual transformation is performed by aimmd.network.rescale_utils.rescale() and is applied in-place to the network output under torch.no_grad().

Typical usage Subclass Rescalable together with a standard nn.Module implementation:

• define forward(…) as usual,

• call set_knots_and_values() to enable rescaling.

Rescaling is disabled by default: buffers are filled with NaNs and only entries up to the first NaN are considered active.

class aimmd.network.rescalable.Rescalable(max_knots=100)

Bases: ABC, Module

Mixin that applies output rescaling at the end of __call__.

This class is intended to be used with multiple inheritance alongside an nn.Module that implements forward(…). By overriding __call__(), it ensures:

1. the usual PyTorch module call path runs first (hooks, autocast, etc.),

2. the resulting output q is optionally rescaled based on stored knots and values.

Rescaling parameters are stored as buffers so they are:

• moved with the module across devices (.to(…)),

• included in module state (unless explicitly filtered),

• not treated as trainable parameters.

Notes

• The rescaling is applied under torch.no_grad() and mutates the output tensor in-place (as implemented by rescale()).

• Rescaling is considered active when rescale_knots contains at least one non-NaN entry.

Parameters:

max_knots – Maximum number of knots/values stored in the internal buffers.

Two buffers of shape (max_knots,) are registered: rescale_knots and rescale_values. Entries are initialized to NaN and only non-NaN entries are used. :type max_knots: int, optional

__init__(max_knots=100)

Parameters:

max_knots – Maximum number of knots/values stored in the internal buffers.

Two buffers of shape (max_knots,) are registered: rescale_knots and rescale_values. Entries are initialized to NaN and only non-NaN entries are used. :type max_knots: int, optional

set_knots_and_values(knots, values)

Set the active rescaling definition.

Parameters:

• knots (Iterable[float]) – Knot positions in the original coordinate. Must be iterable.

• values – Target values at each knot. Must be iterable and have the same

length as knots. :type values: Iterable[float]

Raises:

TypeError – If knots or values are not iterable, or if they do not have the same length.

Notes

• Values are copied into the buffers starting at index 0.

• Only the first len(knots) entries are overwritten; remaining buffer entries are left unchanged. If you need to clear previous knots, call reset_parameters() first (or set fewer knots after a reset).

reset_parameters()

Reset module parameters and clear rescaling buffers.

This cooperates with multiple inheritance by calling reset_parameters on the next class in the MRO if it exists, then clears the rescaling buffers by filling them with NaNs.

Notes

• Clearing the buffers disables rescaling until new knots/values are set.

Rescaling utilities for committor-like coordinates.

This module contains small numerical helpers used to re-map the logit-committor q in order to improve sampling uniformity in the reactive region. It does so by making the rescaled crossing probabilities from both states as close as possible to the theoretical P1->2 = 1 / expit(+q), P2-1 = 1 / expit(-q).

The two main operations are:

• find_knots_and_values(): infer a piecewise-linear remapping in logit space from per-side “extreme” statistics and cumulative crossing-probability measures.

• rescale(): apply that remapping to an array/tensor of q values (performed in-place).

Conventions

• q, extremes1, extremes2, knots, and values are in logit-committor units (i.e. real numbers where expit(q) maps to (0, 1)).

• Masks and interpolation are performed in logit space, but interpolation is carried out on log(expit(…)) to keep monotonic behavior in probability space.

Notes

• This module prints informative messages via aimmd._config.print.

• The logic is heuristic and intentionally conservative: several early returns produce empty knot/value sets when rescaling is not applicable.

aimmd.network.rescale_utils.find_knots_and_values(extremes1, extremes2, xP1, xP2)

Infer rescaling knots and target values in logit space.

This routine constructs a nonlinear coordinate transformation represented as a piecewise-linear map in logit space:

• knots: x-positions (in the original logit coordinate) where slope changes,

• values: y-positions (in the rescaled coordinate) at the corresponding knots.

The transformation is later applied by rescale().

This transformation will make crossing probabilities from both states as close as possible to the theoretical P1->2 = 1 / expit(+q), P2-1 = 1 / expit(-q). In this way, the shooting point selection will be become more uniform in the reactive region, hopefully improving the sampling.

Inputs represent statistics from two “sides” (1 and 2) of an ensemble. Conceptually, the function:

1. normalizes the provided crossing-probability series (xP1, xP2) into codomain curves (N1, N2) in probability space;

2. determines the domain of action (kmin, kmax) from finite extremes;

3. interpolates both sides onto a fine grid q in logit space (101 points);

4. estimates a transition-state shift ts and an overall rescaling factor r;

5. generates a set of non-unique candidate knot locations and associated target values by matching geometric levels between the “actual” and “theoretical” curves;

6. removes redundant / non-monotone knots.

Parameters:

extremes1 – Logit-coordinate “extreme” values for side 1. May contain math.inf values

(treated as state points and removed before the action domain is computed). Converted with np.asarray(extremes1). :type extremes1: array-like of float :param extremes2: Logit-coordinate “extreme” values for side 2. Same conventions as extremes1. Converted with np.asarray(extremes2). :type extremes2: array-like of float :param xP1: Cumulative or aggregated crossing-probability measure for side 1. Must be indexable; if empty, the function falls back to a trivial normalization for that side. Converted with np.asarray(xP1). :type xP1: array-like of float :param xP2: Crossing-probability measure for side 2. Same conventions as xP1. Converted with np.asarray(xP2). :type xP2: array-like of float

Returns:

• knots (numpy.ndarray, dtype=float) – Knot positions (x-coordinates) in the original logit coordinate.

May be empty if no rescaling is applicable.

• values (numpy.ndarray, dtype=float) – Target positions (y-coordinates) in the rescaled logit coordinate,

aligned with knots. Same length as knots.

Notes

• Several “nothing to do” checks return empty arrays:

  • codomain interval invalid (vmin >= vmax),

  • action domain invalid (kmin >= kmax).

• np.isinf(extremes*) entries are removed before computing the domain of action and interpolation.

• The printed diagnostics include the estimated transition-state shift ts and total rescaling factor r.

aimmd.network.rescale_utils.rescale(q, knots, values)

Apply a piecewise-linear rescaling to q (in-place).

The rescaling is defined by matching knots (x-coordinates in the original space) to values (y-coordinates in the rescaled space). Between knots the mapping is linear, and the end segments extrapolate using the slope of the first/last interior segment.

This function mutates q in-place and returns it for convenience.

Parameters:

q – 1D (or broadcastable) array/tensor of values in logit space to be

transformed. If q is a torch.Tensor, bucketization uses torch.bucketize; otherwise NumPy digitization is used. :type q: numpy.ndarray or torch.Tensor :param knots: Knot positions (x-coordinates) in the original space. Typically produced by find_knots_and_values(). :type knots: array-like of float :param values: Target positions (y-coordinates) corresponding to knots. Must have the same length as knots. :type values: array-like of float

Returns:

q – The same object as the input q, modified in-place.

Return type:

numpy.ndarray or torch.Tensor

Notes

• If len(knots) < 1, the function returns q unchanged.

• If len(knots) == 1, the mapping is a pure shift about knots[0]: q[:] = (q - knots[0]) + values[0].

• For len(knots) >= 2, each bucket gets its own affine transform a * (q - x0) + b computed from adjacent knot/value pairs.

Utilities

Utility helpers for AIMMD network components.

This module provides:

• PlaceholderNetwork: a minimal torch.nn.Module used as a safe default when no trained network has been configured.

• placeholder: a module-level instance of PlaceholderNetwork with an additional __source__ attribute used for provenance/logging.

• extract_indices_and_series(): flatten per-path frame indices (and optional per-frame series) into contiguous NumPy arrays, while constructing boolean masks that label the backward and forward portions of each path relative to its shooting point.

Design constraints

• The placeholder network is intentionally stateless: it serializes to an empty state dict and ignores state loading.

• extract_indices_and_series is written against a Path-like interface and only relies on the attributes/methods it explicitly touches.

class aimmd.network.utils.PlaceholderNetwork

Bases: Module

Stateless default network used as a placeholder.

This class exists so that components expecting a torch.nn.Module can run even when the user has not provided a real model.

Behavior

• Forward pass returns only the first channel/feature: x[:, :1].

• The model is stateless with respect to serialization: state_dict() returns {} and load_state_dict() does nothing.

Notes

• A dummy parameter is registered to ensure the module has a parameter tensor that anchors device/dtype when code expects parameters to exist (e.g. for .to(device) semantics).

forward(x)

Compute placeholder output.

Parameters:

x – Input tensor. Expected shape is typically (n_samples, n_features)

so that slicing x[:, :1] yields a tensor of shape (n_samples, 1). The method does not validate rank explicitly; it relies on PyTorch slicing semantics. :type x: torch.Tensor

Returns:

y – A view of the first column of x, i.e. x[:, :1]. For a 2D input

this has shape (n_samples, 1) and shares storage with x. :rtype: torch.Tensor

state_dict(*args, **kwargs)

Return an empty state dictionary.

Parameters:

*args – Accepted for API compatibility with torch.nn.Module.state_dict

but ignored. :param **kwargs: Accepted for API compatibility with torch.nn.Module.state_dict but ignored.

Returns:

state – Always an empty dict {}.

Return type:

dict

load_state_dict(state_dict, strict=True)

Load a state dictionary (no-op).

Parameters:

• state_dict (Mapping[str, Any]) – Ignored.

• strict (bool, optional) – Ignored. Present for compatibility with PyTorch’s API.

Return type:

None

aimmd.network.utils.extract_indices_and_series(paths, key, *names)

Extract and concatenate indices and per-frame series across paths.

The function iterates over a selected subset of paths and produces flat arrays suitable for downstream vectorized analysis/plotting. For each path, it:

1. obtains the internal frame indices via path.internal('indices');

2. builds two boolean masks (backward/forward) in the local index space: - backward: frames from path start through the shooting frame (inclusive); - forward: frames from the shooting frame through the end (inclusive),

but only if the path actually transitions (see below);

3. loads each requested series name via path.get over the global index window defined by the first/last internal index;

4. appends indices, masks, and series to global outputs.

A path is skipped entirely if loading any requested series fails or if a loaded series does not match the number of internal indices.

Path transition rule Forward frames are only marked if path.type[2] != path.type[1]. This is interpreted as “the path ends in a different state than it started” (i.e. it is “going somewhere”); otherwise the forward mask for that path remains all False.

Parameters:

paths – Sequence of Path-like objects.

Each path is expected to provide the following attributes/methods (as used by the current implementation):

• path.type :

Indexable object where elements 1 and 2 encode start/end states.

• path.internal(‘indices’) -> array-like[int] :

Internal frame indices used as the reference axis for series alignment.

• path.shooting_index : int

Global frame index of the shooting point.

• path.indices : array-like[int]

Path-global indices; path.indices[0] is used as the first global index.

• path.get(name, start=int, stop=int, raise_if_missing=bool) :

Returns a per-frame series aligned to the global indices in [start, stop). The return value must have length equal to len(path.internal('indices')). :type paths: Sequence[Path-like] :param key: Path selector.

• If None, all paths are processed in their natural order (equivalent

to selecting range(len(paths))).

• Otherwise, key is used to index np.arange(len(paths)) and then

flattened. This supports boolean masks, integer index arrays, slices, etc., as long as NumPy advanced indexing accepts it. :type key: None or array-like :param *names: Names of additional per-frame series to extract for each path. For each name, the function calls:

path.get(name, start=path_indices[0], stop=path_indices[-1] + 1, raise_if_missing=True)

and requires that the returned object has length equal to the number of internal indices. :type *names: str

Returns:

• indices (numpy.ndarray, dtype=int) – Concatenated internal indices from all successfully processed paths.

• back_mask (numpy.ndarray, dtype=bool) – Boolean mask aligned with indices. True for frames belonging to the

backward segment (start through shooting, inclusive).

• forw_mask (numpy.ndarray, dtype=bool) – Boolean mask aligned with indices. True for frames belonging to the

forward segment (shooting through end) only when the path transitions according to path.type[2] != path.type[1]; otherwise False.

• *series (numpy.ndarray) – For each name in names, a NumPy array concatenating the corresponding

per-frame series across successfully processed paths, aligned with indices. The dtype is inferred by np.array(…) from the appended values.

• n_selected (int) – Number of selected paths after applying key (i.e. len(key) after

normalization). This counts selected candidates, not how many were successfully processed (some may be skipped due to I/O/validation errors).

Notes

• The function is intentionally tolerant to missing/failed series loads: it simply skips problematic paths.

aimmd.network.utils.extract_lsr_pairs(paths, key, lagtime, name)

Extract time-lagged frame pairs from continuous paths for latent space regularization.

For each successfully-loaded path with n frames, this function produces n - lagtime pairs (frame_t, frame_{t+lagtime}), concatenated across all selected paths. These pairs are used by the latent space regularization (LSR) loss in aimmd.network.fit.fit(). The skipping logic (skip on I/O failure) mirrors extract_indices_and_series() exactly: only paths for which name can be fully loaded are included, so pair indices are consistent with the descriptor arrays returned by that function when called with the same key.

Parameters:

• paths (PathEnsemble) – The path ensemble (same interface as for extract_indices_and_series).

• key (array-like or None) – Path selector (same semantics as extract_indices_and_series).

• lagtime – Number of frames to shift between the two sides of each pair. Must

be >= 1. Paths with fewer than lagtime + 1 frames contribute no pairs. :type lagtime: int :param name: Name of the per-frame series to extract (e.g. 'descriptors' or 'coordinates'). Passed to path.get. :type name: str

Returns:

• data_t (numpy.ndarray) – Concatenated frames at time t, shape (n_pairs, *frame_shape).

• data_tau (numpy.ndarray) – Concatenated frames at time t+lagtime, same shape as data_t.

• n_selected (int) – Number of paths selected by key (including skipped ones).

• n_pairs (int) – Total number of valid pairs extracted.

aimmd.network.utils.extract_mar_sequences(paths, key, lagtime, name)

Extract ordered frame sequences from reactive transition paths for MAR.

Unlike extract_lsr_pairs() (which produces flat concatenated pairs), this function preserves path identity: the return value is a list of per-path arrays, one per successfully loaded reactive path. Each array contains every lagtime-th frame (indices 0, lagtime, 2*lagtime, …) of the internal path frames. Paths with fewer than 3 subsampled frames are silently skipped.

The caller is responsible for selecting only reactive path keys (e.g. the union of free1to2, free2to1, shot1to2, shot2to1 type conditions) before passing key.

Parameters:

• paths (PathEnsemble) – The path ensemble (same interface as for extract_lsr_pairs).

• key (array-like or None) – Path selector (same semantics as extract_lsr_pairs).

• lagtime – Subsampling stride. A value of 1 uses every frame; a value of k uses

every k-th frame (indices 0, k, 2k, …). Must be >= 1. :type lagtime: int :param name: Name of the per-frame series to extract (e.g. 'descriptors', 'coordinates', 'filenames', 'locs'). Passed to path.get. :type name: str

Returns:

• sequences (list of numpy.ndarray) – One array per successfully loaded path, shape

(n_subsampled_frames_i, *frame_shape).

• n_selected (int) – Number of paths selected by key (including skipped ones).

• n_sequences (int) – Number of paths that produced a valid sequence.

Graph-neural-network support

The module aimmd.network.graph_utils provides optional support for graph-neural-network committor models (atom-coordinate descriptors, PyG graph construction and SQLite graph caching, and a shared atom_types one-hot encoding for multi-system runs). It requires the optional graphs extras (torch-geometric, torch-cluster) plus mlcolvar and lz4; install them with pip install "aimmd-lab[graphs]" and see Advanced Usage.