Features Module

The features module provides unified feature extraction for molecules and proteins.

MoleculeFeaturizer

MoleculeFeaturizer

MoleculeFeaturizer(featurizer_name: str = 'ecfp', n_jobs: int = 8, device: str = 'auto')

Efficient molecule featurization using molfeat.

Provides batch featurization with SMILES deduplication for efficiency. Supports fingerprints, descriptors, and neural embeddings.

Attributes:

Name	Type	Description
featurizer_name		Name of the featurizer to use.
n_jobs		Number of parallel workers for featurization.
_transformer		Cached molfeat transformer instance.

Examples:

>>> featurizer = MoleculeFeaturizer("ecfp")
>>> features = featurizer.featurize(["CCO", "CCCO", "CCCCO"])
>>> print(features.shape)  # (3, 2048)

>>> # With SMILES deduplication
>>> smiles = ["CCO", "CCCO", "CCO", "CCCCO", "CCCO"]
>>> features = featurizer.featurize_deduplicated(smiles)
>>> print(features.shape)  # (5, 2048) - returns features for all input SMILES

Initialize the molecule featurizer.

Parameters:

Name	Type	Description	Default
featurizer_name	str	Name of the molfeat featurizer to use.	'ecfp'
n_jobs	int	Number of parallel workers for featurization.	8
device	str	Device for neural featurizers ('auto', 'cpu', 'cuda').	'auto'

transformer property

transformer

Get or create the molfeat transformer (lazy initialization).

is_fingerprint property

is_fingerprint: bool

Check if this is a fingerprint-based featurizer (including count fingerprints).

is_neural property

is_neural: bool

Check if this is a neural embedding featurizer.

featurize

featurize(smiles: Union[str, List[str]], ignore_errors: bool = True) -> NDArray[float32]

Featurize one or more SMILES strings.

Parameters:

Name	Type	Description	Default
smiles	Union[str, List[str]]	Single SMILES string or list of SMILES.	required
ignore_errors	bool	If True, return NaN for invalid SMILES.	True

Returns:

Type	Description
NDArray[float32]	Feature array of shape (n_molecules, feature_dim).

featurize_deduplicated

featurize_deduplicated(smiles: List[str], ignore_errors: bool = True) -> NDArray[float32]

Featurize SMILES with deduplication for efficiency.

Unique SMILES are featurized once, then results are mapped back to the original list. This is efficient when there are many duplicate SMILES across datasets.

Parameters:

Name	Type	Description	Default
smiles	List[str]	List of SMILES strings (may contain duplicates).	required
ignore_errors	bool	If True, return NaN for invalid SMILES.	True

Returns:

Type	Description
NDArray[float32]	Feature array of shape (n_molecules, feature_dim) in original order.

featurize_datasets

featurize_datasets(datasets: Dict[str, 'MoleculeDataset'], deduplicate: bool = True) -> Dict[str, NDArray[float32]]

Featurize multiple datasets efficiently.

When deduplicate=True, collects all unique SMILES across all datasets, featurizes them once, then maps back to each dataset.

Parameters:

Name	Type	Description	Default
datasets	Dict[str, 'MoleculeDataset']	Dictionary mapping task IDs to MoleculeDataset instances.	required
deduplicate	bool	If True, deduplicate SMILES across all datasets.	True

Returns:

Type	Description
Dict[str, NDArray[float32]]	Dictionary mapping task IDs to feature arrays.

get_feature_dim

get_feature_dim() -> int

Get the feature dimension for this featurizer.

Returns:

Type	Description
int	Number of features produced by this featurizer.

Available Featurizers

Fingerprints

Featurizer	Description	Dimensions
ecfp	Extended Connectivity Fingerprints	2048
fcfp	Functional Connectivity Fingerprints	2048
maccs	MACCS Structural Keys	167
topological	Topological Fingerprints	2048
avalon	Avalon Fingerprints	512
atompair	Atom Pair Fingerprints	2048
pattern	Pattern Fingerprints	2048
layered	Layered Fingerprints	2048
secfp	SMILES Extended Connectivity Fingerprints	2048
erg	Extended Reduced Graph	varies
estate	E-State Fingerprints	79
rdkit	RDKit Fingerprints	2048

Count variants (ecfp-count, fcfp-count, etc.) are also available.

Descriptors

Featurizer	Description	Dimensions
desc2D	2D Molecular Descriptors	~200
mordred	Mordred Descriptors	~1600
cats2D	CATS2D Pharmacophore	varies
pharm2D	2D Pharmacophore	varies
scaffoldkeys	Scaffold Keys	varies

Neural Embeddings

