Policies

base

Adaptive sampling policy base class and registry.

POLICY_REGISTRY module-attribute

POLICY_REGISTRY: Dict[str, Type['Policy']] = {}

Policy

Bases: ABC

Base class for adaptive sampling policies.

Cluster-based policies implement :meth:select_clusters and set requires_clustering = True. Frame-level policies set requires_clustering = False and implement :meth:select_frames.

select_clusters abstractmethod

select_clusters(cluster_stats: ClusterStats, n_seeds: int) -> List[int]

Select cluster IDs from which to draw seed frames.

Parameters:

Name	Type	Description	Default
cluster_stats	dict	Per-cluster population and frame lists.	required
n_seeds	int	Maximum number of clusters (seeds) to select.	required

Returns:

Type	Description
list of int	Selected cluster IDs.

select_frames

select_frames(dataset: Dataset, n_seeds: int) -> List[SeedResult]

Select seed frames directly from loaded features.

Frame-level policies override this method. Cluster-based policies raise :exc:NotImplementedError.

Parameters:

Name	Type	Description	Default
dataset	Dataset	Loaded feature dataset.	required
n_seeds	int	Number of seed frames to select.	required

Returns:

Type	Description
list of SeedResult	Selected seed frames.

register_policy

register_policy(cls: Type['Policy']) -> Type['Policy']

Register a policy class in :data:POLICY_REGISTRY.

Parameters:

Name	Type	Description	Default
cls	type	Policy subclass with a name class attribute.	required

Returns:

Type	Description
type	The registered policy class (unchanged).

Raises:

Type	Description
ValueError	If the policy name is missing or already registered.

get_policy

get_policy(name: str, **kwargs) -> Policy

Instantiate a registered policy by name.

Parameters:

Name	Type	Description	Default
name	str	Registered policy name.	required
**kwargs		Constructor arguments forwarded to the policy class.	{}

Returns:

Type	Description
Policy	Policy instance.

Raises:

Type	Description
ValueError	If the policy name is unknown.

list_policies

list_policies() -> List[str]

Return names of all registered policies.

Returns:

Type	Description
list of str	Sorted policy names.

least_counts

Least-counts adaptive sampling policy.

LeastCountsPolicy

Bases: Policy

Select clusters with the smallest populations.

Clusters are sorted by ascending population and the first n_seeds cluster IDs are returned (one seed per cluster).

select_clusters

select_clusters(cluster_stats: ClusterStats, n_seeds: int) -> List[int]

Select the least-populated clusters.

Parameters:

Name	Type	Description	Default
cluster_stats	dict	Per-cluster statistics.	required
n_seeds	int	Number of clusters to select.	required

Returns:

Type	Description
list of int	Cluster IDs with smallest populations.

random

Random cluster selection policy.

RandomPolicy

RandomPolicy(random_state: Optional[int] = None)

Bases: Policy

Uniformly sample cluster IDs at random.

Parameters:

Name	Type	Description	Default
random_state	int or None	Seed for the random number generator.	None

select_clusters

select_clusters(cluster_stats: ClusterStats, n_seeds: int) -> List[int]

Randomly sample n_seeds distinct cluster IDs.

Parameters:

Name	Type	Description	Default
cluster_stats	dict	Per-cluster statistics.	required
n_seeds	int	Number of clusters to sample.	required

Returns:

Type	Description
list of int	Randomly selected cluster IDs.

fast

FAST (Fluctuation Amplification of Specific Traits) sampling policy.

FastPolicy

FastPolicy(feature_indices: Sequence[int], directions: Optional[Sequence[Direction]] = None, weights: Optional[Sequence[float]] = None, alpha: float = 1.0)

Bases: Policy

Select clusters by balancing feature-directed exploitation and exploration.

Implements the FAST reward from Zimmerman & Bowman (2015): r(i) = phi_bar(i) + alpha * psi_bar(i), where phi_bar is a feature-scaled directed component and psi_bar favors poorly sampled clusters.

Parameters:

