Spaces:
Running on Zero
Running on Zero
File size: 4,449 Bytes
6fe4366 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """Data utilities for OT diagnostics: dataset context, graph collection, and noise init."""
from __future__ import annotations
from typing import List, Tuple
import torch
from gem import utils
from omegaconf import DictConfig
from gem.datasets.dataset_context import build_dataset_context
from gem.models.extra_features import DummyExtraFeatures, ExtraFeatures
from gem.models.extra_features_molecular import ExtraMolecularFeatures
from gem.flow_matching.noise_distribution import NoiseDistribution
from gem.flow_matching import flow_matching_utils
def build_data_context(cfg: DictConfig):
"""Create datamodule, dataset_infos, and feature builders similar to train/eval."""
datamodule, dataset_infos, dataset_smiles = build_dataset_context(cfg)
extra_features = ExtraFeatures(
cfg.model.extra_features,
cfg.model.rrwp_steps,
dataset_info=dataset_infos,
)
domain_features = (
ExtraMolecularFeatures(dataset_infos=dataset_infos)
if dataset_smiles is not None
else DummyExtraFeatures()
)
dataset_infos.compute_input_output_dims(
datamodule=datamodule,
extra_features=extra_features,
domain_features=domain_features,
)
return datamodule, dataset_infos
def collect_graphs_from_data(datamodule, M: int) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
"""Collect at least M graphs from the training loader as discrete label tensors.
Returns two lists of length >= M:
- node label tensors of shape (n,)
- edge type tensors of shape (n, n)
"""
node_list: List[torch.Tensor] = []
edge_list: List[torch.Tensor] = []
loader = datamodule.train_dataloader()
it = iter(loader)
while len(node_list) < M:
try:
batch = next(it)
except StopIteration:
it = iter(datamodule.train_dataloader())
batch = next(it)
dense_data, node_mask = utils.to_dense(
batch.x, batch.edge_index, batch.edge_attr, batch.batch
)
graphs = dense_data.mask(node_mask, collapse=True).split(node_mask)
for g in graphs:
node_list.append(g.X.long().cpu())
edge_list.append(g.E.long().cpu())
if len(node_list) >= M:
break
return node_list, edge_list
def train_node_count_distribution(datamodule) -> torch.Tensor:
"""Compute distribution of number of nodes per graph over the training dataset."""
max_nodes = 0
for batch in datamodule.train_dataloader():
_, counts = torch.unique(batch.batch, return_counts=True)
if counts.numel() > 0:
max_nodes = max(max_nodes, int(counts.max().item()))
if max_nodes == 0:
return torch.tensor([1.0])
counts_hist = torch.zeros(max_nodes + 1, dtype=torch.float32)
total = 0.0
for batch in datamodule.train_dataloader():
_, counts = torch.unique(batch.batch, return_counts=True)
for c in counts:
n = int(c.item())
counts_hist[n] += 1.0
total += 1.0
if total > 0:
counts_hist = counts_hist / total
return counts_hist
def initialize_random_graphs_with_counts(
counts: List[int],
dataset_info,
device: torch.device = torch.device("cpu"),
transition: str = "marginal",
) -> List[Tuple[torch.Tensor, torch.Tensor]]:
"""Sample random graphs with exact node counts per graph using the noise distribution."""
if not counts:
return []
B = len(counts)
n_max = max(int(c) for c in counts)
if n_max <= 0:
return []
n_nodes = torch.tensor(counts, device=device, dtype=torch.long)
arange = torch.arange(n_max, device=device).unsqueeze(0).expand(B, -1)
node_mask = arange < n_nodes.unsqueeze(1)
noise_dist = NoiseDistribution(transition, dataset_info)
limit_dist = noise_dist.get_limit_dist()
z_T = flow_matching_utils.sample_discrete_feature_noise(
limit_dist=limit_dist, node_mask=node_mask
)
graphs: List[Tuple[torch.Tensor, torch.Tensor]] = []
for i in range(B):
n = int(counts[i])
if n <= 0:
graphs.append((torch.empty(0, dtype=torch.long), torch.empty(0, 0, dtype=torch.long)))
continue
node_types = torch.argmax(z_T.X[i, :n], dim=-1)
edge_types = torch.argmax(z_T.E[i, :n, :n], dim=-1)
graphs.append((node_types.cpu(), edge_types.cpu()))
return graphs
|