Add conditional sweeping

Author: amorehead

.gitignore

@@ -168,6 +168,7 @@ scripts/nautilus/persistent_storage.yaml
 
 bio-diffusion/
 logs/
+outputs/
 
 data/EDM/GEOM
 data/EDM/QM9

configs/mol_gen_eval_conditional_qm9.yaml

@@ -24,4 +24,5 @@ batch_size: 100
 debug_break: false
 sweep_property_values: false
 num_sweeps: 10
+experiment_name: ${.property}-conditioning
 output_dir: ""

src/models/__init__.py

@@ -10,6 +10,7 @@
 
 from functools import partial
 from matplotlib.lines import Line2D
+from omegaconf import DictConfig
 from torch.utils.data import DataLoader
 from torch.distributions.categorical import Categorical
 from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union
@@ -18,6 +19,8 @@
 from torchtyping import TensorType, patch_typeguard
 from typeguard import typechecked
 
+from src.models.components import save_xyz_file, visualize_mol_chain
+
 patch_typeguard()  # use before @typechecked
 
 HALT_FILE_EXTENSION = "done"
@@ -194,6 +197,70 @@ def log_grad_flow_full(
         wandb_run.log({"Gradient flow": plt})
 
 
+@typechecked
+def sample_sweep_conditionally(
+    model: nn.Module,
+    props_distr: object,
+    num_nodes: int = 19,
+    num_frames: int = 100
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    num_nodes_ = torch.tensor([num_nodes] * num_frames, device=model.device)
+
+    context = []
+    for key in props_distr.distributions:
+        min_val, max_val = props_distr.distributions[key][num_nodes]['params']
+        mean, mad = props_distr.normalizer[key]['mean'], props_distr.normalizer[key]['mad']
+        min_val = ((min_val - mean) / (mad)).cpu().numpy()
+        max_val = ((max_val - mean) / (mad)).cpu().numpy()
+        context_row = torch.tensor(np.linspace(min_val, max_val, num_frames)).unsqueeze(1)
+        context.append(context_row)
+    context = torch.cat(context, dim=-1).float().to(model.device)
+
+    x, one_hot, charges, batch_index = model.sample(
+        num_samples=num_frames,
+        num_nodes=num_nodes_,
+        context=context,
+        fix_noise=True
+    )
+    return x, one_hot, charges, batch_index
+
+
+@typechecked
+def save_and_sample_conditionally(
+    cfg: DictConfig,
+    model: nn.Module,
+    props_distr: object,
+    dataset_info: Dict[str, Any],
+    epoch: int = 0,
+    id_from: int = 0
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    x, one_hot, charges, batch_index = sample_sweep_conditionally(
+        model=model,
+        props_distr=props_distr
+    )
+
+    save_xyz_file(
+        path=f"outputs/{cfg.experiment_name}/analysis/run{epoch}/",
+        positions=x,
+        one_hot=one_hot,
+        charges=charges,
+        dataset_info=dataset_info,
+        id_from=id_from,
+        name="conditional",
+        batch_index=batch_index
+    )
+
+    visualize_mol_chain(
+        path=f"outputs/{cfg.experiment_name}/analysis/run{epoch}/",
+        dataset_info=dataset_info,
+        wandb_run=None,
+        spheres_3d=True,
+        mode="conditional"
+    )
+
+    return x, one_hot, charges
+
+
 class NumNodesDistribution(nn.Module):
     """
     Adapted from: https://github.com/ehoogeboom/e3_diffusion_for_molecules

src/models/components/variational_diffusion.py

@@ -829,8 +829,10 @@ def sample_normal(
     ) -> TensorType["batch_num_nodes", "num_x_dims_plus_num_node_scalar_features"]:
         """Sample from a Normal distribution."""
         if fix_noise:
-            raise NotImplementedError("The `fix_noise` option is currently not supported.")
-        eps = self.sample_combined_position_feature_noise(batch_index, node_mask, generate_x_only=generate_x_only)
+            batch_index_ = torch.zeros_like(batch_index)  # broadcast same noise across batch
+            eps = self.sample_combined_position_feature_noise(batch_index_, node_mask, generate_x_only=generate_x_only)
+        else:
+            eps = self.sample_combined_position_feature_noise(batch_index, node_mask, generate_x_only=generate_x_only)
         return mu + sigma[batch_index] * eps
 
     @typechecked
@@ -1317,7 +1319,12 @@ def mol_gen_sample(
             context = context * node_mask.float().unsqueeze(-1)
 
         # sample from the noise distribution (i.e., p(z_T))
-        z = self.sample_combined_position_feature_noise(batch_index, node_mask, generate_x_only=generate_x_only)
+        if fix_noise:
+            batch_index_ = torch.zeros_like(batch_index)  # broadcast same noise across batch
+            z = self.sample_combined_position_feature_noise(batch_index_, node_mask, generate_x_only=generate_x_only)
+        else:
+            z = self.sample_combined_position_feature_noise(batch_index, node_mask, generate_x_only=generate_x_only)
+
         self.assert_mean_zero_with_mask(z[:, :self.num_x_dims], node_mask)
 
         # iteratively sample p(z_s | z_t) for `t = 1, ..., T`, with `s = t - 1`.

src/models/geom_mol_gen_ddpm.py

@@ -589,8 +589,9 @@ def sample(
         num_nodes: Optional[TensorType["batch_size"]] = None,
         node_mask: Optional[TensorType["batch_num_nodes"]] = None,
         context: Optional[TensorType["batch_size", "num_context_features"]] = None,
+        fix_noise: bool = False,
         num_timesteps: Optional[int] = None
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         # node count-conditioning
         if num_nodes is None:
             num_nodes = self.ddpm.num_nodes_distribution.sample(num_samples)
@@ -605,7 +606,7 @@ def sample(
         context = None
 
         # sampling
-        xh, _, _ = self.ddpm.mol_gen_sample(
+        xh, batch_index, _ = self.ddpm.mol_gen_sample(
             num_samples=num_samples,
             num_nodes=num_nodes,
             node_mask=node_mask,
@@ -618,7 +619,7 @@ def sample(
         one_hot = xh[:, self.num_x_dims:-1] if self.include_charges else xh[:, self.num_x_dims:]
         charges = xh[:, -1:] if self.include_charges else torch.zeros(0, device=self.device)
 
-        return x, one_hot, charges
+        return x, one_hot, charges, batch_index
 
     @torch.no_grad()
     @typechecked

src/models/qm9_mol_gen_ddpm.py

@@ -590,8 +590,9 @@ def sample(
         num_nodes: Optional[TensorType["batch_size"]] = None,
         node_mask: Optional[TensorType["batch_num_nodes"]] = None,
         context: Optional[TensorType["batch_size", "num_context_features"]] = None,
+        fix_noise: bool = False,
         num_timesteps: Optional[int] = None
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         # node count-conditioning
         if num_nodes is None:
             num_nodes = self.ddpm.num_nodes_distribution.sample(num_samples)
@@ -610,7 +611,7 @@ def sample(
             context = None
 
         # sampling
-        xh, _, _ = self.ddpm.mol_gen_sample(
+        xh, batch_index, _ = self.ddpm.mol_gen_sample(
             num_samples=num_samples,
             num_nodes=num_nodes,
             node_mask=node_mask,
@@ -623,7 +624,7 @@ def sample(
         one_hot = xh[:, self.num_x_dims:-1] if self.include_charges else xh[:, self.num_x_dims:]
         charges = xh[:, -1:] if self.include_charges else torch.zeros(0, device=self.device)
 
-        return x, one_hot, charges
+        return x, one_hot, charges, batch_index
 
     @torch.no_grad()
     @typechecked

src/mol_gen_eval_conditional_qm9.py

@@ -17,7 +17,7 @@
 
 from src.datamodules.components.edm import get_bond_length_arrays
 from src.datamodules.components.edm.datasets_config import QM9_WITH_H, QM9_WITHOUT_H
-from src.models import NumNodesDistribution, PropertiesDistribution, compute_mean_mad
+from src.models import NumNodesDistribution, PropertiesDistribution, compute_mean_mad, save_and_sample_conditionally
 from src.utils.pylogger import get_pylogger
 
 from src import LR_SCHEDULER_MANUAL_INTERPOLATION_HELPER_CONFIG_ITEMS, LR_SCHEDULER_MANUAL_INTERPOLATION_PRIMARY_CONFIG_ITEMS, get_classifier, test_with_property_classifier, utils
@@ -96,7 +96,7 @@ def __iter__(self):
     def sample(self) -> Dict[str, Any]:
         num_nodes = self.nodes_distr.sample(self.num_samples).to(self.device)
         context = self.props_distr.sample_batch(num_nodes).to(self.device)
-        x, one_hot, _ = self.model.sample(
+        x, one_hot, _, _ = self.model.sample(
             num_samples=self.num_samples,
             num_nodes=num_nodes,
             context=context
@@ -165,14 +165,12 @@ def evaluate(cfg: DictConfig) -> Tuple[dict, dict]:
 
     assert (
         os.path.exists(cfg.generator_model_filepath) and
-        os.path.exists(cfg.classifier_model_dir) and
+        (os.path.exists(cfg.classifier_model_dir) or cfg.sweep_property_values) and
         cfg.property in cfg.generator_model_filepath and
-        cfg.property in cfg.classifier_model_dir
+        (cfg.property in cfg.classifier_model_dir or cfg.sweep_property_values)
     )
 
-    log.info("Loading classifier model!")
     device = f"cuda:{cfg.trainer.devices[0]}" if torch.cuda.is_available() else "cpu"
-    classifier = get_classifier(cfg.classifier_model_dir).to(device)
 
     log.info(f"Instantiating datamodule <{cfg.datamodule._target_}>")
     datamodule: LightningDataModule = hydra.utils.instantiate(cfg.datamodule)
@@ -221,8 +219,6 @@ def evaluate(cfg: DictConfig) -> Tuple[dict, dict]:
             bonds[0], bonds[1], bonds[2]
         )
 
-    log.info("Creating dataloader with generator!")
-
     splits = ["train", "valid", "test"]
     dataloaders = [
         datamodule.train_dataloader(),
@@ -250,8 +246,21 @@ def evaluate(cfg: DictConfig) -> Tuple[dict, dict]:
     nodes_distr = NumNodesDistribution(histogram)
 
     if cfg.sweep_property_values:
-        raise NotImplementedError()
+        log.info(f"Sampling conditionally via a sweep!")
+
+        for i in range(cfg.num_sweeps):
+            log.info(f"Sampling sweep {i + 1}/{cfg.num_sweeps}!")
+            save_and_sample_conditionally(
+                cfg=cfg,
+                model=model,
+                props_distr=props_distr,
+                dataset_info=dataset_info,
+                epoch=i,
+                id_from=0
+            )
     else:
+        log.info("Creating dataloader with generator!")
+        
         conditional_diffusion_dataloader = ConditionalDiffusionDataLoader(
             model=model,
             nodes_distr=nodes_distr,
@@ -261,6 +270,9 @@ def evaluate(cfg: DictConfig) -> Tuple[dict, dict]:
             dataset_info=dataset_info,
             iterations=cfg.iterations
         )
+
+        log.info("Loading classifier model!")
+        classifier = get_classifier(cfg.classifier_model_dir).to(device)
 
         log.info("Evaluating classifier on generator's samples!")
         loss = test_with_property_classifier(