chemprop.data.collate

Classes

BatchMolGraph	A BatchMolGraph represents a batch of individual MolGraphs.
TrainingBatch
BatchMolAtomBondGraph	A BatchMolGraph represents a batch of individual MolGraphs.
MolAtomBondTrainingBatch
MulticomponentTrainingBatch

Functions

collate_batch(batch)
collate_cuik_batch(batch)
collate_mol_atom_bond_batch(batch)
collate_multicomponent(batches)

Module Contents

class chemprop.data.collate.BatchMolGraph

A BatchMolGraph represents a batch of individual MolGraphs.

It has all the attributes of a MolGraph with the addition of the batch attribute. This class is intended for use with data loading, so it uses Tensors to store data

mgs: dataclasses.InitVar[Sequence[chemprop.data.molgraph.MolGraph]]

A list of individual MolGraphs to be batched together

V: torch.Tensor

the atom feature matrix

E: torch.Tensor

the bond feature matrix

edge_index: torch.Tensor

an tensor of shape 2 x E containing the edges of the graph in COO format

rev_edge_index: torch.Tensor

A tensor of shape E that maps from an edge index to the index of the source of the reverse edge in the edge_index attribute.

batch: torch.Tensor

the index of the parent MolGraph in the batched graph

__post_init__(mgs)

Parameters:

mgs (Sequence[chemprop.data.molgraph.MolGraph])

__len__()

the number of individual MolGraphs in this batch

Return type:

int

to(device)

Parameters:

device (str | torch.device)

class chemprop.data.collate.TrainingBatch

Bases: NamedTuple

bmg: BatchMolGraph | chemprop.featurizers.molgraph.molecule.BatchCuikMolGraph

V_d: torch.Tensor | None

X_d: torch.Tensor | None

Y: torch.Tensor | None

w: torch.Tensor

lt_mask: torch.Tensor | None

gt_mask: torch.Tensor | None

chemprop.data.collate.collate_batch(batch)

Parameters:

batch (Iterable[chemprop.data.datasets.Datum])

Return type:

TrainingBatch

chemprop.data.collate.collate_cuik_batch(batch)

Parameters:

batch (chemprop.data.datasets.CuikBatchedDatum)

Return type:

TrainingBatch

class chemprop.data.collate.BatchMolAtomBondGraph

Bases: BatchMolGraph

A BatchMolGraph represents a batch of individual MolGraphs.

It has all the attributes of a MolGraph with the addition of the batch attribute. This class is intended for use with data loading, so it uses Tensors to store data

bond_batch: torch.Tensor

A tensor of indices that show which MolGraph each bond belongs to in the batch

__post_init__(mgs)

Parameters:

mgs (Sequence[chemprop.data.molgraph.MolGraph])

to(device)

class chemprop.data.collate.MolAtomBondTrainingBatch

Bases: NamedTuple

bmg: BatchMolAtomBondGraph

V_d: torch.Tensor | None

E_d: torch.Tensor | None

X_d: torch.Tensor | None

Ys: tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]

w: tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]

lt_masks: tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]

gt_masks: tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]

constraints: tuple[torch.Tensor | None, torch.Tensor | None]

chemprop.data.collate.collate_mol_atom_bond_batch(batch)

Parameters:

batch (Iterable[chemprop.data.datasets.MolAtomBondDatum])

Return type:

MolAtomBondTrainingBatch

class chemprop.data.collate.MulticomponentTrainingBatch

Bases: NamedTuple

bmgs: list[BatchMolGraph]

V_ds: list[torch.Tensor | None]

X_d: torch.Tensor | None

Y: torch.Tensor | None

w: torch.Tensor

lt_mask: torch.Tensor | None

gt_mask: torch.Tensor | None

chemprop.data.collate.collate_multicomponent(batches)

Parameters:

batches (Iterable[Iterable[chemprop.data.datasets.Datum]])

Return type:

MulticomponentTrainingBatch