Featurizer	Description	Dimensions
ChemBERTa-77M-MLM	ChemBERTa masked language model	384
ChemBERTa-77M-MTR	ChemBERTa multi-task regression	384
MolT5	Molecular T5 embeddings	768
Roberta-Zinc480M-102M	RoBERTa trained on ZINC	768
gin_supervised_infomax	GIN infomax embeddings	300
gin_supervised_contextpred	GIN context prediction	300
gin_supervised_edgepred	GIN edge prediction	300
gin_supervised_masking	GIN masking embeddings	300

Run themap list-featurizers to see all available featurizers.

ProteinFeaturizer

ProteinFeaturizer

ProteinFeaturizer(featurizer_name: str = 'esm2_t33_650M_UR50D', layer: Optional[int] = None, device: str = 'auto')

Efficient protein featurization using ESM models.

Provides batch featurization of protein sequences using ESM2 or ESM3 models. Models are cached globally to avoid reloading.

Attributes:

Name	Type	Description
featurizer_name		Name of the ESM model to use.
layer		Which transformer layer to extract embeddings from.
device		Device for computation ('auto', 'cpu', 'cuda').

Examples:

>>> featurizer = ProteinFeaturizer("esm2_t33_650M_UR50D")
>>> sequences = {"P1": "MKTVRQ...", "P2": "MENLNM..."}
>>> features = featurizer.featurize(sequences)
>>> print(features.shape)  # (2, 1280)

Initialize the protein featurizer.

Parameters:

Name	Type	Description	Default
featurizer_name	str	Name of the ESM model to use.	'esm2_t33_650M_UR50D'
layer	Optional[int]	Which transformer layer to extract embeddings from. If None, uses the default for the model.	None
device	str	Device for computation ('auto', 'cpu', 'cuda').	'auto'

is_esm2 property

is_esm2: bool

Check if this is an ESM2 model.

is_esm3 property

is_esm3: bool

Check if this is an ESM3 model.

featurize

featurize(sequences: Union[Dict[str, str], List[str]]) -> NDArray[float32]

Featurize protein sequences.

Parameters:

Name	Type	Description	Default
sequences	Union[Dict[str, str], List[str]]	Either a dictionary mapping protein IDs to sequences, or a list of sequences.	required

Returns:

Type	Description
NDArray[float32]	Feature array of shape (n_proteins, embedding_dim).

featurize_from_fasta

featurize_from_fasta(fasta_path: Union[str, Path]) -> Dict[str, NDArray[float32]]

Featurize proteins from a FASTA file.

Parameters:

Name	Type	Description	Default
fasta_path	Union[str, Path]	Path to the FASTA file.	required

Returns:

Type	Description
Dict[str, NDArray[float32]]	Dictionary mapping protein IDs to feature vectors.

featurize_directory

featurize_directory(directory: Union[str, Path], pattern: str = '*.fasta') -> Dict[str, NDArray[float32]]

Featurize all proteins from FASTA files in a directory.

Each FASTA file is expected to contain one protein sequence. The filename (without extension) is used as the task/protein ID.

Parameters:

Name	Type	Description	Default
directory	Union[str, Path]	Path to directory containing FASTA files.	required
pattern	str	Glob pattern for finding FASTA files.	'*.fasta'

Returns:

Type	Description
Dict[str, NDArray[float32]]	Dictionary mapping task IDs to feature vectors.

get_feature_dim

get_feature_dim() -> int

Get the feature dimension for this featurizer.

Returns:

Type	Description
int	Number of features produced by this featurizer.

Available Models

ESM2 Models

Model	Parameters	Layers	Embedding Dim
esm2_t6_8M_UR50D	8M	6	320
esm2_t12_35M_UR50D	35M	12	480
esm2_t30_150M_UR50D	150M	30	640
esm2_t33_650M_UR50D	650M	33	1280

ESM3 Models

Model	Description
esm3_sm_open_v1	ESM3 small open model

FeatureCache

FeatureCache

FeatureCache(cache_dir: Union[str, Path])

Disk-based feature caching for molecules and proteins.

Caches computed features to NPZ/NPY files for efficient reuse across runs. Supports both molecule features (per-dataset with labels) and protein features (single vector per task).

Attributes:

Name	Type	Description
cache_dir		Root directory for cached features.
molecule_dir		Directory for molecule features.
protein_dir		Directory for protein features.

Directory Structure

cache_dir/ ├── molecule/ │ └── {featurizer}/ │ └── {task_id}.npz # features + labels └── protein/ └── {featurizer}/ └── {task_id}.npy # single vector

Examples:

>>> cache = FeatureCache("./feature_cache")
>>> # Save molecule features
>>> cache.save_molecule_features("CHEMBL123", "ecfp", features, labels)
>>> # Load if exists
>>> features, labels = cache.load_molecule_features("CHEMBL123", "ecfp")
>>> if features is None:
...     # Compute and save
...     cache.save_molecule_features("CHEMBL123", "ecfp", features, labels)