Name	Type	Description	Default
feature_indices	sequence of int	Feature column indices to optimize (required).	required
directions	sequence of str or None	maximize or minimize per feature. Defaults to all maximize.	None
weights	sequence of float or None	Weights per feature. Defaults to equal weights.	None
alpha	float	Relative weight of the exploration term. Default 1.0.	1.0

select_clusters

select_clusters(cluster_stats: ClusterStats, n_seeds: int) -> List[int]

Select clusters with the highest FAST reward scores.

Parameters:

Name	Type	Description	Default
cluster_stats	dict	Per-cluster statistics.	required
n_seeds	int	Number of clusters to select.	required

Returns:

Type	Description
list of int	Cluster IDs with highest rewards.

compute_fast_rewards

compute_fast_rewards(cluster_stats: ClusterStats, feature_indices: Sequence[int], directions: Sequence[Direction], weights: Sequence[float], alpha: float) -> Tuple[Dict[int, float], Dict[int, float], Dict[int, float]]

Compute FAST reward components for all clusters.

Parameters:

Name	Type	Description	Default
cluster_stats	dict	Per-cluster statistics.	required
feature_indices	sequence of int	Feature column indices to use.	required
directions	sequence of str	Optimization direction per feature.	required
weights	sequence of float	Weights per feature.	required
alpha	float	Exploration/exploitation balance parameter.	required

Returns:

Type	Description
tuple of dict	(directed_scores, exploration_scores, rewards) keyed by cluster ID.

feature_scale

feature_scale(values: Dict[int, float], direction: Direction) -> Dict[int, float]

Min-max scale cluster descriptor values to [0, 1].

Parameters:

Name	Type	Description	Default
values	dict	Mapping from cluster ID to raw descriptor value.	required
direction	str	maximize or minimize.	required

Returns:

Type	Description
dict	Scaled values in [0, 1]. Returns zeros when all values are equal.

knn_as

k-nearest neighbors adaptive sampling policy.

KnnAsPolicy

KnnAsPolicy(k: int = 5, scoring: ScoringMode = 'vectorsum', cluster_centers: Optional[ndarray] = None)

Bases: Policy

Select clusters using k-nearest neighbors adaptive sampling.

The original kNN-AS algorithm ranks states by local-neighborhood geometry. AdaptivePy policies select clusters, so this implementation applies the same ranking to cluster representative vectors and lets the seed-selection layer choose a frame from each selected cluster.

Parameters:

Name	Type	Description	Default
k	int	Number of nearest-neighbor records requested. Includes the query point itself when returned by scikit-learn. Default 5.	5
scoring	str	vectorsum matches the upstream vector-sum magnitude mode, while distance uses mean neighbor distance.	'vectorsum'
cluster_centers	ndarray or None	Optional cluster centers, shape (n_clusters, n_features).	None

select_clusters

select_clusters(cluster_stats: ClusterStats, n_seeds: int) -> List[int]

Select clusters with the highest kNN-AS scores.

compute_knn_as_scores

compute_knn_as_scores(vectors: ndarray, k: int, scoring: ScoringMode = 'vectorsum') -> Tuple[np.ndarray, int]

Compute kNN-AS scores for representative feature vectors.

Parameters:

Name	Type	Description	Default
vectors	ndarray	Representative feature matrix, shape (n_states, n_features).	required
k	int	Number of nearest-neighbor records requested from scikit-learn. This includes the point itself when present, matching the upstream algorithm.	required
scoring	str	vectorsum for the magnitude of summed neighbor displacement vectors, or distance for mean Euclidean neighbor distance.	'vectorsum'

Returns:

Type	Description
tuple	(scores, effective_k) where scores has one value per row and effective_k is the clamped neighbor count used for fitting.

ma_reap

Multiagent REAP (MA-REAP) adaptive sampling policy.

MaReapPolicy

