Data Utilities

This page documents the data loading and processing utilities.

Dataset Classes

MolecularDataset

molax.utils.data.MolecularDataset

Dataset class for molecular graphs using jraph format.

Attributes:

Name	Type	Description
graphs	List[GraphsTuple]	List of jraph.GraphsTuple objects
labels	List[float]	Array of property labels
n_node_features		Number of node features

__init__

__init__(
    data: Union[DataFrame, str, Path],
    smiles_col: str = "smiles",
    label_col: str = "property",
    features: str = DEFAULT_FEATURIZER,
)

Initialize dataset from DataFrame or CSV file.

Parameters:

Name	Type	Description	Default
data	Union[DataFrame, str, Path]	DataFrame or path to CSV file	required
smiles_col	str	Column name for SMILES strings	'smiles'
label_col	str	Column name for property labels	'property'
features	str	Atom featurizer name, see :mod:molax.utils.featurizers	DEFAULT_FEATURIZER

get_batched

get_batched(
    indices: Optional[List[int]] = None,
    pad_to_nodes: Optional[int] = None,
    pad_to_edges: Optional[int] = None,
    pad_to_graphs: Optional[int] = None,
) -> Tuple[jraph.GraphsTuple, jnp.ndarray]

Get a batched GraphsTuple for the specified indices.

Parameters:

Name	Type	Description	Default
indices	Optional[List[int]]	List of indices to include. If None, returns all data.	None
pad_to_nodes	Optional[int]	Pad to this many nodes for consistent JIT shapes	None
pad_to_edges	Optional[int]	Pad to this many edges	None
pad_to_graphs	Optional[int]	Pad to this many graphs	None

Returns:

Type	Description
Tuple[GraphsTuple, ndarray]	Tuple of (batched_graphs, labels)

compute_padding_sizes

compute_padding_sizes(batch_size: int) -> Tuple[int, int, int]

Compute fixed padding sizes for efficient JIT compilation.

Parameters:

Name	Type	Description	Default
batch_size	int	Maximum batch size	required

Returns:

Type	Description
Tuple[int, int, int]	Tuple of (max_nodes, max_edges, n_graphs) for padding

split

split(
    test_size: float = 0.2, seed: Optional[int] = None
) -> Tuple[MolecularDataset, MolecularDataset]

Split dataset into train and test sets.

Parameters:

Name	Type	Description	Default
test_size	float	Fraction for test set	0.2
seed	Optional[int]	Random seed for reproducibility	None

Returns:

Type	Description
Tuple[MolecularDataset, MolecularDataset]	Tuple of (train_dataset, test_dataset)

---

Graph Conversion

Functions for converting molecular representations to graph format.

smiles_to_jraph

molax.utils.data.smiles_to_jraph

smiles_to_jraph(
    smiles: str, features: str = DEFAULT_FEATURIZER
) -> jraph.GraphsTuple

Convert SMILES string to jraph GraphsTuple format.

Parameters:

Name	Type	Description	Default
smiles	str	SMILES string representing the molecule	required
features	str	Atom featurizer name. "basic" (default) gives the six raw descriptors; "rich" gives a 29-dimensional one-hot encoding that trains substantially better. See :mod:molax.utils.featurizers.	DEFAULT_FEATURIZER

Returns:

Type	Description
GraphsTuple	jraph.GraphsTuple containing the molecular graph

Raises:

Type	Description
ValueError	If the SMILES string is invalid or the featurizer is unknown

batch_graphs

molax.utils.data.batch_graphs

batch_graphs(
    graphs: List[GraphsTuple],
    pad_to_nodes: Optional[int] = None,
    pad_to_edges: Optional[int] = None,
    pad_to_graphs: Optional[int] = None,
) -> jraph.GraphsTuple

Batch multiple graphs into a single padded GraphsTuple.

Padding ensures consistent shapes for JIT compilation efficiency.

Parameters:

Name	Type	Description	Default
graphs	List[GraphsTuple]	List of individual GraphsTuple objects	required
pad_to_nodes	Optional[int]	Pad total nodes to this number (default: auto)	None
pad_to_edges	Optional[int]	Pad total edges to this number (default: auto)	None
pad_to_graphs	Optional[int]	Pad to this many graphs (default: len(graphs) + 1)	None

Returns:

Type	Description
GraphsTuple	Single batched and padded GraphsTuple

unbatch_graphs

molax.utils.data.unbatch_graphs

unbatch_graphs(batched: GraphsTuple) -> List[jraph.GraphsTuple]

Unbatch a batched GraphsTuple back to individual graphs.

Parameters:

Name	Type	Description	Default
batched	GraphsTuple	Batched GraphsTuple	required

Returns:

Type	Description
List[GraphsTuple]	List of individual GraphsTuple objects

---

Atom Featurizers

Node features are produced by a named atom featurizer. "basic" is the default and emits the original six raw descriptors; "rich" emits a 29-dimensional one-hot encoding that trains substantially better.

On ESOL, with the same architecture and training budget for both (UncertaintyGCN, hidden [128, 128, 128], 1500 epochs, 80/20 split, seed 42):

Featurizer	Dims	Test RMSE
"basic"	6	1.33
"rich"	29	0.92
predict the training mean	—	2.13

from molax import ATOM_FEATURIZERS, MolecularDataset, smiles_to_jraph

graph = smiles_to_jraph("CCO", features="rich")
dataset = MolecularDataset("datasets/bace.csv", features="rich")

n_features = ATOM_FEATURIZERS["rich"].dim  # 29 — read it, don't hardcode it

AtomFeaturizer

molax.utils.featurizers.AtomFeaturizer dataclass

A named atom featurizer with a fixed output width.

Attributes:

Name	Type	Description
name	str	Registry key for this featurizer
dim	int	Number of features produced per atom
fn	Callable[[Atom], List[float]]	Callable mapping an RDKit atom to a list of floats of length dim

get_atom_featurizer

molax.utils.featurizers.get_atom_featurizer

get_atom_featurizer(features: str = DEFAULT_FEATURIZER) -> AtomFeaturizer

Look up a registered atom featurizer by name.

Parameters:

Name	Type	Description	Default
features	str	Registry key, one of ATOM_FEATURIZERS	DEFAULT_FEATURIZER

Returns:

Type	Description
AtomFeaturizer	The matching AtomFeaturizer

Raises:

Type	Description
ValueError	If no featurizer is registered under that name