Initialize the feature cache.

Parameters:

Name	Type	Description	Default
cache_dir	Union[str, Path]	Root directory for cached features.	required

has_molecule_features

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

Check if molecule features are cached.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID to check.	required
featurizer	str	Name of the featurizer.	required

Returns:

Type	Description
bool	True if features are cached.

has_protein_features

has_protein_features(task_id: str, featurizer: str) -> bool

Check if protein features are cached.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID to check.	required
featurizer	str	Name of the featurizer.	required

Returns:

Type	Description
bool	True if features are cached.

save_molecule_features

save_molecule_features(task_id: str, featurizer: str, features: NDArray[float32], labels: NDArray[int32], metadata: Optional[Dict[str, Any]] = None) -> Path

Save molecule features to cache.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID for the dataset.	required
featurizer	str	Name of the featurizer used.	required
features	NDArray[float32]	Feature matrix of shape (n_molecules, feature_dim).	required
labels	NDArray[int32]	Binary labels of shape (n_molecules,).	required
metadata	Optional[Dict[str, Any]]	Optional metadata dictionary.	None

Returns:

Type	Description
Path	Path to the saved file.

load_molecule_features

load_molecule_features(task_id: str, featurizer: str) -> Tuple[Optional[NDArray[float32]], Optional[NDArray[int32]]]

Load molecule features from cache.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID for the dataset.	required
featurizer	str	Name of the featurizer.	required

Returns:

Type	Description
Tuple[Optional[NDArray[float32]], Optional[NDArray[int32]]]	Tuple of (features, labels), or (None, None) if not cached.

save_protein_features

save_protein_features(task_id: str, featurizer: str, features: NDArray[float32]) -> Path

Save protein features to cache.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID for the protein.	required
featurizer	str	Name of the featurizer used.	required
features	NDArray[float32]	Feature vector of shape (feature_dim,).	required

Returns:

Type	Description
Path	Path to the saved file.

load_protein_features

load_protein_features(task_id: str, featurizer: str) -> Optional[NDArray[float32]]

Load protein features from cache.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID for the protein.	required
featurizer	str	Name of the featurizer.	required

Returns:

Type	Description
Optional[NDArray[float32]]	Feature vector, or None if not cached.

save_all_molecule_features

save_all_molecule_features(features_dict: Dict[str, NDArray[float32]], labels_dict: Dict[str, NDArray[int32]], featurizer: str) -> List[Path]

Save molecule features for multiple datasets.

Parameters:

Name	Type	Description	Default
features_dict	Dict[str, NDArray[float32]]	Dictionary mapping task IDs to feature matrices.	required
labels_dict	Dict[str, NDArray[int32]]	Dictionary mapping task IDs to label arrays.	required
featurizer	str	Name of the featurizer used.	required

Returns:

Type	Description
List[Path]	List of paths to saved files.

save_all_protein_features

save_all_protein_features(features_dict: Dict[str, NDArray[float32]], featurizer: str) -> List[Path]

Save protein features for multiple tasks.

Parameters:

Name	Type	Description	Default
features_dict	Dict[str, NDArray[float32]]	Dictionary mapping task IDs to feature vectors.	required
featurizer	str	Name of the featurizer used.	required

Returns:

Type	Description
List[Path]	List of paths to saved files.

load_all_molecule_features

load_all_molecule_features(task_ids: List[str], featurizer: str) -> Tuple[Dict[str, NDArray[float32]], Dict[str, NDArray[int32]], List[str]]

Load molecule features for multiple datasets.

Parameters:

Name	Type	Description	Default
task_ids	List[str]	List of task IDs to load.	required
featurizer	str	Name of the featurizer.	required

Returns:

Type	Description
Dict[str, NDArray[float32]]	Tuple of (features_dict, labels_dict, missing_ids).
Dict[str, NDArray[int32]]	missing_ids contains task IDs that were not found in cache.

load_all_protein_features

load_all_protein_features(task_ids: List[str], featurizer: str) -> Tuple[Dict[str, NDArray[float32]], List[str]]

Load protein features for multiple tasks.

Parameters:

Name	Type	Description	Default
task_ids	List[str]	List of task IDs to load.	required
featurizer	str	Name of the featurizer.	required

Returns:

Type	Description
Dict[str, NDArray[float32]]	Tuple of (features_dict, missing_ids).
List[str]	missing_ids contains task IDs that were not found in cache.

clear

clear(featurizer: Optional[str] = None) -> int

Clear cached features.

Parameters:

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

Returns:

Type	Description
int	Number of files deleted.

get_statistics

get_statistics() -> Dict[str, Any]

Get statistics about cached features.

Returns:

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

See the Performance Optimization tutorial for caching and parallelism tips.