Pipeline Module

The pipeline module provides a high-level API for running distance computation workflows.

Entry Points

quick_distance

quick_distance

quick_distance(data_dir: str, output_dir: str = 'output', molecule_featurizer: str = 'ecfp', molecule_method: str = 'euclidean', n_jobs: int = 8, device: str = 'auto') -> Dict[str, DistanceMatrix]

Quick distance computation with minimal configuration.

Parameters:

Name	Type	Description	Default
data_dir	str	Path to data directory with train/test folders.	required
output_dir	str	Path to output directory.	'output'
molecule_featurizer	str	Molecule featurizer name.	'ecfp'
molecule_method	str	Distance method.	'euclidean'
n_jobs	int	Number of parallel jobs.	8
device	str	Device for OTDD ("auto" | "cpu" | "cuda"). Ignored by Euclidean/Cosine.	'auto'

Returns:

Type	Description
Dict[str, DistanceMatrix]	Dictionary of distance matrices.

run_pipeline

run_pipeline

run_pipeline(config_path: str) -> Dict[str, DistanceMatrix]

Convenience function to run pipeline from config file.

Parameters:

Name	Type	Description	Default
config_path	str	Path to YAML config file.	required

Returns:

Type	Description
Dict[str, DistanceMatrix]	Dictionary of distance matrices.

Pipeline

Pipeline

Pipeline(config: PipelineConfig)

Main pipeline orchestrator for THEMAP distance computation.

Orchestrates the complete workflow: 1. Load datasets from train/test directories 2. Compute molecule features (with caching) 3. Compute protein features (with caching, if enabled) 4. Compute distance matrices 5. Combine matrices (if needed) 6. Save results to output directory

Attributes:

Name	Type	Description
config		Pipeline configuration
loader		Dataset loader
cache	Optional[FeatureCache]	Feature cache (if save_features is enabled)
mol_featurizer	MoleculeFeaturizer	Molecule featurizer (if molecule distance enabled)
prot_featurizer	ProteinFeaturizer	Protein featurizer (if protein distance enabled)

Examples:

>>> config = PipelineConfig.from_yaml("config.yaml")
>>> pipeline = Pipeline(config)
>>> results = pipeline.run()
>>> print(results["molecule"])  # molecule distance matrix

Initialize the pipeline.

Parameters:

Name	Type	Description	Default
config	PipelineConfig	Pipeline configuration	required

mol_featurizer property

mol_featurizer: MoleculeFeaturizer

Get molecule featurizer (lazy initialization).

prot_featurizer property

prot_featurizer: ProteinFeaturizer

Get protein featurizer (lazy initialization).

run

run() -> Dict[str, DistanceMatrix]

Run the complete pipeline.

Returns:

Type	Description
Dict[str, DistanceMatrix]	Dictionary mapping distance type to distance matrix:
Dict[str, DistanceMatrix]	"molecule": molecule distance matrix (if enabled)
Dict[str, DistanceMatrix]	"protein": protein distance matrix (if enabled)
Dict[str, DistanceMatrix]	"combined": combined matrix (if combination != "separate")

Raises:

Type	Description
ValueError	If no distance type is enabled

Configuration Classes

PipelineConfig

PipelineConfig dataclass

PipelineConfig(data: DataConfig = DataConfig(), molecule: MoleculeDistanceConfig = MoleculeDistanceConfig(), protein: ProteinDistanceConfig = ProteinDistanceConfig(), combination: CombinationConfig = CombinationConfig(), output: OutputConfig = OutputConfig(), compute: ComputeConfig = ComputeConfig())

Main configuration for the THEMAP pipeline.

Example YAML:

data:
  directory: "datasets/TDC"
  task_list: "tasks.json"  # Optional

distances:
  molecule:
    enabled: true
    featurizer: "ecfp"
    method: "euclidean"
  protein:
    enabled: false
    featurizer: "esm2_t33_650M_UR50D"
    method: "cosine"

combination:
  strategy: "weighted_average"
  weights:
    molecule: 0.7
    protein: 0.3

output:
  directory: "output/"
  save_features: true

compute:
  n_jobs: 8
  batch_size: 1000

from_yaml classmethod

from_yaml(path: Union[str, Path]) -> PipelineConfig

Load configuration from a YAML file.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Path to the YAML configuration file.	required

Returns:

Type	Description
PipelineConfig	PipelineConfig instance.

from_dict classmethod

from_dict(config_dict: Dict[str, Any]) -> PipelineConfig

Create configuration from a dictionary.

Parameters:

Name	Type	Description	Default
config_dict	Dict[str, Any]	Dictionary with configuration values.	required

Returns:

Type	Description
PipelineConfig	PipelineConfig instance.

to_dict

to_dict() -> Dict[str, Any]

Convert configuration to dictionary.

to_yaml

to_yaml(path: Union[str, Path]) -> None

