Acquisition Functions

This page documents the acquisition functions for active learning sample selection.

Uncertainty-based Acquisition

Functions that select samples based on model uncertainty.

uncertainty_sampling

molax.acquisition.uncertainty.uncertainty_sampling

uncertainty_sampling(
    model: UncertaintyGCN, pool_graphs: List[GraphsTuple], n_samples: int = 10
) -> jnp.ndarray

Compute uncertainty scores for pool samples using MC dropout.

Parameters:

Name	Type	Description	Default
model	UncertaintyGCN	UncertaintyGCN model	required
pool_graphs	List[GraphsTuple]	List of jraph.GraphsTuple for pool samples	required
n_samples	int	Number of MC dropout samples	10

Returns:

Type	Description
ndarray	Array of uncertainty scores for each pool sample

ensemble_uncertainty_sampling

molax.acquisition.uncertainty.ensemble_uncertainty_sampling

ensemble_uncertainty_sampling(
    ensemble: DeepEnsemble,
    pool_graphs: List[GraphsTuple],
    use_epistemic: bool = True,
) -> jnp.ndarray

Compute uncertainty scores for pool samples using ensemble disagreement.

Unlike MC Dropout, ensembles provide uncertainty in a single forward pass by measuring disagreement between independently trained models.

Parameters:

Name	Type	Description	Default
ensemble	DeepEnsemble	DeepEnsemble model	required
pool_graphs	List[GraphsTuple]	List of jraph.GraphsTuple for pool samples	required
use_epistemic	bool	If True, use epistemic uncertainty (model disagreement). If False, use total uncertainty (epistemic + aleatoric).	True

Returns:

Type	Description
ndarray	Array of uncertainty scores for each pool sample

evidential_uncertainty_sampling

molax.acquisition.uncertainty.evidential_uncertainty_sampling

evidential_uncertainty_sampling(
    model: EvidentialGCN,
    pool_graphs: List[GraphsTuple],
    use_epistemic: bool = True,
) -> jnp.ndarray

Compute uncertainty scores for pool samples using evidential prediction.

Unlike MC Dropout, evidential models provide uncertainty in a single forward pass by predicting the parameters of a higher-order distribution.

Parameters:

Name	Type	Description	Default
model	EvidentialGCN	EvidentialGCN model	required
pool_graphs	List[GraphsTuple]	List of jraph.GraphsTuple for pool samples	required
use_epistemic	bool	If True, use epistemic uncertainty (model uncertainty). If False, use total uncertainty (epistemic + aleatoric).	True

Returns:

Type	Description
ndarray	Array of uncertainty scores for each pool sample

---

Diversity-based Acquisition

Functions that select diverse samples in feature space.

diversity_sampling

molax.acquisition.uncertainty.diversity_sampling

diversity_sampling(
    pool_graphs: List[GraphsTuple],
    labeled_graphs: List[GraphsTuple],
    n_select: int,
) -> List[int]

Select diverse samples using greedy farthest point sampling.

Uses mean node features as molecular fingerprints.

Parameters:

Name	Type	Description	Default
pool_graphs	List[GraphsTuple]	List of pool sample graphs	required
labeled_graphs	List[GraphsTuple]	List of labeled sample graphs	required
n_select	int	Number of samples to select	required

Returns:

Type	Description
List[int]	List of selected indices into pool_graphs

---

Combined Acquisition

Functions that combine multiple acquisition strategies.

combined_acquisition

molax.acquisition.uncertainty.combined_acquisition

combined_acquisition(
    model: UncertaintyGCN,
    pool_graphs: List[GraphsTuple],
    labeled_graphs: List[GraphsTuple],
    n_select: int,
    uncertainty_weight: float = 0.7,
    n_mc_samples: int = 10,
) -> List[int]

Combined uncertainty and diversity acquisition.

Parameters:

Name	Type	Description	Default
model	UncertaintyGCN	UncertaintyGCN model	required
pool_graphs	List[GraphsTuple]	List of pool sample graphs	required
labeled_graphs	List[GraphsTuple]	List of labeled sample graphs	required
n_select	int	Number of samples to select	required
uncertainty_weight	float	Weight for uncertainty vs diversity (0-1)	0.7
n_mc_samples	int	Number of MC samples for uncertainty estimation	10

Returns:

Type	Description
List[int]	List of selected indices into pool_graphs

combined_ensemble_acquisition

molax.acquisition.uncertainty.combined_ensemble_acquisition

combined_ensemble_acquisition(
    ensemble: DeepEnsemble,
    pool_graphs: List[GraphsTuple],
    labeled_graphs: List[GraphsTuple],
    n_select: int,
    uncertainty_weight: float = 0.7,
    use_epistemic: bool = True,
) -> List[int]

Combined uncertainty and diversity acquisition using ensemble.

Parameters:

Name	Type	Description	Default
ensemble	DeepEnsemble	DeepEnsemble model	required
pool_graphs	List[GraphsTuple]	List of pool sample graphs	required
labeled_graphs	List[GraphsTuple]	List of labeled sample graphs	required
n_select	int	Number of samples to select	required
uncertainty_weight	float	Weight for uncertainty vs diversity (0-1)	0.7
use_epistemic	bool	If True, use epistemic uncertainty for acquisition	True

Returns:

Type	Description
List[int]	List of selected indices into pool_graphs

combined_evidential_acquisition

molax.acquisition.uncertainty.combined_evidential_acquisition

combined_evidential_acquisition(
    model: EvidentialGCN,
    pool_graphs: List[GraphsTuple],
    labeled_graphs: List[GraphsTuple],
    n_select: int,
    uncertainty_weight: float = 0.7,
    use_epistemic: bool = True,
) -> List[int]

Combined uncertainty and diversity acquisition using evidential model.

Parameters:

Name	Type	Description	Default
model	EvidentialGCN	EvidentialGCN model	required
pool_graphs	List[GraphsTuple]	List of pool sample graphs	required
labeled_graphs	List[GraphsTuple]	List of labeled sample graphs	required
n_select	int	Number of samples to select	required
uncertainty_weight	float	Weight for uncertainty vs diversity (0-1)	0.7
use_epistemic	bool	If True, use epistemic uncertainty for acquisition	True

Returns:

Type	Description
List[int]	List of selected indices into pool_graphs