Spaces:
Running on Zero
Running on Zero
| """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 | |