Save configuration to a YAML file.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Path to save the YAML file.	required

validate

validate() -> List[str]

Validate configuration and return list of warnings/errors.

Returns:

Type	Description
List[str]	List of warning/error messages.

DataConfig

DataConfig dataclass

DataConfig(directory: Path = (lambda: Path('datasets'))(), task_list: Optional[str] = None)

Configuration for data loading.

MoleculeDistanceConfig

MoleculeDistanceConfig dataclass

MoleculeDistanceConfig(enabled: bool = True, featurizer: str = 'ecfp', method: str = 'euclidean')

Configuration for molecule-based distance computation.

ProteinDistanceConfig

ProteinDistanceConfig dataclass

ProteinDistanceConfig(enabled: bool = False, featurizer: str = 'esm2_t33_650M_UR50D', method: str = 'cosine', layer: Optional[int] = None)

Configuration for protein-based distance computation.

OutputConfig

OutputConfig dataclass

OutputConfig(directory: Path = (lambda: Path('output'))(), format: str = 'csv', save_features: bool = True)

Configuration for output files.

ComputeConfig

ComputeConfig dataclass

ComputeConfig(n_jobs: int = 8, batch_size: int = 1000, device: str = 'auto')

Configuration for computation settings.

CombinationConfig

CombinationConfig dataclass

CombinationConfig(strategy: str = 'weighted_average', weights: Dict[str, float] = (lambda: {'molecule': 0.5, 'protein': 0.5})())

Configuration for combining multiple distance matrices.

Featurization Pipeline

FeatureStore

FeatureStore dataclass

FeatureStore(cache_dir: Path)

Disk-based feature storage for datasets.

File structure

cache_dir/ molecules/ {featurizer_name}/ {task_id}.npz # Contains 'features' and 'labels' protein/ {featurizer_name}/ {task_id}.npy # Single vector per task metadata/ {metadata_type}/ {featurizer_name}/ {task_id}.npy # Single vector per task

Attributes:

Name	Type	Description
cache_dir	Path	Root directory for feature storage

save_molecule_features

save_molecule_features(task_id: str, features: NDArray[float32], labels: NDArray[int32], featurizer_name: str) -> None

Save molecule dataset features.

Parameters:

Name	Type	Description	Default
task_id	str	Task identifier	required
features	NDArray[float32]	Feature matrix of shape (n_molecules, feature_dim)	required
labels	NDArray[int32]	Label array of shape (n_molecules,)	required
featurizer_name	str	Name of the featurizer used	required

load_molecule_features

load_molecule_features(task_id: str, featurizer_name: str) -> Optional[Dict[str, NDArray]]

Load molecule dataset features.

Parameters:

Name	Type	Description	Default
task_id	str	Task identifier	required
featurizer_name	str	Name of the featurizer used	required

Returns:

Type	Description
Optional[Dict[str, NDArray]]	Dictionary with 'features' and 'labels' arrays, or None if not found

save_metadata_features

save_metadata_features(task_id: str, features: NDArray[float32], metadata_type: str, featurizer_name: str) -> None

Save metadata features (single vector per task).

Parameters:

Name	Type	Description	Default
task_id	str	Task identifier	required
features	NDArray[float32]	Feature vector of shape (feature_dim,)	required
metadata_type	str	Type of metadata (e.g., 'protein', 'description')	required
featurizer_name	str	Name of the featurizer used	required

load_metadata_features

load_metadata_features(task_id: str, metadata_type: str, featurizer_name: str) -> Optional[NDArray[float32]]

Load metadata features.

Parameters:

Name	Type	Description	Default
task_id	str	Task identifier	required
metadata_type	str	Type of metadata	required
featurizer_name	str	Name of the featurizer used	required

Returns:

Type	Description
Optional[NDArray[float32]]	Feature vector or None if not found

has_molecule_features

has_molecule_features(task_id: str, featurizer_name: str) -> bool

Check if molecule features exist for a task.

has_metadata_features

has_metadata_features(task_id: str, metadata_type: str, featurizer_name: str) -> bool

Check if metadata features exist for a task.

get_cached_task_ids

get_cached_task_ids(featurizer_name: str, feature_type: str = 'molecules') -> List[str]

Get list of task IDs that have cached features.

Parameters:

Name	Type	Description	Default
featurizer_name	str	Name of the featurizer	required
feature_type	str	'molecules' or metadata type name	'molecules'

Returns:

Type	Description
List[str]	List of task IDs with cached features

clear_cache

clear_cache(featurizer_name: Optional[str] = None) -> None

Clear cached features.

Parameters:

Name	Type	Description	Default
featurizer_name	Optional[str]	If provided, only clear features for this featurizer. If None, clear all cached features.	None

FeaturizationPipeline

FeaturizationPipeline

