constrained_generation_template.py

import argparse

import ase
from ase import Atoms
import numpy as np
from pathlib import Path
from glob import glob
from tqdm import tqdm
from datetime import datetime
from omegaconf import DictConfig, OmegaConf

import torch
from torch.utils.data import Dataset
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader

from dosmatgen.diffusion.diffusion_cfg import CSPDiffusion
from dosmatgen.utils.constants import cdvae_train_num_elements_distribution
from dosmatgen.utils.utils import decode
from dosmatgen.dataset.datamodule import CrystalDataModule, worker_init_fn
from dosmatgen.dataset.dataset import CrystalDataset


def build_first_atom_mask(batch, device):
    """Returns a (num_atoms, 1) mask that conditions only the first atom of each structure."""
    mask = torch.zeros(batch.num_nodes, 1, device=device)
    offset = 0
    for n in batch.num_atoms:
        mask[offset] = 1.0
        offset += n
    return mask


def find_batch_by_key(data_loader, key):
    """Search data_loader for the first batch containing the given structure_id."""
    for batch in data_loader:
        if key in batch.structure_id:
            return batch
    return None


def get_loader(data_module, split):
    if split == "train":
        return data_module.train_dataloader()
    elif split == "val":
        return data_module.val_dataloader()
    elif split == "test":
        return data_module.test_dataloader()
    else:
        raise ValueError(f"Unknown split '{split}'. Choose from train, val, test.")


def diffuse(
    start_loader,
    target_loader,
    model,
    start_key,
    target_key,
    n_candidates,
    step_lr,
    diff_ratio,
    w,
):
    # find target DOS
    print(f"Searching for target DOS (key={target_key})...")
    target_batch = find_batch_by_key(target_loader, target_key)
    if target_batch is None:
        raise ValueError(f"Target key '{target_key}' not found in target loader.")

    target_idx = list(target_batch.structure_id).index(target_key)
    offset = sum(target_batch.num_atoms[:target_idx].tolist())
    # DOS of the first atom of the target structure: shape [dos_dim]
    target_y = target_batch.y[offset].clone()

    # find template structure
    print(f"Searching for template structure (key={start_key})...")
    start_batch = find_batch_by_key(start_loader, start_key)
    if start_batch is None:
        raise ValueError(f"Start key '{start_key}' not found in start loader.")

    # patch batch.y: first atom of each structure in the template batch gets
    # the target DOS; all other atoms are zeroed out (unconditioned)
    patched_y = torch.zeros_like(start_batch.y)
    atom_offset = 0
    for n in start_batch.num_atoms:
        patched_y[atom_offset] = target_y
        atom_offset += n
    start_batch.y = patched_y

    start_batch = start_batch.to('cuda')
    mask = build_first_atom_mask(start_batch, device='cuda')

    batch_outputs = []
    for i in range(n_candidates):
        outputs, _ = model.masked_cfg_sample(
            start_batch,
            mask,
            step_lr=step_lr,
            diff_ratio=diff_ratio,
            w=w,
        )

        outputs = {
            'structure_id': start_batch.structure_id,
            'num_atoms': outputs['num_atoms'].detach().cpu(),
            'atom_types': outputs['atom_types'].detach().cpu(),
            'frac_coords': outputs['frac_coords'].detach().cpu(),
            'lattices': outputs['lattices'].detach().cpu(),
        }
        batch_outputs.append(outputs)

    all_outputs = {
        0: {
            "batch": start_batch,
            "outputs": batch_outputs,
        }
    }

    return all_outputs


def main(args):
    root_path = Path(args.root_path)

    now = datetime.now()
    formatted_time = now.strftime("%d%m%Y_%H%M%S")
    save_path = Path(args.save_path) / formatted_time

    # load config
    print("Loading model...")
    config_path = root_path / 'hparams.yaml'
    config = OmegaConf.load(config_path)

    # load checkpoint
    ckpt_path = glob(str(root_path / '*.ckpt'))
    if len(ckpt_path) == 0:
        raise ValueError("No checkpoint file found.")
    elif len(ckpt_path) > 1:
        raise ValueError("Multiple checkpoint files found.")
    ckpt_path = ckpt_path[0]

    model = CSPDiffusion.load_from_checkpoint(ckpt_path, config=config)
    model.to('cuda')

    # load data module
    print("Loading data loaders...")
    config.datamodule.batch_size.test = args.batch_size
    config.datamodule.batch_size.train = args.batch_size
    config.datamodule.batch_size.val = args.batch_size

    data_module = CrystalDataModule(config, scaler_path=str(root_path))

    needs_fit = args.start_split in ("train", "val") or args.target_split in ("train", "val")
    if needs_fit:
        data_module.setup(stage="fit")
    data_module.setup(stage="test")

    start_loader = get_loader(data_module, args.start_split)
    target_loader = get_loader(data_module, args.target_split)

    # diffuse
    print("Denoising for generation...")
    all_outputs = diffuse(
        start_loader,
        target_loader,
        model,
        args.start_key,
        args.target_key,
        args.n_candidates,
        args.step_lr,
        args.diff_ratio,
        args.w,
    )

    # decode
    print("Decoding to ase.Atoms objects...")
    atoms_list = {}

    for j in tqdm(all_outputs):
        output_dict = all_outputs[j]['outputs'][0]

        start_idx = 0
        for i in range(len(output_dict['structure_id'])):
            sid = output_dict['structure_id'][i]
            num_atoms = output_dict['num_atoms'][i]
            lattices = output_dict['lattices'][i]

            end_idx = start_idx + num_atoms
            atom_types = output_dict['atom_types'][start_idx:end_idx]
            frac_coords = output_dict['frac_coords'][start_idx:end_idx]

            assert len(atom_types) == num_atoms.item() == frac_coords.shape[0]

            atoms = Atoms(
                cell=lattices.cpu().numpy(),
                scaled_positions=frac_coords.cpu().numpy(),
                numbers=atom_types
            )

            atoms_list[sid] = atoms

            start_idx = end_idx

    save_path.mkdir(exist_ok=True)
    for k, v in atoms_list.items():
        ase.io.write(str(save_path / f"{k}.cif"), v)

    print("Done!")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--root_path', type=str, required=True, help="Path to the folder containing model.ckpt, hparams.yaml, and scaler.pt")
    parser.add_argument('--save_path', type=str, default="structures/", help="Path to save the generated structures")
    parser.add_argument('--batch_size', type=int, default=1)
    parser.add_argument('--step_lr', type=float, default=5e-6)
    parser.add_argument('--n_candidates', type=int, default=1, help="Number of candidate structures to generate per template")
    parser.add_argument('--diff_ratio', type=float, default=0.5, help="Timestep fraction at which denoising starts; (0, 1]. Use < 1 to initialize from template.")
    parser.add_argument('--w', type=float, default=1.0, help="CFG guidance weight; higher = stronger conditioning")
    parser.add_argument('--start_key', type=str, required=True, help="Structure ID of the template structure to initialize diffusion from")
    parser.add_argument('--target_key', type=str, required=True, help="Structure ID whose DOS is used for conditioning")
    parser.add_argument('--start_split', type=str, default="test", choices=["train", "val", "test"], help="Dataset split to search for the template structure")
    parser.add_argument('--target_split', type=str, default="test", choices=["train", "val", "test"], help="Dataset split to search for the target DOS")
    args = parser.parse_args()

    main(args)