"""Copyright (c) Meta Platforms, Inc. and affiliates.""" from __future__ import annotations from copy import deepcopy from pathlib import Path from typing import Any, Literal, Sequence import click import pytorch_lightning as pl import torch from pytorch_lightning.callbacks import BasePredictionWriter from pytorch_lightning.loggers.wandb import WandbLogger from torch_geometric.data import Batch, Data, DataLoader import wandb from diffcsp.script_utils import GenDataset from flowmm.model.eval_utils import ( CSPDataset, get_loaders, load_cfg, load_date_from_wandb, load_group_from_wandb, load_id_from_wandb, load_model, load_project_from_wandb, register_omega_conf_resolvers, ) from flowmm.old_eval.generation_metrics import compute_generation_metrics from flowmm.old_eval.lattice_metrics import compute_lattice_metrics from flowmm.old_eval.reconstruction_metrics import compute_reconstruction_metrics TASKS_TYPE = Literal[ "reconstruct", "recon_trajectory", "generate", "gen_trajectory", "pred" ] TASKS = deepcopy(TASKS_TYPE.__args__) STAGE_TYPE = Literal["train", "val", "test"] STAGES = deepcopy(STAGE_TYPE.__args__) register_omega_conf_resolvers() class TorchPredictionWriter(BasePredictionWriter): def __init__( self, output_dir: Path | str, write_interval: Literal["batch", "epoch", "batch_and_epoch"] = "epoch", ): super().__init__(write_interval) self.output_dir = Path(output_dir) self.output_dir.mkdir(exist_ok=True, parents=True) def write_on_epoch_end( self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", predictions: Sequence[Any], batch_indices: Sequence[Any] | None, ) -> None: # flatten extra nesting level from single_gpu mode (predictions/batch_indices # may be wrapped as [[batch1, ...]] instead of [batch1, ...]) if predictions and isinstance(predictions[0], list): predictions = [item for sublist in predictions for item in sublist] if batch_indices and isinstance(batch_indices[0], list): batch_indices = [item for sublist in batch_indices for item in sublist] # this will create N (num processes) files in `output_dir` each containing # the predictions of it's respective rank torch.save( predictions, self.output_dir / f"predictions_{trainer.global_rank:02d}.pt" ) # optionally, you can also save `batch_indices` to get the information about the data index # from your prediction data torch.save( batch_indices, self.output_dir / f"batch_indices_{trainer.global_rank:02d}.pt", ) @click.group() def cli(): pass @cli.command() @click.argument("checkpoint", type=Path) @click.option("--stage", type=click.Choice(STAGES, case_sensitive=False), default="val") @click.option("--batch_size", type=int, default=16384) @click.option("--num_evals", type=int, default=1) @click.option("--limit_predict_batches", type=str, default="1.") @click.option("--num_steps", type=int, default=None) @click.option( "--div_mode", type=click.Choice(["exact", "rademacher"], case_sensitive=False), default=None, ) @click.option( "--single_gpu/--multi_gpu", is_flag=True, show_default=True, default=False, help="use one gpu, not ddp", ) def nll( checkpoint: Path, stage: STAGE_TYPE, batch_size: int, num_evals: int, limit_predict_batches: str, num_steps: int, div_mode: bool, single_gpu: bool, ) -> None: raise NotImplementedError( "there are currently base distributions which make this unappealing." ) def get_target_dir(checkpoint: Path, subdir: bool) -> Path: if subdir: target_dir = checkpoint.parent / subdir else: target_dir = checkpoint.parent return target_dir.resolve() import random import os def subsample_test_set(file_path, ratio): # 加载 .pt 文件 data = torch.load(file_path) if not isinstance(data, list): raise ValueError("Loaded data is not a list.") # 计算要抽取的元素数量 num_elements = len(data) num_samples = int(num_elements * ratio) if num_samples <= 0: raise ValueError("Ratio too small to sample any elements.") # 随机抽样 sampled_data = random.sample(data, num_samples) # 构造新文件路径 dir_name = os.path.dirname(file_path) base_name = os.path.basename(file_path) name_part, ext = os.path.splitext(base_name) new_file_path = os.path.join(dir_name, f"{name_part}_sub_{ratio}{ext}") # 保存抽样后的数据 torch.save(sampled_data, new_file_path) print(f"Sampled {num_samples} elements and saved to: {new_file_path}") @cli.command() @click.argument("checkpoint", type=Path) @click.option("--stage", type=click.Choice(STAGES, case_sensitive=False), default="val") @click.option("--batch_size", type=int, default=None) @click.option("--num_evals", type=int, default=1) @click.option("--limit_predict_batches", type=str, default="1.") @click.option("--num_steps", type=int, default=None) @click.option( "--single_gpu/--multi_gpu", is_flag=True, show_default=True, default=False, help="use one gpu, not ddp", ) @click.option( "--subdir", type=str, default="", help="subdir name at level of checkpoint" ) @click.option( "--inference_anneal_slope", type=float, default=None, ) @click.option( "--inference_anneal_offset", type=float, default=None, ) @click.option( "--inference_anneal_types/--no-inference_anneal_types", is_flag=True, show_default=True, default=False, ) @click.option( "--inference_anneal_coords/--no-inference_anneal_coords", is_flag=True, show_default=True, default=True, ) @click.option( "--inference_anneal_lattice/--no-inference_anneal_lattice", is_flag=True, show_default=True, default=False, ) @click.option( "--compute_traj_velo_norms", is_flag=True, show_default=True, default=False, ) @click.option('--sample_way', type = str, required = True) @click.option('--sample_lbfgs_lr', type = float, default = 0.00005) @click.option('--sample_lbfgs_iter', type = int, default = 20) @click.option('--sample_md_step', type = int, default = 100) @click.option('--sample_md_grad', type = str, default = "manifold") @click.option('--use_subset', type = bool, default = False) @click.option('--subset_ratio', type = float, default = 0.1) def reconstruct( checkpoint: Path, stage: STAGE_TYPE, batch_size: int | None, num_evals: int, limit_predict_batches: str, num_steps: int | None, single_gpu: bool, subdir: str, inference_anneal_slope: float | None, inference_anneal_offset: float | None, inference_anneal_types: bool, inference_anneal_coords: bool, inference_anneal_lattice: bool, compute_traj_velo_norms: bool | None, sample_way: str, sample_lbfgs_lr: float, sample_lbfgs_iter: int, sample_md_step: int, sample_md_grad: str, use_subset: bool, subset_ratio: float ) -> None: cfg, model = load_model(checkpoint) model.set_sample_param( sample_way = sample_way, sample_lbfgs_lr = sample_lbfgs_lr, sample_lbfgs_iter = sample_lbfgs_iter, sample_md_step = sample_md_step, sample_md_grad = sample_md_grad ) if "null" not in cfg.model.manifold_getter.atom_type_manifold: raise ValueError( f"you cannot do reconstruction with an unconditional atom_type_manifold {cfg.model.manifold_getter.atom_type_manifold=}" ) stage = stage.lower() if batch_size is None: # this must be explicitly set since the default is int batch_size = getattr(cfg.data.datamodule.batch_size, stage) print(f"Using {batch_size=} from default cfg") else: setattr(cfg.data.datamodule.batch_size, stage, batch_size) print(f"Using custom {batch_size=}") # update cfg if num_steps is not None: cfg.integrate.num_steps = num_steps if inference_anneal_slope is not None: cfg.integrate.inference_anneal_slope = inference_anneal_slope if inference_anneal_offset is not None: assert (0 <= inference_anneal_offset) and (inference_anneal_offset < 1) cfg.integrate.inference_anneal_offset = inference_anneal_offset if compute_traj_velo_norms: cfg.integrate.compute_traj_velo_norms = compute_traj_velo_norms cfg.integrate.inference_anneal_types = inference_anneal_types cfg.integrate.inference_anneal_coords = inference_anneal_coords cfg.integrate.inference_anneal_lattice = inference_anneal_lattice if use_subset: origin_test_path = cfg.data.datamodule.datasets.test[0].save_path test_subset_path = f"{origin_test_path[:-3]}_sub_{subset_ratio}.pt" cfg.data.datamodule.datasets.test[0].save_path = test_subset_path if not os.path.exists(test_subset_path): subsample_test_set(origin_test_path, subset_ratio) loaders = get_loaders(cfg) loader = loaders[STAGES.index(stage)] target_dir = get_target_dir(checkpoint, subdir) assert num_evals > 0 directories = [f"reconstruct_{i:02d}" for i in range(num_evals)] for directory in directories: pred_writer = TorchPredictionWriter( output_dir=target_dir / directory, write_interval="epoch", ) # save num_steps (target_dir / directory / "num_steps.txt").write_text( str(cfg.integrate.num_steps) ) if single_gpu: trainer = pl.Trainer( accelerator="gpu", devices=1, callbacks=[pred_writer], limit_predict_batches=eval(limit_predict_batches), inference_mode=True if sample_way == "flow" else False, logger=False, ) else: trainer = pl.Trainer( accelerator="gpu", strategy="ddp", devices="auto", callbacks=[pred_writer], limit_predict_batches=eval(limit_predict_batches), inference_mode=True if sample_way == "flow" else False, logger=False, ) trainer.predict( model, dataloaders=loader, return_predictions=False, ckpt_path=checkpoint, ) @cli.command(name="recon_trajectory") @click.argument("checkpoint", type=Path) @click.option("--stage", type=click.Choice(STAGES, case_sensitive=False), default="val") @click.option("--batch_size", type=int, default=16384) @click.option("--num_evals", type=int, default=1) @click.option("--limit_predict_batches", type=str, default="1.") @click.option("--num_steps", type=int, default=None) @click.option( "--single_gpu/--multi_gpu", is_flag=True, show_default=True, default=False, help="use one gpu, not ddp", ) @click.option( "--subdir", type=str, default="", help="subdir name at level of checkpoint" ) @click.option( "--inference_anneal_slope", type=float, default=None, ) @click.option( "--inference_anneal_offset", type=float, default=None, ) @click.option( "--inference_anneal_types/--no-inference_anneal_types", is_flag=True, show_default=True, default=False, ) @click.option( "--inference_anneal_coords/--no-inference_anneal_coords", is_flag=True, show_default=True, default=True, ) @click.option( "--inference_anneal_lattice/--no-inference_anneal_lattice", is_flag=True, show_default=True, default=False, ) @click.option( "--compute_traj_velo_norms", is_flag=True, show_default=True, default=False, ) def recon_trajectory( checkpoint: Path, stage: STAGE_TYPE, batch_size: int | None, num_evals: int, limit_predict_batches: str, num_steps: int | None, single_gpu: bool, subdir: str, inference_anneal_slope: float | None, inference_anneal_offset: float | None, inference_anneal_types: bool, inference_anneal_coords: bool, inference_anneal_lattice: bool, compute_traj_velo_norms: bool | None, ) -> None: cfg, model = load_model(checkpoint) if "null" not in cfg.model.manifold_getter.atom_type_manifold: raise ValueError( f"you cannot do reconstruction with an unconditional atom_type_manifold {cfg.model.manifold_getter.atom_type_manifold=}" ) stage = stage.lower() if batch_size is None: # this must be explicitly set since the default is int batch_size = getattr(cfg.data.datamodule.batch_size, stage) print(f"Using {batch_size=} from default cfg") else: setattr(cfg.data.datamodule.batch_size, stage, batch_size) print(f"Using custom {batch_size=}") # update cfg if num_steps is not None: cfg.integrate.num_steps = num_steps if inference_anneal_slope is not None: cfg.integrate.inference_anneal_slope = inference_anneal_slope if inference_anneal_offset is not None: assert (0 <= inference_anneal_offset) and (inference_anneal_offset < 1) cfg.integrate.inference_anneal_offset = inference_anneal_offset if compute_traj_velo_norms: cfg.integrate.compute_traj_velo_norms = compute_traj_velo_norms cfg.integrate.inference_anneal_types = inference_anneal_types cfg.integrate.inference_anneal_coords = inference_anneal_coords cfg.integrate.inference_anneal_lattice = inference_anneal_lattice # THIS ADDS A NEW FIELD TO CFG cfg.integrate.entire_traj = True loaders = get_loaders(cfg) loader = loaders[STAGES.index(stage)] target_dir = get_target_dir(checkpoint, subdir) assert num_evals > 0 directories = [f"recon_trajectory_{i:02d}" for i in range(num_evals)] for directory in directories: pred_writer = TorchPredictionWriter( output_dir=target_dir / directory, write_interval="epoch", ) # save num_steps (target_dir / directory / "num_steps.txt").write_text( str(cfg.integrate.num_steps) ) if single_gpu: trainer = pl.Trainer( accelerator="gpu", devices=1, callbacks=[pred_writer], limit_predict_batches=eval(limit_predict_batches), ) else: trainer = pl.Trainer( accelerator="gpu", strategy="ddp", devices="auto", callbacks=[pred_writer], limit_predict_batches=eval(limit_predict_batches), ) trainer.predict( model, dataloaders=loader, return_predictions=False, ckpt_path=checkpoint, ) @cli.command() @click.argument("checkpoint", type=Path) @click.option("--num_samples", type=int, default=10_000) @click.option("--batch_size", type=int, default=16384) @click.option("--num_steps", type=int, default=None) @click.option( "--single_gpu/--multi_gpu", is_flag=True, show_default=True, default=False, help="use one gpu, not ddp", ) @click.option( "--subdir", type=str, default="", help="subdir name at level of checkpoint" ) @click.option("--gen_id", type=int, default=0, help=r"folder name is generate_{gen_id}") @click.option( "--inference_anneal_slope", type=float, default=None, ) @click.option( "--inference_anneal_offset", type=float, default=None, ) @click.option( "--inference_anneal_types/--no-inference_anneal_types", is_flag=True, show_default=True, default=False, ) @click.option( "--inference_anneal_coords/--no-inference_anneal_coords", is_flag=True, show_default=True, default=True, ) @click.option( "--inference_anneal_lattice/--no-inference_anneal_lattice", is_flag=True, show_default=True, default=False, ) @click.option( "--compute_traj_velo_norms", is_flag=True, show_default=True, default=False, ) @click.option('--sample_way', type = str, required = True) @click.option('--sample_lbfgs_lr', type = float, default = 0.00005) @click.option('--sample_lbfgs_iter', type = int, default = 20) @click.option('--sample_md_step', type = int, default = 100) @click.option('--sample_md_grad', type = str, default = "manifold") def generate( checkpoint: Path, num_samples: int, batch_size: int | None, num_steps: int | None, single_gpu: bool, subdir: str, gen_id: int, inference_anneal_slope: float | None, inference_anneal_offset: float | None, inference_anneal_types: bool, inference_anneal_coords: bool, inference_anneal_lattice: bool, compute_traj_velo_norms: bool | None, sample_way: str, sample_lbfgs_lr: float, sample_lbfgs_iter: int, sample_md_step: int, sample_md_grad: str, ) -> None: cfg, model = load_model(checkpoint) model.set_sample_param( sample_way = sample_way, sample_lbfgs_lr = sample_lbfgs_lr, sample_lbfgs_iter = sample_lbfgs_iter, sample_md_step = sample_md_step, sample_md_grad = sample_md_grad ) if "null" in cfg.model.manifold_getter.atom_type_manifold: raise ValueError( f"you cannot do generation with a conditional atom_type_manifold {cfg.model.manifold_getter.atom_type_manifold=}" ) # update cfg if num_steps is not None: cfg.integrate.num_steps = num_steps if inference_anneal_slope is not None: cfg.integrate.inference_anneal_slope = inference_anneal_slope if inference_anneal_offset is not None: assert (0 <= inference_anneal_offset) and (inference_anneal_offset < 1) cfg.integrate.inference_anneal_offset = inference_anneal_offset if compute_traj_velo_norms: cfg.integrate.compute_traj_velo_norms = compute_traj_velo_norms cfg.integrate.inference_anneal_types = inference_anneal_types cfg.integrate.inference_anneal_coords = inference_anneal_coords cfg.integrate.inference_anneal_lattice = inference_anneal_lattice sample_set = GenDataset(dataset=cfg.data.dataset_name, total_num=num_samples) loader = DataLoader(sample_set, batch_size=batch_size) target_dir = get_target_dir(checkpoint, subdir) directories = [f"generate_{gen_id:02d}"] for directory in directories: pred_writer = TorchPredictionWriter( output_dir=target_dir / directory, write_interval="epoch", ) # save num_steps (target_dir / directory / "num_steps.txt").write_text( str(cfg.integrate.num_steps) ) if single_gpu: trainer = pl.Trainer( accelerator="gpu", devices=1, callbacks=[pred_writer], ) else: trainer = pl.Trainer( accelerator="gpu", strategy="ddp", devices="auto", callbacks=[pred_writer], ) trainer.predict( model, dataloaders=loader, return_predictions=False, ckpt_path=checkpoint, ) @cli.command(name="gen_trajectory") @click.argument("checkpoint", type=Path) @click.option("--num_samples", type=int, default=256) @click.option("--batch_size", type=int, default=256) @click.option("--num_steps", type=int, default=None) @click.option( "--single_gpu/--multi_gpu", is_flag=True, show_default=True, default=False, help="use one gpu, not ddp", ) @click.option( "--subdir", type=str, default="", help="subdir name at level of checkpoint" ) @click.option("--gen_id", type=int, default=0, help=r"folder name is generate_{gen_id}") @click.option( "--inference_anneal_slope", type=float, default=None, ) @click.option( "--inference_anneal_offset", type=float, default=None, ) @click.option( "--inference_anneal_types/--no-inference_anneal_types", is_flag=True, show_default=True, default=False, ) @click.option( "--inference_anneal_coords/--no-inference_anneal_coords", is_flag=True, show_default=True, default=True, ) @click.option( "--inference_anneal_lattice/--no-inference_anneal_lattice", is_flag=True, show_default=True, default=False, ) @click.option( "--compute_traj_velo_norms", is_flag=True, show_default=True, default=False, ) def gen_trajectory( checkpoint: Path, num_samples: int, batch_size: int | None, num_steps: int | None, single_gpu: bool, subdir: str, gen_id: int, inference_anneal_slope: float | None, inference_anneal_offset: float | None, inference_anneal_types: bool, inference_anneal_coords: bool, inference_anneal_lattice: bool, compute_traj_velo_norms: bool | None, ) -> None: cfg, model = load_model(checkpoint) if "null" in cfg.model.manifold_getter.atom_type_manifold: raise ValueError( f"you cannot do generation with a conditional atom_type_manifold {cfg.model.manifold_getter.atom_type_manifold=}" ) # update cfg if num_steps is not None: cfg.integrate.num_steps = num_steps if inference_anneal_slope is not None: cfg.integrate.inference_anneal_slope = inference_anneal_slope if inference_anneal_offset is not None: assert (0 <= inference_anneal_offset) and (inference_anneal_offset < 1) cfg.integrate.inference_anneal_offset = inference_anneal_offset if compute_traj_velo_norms: cfg.integrate.compute_traj_velo_norms = compute_traj_velo_norms cfg.integrate.inference_anneal_types = inference_anneal_types cfg.integrate.inference_anneal_coords = inference_anneal_coords cfg.integrate.inference_anneal_lattice = inference_anneal_lattice # THIS ADDS A NEW FIELD TO CFG cfg.integrate.entire_traj = True sample_set = GenDataset(dataset=cfg.data.dataset_name, total_num=num_samples) loader = DataLoader(sample_set, batch_size=batch_size) target_dir = get_target_dir(checkpoint, subdir) directories = [f"gen_trajectory_{gen_id:02d}"] for directory in directories: pred_writer = TorchPredictionWriter( output_dir=target_dir / directory, write_interval="epoch", ) # save num_steps (target_dir / directory / "num_steps.txt").write_text( str(cfg.integrate.num_steps) ) if single_gpu: trainer = pl.Trainer( accelerator="gpu", devices=1, callbacks=[pred_writer], ) else: trainer = pl.Trainer( accelerator="gpu", strategy="ddp", devices="auto", callbacks=[pred_writer], ) trainer.predict( model, dataloaders=loader, return_predictions=False, ckpt_path=checkpoint, ) @cli.command() @click.argument("checkpoint", type=Path) @click.argument("atom_types_path", type=Path) @click.option("--batch_size", type=int, default=16384) @click.option("--num_steps", type=int, default=None) @click.option( "--single_gpu/--multi_gpu", is_flag=True, show_default=True, default=False, help="use one gpu, not ddp", ) @click.option( "--subdir", type=str, default="", help="subdir name at level of checkpoint" ) @click.option("--pred_id", type=int, default=0, help=r"folder name is pred_{pred_id}") @click.option( "--inference_anneal_slope", type=float, default=None, ) @click.option( "--inference_anneal_offset", type=float, default=None, ) @click.option( "--inference_anneal_types/--no-inference_anneal_types", is_flag=True, show_default=True, default=False, ) @click.option( "--inference_anneal_coords/--no-inference_anneal_coords", is_flag=True, show_default=True, default=True, ) @click.option( "--inference_anneal_lattice/--no-inference_anneal_lattice", is_flag=True, show_default=True, default=False, ) def predict( checkpoint: Path, atom_types_path: Path, batch_size: int | None, num_steps: int | None, single_gpu: bool, subdir: str, pred_id: int, inference_anneal_slope: float | None, inference_anneal_offset: float | None, inference_anneal_types: bool, inference_anneal_coords: bool, inference_anneal_lattice: bool, ) -> None: cfg, model = load_model(checkpoint) # update cfg if num_steps is not None: cfg.integrate.num_steps = num_steps if inference_anneal_slope is not None: cfg.integrate.inference_anneal_slope = inference_anneal_slope if inference_anneal_offset is not None: assert (0 <= inference_anneal_offset) and (inference_anneal_offset < 1) cfg.integrate.inference_anneal_offset = inference_anneal_offset cfg.integrate.inference_anneal_types = inference_anneal_types cfg.integrate.inference_anneal_coords = inference_anneal_coords cfg.integrate.inference_anneal_lattice = inference_anneal_lattice with open(atom_types_path, "r") as f: atom_types = f.read() atom_types = eval(atom_types) dataset = CSPDataset(atom_types) loader = DataLoader(dataset, batch_size=batch_size) target_dir = get_target_dir(checkpoint, subdir) directories = [f"pred_{pred_id:02d}"] for directory in directories: pred_writer = TorchPredictionWriter( output_dir=target_dir / directory, write_interval="epoch", ) # save num_steps (target_dir / directory / "num_steps.txt").write_text( str(cfg.integrate.num_steps) ) if single_gpu: trainer = pl.Trainer( accelerator="gpu", devices=1, callbacks=[pred_writer], ) else: trainer = pl.Trainer( accelerator="gpu", strategy="ddp", devices="auto", callbacks=[pred_writer], ) trainer.predict( model, dataloaders=loader, return_predictions=False, ckpt_path=checkpoint, ) def _get_consolidated_path(directory: Path, task: str) -> str: return directory / f"consolidated_{task}.pt" def _list_of_dicts_to_dict_of_lists( lod: list[dict[str, torch.Tensor | Batch]], keys_to_ignore: tuple[str] = (), ) -> dict[str, list[torch.Tensor] | list[Data]]: out = {k: [] for k in lod[0].keys()} for key, val in out.items(): for d in lod: if key in keys_to_ignore: continue elif isinstance(d[key], Batch): val.append(d[key].to_data_list()) else: val.append(d[key]) return out def _consolidate( target_dir: Path, task: TASKS_TYPE ) -> dict[str, list[dict[str, torch.Tensor | list[Data] | list[int]]]]: pattern = f"{task}_??" directories = sorted(list(target_dir.glob(pattern))) if not directories: print(f"no! directories found with the pattern {pattern}") return None else: print(f"yes directories found with the pattern {pattern}") # to a single array num_stepss = [] out_by_eval = [] for _, directory in enumerate(directories): preds = sorted(list(directory.glob("predictions_??.pt"))) batches = sorted(list(directory.glob("batch_indices_??.pt"))) assert len(preds) == len(batches) # get and empty out first_pred = torch.load(preds[0], map_location="cpu") # Handle both [dicts] and [[dicts]] formats if first_pred and isinstance(first_pred[0], list): keys = first_pred[0][0].keys() else: keys = first_pred[0].keys() out = {k: [] for k in keys} order = [] # collect from output, keeping track of order for pred, batch in zip(preds, batches): pred = torch.load(pred, map_location="cpu") batch = torch.load(batch, map_location="cpu") # Unwrap extra nesting if present (from single_gpu mode with re-wrapping) if pred and isinstance(pred[0], list): pred = pred[0] if batch and isinstance(batch[0], list): batch = batch[0] for pp, bb in zip(pred, batch): for k, v in pp.items(): if isinstance(v, torch.Tensor): out[k].append(v) elif isinstance(v, Data): out[k].extend(v.to_data_list()) else: raise TypeError("don't know what to do with that type.") if isinstance(bb, list): order.extend(bb) else: order.append(bb) # organize out into our datatype for k, v in out.items(): if k == "input_data_batch": out[k] = Batch.from_data_list(v) elif k in ["atom_types", "frac_coords", "lattices", "lengths", "angles"]: if "trajectory" in task: out[k] = torch.concat(v, dim=1) else: out[k] = torch.concat(v, dim=0) elif k == "num_atoms": out[k] = torch.concat(v, dim=0) else: raise ValueError(f"don't know what to do with {k=}") num_stepss.append(eval((target_dir / directory / "num_steps.txt").read_text())) out["batch_indices"] = torch.tensor(order) out_by_eval.append(out) # make sure all num_steps are the same assert all( [num_stepss[0] == i for i in num_stepss] ), f"not all num_steps agreed, got {num_stepss=}" num_steps = num_stepss[0] # save num_steps (_get_consolidated_path(target_dir, task).parent / "num_steps.txt").write_text( str(num_steps) ) out_by_eval = _list_of_dicts_to_dict_of_lists(out_by_eval) torch.save(out_by_eval, _get_consolidated_path(target_dir, task)) return out_by_eval def _create_eval_pt( consolidated: dict[str, torch.Tensor] | None, target_dir: Path, filename: str, ) -> Path: # consolidated = {k: v.reshape(-1, v.shape[-1]) if k != "lattices" else v.reshape(-1, *v.shape[-2:]) for k, v in r.items()} if consolidated is None: raise ValueError( f"you cannot try to save an eval_pt with no data, {consolidated=}" ) consolidated = {k: v[0] for k, v in consolidated.items()} consolidated["eval_setting"] = None path = target_dir / filename torch.save(consolidated, path) return path @cli.command() @click.argument("checkpoint", type=Path) @click.option( "--subdir", type=str, default="", help="subdir name at level of checkpoint" ) @click.option( "--path_eval_pt", show_default=True, default=None, help="select a path to save the eval_pt, otherwise it is not saved", ) @click.option( "--task_to_save", default=None, help="if it is ambiguous, you can select the task for path_eval_pt", ) def consolidate( checkpoint: Path, subdir: str, path_eval_pt: str | None, task_to_save: TASKS_TYPE | None, ) -> None: target_dir = get_target_dir(checkpoint, subdir) r = _consolidate(target_dir, "reconstruct") rt = _consolidate(target_dir, "recon_trajectory") g = _consolidate(target_dir, "generate") gt = _consolidate(target_dir, "gen_trajectory") p = _consolidate(target_dir, "pred") consolidations = {k: v for k, v in zip(TASKS, [r, rt, g, gt, p])} did_consolidate = {k: v != None for k, v in consolidations.items()} if any(did_consolidate.values()): if task_to_save is not None: print(f"consolidating {task_to_save}") consolidated = consolidations[task_to_save] elif sum(did_consolidate.values()) == 1: consolidated_task = list(consolidations.keys())[ list((did_consolidate.values())).index(True) ] print(f"only {consolidated_task} was consolidated") consolidated = consolidations[consolidated_task] else: raise ValueError( f"more than one task was consolidated, so you must specify which one to save as eval_pt. FYI: {did_consolidate=}" ) else: raise ValueError( "nothing was consolidated, so the program cannot print an eval for dft" ) if sum(did_consolidate.values()) > 0 and path_eval_pt is not None: path = _create_eval_pt(consolidated, target_dir, path_eval_pt) print("eval_pt:") print(path) def _reconstruction_metrics_wandb( target_dir: Path, consolidated_reconstruction_path: Path, global_step: int, stage: STAGE_TYPE, ) -> dict[str, float]: recon_metrics = {} if consolidated_reconstruction_path.exists(): # should be 20 evals tmp, num_evals = compute_reconstruction_metrics( consolidated_reconstruction_path, multi_eval=True, metrics_path=target_dir / f"old_eval_metrics_reconstruct_multi.json", ground_truth_path=None, # unnecessary since we save this when we consolidate ) recon_metrics.update( {f"{stage}/" + k + f"_{num_evals:02d}": v for k, v in tmp.items()} ) # should be 01 eval tmp, num_evals = compute_reconstruction_metrics( consolidated_reconstruction_path, multi_eval=False, metrics_path=target_dir / f"old_eval_metrics_reconstruct_single.json", ground_truth_path=None, # unnecessary since we save this when we consolidate ) recon_metrics.update( {f"{stage}/" + k + f"_{num_evals:02d}": v for k, v in tmp.items()} ) # this will be 01 evals recon_num_steps = eval( (consolidated_reconstruction_path.parent / "num_steps.txt").read_text() ) recon_metrics.update({f"{stage}/recon_num_steps": recon_num_steps}) recon_metrics.update({"trainer/global_step": global_step}) return recon_metrics def _generation_metrics_wandb( target_dir: Path, consolidated_generation_path: Path, gt_dataset_path: Path, global_step: int, eval_model_name: Literal["carbon", "mp20", "perovskite"], n_subsamples: int, stage: STAGE_TYPE, ) -> dict[str, float]: gen_metrics = {} if consolidated_generation_path.exists(): tmp = compute_generation_metrics( path=consolidated_generation_path, metrics_path=target_dir / f"old_eval_metrics_generate.json", ground_truth_path=gt_dataset_path, eval_model_name=eval_model_name, n_subsamples=n_subsamples, ) gen_metrics.update({f"{stage}/" + k: v for k, v in tmp.items()}) gen_num_steps = eval( (consolidated_generation_path.parent / "num_steps.txt").read_text() ) gen_metrics.update({f"{stage}/gen_n_subsamples": n_subsamples}) gen_metrics.update({f"{stage}/gen_num_steps": gen_num_steps}) gen_metrics.update({"trainer/global_step": global_step}) return gen_metrics @cli.command(name="old_eval_metrics") @click.argument("checkpoint", type=Path) @click.option( "--do_not_log_wandb", is_flag=True, show_default=True, default=False, help="do not log results in the wandb training run", ) @click.option( "--subdir", type=str, default="", help="subdir name at level of checkpoint" ) @click.option( "--gen_subsamples", type=int, default=1_000, help="gen metrics are on this many subsamples", ) @click.option("--stage", type=click.Choice(STAGES, case_sensitive=False), default="val") def old_eval_metrics( checkpoint: Path, do_not_log_wandb: bool, subdir: str, gen_subsamples: int, stage: STAGE_TYPE, ) -> None: log_wandb = not do_not_log_wandb target_dir = get_target_dir(checkpoint, subdir) chkp = torch.load(checkpoint) global_step = chkp["global_step"] print(f"") print(f"======= reconstruction =======") print(f"") consolidated_reconstruction_path = target_dir / _get_consolidated_path( target_dir, "reconstruct" ) if consolidated_reconstruction_path.exists(): recon_metrics = _reconstruction_metrics_wandb( target_dir, consolidated_reconstruction_path, global_step, stage ) else: recon_metrics = {} print(f"{consolidated_reconstruction_path=} not found") print(f"Not performing reconstruction metrics!") print(f"") print(f"======= generation =======") print(f"") consolidated_generation_path = target_dir / _get_consolidated_path( target_dir, "generate" ) if consolidated_generation_path.exists(): cfg = load_cfg(checkpoint) gen_metrics = _generation_metrics_wandb( target_dir=target_dir, consolidated_generation_path=consolidated_generation_path, gt_dataset_path=cfg.data.datamodule.datasets[stage][0].save_path, global_step=global_step, eval_model_name=cfg.data.eval_model_name, n_subsamples=gen_subsamples, stage=stage, ) else: gen_metrics = {} print(f"{consolidated_generation_path=} not found") print(f"Not performing generation metrics!") print(f"") if ( not consolidated_reconstruction_path.exists() and not consolidated_generation_path.exists() ): raise FileNotFoundError( f"Neither {consolidated_reconstruction_path=} nor {consolidated_generation_path=} exists." ) if log_wandb: cfg = load_cfg(checkpoint) wandb_config = cfg.logging.wandb wandb_config.project = load_project_from_wandb(checkpoint) wandb_config.group = load_group_from_wandb(checkpoint) wandb_config.job_type = "cdvae_metrics" wandb_config.tags = [ load_date_from_wandb(checkpoint), load_id_from_wandb(checkpoint), ] wandb_config = dict(wandb_config) del wandb_config["log_model"] wandb.init(**wandb_config) wandb.log(recon_metrics, global_step) wandb.log(gen_metrics, global_step) wandb.finish() @cli.command(name="lattice_metrics") @click.argument("checkpoint", type=Path) # @click.option( # "--do_not_log_wandb", # is_flag=True, # show_default=True, # default=False, # help="do not log results in the wandb training run", # ) @click.option( "--subdir", type=str, default="", help="subdir name at level of checkpoint" ) @click.option("--stage", type=click.Choice(STAGES, case_sensitive=False), default="val") def lattice_metrics( checkpoint: Path, # do_not_log_wandb: bool, subdir: str, stage: STAGE_TYPE, ) -> None: # log_wandb = not do_not_log_wandb target_dir = get_target_dir(checkpoint, subdir) chkp = torch.load(checkpoint) global_step = chkp["global_step"] consolidated_reconstruction_path = target_dir / _get_consolidated_path( target_dir, "reconstruct" ) if consolidated_reconstruction_path.exists(): compute_lattice_metrics( consolidated_reconstruction_path, metrics_path=target_dir / f"lattice_metrics_reconstruct_single.json", ) # right now this returns nothing since it just plots the distribution consolidated_generation_path = target_dir / _get_consolidated_path( target_dir, "generate" ) if consolidated_generation_path.exists(): cfg = load_cfg(checkpoint) compute_lattice_metrics( consolidated_generation_path, metrics_path=target_dir / f"lattice_metrics_generate.json", ground_truth_path=cfg.data.datamodule.datasets[stage][0].save_path, ) # right now this returns nothing since it just plots the distribution if ( not consolidated_reconstruction_path.exists() and not consolidated_generation_path.exists() ): raise FileNotFoundError( f"Neither {consolidated_reconstruction_path=} nor {consolidated_generation_path=} exist." ) if __name__ == "__main__": cli()