MaReapPolicy(agent_assignments: Dict[str, Sequence[str]], traj_names: Sequence[str], cluster_centers: Optional[ndarray] = None, n_candidates: int = 10, initial_weights: Optional[ndarray] = None, delta: float = 0.05, stakes_method: StakesMethod = 'percentage', stakes_k: Optional[float] = None, regime: Regime = 'collaborative')

Bases: Policy

Multiagent REAP cluster selection policy.

Implements Kleiman & Shukla (2022): least-counts candidates, per-agent stakes, learned CV weights, and multiagent reward aggregation.

Parameters:

Name	Type	Description	Default
agent_assignments	dict	Maps agent names to feature file stems.	required
traj_names	list of str	Ordered trajectory stems from the loaded dataset.	required
cluster_centers	ndarray or None	Cluster centroids, shape (n_clusters, n_features).	None
n_candidates	int	Number of least-count clusters to consider.	10
initial_weights	ndarray or None	Starting CV weights per agent or shared across agents.	None
delta	float	Maximum per-feature weight change per round.	0.05
stakes_method	str	percentage, equal, max, or logistic.	'percentage'
stakes_k	float or None	Logistic steepness when stakes_method='logistic'.	None
regime	str	collaborative, noncollaborative, or competitive.	'collaborative'

select_clusters

select_clusters(cluster_stats: ClusterStats, n_seeds: int) -> List[int]

Select clusters using the MA-REAP reward pipeline.

aggregate_agent_scores

aggregate_agent_scores(scores: ndarray, regime: Regime) -> np.ndarray

Combine per-agent scores into global candidate scores (eqs. 7-9).

apply_stakes_method

apply_stakes_method(raw_counts: ndarray, method: StakesMethod, stakes_k: Optional[float] = None) -> np.ndarray

Convert raw frame counts to normalized stakes per candidate column.

compute_agent_scores

compute_agent_scores(means: ndarray, stdev: ndarray, stakes_agent: ndarray, candidate_features: ndarray, weights: ndarray) -> np.ndarray

Per-candidate reward for one agent (eq. 2 in Kleiman & Shukla 2022).

maxent_vampnet

Maximum-entropy VAMPNet adaptive sampling policy.

MaxEntVampNetPolicy

MaxEntVampNetPolicy(n_features: int, output_states: Optional[int] = None, n_states: Optional[int] = None, lagtime: int = DEFAULT_LAGTIME, hidden_layers: Optional[Sequence[int]] = None, learning_rate: float = DEFAULT_LEARNING_RATE, batch_size: int = DEFAULT_BATCH_SIZE, epochs: int = DEFAULT_EPOCHS, device: str = DEFAULT_DEVICE, num_threads: int = DEFAULT_NUM_THREADS, epsilon: float = VAMPNET_EPSILON, estimator: Optional[Any] = None)

Bases: Policy

Select frames by Shannon entropy of VAMPNet soft state assignments.

Implements the entropy-only MaxEnt VAMPNet acquisition function from Kleiman & Shukla (2023). Features are passed directly to a VAMPNet trained on lagged trajectory pairs; frames with the highest entropy of softmax state probabilities are selected as seeds. This policy does not require clustering.

Parameters:

Name	Type	Description	Default
n_states	int or None	Number of softmax output nodes. Defaults to the feature dimensionality.	None
output_states	int or None	Backward-compatible alias for n_states.	None
lagtime	int	Lag time in frames for VAMPNet training.	DEFAULT_LAGTIME
hidden_layers	sequence of int or None	Hidden MLP layer widths. Defaults to the author repository pattern.	None
learning_rate	float	VAMPNet learning rate.	DEFAULT_LEARNING_RATE
batch_size	int	Training batch size.	DEFAULT_BATCH_SIZE
epochs	int	Training epochs per policy invocation.	DEFAULT_EPOCHS
device	str	PyTorch device name, typically cpu or cuda.	DEFAULT_DEVICE
num_threads	int	CPU threads used by PyTorch during training.	DEFAULT_NUM_THREADS
epsilon	float	Numerical regularization constant passed to deeptime VAMPNet.	VAMPNET_EPSILON
estimator	object or None	Optional pre-fitted estimator for testing. When provided, training is skipped and this estimator is used for scoring.	None

