MolecularDiffusion.modules.models.syncogen.data.dataloader

Classes

FaultTolerantDistributedSampler
RandomFaultTolerantSampler
SyncogenDataManager
SyncogenDataset

Functions

expand_edge_indices_to_matrix(edge_index, edge_attr, ...)

Expand sparse edge indices to full (N, N) reaction index matrix.

Module Contents

class MolecularDiffusion.modules.models.syncogen.data.dataloader.FaultTolerantDistributedSampler(*args, **kwargs)

Bases: torch.utils.data.DistributedSampler

load_state_dict(state_dict)

state_dict()

counter = 0

restarting = False

class MolecularDiffusion.modules.models.syncogen.data.dataloader.RandomFaultTolerantSampler(*args, generator=None, **kwargs)

Bases: torch.utils.data.RandomSampler

load_state_dict(state_dict)

state_dict()

counter = 0

restarting = False

class MolecularDiffusion.modules.models.syncogen.data.dataloader.SyncogenDataManager(*, graphs_path: str | pathlib.Path, conformers_path: str | pathlib.Path | None = None, pharmacophore_path: str | pathlib.Path | None = None, train_size: float = 0.9, validation_size: float = 0.1, test_size: float = 0.0, overfit: bool = False, n_overfit: int | None = None, batch_size: int = 32, eval_batch_size: int = 64, num_workers: int = 4, pin_memory: bool = True, shuffle_train: bool = True, sample_conformer: bool = False, load_pharmacophores: bool = False, load_bonds: bool = False, coord_mask_value: float = 0.0, valid_seed: int | None = None, max_bbs=5)

ensure_train_lengths_loaded()

Loads the train lengths/probs from the lengths_path. Raises if path not present.

get_dataloaders() → Tuple[torch_geometric.loader.DataLoader, torch_geometric.loader.DataLoader]

get_graph_data_splits() → Dict[str, List[torch_geometric.data.Data]]

get_lengths_path() → pathlib.Path

get_split_cache_dir() → pathlib.Path

get_split_path(split_name: str) → pathlib.Path

sample_n_nodes(batch_size: int) → torch.Tensor | None

batch_size = 32

conformers_path = None

coord_mask_value = 0.0

eval_batch_size = 64

graphs_path

load_bonds = False

load_pharmacophores = False

max_atoms

max_bbs = 5

n_overfit = None

num_workers = 4

overfit = False

pharmacophore_path = None

pin_memory = True

sample_conformer = False

shuffle_train = True

test_size = 0.0

train_length_probs: torch.Tensor | None = None

train_length_values: torch.Tensor | None = None

train_size = 0.9

valid_seed = None

validation_size = 0.1

class MolecularDiffusion.modules.models.syncogen.data.dataloader.SyncogenDataset(conformers_path: str, data_list: List[torch_geometric.data.Data], pharmacophore_path: str | None = None, sample_conformer: bool = False, coord_mask_value: float = 0.0, load_pharmacophores: bool = False, load_bonds: bool = False)

Bases: torch_geometric.data.Dataset

get(idx: int)

len()

conformers_path

coord_mask_value = 0.0

data_list

load_bonds = False

load_pharmacophores = False

pharmacophore_path = None

sample_conformer = False

MolecularDiffusion.modules.models.syncogen.data.dataloader.expand_edge_indices_to_matrix(edge_index, edge_attr, n_nodes, device)

Expand sparse edge indices to full (N, N) reaction index matrix.