import torch import torch.nn as nn import torch.nn.functional as F import pytorch_lightning as pl import time import wandb import os import random import pickle import numpy as np from src.models.transformer_model import GraphTransformer from src.diffusion.noise_schedule import DiscreteUniformTransition, PredefinedNoiseScheduleDiscrete, \ MarginalUniformTransition from src.diffusion import diffusion_utils from src.metrics.train_metrics import TrainLossDiscrete from src.metrics.abstract_metrics import SumExceptBatchMetric, SumExceptBatchKL, NLL import src.utils from torch_geometric.utils import to_dense_batch from src.datasets.schenker_dataset import SchenkerDiffHeteroGraphData from src.rule_guidance import ParallelChecker, DissonanceChecker from src.schenker_gnn.for_diffusion.infer_structure_from_rhythm import load_score, extract_structure_sparse from src.schenker_gnn.config import DEVICE from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts class CustomUnpickler(pickle.Unpickler): def persistent_load(self, pid): # Here you can implement logic to handle the persistent id. # For instance, if you don't need to resolve external references, # you can simply return the pid or raise an informative error. # In this example, we'll simply return pid. return None def load_pickle_with_persistent(path): with open(path, 'rb') as f: return CustomUnpickler(f).load() class DiscreteDenoisingDiffusion(pl.LightningModule): def __init__(self, cfg, dataset_infos, train_metrics, sampling_metrics, visualization_tools, extra_features, domain_features): super().__init__() input_dims = dataset_infos.input_dims output_dims = dataset_infos.output_dims nodes_dist = dataset_infos.nodes_dist self.cfg = cfg self.name = cfg.general.name self.model_dtype = torch.float32 self.T = cfg.model.diffusion_steps self.Xdim = input_dims['X'] self.Edim = input_dims['E'] self.ydim = input_dims['y'] self.rdim = input_dims['r'] self.Xdim_output = output_dims['X'] self.Edim_output = output_dims['E'] self.ydim_output = output_dims['y'] self.rdim_output = output_dims['r'] self.node_dist = nodes_dist self.dataset_info = dataset_infos self.train_loss = TrainLossDiscrete(self.cfg.model.lambda_train) self.val_nll = NLL() self.val_X_kl = SumExceptBatchKL() self.val_E_kl = SumExceptBatchKL() self.val_X_logp = SumExceptBatchMetric() self.val_E_logp = SumExceptBatchMetric() self.test_nll = NLL() self.test_X_kl = SumExceptBatchKL() self.test_E_kl = SumExceptBatchKL() self.test_X_logp = SumExceptBatchMetric() self.test_E_logp = SumExceptBatchMetric() self.train_metrics = train_metrics self.sampling_metrics = sampling_metrics self.visualization_tools = visualization_tools self.extra_features = extra_features self.domain_features = domain_features self.model = GraphTransformer(n_layers=cfg.model.n_layers, input_dims=input_dims, hidden_mlp_dims=cfg.model.hidden_mlp_dims, hidden_dims=cfg.model.hidden_dims, output_dims=output_dims, act_fn_in=nn.ReLU(), act_fn_out=nn.ReLU()) self.noise_schedule = PredefinedNoiseScheduleDiscrete(cfg.model.diffusion_noise_schedule, timesteps=cfg.model.diffusion_steps) if cfg.model.transition == 'uniform': self.transition_model = DiscreteUniformTransition(x_classes=self.Xdim_output, e_classes=self.Edim_output, y_classes=self.ydim_output) x_limit = torch.ones(self.Xdim_output) / self.Xdim_output e_limit = torch.ones(self.Edim_output) / self.Edim_output y_limit = torch.ones(self.ydim_output) / self.ydim_output self.limit_dist = src.utils.PlaceHolder(X=x_limit, E=e_limit, y=y_limit) elif cfg.model.transition == 'marginal': node_types = self.dataset_info.node_types.float() + 1e-6 # adding small constant to avoid zeros x_marginals = node_types / torch.sum(node_types) # x_marginals = torch.softmax(x_marginals, dim = 1) # making sure everything still adds up to 1 edge_types = self.dataset_info.edge_types.float() e_marginals = edge_types / torch.sum(edge_types) print(f"Marginal distribution of the classes: {x_marginals} for nodes, {e_marginals} for edges") self.transition_model = MarginalUniformTransition(x_marginals=x_marginals, e_marginals=e_marginals, y_classes=self.ydim_output) self.limit_dist = src.utils.PlaceHolder(X=x_marginals, E=e_marginals, y=torch.ones(self.ydim_output) / self.ydim_output) self.save_hyperparameters(ignore=['train_metrics', 'sampling_metrics']) self.start_epoch_time = None self.train_iterations = None self.val_iterations = None self.log_every_steps = cfg.general.log_every_steps self.number_chain_steps = cfg.general.number_chain_steps self.best_val_nll = 1e8 self.val_counter = 0 def training_step(self, data, i): if data.edge_index.numel() == 0: self.print("Found a batch with no edges. Skipping.") return dense_data, node_mask = src.utils.to_dense(data.x, data.edge_index, data.edge_attr, data.batch) dense_data = dense_data.mask(node_mask) X, E = dense_data.X, dense_data.E noisy_data = self.apply_noise(X, E, data.y, node_mask) extra_data = self.compute_extra_data(noisy_data) # extra_data = src.utils.PlaceHolder(X=data.r, E=data.s, y=data.y) dense_r, r_mask = to_dense_batch(data.r, data.batch) # extra_data.X = dense_r pred = self.forward(noisy_data, extra_data, dense_r, node_mask) loss = self.train_loss(masked_pred_X=pred.X, masked_pred_E=pred.E, pred_y=pred.y, true_X=X, true_E=E, true_y=data.y, log=i % self.log_every_steps == 0) self.train_metrics(masked_pred_X=pred.X, masked_pred_E=pred.E, true_X=X, true_E=E, log=i % self.log_every_steps == 0) return {'loss': loss} # def configure_optimizers(self): # return torch.optim.AdamW(self.parameters(), lr=self.cfg.train.lr, amsgrad=True, # weight_decay=self.cfg.train.weight_decay) def configure_optimizers(self): # 1) build your optimizer as before optimizer = torch.optim.AdamW( self.parameters(), lr=self.cfg.train.lr, amsgrad=True, weight_decay=self.cfg.train.weight_decay ) # 2) wrap it in a CosineAnnealingWarmRestarts scheduler scheduler = CosineAnnealingWarmRestarts( optimizer, T_0=self.cfg.train.T_0, # epochs until first restart T_mult=self.cfg.train.T_mult, # multiply T_i by this after each restart eta_min=getattr(self.cfg.train, 'min_lr', 0.0) # verbose='deprecated' ) # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( # optimizer, # T_max=self.cfg.train.T_max, # eta_min=0, # verbose='deprecated' # ) # scheduler = torch.optim.lr_scheduler.OneCycleLR( # optimizer, # max_lr=self.cfg.train.lr, # epochs=self.cfg.train.n_epochs, # steps_per_epoch=self.trainer.estimated_stepping_batches # ) # 3) return both in Lightning’s dict format return { 'optimizer': optimizer, 'lr_scheduler': { 'scheduler': scheduler, 'interval': 'epoch', # calls scheduler.step() every epoch 'frequency': 1, 'name': 'cosine_restart' } } def on_fit_start(self) -> None: self.train_iterations = len(self.trainer.datamodule.train_dataloader()) self.print("Size of the input features", self.Xdim, self.Edim, self.ydim) if self.local_rank == 0: src.utils.setup_wandb(self.cfg) def on_train_epoch_start(self) -> None: self.print("Starting train epoch...") self.start_epoch_time = time.time() # self.train_loss.reset() # self.train_metrics.reset() def on_train_epoch_end(self) -> None: to_log = self.train_loss.log_epoch_metrics() self.print(f"Epoch {self.current_epoch}: X_CE: {to_log['train_epoch/x_CE'] :.3f}" f" -- E_CE: {to_log['train_epoch/E_CE'] :.3f} --" # f" y_CE: {to_log['train_epoch/y_CE'] :.3f}" f" -- {time.time() - self.start_epoch_time:.1f}s ") # epoch_at_metrics, epoch_bond_metrics = self.train_metrics.log_epoch_metrics() # self.print(f"Epoch {self.current_epoch}: {epoch_at_metrics} -- {epoch_bond_metrics}") if torch.cuda.is_available(): # print(torch.cuda.memory_summary()) pass else: print("CUDA is not available. Skipping memory summary.") def on_validation_epoch_start(self) -> None: self.val_nll.reset() self.val_X_kl.reset() self.val_E_kl.reset() self.val_X_logp.reset() self.val_E_logp.reset() self.sampling_metrics.reset() def validation_step(self, data, i): dense_data, node_mask = src.utils.to_dense(data.x, data.edge_index, data.edge_attr, data.batch) dense_data = dense_data.mask(node_mask) noisy_data = self.apply_noise(dense_data.X, dense_data.E, data.y, node_mask) extra_data = self.compute_extra_data(noisy_data) # extra_data = src.utils.PlaceHolder(X=data.r, E=data.s, y=data.y) dense_r, r_mask = to_dense_batch(data.r, data.batch) pred = self.forward(noisy_data, extra_data, dense_r, node_mask) nll = self.compute_val_loss(pred, noisy_data, dense_data.X, dense_data.E, dense_r, data.y, node_mask, test=False) return {'loss': nll} def on_validation_epoch_end(self) -> None: metrics = [self.val_nll.compute(), self.val_X_kl.compute() * self.T, self.val_E_kl.compute() * self.T, self.val_X_logp.compute(), self.val_E_logp.compute()] if wandb.run: wandb.log({"val/epoch_NLL": metrics[0], "val/X_kl": metrics[1], "val/E_kl": metrics[2], "val/X_logp": metrics[3], "val/E_logp": metrics[4]}, commit=False) self.print(f"Epoch {self.current_epoch}: Val NLL {metrics[0] :.2f} -- Val Atom type KL {metrics[1] :.2f} -- ", f"Val Edge type KL: {metrics[2] :.2f}") # Log val nll with default Lightning logger, so it can be monitored by checkpoint callback val_nll = metrics[0] self.log("val/epoch_NLL", val_nll, sync_dist=True) if val_nll < self.best_val_nll: self.best_val_nll = val_nll self.print('Val loss: %.4f \t Best val loss: %.4f\n' % (val_nll, self.best_val_nll)) self.val_counter += 1 if self.val_counter % self.cfg.general.sample_every_val == 0: start = time.time() samples_left_to_generate = self.cfg.general.samples_to_generate samples_left_to_save = self.cfg.general.samples_to_save chains_left_to_save = self.cfg.general.chains_to_save samples = [] ident = 0 while samples_left_to_generate > 0: bs = 2 * self.cfg.train.batch_size to_generate = min(samples_left_to_generate, bs) to_save = min(samples_left_to_save, bs) chains_save = min(chains_left_to_save, bs) samples.extend(self.sample_batch(batch_id=ident, batch_size=to_generate, num_nodes=None, save_final=to_save, keep_chain=chains_save, number_chain_steps=self.number_chain_steps)) ident += to_generate samples_left_to_save -= to_save samples_left_to_generate -= to_generate chains_left_to_save -= chains_save self.print("Computing sampling metrics...") self.sampling_metrics.forward(samples, self.name, self.current_epoch, val_counter=-1, test=False, local_rank=self.local_rank) self.print(f'Done. Sampling took {time.time() - start:.2f} seconds\n') print("Validation epoch end ends...") def on_test_epoch_start(self) -> None: self.print("Starting test...") self.test_nll.reset() self.test_X_kl.reset() self.test_E_kl.reset() self.test_X_logp.reset() self.test_E_logp.reset() if self.local_rank == 0: src.utils.setup_wandb(self.cfg) def test_step(self, data, i): dense_data, node_mask = src.utils.to_dense(data.x, data.edge_index, data.edge_attr, data.batch) dense_data = dense_data.mask(node_mask) noisy_data = self.apply_noise(dense_data.X, dense_data.E, data.y, node_mask) extra_data = self.compute_extra_data(noisy_data) dense_r, r_mask = to_dense_batch(data.r, data.batch) # extra_data.X = dense_r pred = self.forward(noisy_data, extra_data, dense_r, node_mask) nll = self.compute_val_loss(pred, noisy_data, dense_data.X, dense_data.E, dense_r, data.y, node_mask, test=True) return {'loss': nll} def on_test_epoch_end(self) -> None: """ Measure likelihood on a test set and compute stability metrics. """ metrics = [self.test_nll.compute(), self.test_X_kl.compute(), self.test_E_kl.compute(), self.test_X_logp.compute(), self.test_E_logp.compute()] if wandb.run: wandb.log({"test/epoch_NLL": metrics[0], "test/X_kl": metrics[1], "test/E_kl": metrics[2], "test/X_logp": metrics[3], "test/E_logp": metrics[4]}, commit=False) self.print(f"Epoch {self.current_epoch}: Test NLL {metrics[0] :.2f} -- Test Atom type KL {metrics[1] :.2f} -- ", f"Test Edge type KL: {metrics[2] :.2f}") test_nll = metrics[0] if wandb.run: wandb.log({"test/epoch_NLL": test_nll}, commit=False) self.print(f'Test loss: {test_nll :.4f}') samples_left_to_generate = self.cfg.general.final_model_samples_to_generate samples_left_to_save = self.cfg.general.final_model_samples_to_save chains_left_to_save = self.cfg.general.final_model_chains_to_save samples = [] id = 0 while samples_left_to_generate > 0: self.print(f'Samples left to generate: {samples_left_to_generate}/' f'{self.cfg.general.final_model_samples_to_generate}', end='', flush=True) bs = 2 * self.cfg.train.batch_size to_generate = min(samples_left_to_generate, bs) to_save = min(samples_left_to_save, bs) chains_save = min(chains_left_to_save, bs) samples.extend(self.sample_batch(id, to_generate, num_nodes=None, save_final=to_save, keep_chain=chains_save, number_chain_steps=self.number_chain_steps)) id += to_generate samples_left_to_save -= to_save samples_left_to_generate -= to_generate chains_left_to_save -= chains_save self.print("Saving the generated graphs") filename = f'generated_samples1.txt' for i in range(2, 10): if os.path.exists(filename): filename = f'generated_samples{i}.txt' else: break with open(filename, 'w') as f: for item in samples: f.write(f"N={item[0].shape[0]}\n") atoms = item[0].tolist() f.write("X: \n") for at in atoms: f.write(f"{at} ") f.write("\n") f.write("E: \n") for bond_list in item[1]: for bond in bond_list: f.write(f"{int(bond)} ") f.write("\n") f.write("R: \n") for r_list in item[2]: for r in r_list: f.write(f"{r} ") f.write("\n") for name in item[3]: f.write(name) f.write("\n") f.write("\n") self.print("Generated graphs Saved. Computing sampling metrics...") self.sampling_metrics(samples, self.name, self.current_epoch, self.val_counter, test=True, local_rank=self.local_rank) self.print("Done testing.") def kl_prior(self, X, E, node_mask): """Computes the KL between q(z1 | x) and the prior p(z1) = Normal(0, 1). This is essentially a lot of work for something that is in practice negligible in the loss. However, you compute it so that you see it when you've made a mistake in your noise schedule. """ # Compute the last alpha value, alpha_T. ones = torch.ones((X.size(0), 1), device=X.device) Ts = self.T * ones alpha_t_bar = self.noise_schedule.get_alpha_bar(t_int=Ts) # (bs, 1) Qtb = self.transition_model.get_Qt_bar(alpha_t_bar, self.device) # Compute transition probabilities probX = X @ Qtb.X # (bs, n, dx_out) probE = E @ Qtb.E.unsqueeze(1) # (bs, n, n, de_out) assert probX.shape == X.shape bs, n, _ = probX.shape limit_X = self.limit_dist.X[None, None, :].expand(bs, n, -1).type_as(probX) limit_E = self.limit_dist.E[None, None, None, :].expand(bs, n, n, -1).type_as(probE) # Make sure that masked rows do not contribute to the loss limit_dist_X, limit_dist_E, probX, probE = diffusion_utils.mask_distributions(true_X=limit_X.clone(), true_E=limit_E.clone(), pred_X=probX, pred_E=probE, node_mask=node_mask) kl_distance_X = F.kl_div(input=probX.log(), target=limit_dist_X, reduction='none') # kl_distance_E = F.kl_div(input=probE.log(), target=limit_dist_E, reduction='none') return diffusion_utils.sum_except_batch(kl_distance_X) #+ \ # diffusion_utils.sum_except_batch(kl_distance_E) def compute_Lt(self, X, E, y, pred, noisy_data, node_mask, test): pred_probs_X = F.softmax(pred.X, dim=-1) pred_probs_E = F.softmax(pred.E, dim=-1) pred_probs_y = F.softmax(pred.y, dim=-1) Qtb = self.transition_model.get_Qt_bar(noisy_data['alpha_t_bar'], self.device) Qsb = self.transition_model.get_Qt_bar(noisy_data['alpha_s_bar'], self.device) Qt = self.transition_model.get_Qt(noisy_data['beta_t'], self.device) # Compute distributions to compare with KL bs, n, d = X.shape prob_true = diffusion_utils.posterior_distributions(X=X, E=E, y=y, X_t=noisy_data['X_t'], E_t=noisy_data['E_t'], y_t=noisy_data['y_t'], Qt=Qt, Qsb=Qsb, Qtb=Qtb) prob_true.E = prob_true.E.reshape((bs, n, n, -1)) prob_pred = diffusion_utils.posterior_distributions(X=pred_probs_X, E=pred_probs_E, y=pred_probs_y, X_t=noisy_data['X_t'], E_t=noisy_data['E_t'], y_t=noisy_data['y_t'], Qt=Qt, Qsb=Qsb, Qtb=Qtb) prob_pred.E = prob_pred.E.reshape((bs, n, n, -1)) # Reshape and filter masked rows prob_true_X, prob_true_E, prob_pred.X, prob_pred.E = diffusion_utils.mask_distributions(true_X=prob_true.X, true_E=prob_true.E, pred_X=prob_pred.X, pred_E=prob_pred.E, node_mask=node_mask) kl_x = (self.test_X_kl if test else self.val_X_kl)(prob_true.X, torch.log(prob_pred.X)) # kl_e = (self.test_E_kl if test else self.val_E_kl)(prob_true.E, torch.log(prob_pred.E)) # kl_e = 0 return self.T * (kl_x ) #+ kl_e) def reconstruction_logp(self, t, X, E, r, node_mask): # Compute noise values for t = 0. t_zeros = torch.zeros_like(t) beta_0 = self.noise_schedule(t_zeros) Q0 = self.transition_model.get_Qt(beta_t=beta_0, device=self.device) probX0 = X @ Q0.X # (bs, n, dx_out) probE0 = E @ Q0.E.unsqueeze(1) # (bs, n, n, de_out) sampled0 = diffusion_utils.sample_discrete_features(probX=probX0, probE=probE0, node_mask=node_mask) X0 = F.one_hot(sampled0.X, num_classes=self.Xdim_output).float() E0 = F.one_hot(sampled0.E, num_classes=self.Edim_output).float() y0 = sampled0.y assert (X.shape == X0.shape) and (E.shape == E0.shape) sampled_0 = src.utils.PlaceHolder(X=X0, E=E0, y=y0).mask(node_mask) # Predictions noisy_data = {'X_t': sampled_0.X, 'E_t': sampled_0.E, 'y_t': sampled_0.y, 'node_mask': node_mask, 't': torch.zeros(X0.shape[0], 1).type_as(y0)} extra_data = self.compute_extra_data(noisy_data) pred0 = self.forward(noisy_data, extra_data, r, node_mask) # Normalize predictions probX0 = F.softmax(pred0.X, dim=-1) probE0 = F.softmax(pred0.E, dim=-1) proby0 = F.softmax(pred0.y, dim=-1) # Set masked rows to arbitrary values that don't contribute to loss probX0[~node_mask] = torch.ones(self.Xdim_output).type_as(probX0) probE0[~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2))] = torch.ones(self.Edim_output).type_as(probE0) diag_mask = torch.eye(probE0.size(1)).type_as(probE0).bool() diag_mask = diag_mask.unsqueeze(0).expand(probE0.size(0), -1, -1) probE0[diag_mask] = torch.ones(self.Edim_output).type_as(probE0) return src.utils.PlaceHolder(X=probX0, E=probE0, y=proby0) def apply_noise(self, X, E, y, node_mask): """ Sample noise and apply it to the data. """ # Sample a timestep t. # When evaluating, the loss for t=0 is computed separately lowest_t = 0 if self.training else 1 t_int = torch.randint(lowest_t, self.T + 1, size=(X.size(0), 1), device=X.device).float() # (bs, 1) s_int = t_int - 1 t_float = t_int / self.T s_float = s_int / self.T # beta_t and alpha_s_bar are used for denoising/loss computation beta_t = self.noise_schedule(t_normalized=t_float) # (bs, 1) alpha_s_bar = self.noise_schedule.get_alpha_bar(t_normalized=s_float) # (bs, 1) alpha_t_bar = self.noise_schedule.get_alpha_bar(t_normalized=t_float) # (bs, 1) Qtb = self.transition_model.get_Qt_bar(alpha_t_bar, device=self.device) # (bs, dx_in, dx_out), (bs, de_in, de_out) assert (abs(Qtb.X.sum(dim=2) - 1.) < 1e-4).all(), Qtb.X.sum(dim=2) - 1 assert (abs(Qtb.E.sum(dim=2) - 1.) < 1e-4).all() # Compute transition probabilities probX = X @ Qtb.X # (bs, n, dx_out) probE = E @ Qtb.E.unsqueeze(1) # (bs, n, n, de_out) sampled_t = diffusion_utils.sample_discrete_features(probX=probX, probE=probE, node_mask=node_mask) X_t = F.one_hot(sampled_t.X, num_classes=self.Xdim_output) E_t = F.one_hot(sampled_t.E, num_classes=self.Edim_output) assert (X.shape == X_t.shape) and (E.shape == E_t.shape) z_t = src.utils.PlaceHolder(X=X_t, E=E_t, y=y).type_as(X_t).mask(node_mask) noisy_data = {'t_int': t_int, 't': t_float, 'beta_t': beta_t, 'alpha_s_bar': alpha_s_bar, 'alpha_t_bar': alpha_t_bar, 'X_t': z_t.X, 'E_t': z_t.E, 'y_t': z_t.y, 'node_mask': node_mask} return noisy_data def compute_val_loss(self, pred, noisy_data, X, E, r, y, node_mask, test=False): """Computes an estimator for the variational lower bound. pred: (batch_size, n, total_features) noisy_data: dict X, E, r, y : (bs, n, dx), (bs, n, n, de), (bs, n, dr), (bs, dy) node_mask : (bs, n) Output: nll (size 1) """ t = noisy_data['t'] # 1. N = node_mask.sum(1).long() log_pN = self.node_dist.log_prob(N) # 2. The KL between q(z_T | x) and p(z_T) = Uniform(1/num_classes). Should be close to zero. kl_prior = self.kl_prior(X, E, node_mask) # 3. Diffusion loss loss_all_t = self.compute_Lt(X, E, y, pred, noisy_data, node_mask, test) # 4. Reconstruction loss # Compute L0 term : -log p (X, E, y | z_0) = reconstruction loss prob0 = self.reconstruction_logp(t, X, E, r, node_mask) loss_term_0 = self.val_X_logp(X * prob0.X.log()) #+ self.val_E_logp(E * prob0.E.log()) # Combine terms nlls = - log_pN + kl_prior + loss_all_t - loss_term_0 assert len(nlls.shape) == 1, f'{nlls.shape} has more than only batch dim.' # Update NLL metric object and return batch nll nll = (self.test_nll if test else self.val_nll)(nlls) # Average over the batch accuracy_X = (pred.X == X).float().mean() if wandb.run: wandb.log({"kl prior": kl_prior.mean(), "Estimator loss terms": loss_all_t.mean(), "log_pn": log_pN.mean(), "loss_term_0": loss_term_0, "X accuracy": accuracy_X, 'batch_test_nll' if test else 'val_nll': nll}, commit=False) return nll def forward(self, noisy_data, extra_data, rhythmic_data, node_mask): X = torch.cat((noisy_data['X_t'], extra_data.X), dim=2).float() E = torch.cat((noisy_data['E_t'], extra_data.E), dim=3).float() y = torch.hstack((noisy_data['y_t'], extra_data.y)).float() return self.model(X, E, rhythmic_data, y, node_mask) @torch.no_grad() def sample_batch(self, batch_id: int, batch_size: int, keep_chain: int, number_chain_steps: int, save_final: int, num_nodes=None, use_rules=True): """ :param batch_id: int :param batch_size: int :param num_nodes: int, tensor (batch_size) (optional) for specifying number of nodes :param save_final: int: number of predictions to save to file :param keep_chain: int: number of chains to save to file :param keep_chain_steps: number of timesteps to save for each chain :return: molecule_list. Each element of this list is a tuple (atom_types, charges, positions) """ E, r, names, n_nodes_list = self.sample_r_E(batch_size) num_nodes = torch.tensor([int(x) for x in n_nodes_list]).to(self.device) if num_nodes is None: n_nodes = self.node_dist.sample_n(batch_size, self.device) elif type(num_nodes) == int: n_nodes = num_nodes * torch.ones(batch_size, device=self.device, dtype=torch.int) else: assert isinstance(num_nodes, torch.Tensor) n_nodes = num_nodes n_max = torch.max(n_nodes).item() # Build the masks arange = torch.arange(n_max, device=self.device).unsqueeze(0).expand(batch_size, -1) node_mask = arange < n_nodes.unsqueeze(1) # Sample a piece, and use the R matrix from that # Get a random sample from the data # pass through Stephen's script to get the S matrix, and the R matrix through the data processing (process_file_for_GUI) # Sample noise -- z has size (n_samples, n_nodes, n_features) z_T = diffusion_utils.sample_discrete_feature_noise(limit_dist=self.limit_dist, node_mask=node_mask) X, _, y = z_T.X, z_T.E, z_T.y E_transpose = E.permute(0, 2, 1, 3) # Shape remains (bs, n_nodes, n_nodes, num_edges) # Symmetrize using max operation (ensures strongest connection remains) # E = torch.maximum(E, E_transpose).to(DEVICE) r = r.to(DEVICE) E = E.to(DEVICE) # assert (E == torch.transpose(E, 1, 2)).all() assert number_chain_steps < self.T chain_X_size = torch.Size((number_chain_steps, keep_chain, X.size(1))) chain_E_size = torch.Size((number_chain_steps, keep_chain, E.size(1), E.size(2))) chain_X = torch.zeros(chain_X_size) chain_E = torch.zeros(chain_E_size) # Iteratively sample p(z_s | z_t) for t = 1, ..., T, with s = t - 1. for s_int in reversed(range(0, self.T)): s_array = s_int * torch.ones((batch_size, 1)).type_as(y) t_array = s_array + 1 s_norm = s_array / self.T t_norm = t_array / self.T # Sample z_s if use_rules: sampled_s, discrete_sampled_s = self.sample_p_zs_given_zt_with_rules(s_norm, t_norm, X, E, r, y, node_mask) else: sampled_s, discrete_sampled_s = self.sample_p_zs_given_zt(s_norm, t_norm, X, E, r, y, node_mask) X, _, y = sampled_s.X, sampled_s.E, sampled_s.y discrete_sampled_s_E, _ = self.apply_node_mask_E_r(E, r, node_mask) # Save the first keep_chain graphs write_index = (s_int * number_chain_steps) // self.T chain_X[write_index] = discrete_sampled_s.X[:keep_chain] chain_E[write_index] = discrete_sampled_s_E[:keep_chain] # Sample sampled_s = sampled_s.mask(node_mask, collapse=True) X, _, y = sampled_s.X, sampled_s.E, sampled_s.y E, _ = self.apply_node_mask_E_r(E, r, node_mask) # Prepare the chain for saving if keep_chain > 0: final_X_chain = X[:keep_chain] final_E_chain = E[:keep_chain] chain_X[0] = final_X_chain # Overwrite last frame with the resulting X, E chain_E[0] = final_E_chain chain_X = diffusion_utils.reverse_tensor(chain_X) chain_E = diffusion_utils.reverse_tensor(chain_E) # Repeat last frame to see final sample better chain_X = torch.cat([chain_X, chain_X[-1:].repeat(10, 1, 1)], dim=0) chain_E = torch.cat([chain_E, chain_E[-1:].repeat(10, 1, 1, 1)], dim=0) assert chain_X.size(0) == (number_chain_steps + 10) molecule_list = [] for i in range(batch_size): n = n_nodes[i] atom_types = X[i, :n].cpu() edge_types = E[i, :n, :n].cpu() rhythm_types = r[i, :n, :].cpu() sample_names = names[i] molecule_list.append([atom_types, edge_types, rhythm_types, sample_names]) # Visualize chains if self.visualization_tools is not None: self.print('Visualizing chains...') current_path = os.getcwd() num_molecules = chain_X.size(1) # number of molecules for i in range(num_molecules): result_path = os.path.join(current_path, f'chains/{self.cfg.general.name}/' f'epoch{self.current_epoch}/' f'chains/molecule_{batch_id + i}') if not os.path.exists(result_path): os.makedirs(result_path) _ = self.visualization_tools.visualize_chain(result_path, chain_X[:, i, :].numpy(), chain_E[:, i, :].numpy()) self.print('\r{}/{} complete'.format(i + 1, num_molecules), end='', flush=True) self.print('\nVisualizing molecules...') # Visualize the final molecules result_path = os.path.join(current_path, f'graphs/{self.name}/epoch{self.current_epoch}_b{batch_id}/') self.visualization_tools.visualize(result_path, molecule_list, save_final) self.print("Done.") return molecule_list @staticmethod def apply_node_mask_E_r(E, r, node_mask): """ Applies a node mask to both the adjacency tensor E and the node feature tensor r. Parameters: E (torch.Tensor): Adjacency tensor of shape (batch_size, n_nodes, n_nodes, 2). Represents edge attributes between nodes. r (torch.Tensor): Node feature tensor of shape (batch_size, n_nodes, dr). node_mask (torch.Tensor): Boolean tensor of shape (batch_size, n_nodes) where True indicates a valid node and False an invalid/padded node. Returns: E_masked (torch.Tensor): Masked adjacency tensor where an edge is retained only if both its source and target nodes are valid. Shape remains (batch_size, n_nodes, n_nodes, 2). r_masked (torch.Tensor): Masked node features where invalid nodes are zeroed out. Shape remains (batch_size, n_nodes, dr). How it works: - For r, the mask is expanded to match the last dimension (dr) and multiplied elementwise. - For E, the mask is applied to both the row and column dimensions. That is, for each batch sample, an edge (i,j) is kept only if both node i and node j are valid (i.e., their mask values are True). """ r_mask = node_mask.unsqueeze(-1) # bs, n, 1 e_mask1 = r_mask.unsqueeze(2) # bs, n, 1, 1 e_mask2 = r_mask.unsqueeze(1) r_out = torch.argmax(r, dim=-1) E_out = torch.argmax(E, dim=-1) r_out[node_mask == 0] = - 1 E_out[(e_mask1 * e_mask2).squeeze(-1) == 0] = - 1 return E_out, r_out @staticmethod def sample_r_E(batch_size): """ Samples `batch_size` random pickle files from the directory where i is a random integer between 0 and 1080. Each pickle file contains a dictionary that is converted to a PyG Data object using the pre-defined function `hetero_to_data`. The PyG Data object is expected to have at least the following attributes: - x: Tensor of node features with shape (num_nodes, feature_dim) - edge_index: LongTensor with shape (2, num_edges) - edge_attr: Tensor with shape (num_edges, 2) representing edge attributes - r: Tensor with shape (num_nodes, dr) representing additional node-level features For each sample, the function creates: - An adjacency tensor E_sample of shape (n_nodes, n_nodes, 2) where each edge's attribute is placed at the corresponding (u, v) location. If the original graph has fewer than `n_nodes` nodes, the tensors are padded with zeros; if it has more, they are truncated. - A node feature tensor r_sample of shape (n_nodes, dr) similarly padded or truncated. Finally, the function stacks these into: - E_tensor: Tensor of shape (batch_size, n_nodes, n_nodes, 2) - r_tensor: Tensor of shape (batch_size, n_nodes, dr) Returns: E_tensor, r_tensor """ E_list = [] r_list = [] name_list = [] node_sizes = [] # get samples from OOS distribution # [TODO]: take these from cfg files instead having magic numbers np.random.seed(42) n_samples = 350 # Randomly select 90 indices for the test set test_indices = np.random.choice(n_samples, 20, replace=False) for _ in range(batch_size): # Select a random index between 0 and 1080 (inclusive) idx = random.choice(test_indices) file_path = f"../../../SchenkerDiff/data/schenker/processed/heterdatacleaned/processed/{idx}_processed.pt" # Load the pickle file containing a dictionary data_dict = torch.load(file_path, weights_only=False) # Convert dictionary to a PyG Data object using the provided function data = SchenkerDiffHeteroGraphData.hetero_to_data(data_dict) # Determine the actual number of nodes in the current sample m = data.x.shape[0] # Initialize an adjacency tensor for this sample E_sample = torch.zeros((m, m, 30)) # Fill in the edge attributes: iterate over each edge for i in range(data.edge_index.shape[1]): u = data.edge_index[0, i].item() v = data.edge_index[1, i].item() # Only consider nodes within the allowed range (pad/truncate as needed) if u < m and v < m: E_sample[u, v, :] = data.edge_attr[i, :] # Process the r tensor (node-level additional features) dr = data.r.shape[1] # feature dimension of r r_sample = torch.zeros((m, dr)) # Copy available node features; pad with zeros if necessary or truncate if too many nodes r_sample[:m, :] = data.r[:m, :] # Append this sample's results to the lists E_list.append(E_sample) r_list.append(r_sample) name_list.append(data_dict['name']) node_sizes.append(m) # Stack all samples to form the batch tensors # Determine the maximum number of nodes in the batch max_nodes = max(tensor.shape[0] for tensor in r_list) # Pad the E_list tensors to shape (max_nodes, max_nodes, 3) E_padded = [] for e in E_list: n = e.shape[0] # F.pad expects pad in the format: (pad_last_dim_left, pad_last_dim_right, # pad_second_last_dim_left, pad_second_last_dim_right, ...) # For a tensor of shape (n, n, 3): pad last dimension (3) with (0,0), # second dimension with (0, max_nodes-n), and first dimension with (0, max_nodes-n). pad_amount = (0, 0, 0, max_nodes - n, 0, max_nodes - n) E_padded.append(F.pad(e, pad_amount)) # Stack the padded tensors along a new batch dimension E_tensor = torch.stack(E_padded, dim=0) # Shape: (batch_size, max_nodes, max_nodes, 3) # Pad the r_list tensors to shape (max_nodes, dr) r_padded = [] for r in r_list: n = r.shape[0] # For a tensor of shape (n, dr), pad the first dimension with (0, max_nodes-n) pad_amount = (0, 0, 0, max_nodes - n) r_padded.append(F.pad(r, pad_amount)) # Stack the padded tensors along the batch dimension r_tensor = torch.stack(r_padded, dim=0) # Shape: (batch_size, max_nodes, dr) return E_tensor, r_tensor, name_list, node_sizes def determine_prob_X_prob_E(self, s, t, X_t, E_t, r, y_t, node_mask): bs, n, dxs = X_t.shape beta_t = self.noise_schedule(t_normalized=t) # (bs, 1) alpha_s_bar = self.noise_schedule.get_alpha_bar(t_normalized=s) alpha_t_bar = self.noise_schedule.get_alpha_bar(t_normalized=t) # Retrieve transitions matrix Qtb = self.transition_model.get_Qt_bar(alpha_t_bar, self.device) Qsb = self.transition_model.get_Qt_bar(alpha_s_bar, self.device) Qt = self.transition_model.get_Qt(beta_t, self.device) # Neural net predictions noisy_data = {'X_t': X_t, 'E_t': E_t, 'y_t': y_t, 't': t, 'node_mask': node_mask} extra_data = self.compute_extra_data(noisy_data) pred = self.forward(noisy_data, extra_data, r, node_mask) # pred is the clean graph estimate (G_0) # Normalize predictions pred_X = F.softmax(pred.X, dim=-1) # bs, n, d0 pred_E = F.softmax(pred.E, dim=-1) # bs, n, n, d0 # This is the forward posterior q(G_t-1 | G_t, G_0) p_s_and_t_given_0_X = diffusion_utils.compute_batched_over0_posterior_distribution(X_t=X_t, Qt=Qt.X, Qsb=Qsb.X, Qtb=Qtb.X) p_s_and_t_given_0_E = diffusion_utils.compute_batched_over0_posterior_distribution(X_t=E_t, Qt=Qt.E, Qsb=Qsb.E, Qtb=Qtb.E) # Dim of these two tensors: bs, N, d0, d_t-1 weighted_X = pred_X.unsqueeze(-1) * p_s_and_t_given_0_X # bs, n, d0, d_t-1 unnormalized_prob_X = weighted_X.sum(dim=2) # bs, n, d_t-1 unnormalized_prob_X[torch.sum(unnormalized_prob_X, dim=-1) == 0] = 1e-5 # "slightly denoised sample probability" p_theta(G_t-1 | G_t) prob_X = unnormalized_prob_X / torch.sum(unnormalized_prob_X, dim=-1, keepdim=True) # bs, n, d_t-1 pred_E = pred_E.reshape((bs, -1, pred_E.shape[-1])) weighted_E = pred_E.unsqueeze(-1) * p_s_and_t_given_0_E # bs, N, d0, d_t-1 unnormalized_prob_E = weighted_E.sum(dim=-2) unnormalized_prob_E[torch.sum(unnormalized_prob_E, dim=-1) == 0] = 1e-5 prob_E = unnormalized_prob_E / torch.sum(unnormalized_prob_E, dim=-1, keepdim=True) prob_E = prob_E.reshape(bs, n, n, pred_E.shape[-1]) assert ((prob_X.sum(dim=-1) - 1).abs() < 1e-4).all() assert ((prob_E.sum(dim=-1) - 1).abs() < 1e-4).all() return prob_X, prob_E def sample_p_zs_given_zt_with_rules(self, s, t, X_t, E_t, r, y_t, node_mask, scg_kwargs=None): if scg_kwargs is None: scg_kwargs = { "num_samples": 8, "rules": [ParallelChecker, DissonanceChecker], "disallowed_intervals": [1, 2, 11, 5], "disallowed_parallels": [0, 7, 1, 2, 10, 11, 5] } # find distribution for p_theta(G^t-1 | G^t) prob_X, prob_E = self.determine_prob_X_prob_E(s, t, X_t, E_t, r, y_t, node_mask) candidates = [] scores = [] # Sample several slightly denoised graphs for i in range(scg_kwargs["num_samples"]): # Sample a slightly denoised graph sampled_s = diffusion_utils.sample_discrete_features(prob_X, prob_E, node_mask=node_mask) X_s = F.one_hot(sampled_s.X, num_classes=self.Xdim_output).float() E_s = F.one_hot(sampled_s.E, num_classes=self.Edim_output).float() # Estimate clean graph from slightly denoised graph noisy_data_i = {'X_t': X_s, 'E_t': E_s, 'y_t': y_t, 't': t - 1, 'node_mask': node_mask} extra_data_i = self.compute_extra_data(noisy_data_i) pred_i = self.forward(noisy_data_i, extra_data_i, r, node_mask) # Convert to probability distribution pred_i_X = F.softmax(pred_i.X, dim=-1) pred_i_E = F.softmax(pred_i.E, dim=-1) # Determine score based on all provided score classes. # Note that the rule classes are provided the probability distribution, not a sampled graph score = self.determine_overall_rule_score(pred_i_X, pred_i_E, r, scg_kwargs, num_X_classes=self.Xdim_output, num_E_classes=self.Edim_output) candidates.append((X_s, E_s)) scores.append(score) print(scores) best_idx = torch.argmax(torch.tensor(scores)) X_best, E_best = candidates[best_idx] out_one_hot = src.utils.PlaceHolder(X=X_best, E=E_best, y=torch.zeros(y_t.shape[0], 0)) out_discrete = src.utils.PlaceHolder(X=X_best, E=E_best, y=torch.zeros(y_t.shape[0], 0)) return out_one_hot.mask(node_mask).type_as(y_t), out_discrete.mask(node_mask, collapse=True).type_as(y_t) @staticmethod def determine_overall_rule_score(pred_i_X, pred_i_E, r, scg_kwargs, num_X_classes, num_E_classes): scores = [] for rule_class in scg_kwargs["rules"]: rule_instance = rule_class(pred_i_X, pred_i_E, r, num_X_classes, num_E_classes, scg_kwargs) scores.append(rule_instance.calculate_score()) sum_scores = torch.tensor(scores, dtype=torch.float32).sum() return torch.round(sum_scores).to(torch.int64).item() def sample_p_zs_given_zt(self, s, t, X_t, E_t, r, y_t, node_mask): """Samples from zs ~ p(zs | zt). Only used during sampling. if last_step, return the graph prediction as well""" prob_X, prob_E = self.determine_prob_X_prob_E(s, t, X_t, E_t, r, y_t, node_mask) sampled_s = diffusion_utils.sample_discrete_features(prob_X, prob_E, node_mask=node_mask) X_s = F.one_hot(sampled_s.X, num_classes=self.Xdim_output).float() E_s = F.one_hot(sampled_s.E, num_classes=self.Edim_output).float() assert (E_s == torch.transpose(E_s, 1, 2)).all() assert (X_t.shape == X_s.shape) and (E_t.shape == E_s.shape) out_one_hot = src.utils.PlaceHolder(X=X_s, E=E_s, y=torch.zeros(y_t.shape[0], 0)) out_discrete = src.utils.PlaceHolder(X=X_s, E=E_s, y=torch.zeros(y_t.shape[0], 0)) return out_one_hot.mask(node_mask).type_as(y_t), out_discrete.mask(node_mask, collapse=True).type_as(y_t) def compute_extra_data(self, noisy_data): """ At every training step (after adding noise) and step in sampling, compute extra information and append to the network input. """ extra_features = self.extra_features(noisy_data) extra_molecular_features = self.domain_features(noisy_data) extra_X = torch.cat((extra_features.X, extra_molecular_features.X), dim=-1) extra_E = torch.cat((extra_features.E, extra_molecular_features.E), dim=-1) extra_y = torch.cat((extra_features.y, extra_molecular_features.y), dim=-1) t = noisy_data['t'] extra_y = torch.cat((extra_y, t), dim=-1) return src.utils.PlaceHolder(X=extra_X, E=extra_E, y=extra_y)