select_clusters

select_clusters(cluster_stats: ClusterStats, n_seeds: int) -> List[int]

Not used by this frame-level policy.

select_frames

select_frames(dataset: Dataset, n_seeds: int) -> List[SeedResult]

Train or reuse VAMPNet, score frames by entropy, and select seeds.

compute_shannon_entropy

compute_shannon_entropy(probabilities: ndarray) -> np.ndarray

Compute per-row Shannon entropy from softmax probabilities.

Uses :func:scipy.stats.entropy with natural logarithm, matching the author implementation.

Parameters:

Name	Type	Description	Default
probabilities	ndarray	Softmax probabilities, shape (n_frames, n_states).	required

Returns:

Type	Description
ndarray	Entropy values with shape (n_frames,).

rank_frames_by_entropy

rank_frames_by_entropy(entropy_scores: ndarray, global_indices: Sequence[int], n_seeds: int) -> List[int]

Rank frame row indices by descending entropy.

Ties are broken by ascending global_index for deterministic ordering.

Parameters:

Name	Type	Description	Default
entropy_scores	ndarray	Entropy value per candidate frame.	required
global_indices	sequence of int	Global frame index for each entropy score.	required
n_seeds	int	Number of frames to select.	required

Returns:

Type	Description
list of int	Selected row indices into entropy_scores.

split_trajectories_from_dataset

split_trajectories_from_dataset(dataset: Dataset) -> List[np.ndarray]

Split a dataset feature matrix into per-trajectory arrays.

Parameters:

Name	Type	Description	Default
dataset	Dataset	Loaded dataset with feature_matrix and traj_index_map.	required

Returns:

Type	Description
list of np.ndarray	Feature arrays ordered by trajectory ID.

ts_dar

TS-DAR adaptive sampling policy.

TsDarPolicy

TsDarPolicy(n_features: int, n_states: Optional[int] = None, latent_dim: Optional[int] = None, hidden_layers: Optional[Sequence[int]] = None, encoder_sizes: Optional[Sequence[int]] = None, lagtime: int = DEFAULT_LAGTIME, learning_rate: float = DEFAULT_LEARNING_RATE, batch_size: int = DEFAULT_BATCH_SIZE, epochs: int = DEFAULT_EPOCHS, pretrain: int = DEFAULT_PRETRAIN, beta: float = DEFAULT_BETA, gamma: float = DEFAULT_GAMMA, scaling_temperature: float = DEFAULT_SCALING_TEMPERATURE, proto_update_factor: float = DEFAULT_PROTO_UPDATE_FACTOR, optimizer: str = DEFAULT_OPTIMIZER, device: str = DEFAULT_DEVICE, num_threads: int = DEFAULT_NUM_THREADS, train_split: float = DEFAULT_TRAIN_SPLIT, epsilon: float = TSDAR_EPSILON, random_state: Optional[int] = None, estimator: Optional[Any] = None)

Bases: Policy

Select frames with high TS-DAR out-of-distribution scores.

select_clusters

select_clusters(cluster_stats: ClusterStats, n_seeds: int) -> List[int]

Not used by this frame-level policy.

select_frames

select_frames(dataset: Dataset, n_seeds: int) -> List[SeedResult]

Train or reuse TS-DAR, score frames by OOD score, and select seeds.

compute_ood_scores

compute_ood_scores(embeddings: ndarray, state_centers: ndarray, epsilon: float = TSDAR_EPSILON) -> np.ndarray

Compute TS-DAR OOD scores from cosine distance to nearest state center.

compute_state_centers

compute_state_centers(embeddings: ndarray, states: ndarray, n_states: int) -> np.ndarray

Compute normalized state-center vectors from hyperspherical embeddings.

rank_frames_by_ood

rank_frames_by_ood(ood_scores: ndarray, global_indices: Sequence[int], n_seeds: int) -> List[int]

Rank frame row indices by descending OOD score.