FeaturizationPipeline(cache_dir: Union[str, Path], molecule_featurizer: str = 'ecfp', protein_featurizer: str = 'esm2_t33_650M_UR50D')

Pipeline for batch featurization of molecule datasets.

This pipeline provides efficient batch featurization with: - Global SMILES deduplication across all datasets - Disk-based caching for reuse across sessions - Parallel computation support

Attributes:

Name	Type	Description
store		FeatureStore for disk caching
molecule_featurizer		Name of molecular featurizer to use
protein_featurizer		Name of protein featurizer to use

Examples:

>>> pipeline = FeaturizationPipeline(
...     cache_dir="./feature_cache",
...     molecule_featurizer="ecfp"
... )
>>> # Featurize all datasets
>>> pipeline.featurize_all_datasets(datasets)
>>> # Load features for distance computation
>>> features, labels, ids = pipeline.load_dataset_features(datasets, names)

Initialize the featurization pipeline.

Parameters:

Name	Type	Description	Default
cache_dir	Union[str, Path]	Directory for storing cached features	required
molecule_featurizer	str	Name of the molecular featurizer	'ecfp'
protein_featurizer	str	Name of the protein featurizer	'esm2_t33_650M_UR50D'

featurize_all_datasets

featurize_all_datasets(datasets: List[MoleculeDataset], n_jobs: int = 8, batch_size: int = 1000, force_recompute: bool = False) -> Dict[str, bool]

Featurize all datasets and save features to disk.

This method: 1. Collects all unique SMILES across datasets 2. Batch computes features for unique SMILES 3. Distributes and saves features per dataset

Parameters:

Name	Type	Description	Default
datasets	List[MoleculeDataset]	List of MoleculeDataset objects to featurize	required
n_jobs	int	Number of parallel jobs for featurization	8
batch_size	int	Batch size for featurizer	1000
force_recompute	bool	If True, recompute even if cached	False

Returns:

Type	Description
Dict[str, bool]	Dictionary mapping task_id to success status

load_dataset_features

load_dataset_features(datasets: List[MoleculeDataset], dataset_names: Optional[List[str]] = None) -> Tuple[List[NDArray[float32]], List[NDArray[int32]], List[str]]

Load features for a list of datasets.

Parameters:

Name	Type	Description	Default
datasets	List[MoleculeDataset]	List of MoleculeDataset objects	required
dataset_names	Optional[List[str]]	Optional list of names (uses task_id if not provided)	None

Returns:

Type	Description
List[NDArray[float32]]	Tuple of (features_list, labels_list, valid_names)
List[NDArray[int32]]	features_list: List of feature matrices, one per dataset
List[str]	labels_list: List of label arrays, one per dataset
Tuple[List[NDArray[float32]], List[NDArray[int32]], List[str]]	valid_names: List of dataset names that were successfully loaded

load_features_for_distance

load_features_for_distance(source_datasets: List[MoleculeDataset], target_datasets: List[MoleculeDataset], source_names: Optional[List[str]] = None, target_names: Optional[List[str]] = None) -> Tuple[List[NDArray[float32]], List[NDArray[int32]], List[str], List[NDArray[float32]], List[NDArray[int32]], List[str]]

Load features organized for N×M distance computation.

Parameters:

Name	Type	Description	Default
source_datasets	List[MoleculeDataset]	List of source MoleculeDataset objects (N)	required
target_datasets	List[MoleculeDataset]	List of target MoleculeDataset objects (M)	required
source_names	Optional[List[str]]	Optional names for source datasets	None
target_names	Optional[List[str]]	Optional names for target datasets	None

Returns:

Type	Description
List[NDArray[float32]]	Tuple containing:
List[NDArray[int32]]	source_features: List of N feature matrices
List[str]	source_labels: List of N label arrays
List[NDArray[float32]]	source_ids: List of N dataset names
List[NDArray[int32]]	target_features: List of M feature matrices
List[str]	target_labels: List of M label arrays
Tuple[List[NDArray[float32]], List[NDArray[int32]], List[str], List[NDArray[float32]], List[NDArray[int32]], List[str]]	target_ids: List of M dataset names

get_cache_stats

get_cache_stats() -> Dict[str, Any]

Get statistics about cached features.

Returns:

Type	Description
Dict[str, Any]	Dictionary with cache statistics

Configuration File Format

data:
  directory: "datasets"

distances:
  molecule:
    enabled: true
    featurizer: "ecfp"
    method: "euclidean"
  protein:
    enabled: false
    featurizer: "esm2_t33_650M_UR50D"
    method: "euclidean"

output:
  directory: "output"
  format: "csv"

compute:
  n_jobs: 8
  device: "auto"

Output Files

output/
├── molecule_distances.csv
├── features/
│   ├── ecfp_source.npz
│   └── ecfp_target.npz
└── pipeline_summary.json

See Getting Started for usage examples and CLI Reference for command-line usage.