"""Copyright (c) Meta Platforms, Inc. and affiliates.""" from __future__ import annotations import warnings from functools import partial from typing import Any, Literal import hydra import pytorch_lightning as pl import torch from geoopt.manifolds.euclidean import Euclidean from geoopt.manifolds.product import ProductManifold from omegaconf import DictConfig from pytorch_lightning.utilities.warnings import PossibleUserWarning from torch.func import jvp, vjp from torch_geometric.data import Data from torchmetrics import MeanMetric, MinMetric from diffcsp.common.data_utils import lattices_to_params_shape from manifm.ema import EMA from manifm.model_pl import ManifoldFMLitModule, div_fn, output_and_div from manifm.solvers import projx_integrator, projx_integrator_return_last from flowmm.model.arch import CSPNet, ProjectedConjugatedCSPNet from flowmm.model.solvers import ( projx_cond_integrator_return_last, projx_integrate_xt_to_x1, ) from flowmm.model.standardize import get_affine_stats from flowmm.rfm.manifold_getter import Dims, ManifoldGetter from flowmm.rfm.manifolds.spd import SPDGivenN, SPDNonIsotropicRandom from flowmm.rfm.vmap import VMapManifolds from omegaconf import OmegaConf def output_and_div( vecfield: callable[[torch.Tensor], torch.Tensor], x: torch.Tensor, v: torch.Tensor | None = None, div_mode: Literal["exact", "rademacher"] = "exact", ) -> tuple[torch.Tensor, torch.Tensor]: if div_mode == "exact": dx = vecfield(x) div = div_fn(vecfield)(x) else: dx, vjpfunc = vjp(vecfield, x) vJ = vjpfunc(v)[0] div = torch.sum(vJ * v, dim=-1) return dx, div class MaterialsRFMLitModule(ManifoldFMLitModule): def __init__(self, cfg: DictConfig): pl.LightningModule.__init__(self) self.cfg = cfg self.save_hyperparameters() self.manifold_getter = ManifoldGetter( atom_type_manifold=cfg.model.manifold_getter.atom_type_manifold, coord_manifold=cfg.model.manifold_getter.coord_manifold, lattice_manifold=cfg.model.manifold_getter.lattice_manifold, dataset=cfg.data.dataset_name, analog_bits_scale=cfg.model.manifold_getter.get("analog_bits_scale", None), length_inner_coef=cfg.model.manifold_getter.get("length_inner_coef", None), ) if "null" in cfg.model.manifold_getter.atom_type_manifold: cost_type = 0.0 cost_cross_ent = 0.0 warnings.warn( f"since {cfg.model.manifold_getter.atom_type_manifold=}, set {cost_type=} & {cost_cross_ent=}", PossibleUserWarning, ) else: cost_type = cfg.model.cost_type cost_cross_ent = cfg.model.cost_cross_ent self.costs = { "loss_a": cost_type, "loss_f": cfg.model.cost_coord, "loss_l": cfg.model.cost_lattice, "loss_ce": cost_cross_ent, } if cfg.model.affine_combine_costs: total_cost = sum([v for v in self.costs.values()]) self.costs = {k: v / total_cost for k, v in self.costs.items()} model: CSPNet = hydra.utils.instantiate( self.cfg.vectorfield, _convert_="partial" ) # Model of the vector field. cspnet = ProjectedConjugatedCSPNet( cspnet=model, manifold_getter=self.manifold_getter, lattice_affine_stats=get_affine_stats( cfg.data.dataset_name, self.manifold_getter.lattice_manifold ), coord_affine_stats=get_affine_stats( cfg.data.dataset_name, self.manifold_getter.coord_manifold ), ) if cfg.optim.get("ema_decay", None) is None: self.model = cspnet else: self.model = EMA( cspnet, cfg.optim.ema_decay, ) # use separate metric instance for train, val and test step # to ensure a proper reduction over the epoch self.train_metrics = { "loss": MeanMetric(), "loss_a": MeanMetric(), "loss_f": MeanMetric(), "loss_l": MeanMetric(), "loss_ce": MeanMetric(), "unscaled/loss_a": MeanMetric(), "unscaled/loss_f": MeanMetric(), "unscaled/loss_l": MeanMetric(), "unscaled/loss_ce": MeanMetric(), } self.val_metrics = { "loss": MeanMetric(), "loss_a": MeanMetric(), "loss_f": MeanMetric(), "loss_l": MeanMetric(), "loss_ce": MeanMetric(), "unscaled/loss_a": MeanMetric(), "unscaled/loss_f": MeanMetric(), "unscaled/loss_l": MeanMetric(), "unscaled/loss_ce": MeanMetric(), } self.test_metrics = { "loss": MeanMetric(), "loss_a": MeanMetric(), "loss_f": MeanMetric(), "loss_l": MeanMetric(), "loss_ce": MeanMetric(), "unscaled/loss_a": MeanMetric(), "unscaled/loss_f": MeanMetric(), "unscaled/loss_l": MeanMetric(), "unscaled/loss_ce": MeanMetric(), "nll": MeanMetric(), } # for logging best so far validation accuracy self.val_metrics_best = { "loss": MinMetric(), "loss_a": MinMetric(), "loss_f": MinMetric(), "loss_l": MinMetric(), "loss_ce": MinMetric(), "unscaled/loss_a": MinMetric(), "unscaled/loss_f": MinMetric(), "unscaled/loss_l": MinMetric(), "unscaled/loss_ce": MinMetric(), } if self.cfg.val.compute_nll: self.val_metrics["nll"] = MeanMetric() self.val_metrics_best["nll"] = MinMetric() def on_fit_start(self): super().on_fit_start() for stage in ["train", "val", "test"]: metrics = getattr(self, f"{stage}_metrics") for v in metrics.values(): v = v.to(self.device) @staticmethod def _annealing_schedule( t: torch.Tensor, slope: float, intercept: float ) -> torch.Tensor: return slope * torch.nn.functional.relu(t - intercept) + 1.0 @torch.no_grad() def sample( self, batch: Data, x0: torch.Tensor = None, num_steps: int = 1_000, entire_traj: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ( x1, manifold, a_manifold, f_manifold, l_manifold, dims, mask_a_or_f, ) = self.manifold_getter( batch.batch, batch.atom_types, batch.frac_coords, batch.lengths, batch.angles, split_manifold=True, ) if x0 is None: if self.cfg.base_distribution_from_data: x0 = self.manifold_getter( batch.batch, batch.atom_types_initial, batch.frac_coords_initial, batch.lengths_initial, batch.angles_initial, split_manifold=True, )[0] else: x0 = manifold.random(*x1.shape, dtype=x1.dtype, device=x1.device) else: x0 = x0.to(x1) return self.finish_sampling( x0=x0, manifold=manifold, a_manifold=a_manifold, f_manifold=f_manifold, l_manifold=l_manifold, dims=dims, num_atoms=batch.num_atoms, node2graph=batch.batch, mask_a_or_f=mask_a_or_f, num_steps=num_steps, entire_traj=entire_traj, ) @torch.no_grad() def gen_sample( self, node2graph: torch.LongTensor, dim_coords: int, x0: torch.Tensor = None, num_steps: int = 1_000, entire_traj: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ( shape, manifold, a_manifold, f_manifold, l_manifold, dims, mask_a_or_f, ) = self.manifold_getter.from_empty_batch( node2graph, dim_coords, split_manifold=True ) num_atoms = self.manifold_getter._get_num_atoms(mask_a_or_f) if x0 is None: assert not self.cfg.base_distribution_from_data, "Need to sample from the base distribution" x0 = manifold.random(*shape, device=node2graph.device) else: x0 = x0.to(device=node2graph.device) return self.finish_sampling( x0=x0, manifold=manifold, a_manifold=a_manifold, f_manifold=f_manifold, l_manifold=l_manifold, dims=dims, num_atoms=num_atoms, node2graph=node2graph, mask_a_or_f=mask_a_or_f, num_steps=num_steps, entire_traj=entire_traj, ) @torch.no_grad() def pred_sample( self, node2graph: torch.LongTensor, atom_types: torch.LongTensor, dim_coords: int, x0: torch.Tensor = None, num_steps: int = 1_000, entire_traj: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ( shape, manifold, a_manifold, f_manifold, l_manifold, dims, mask_a_or_f, ) = self.manifold_getter.from_only_atom_types( node2graph, atom_types, dim_coords, split_manifold=True ) num_atoms = self.manifold_getter._get_num_atoms(mask_a_or_f) if x0 is None: assert not self.cfg.base_distribution_from_data, "Need to sample from the base distribution" x0 = manifold.random(*shape, device=node2graph.device) else: x0 = x0.to(device=node2graph.device) return self.finish_sampling( x0=x0, manifold=manifold, a_manifold=a_manifold, f_manifold=f_manifold, l_manifold=l_manifold, dims=dims, num_atoms=num_atoms, node2graph=node2graph, mask_a_or_f=mask_a_or_f, num_steps=num_steps, entire_traj=entire_traj, ) @torch.no_grad() def finish_sampling( self, x0: torch.Tensor, manifold: VMapManifolds, a_manifold: VMapManifolds, f_manifold: VMapManifolds, l_manifold: VMapManifolds, dims: Dims, num_atoms: torch.LongTensor, node2graph: torch.LongTensor, # aka batch.batch mask_a_or_f: torch.BoolTensor, num_steps: int, entire_traj: bool, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: vecfield = partial( self.vecfield, num_atoms=num_atoms, node2graph=node2graph, dims=dims, mask_a_or_f=mask_a_or_f, ) compute_traj_velo_norms = self.cfg.integrate.get( "compute_traj_velo_norms", False ) c = self.cfg.integrate.get("inference_anneal_slope", 0.0) b = self.cfg.integrate.get("inference_anneal_offset", 0.0) anneal_types = self.cfg.integrate.get("inference_anneal_types", False) anneal_coords = self.cfg.integrate.get("inference_anneal_coords", True) anneal_lattice = self.cfg.integrate.get("inference_anneal_lattice", False) print( "Anneal Types:", anneal_types, "Anneal Coords:", anneal_coords, "Anneal Lattice:", anneal_lattice, ) def scheduled_fn_to_integrate( t: torch.Tensor, x: torch.Tensor, cond: torch.Tensor | None = None ) -> torch.Tensor: anneal_factor = self._annealing_schedule(t, c, b) out = vecfield( t=torch.atleast_2d(t), x=torch.atleast_2d(x), manifold=manifold, cond=torch.atleast_2d(cond) if isinstance(cond, torch.Tensor) else cond, ) if anneal_types: out[:, : dims.a].mul_(anneal_factor) if anneal_coords: out[:, dims.a : -dims.l].mul_(anneal_factor) if anneal_lattice: out[:, -dims.l :].mul_(anneal_factor) return out if self.cfg.model.get("self_cond", False): x1 = projx_cond_integrator_return_last( manifold, scheduled_fn_to_integrate, x0, t=torch.linspace(0, 1, num_steps + 1).to(x0.device), method=self.cfg.integrate.get("method", "euler"), projx=True, local_coords=False, pbar=True, ) return x1 elif entire_traj or compute_traj_velo_norms: xs, vs = projx_integrator( manifold, scheduled_fn_to_integrate, x0, t=torch.linspace(0, 1, num_steps + 1).to(x0.device), method=self.cfg.integrate.get("method", "euler"), projx=True, pbar=True, ) else: x1 = projx_integrator_return_last( manifold, scheduled_fn_to_integrate, x0, t=torch.linspace(0, 1, num_steps + 1).to(x0.device), method=self.cfg.integrate.get("method", "euler"), projx=True, local_coords=False, pbar=True, ) return x1 if compute_traj_velo_norms: s = 0 e = dims.a norm_a = a_manifold.inner( xs[..., s:e], vs[..., s:e], vs[..., s:e], data_in_dim=1 ) s = e e += dims.f norm_f = f_manifold.inner( xs[..., s:e], vs[..., s:e], vs[..., s:e], data_in_dim=1 ) s = e e += dims.l norm_l = l_manifold.inner( xs[..., s:e], vs[..., s:e], vs[..., s:e], data_in_dim=1 ) if entire_traj and compute_traj_velo_norms: return xs, norm_a, norm_f, norm_l elif entire_traj and not compute_traj_velo_norms: return xs elif not entire_traj and compute_traj_velo_norms: return xs[0], norm_a, norm_f, norm_l else: # this should happen due to logic above return xs[0] @torch.no_grad() def compute_exact_loglikelihood( self, batch: Data, stage: str, t1: float = 1.0, return_projx_error: bool = False, num_steps: int = 1_000, ): """Computes the negative log-likelihood of a batch of data.""" x1, manifold, dims, mask_a_or_f = self.manifold_getter( batch.batch, batch.atom_types, batch.frac_coords, batch.lengths, batch.angles, split_manifold=False, ) dim = sum(dims) nfe = [0] div_mode = self.cfg.integrate.get( "div_mode", "rademacher" ) # alternative: exact v = None if div_mode == "rademacher": v = torch.randint(low=0, high=2, size=x1.shape).to(x1) * 2 - 1 vecfield = partial( self.vecfield, num_atoms=batch.num_atoms, node2graph=batch.batch, dims=dims, mask_a_or_f=mask_a_or_f, ) def odefunc(t, tensor): nfe[0] += 1 t = t.to(tensor) x = tensor[..., :dim] l_vecfield = lambda x: vecfield( t=torch.atleast_2d(t), x=torch.atleast_2d(x), manifold=manifold ) dx, div = output_and_div(l_vecfield, x, v=v, div_mode=div_mode) if hasattr(manifold, "logdetG"): corr = jvp(manifold.logdetG, (x,), (dx,))[1] div = div + 0.5 * corr.to(div) div = div.reshape(-1, 1) del t, x return torch.cat([dx, div], dim=-1) # Solve ODE on the product manifold of data manifold x euclidean. prod_manis = [ ProductManifold((m, dim), (Euclidean(), 1)) for m in manifold.manifolds ] product_man = VMapManifolds(prod_manis) state1 = torch.cat([x1, torch.zeros_like(x1[..., :1])], dim=-1) with torch.no_grad(): state0 = projx_integrator_return_last( product_man, odefunc, state1, t=torch.linspace(t1, 0, num_steps + 1).to(x1), method=self.cfg.integrate.get("method", "euler"), projx=True, local_coords=False, pbar=True, ) # log number of function evaluations # self.log( # f"{stage}/nfe", # float(nfe[0]), # prog_bar=True, # logger=True, # batch_size=batch.batch_size, # ) x0, logdetjac = state0[..., :dim], state0[..., -1] x0_ = x0 x0 = manifold.projx(x0) # log how close the final solution is to the manifold. integ_error = (x0[..., :dim] - x0_[..., :dim]).abs().max() # self.log(f"{stage}/integ_error", integ_error, batch_size=batch.batch_size) logp0 = manifold.base_logprob(x0) logp1 = logp0 + logdetjac if self.cfg.get("normalize_loglik", True): logp1 = logp1 / batch.num_atoms # Mask out those that left the manifold masked_logp1 = logp1 if isinstance(manifold, (SPDGivenN, SPDNonIsotropicRandom)): mask = integ_error < 1e-5 # self.log( # f"{stage}/frac_within_manifold", # mask.sum() / mask.nelement(), # batch_size=batch.batch_size, # ) masked_logp1 = logp1[mask] if return_projx_error: return logp1, integ_error else: return masked_logp1 def loss_fn(self, batch: Data, *args, **kwargs) -> dict[str, torch.Tensor]: return self.rfm_loss_fn(batch, *args, **kwargs) def rfm_loss_fn( self, batch: Data, x0: torch.Tensor = None ) -> dict[str, torch.Tensor]: ( x1, manifold, a_manifold, f_manifold, l_manifold, dims, mask_a_or_f, ) = self.manifold_getter( batch.batch, batch.atom_types, batch.frac_coords, batch.lengths, batch.angles, split_manifold=True, ) if x0 is None: if self.cfg.base_distribution_from_data: x0 = self.manifold_getter( batch.batch, batch.atom_types_initial, batch.frac_coords_initial, batch.lengths_initial, batch.angles_initial, split_manifold=True, )[0] else: x0 = manifold.random(*x1.shape, dtype=x1.dtype, device=x1.device) vecfield = partial( self.vecfield, num_atoms=batch.num_atoms, node2graph=batch.batch, dims=dims, mask_a_or_f=mask_a_or_f, ) N = x1.shape[0] t = torch.rand(N, dtype=x1.dtype, device=x1.device).reshape(-1, 1) x1 = manifold.projx(x1) x_t, u_t = manifold.cond_u(x0, x1, t) x_t = x_t.reshape(N, x0.shape[-1]) u_t = u_t.reshape(N, x0.shape[-1]) # this projects out the mean from the tangent vectors # our model cannot predict it, so keeping it in inflates the loss u_t = manifold.proju(x_t, u_t) cond = None if self.cfg.model.self_cond: with torch.no_grad(): if torch.rand((1)) < 0.5: cond = projx_integrate_xt_to_x1( manifold, lambda t, x: vecfield( t=torch.atleast_2d(t), x=torch.atleast_2d(x), manifold=manifold, ), x_t, t, ).detach_() u_t_pred = vecfield(t=t, x=x_t, manifold=manifold, cond=cond) diff = u_t_pred - u_t max_num_atoms = mask_a_or_f.size(-1) dim_a_per_atom = dims.a / max_num_atoms dim_f_per_atom = dims.f / max_num_atoms if self.costs["loss_ce"] > 0.0: x1_pred_t = projx_integrate_xt_to_x1( manifold, None, # already computed velocity x_t, t, u_t_pred, ) a1_pred_t = x1_pred_t[:, : dims.a].reshape(N * max_num_atoms, -1) a1 = x1[:, : dims.a].reshape(N * max_num_atoms, -1) if self.manifold_getter.atom_type_manifold == "analog_bits": loss_ce = torch.einsum("bi,bi->b", a1_pred_t, a1).reshape( N, max_num_atoms ) loss_ce = torch.nn.functional.logsigmoid(loss_ce) * mask_a_or_f elif self.manifold_getter.atom_type_manifold == "simplex": a1 = self.manifold_getter._inverse_atomic_one_hot(a1) loss_ce = torch.nn.functional.cross_entropy( a1_pred_t, a1, reduce=False, ).reshape(N, max_num_atoms) loss_ce = loss_ce * mask_a_or_f else: raise ValueError( f"{self.manifold_getter.atom_type_manifold=} cannot do cross entropy" ) loss_ce = (loss_ce.sum(dim=-1) / mask_a_or_f.sum(dim=-1)).mean() else: loss_ce = torch.Tensor([0.0]).squeeze().to(diff) s = 0 e = dims.a loss_a = ( a_manifold.inner(x_t[:, s:e], diff[:, s:e], diff[:, s:e]).mean() / dim_a_per_atom ) # per dim, already per atom s = e e += dims.f loss_f = ( f_manifold.inner(x_t[:, s:e], diff[:, s:e], diff[:, s:e]).mean() / dim_f_per_atom ) # per dim, already per atom s = e e += dims.l loss_l = ( l_manifold.inner(x_t[:, s:e], diff[:, s:e], diff[:, s:e]).mean() / dims.l ) # per dim loss = ( self.costs["loss_a"] * loss_a + self.costs["loss_f"] * loss_f + self.costs["loss_l"] * loss_l + self.costs["loss_ce"] * loss_ce ) return { "loss": loss, "loss_a": self.costs["loss_a"] * loss_a, "loss_f": self.costs["loss_f"] * loss_f, "loss_l": self.costs["loss_l"] * loss_l, "loss_ce": self.costs["loss_ce"] * loss_ce, "unscaled/loss_a": loss_a, "unscaled/loss_f": loss_f, "unscaled/loss_l": loss_l, "unscaled/loss_ce": loss_ce, } def training_step(self, batch: Data, batch_idx: int): loss_dict = self.loss_fn(batch) if torch.isfinite(loss_dict["loss"]): # log train metrics for k, v in loss_dict.items(): self.log( f"train/{k}", v, batch_size=batch.batch_size, sync_dist = True ) self.train_metrics[k].update(v) else: # skip step if loss is NaN. print(f"Skipping iteration because loss is {loss_dict['loss'].item()}.") return None # we can return here dict with any tensors # and then read it in some callback or in `training_epoch_end()` below # remember to always return loss from `training_step()` or else backpropagation will fail! return loss_dict def training_epoch_end(self, outputs: list[Any]): # `outputs` is a list of dicts returned from `training_step()` for train_metric in self.train_metrics.values(): train_metric.reset() def shared_eval_step( self, batch: Data, batch_idx: int, stage: Literal["val", "test"], compute_loss: bool, compute_nll: bool, ) -> dict[str, torch.Tensor]: if stage not in ["val", "test"]: raise ValueError("stage must be 'val' or 'test'.") metrics = getattr(self, f"{stage}_metrics") out = {} if compute_loss: loss_dict = self.loss_fn(batch) for k, v in loss_dict.items(): self.log( f"{stage}/{k}", v, on_epoch=True, prog_bar=True, batch_size=batch.batch_size, sync_dist = True ) metrics[k].update(v) out.update(loss_dict) if compute_nll: nll_dict = {} logprob = self.compute_exact_loglikelihood( batch, stage, num_steps=self.cfg.integrate.get("num_steps", 1_000), ) nll = -logprob.mean() nll_dict[f"{stage}/nll"] = nll nll_dict[f"{stage}/nll_num_steps"] = self.cfg.integrate.num_steps # self.log( # f"{stage}/nll", # nll, # prog_bar=True, # batch_size=batch.batch_size, # ) # self.log( # f"{stage}/nll_num_steps", # self.cfg.integrate.num_steps, # prog_bar=True, # batch_size=batch.batch_size, # ) self.logger.log_metrics(nll_dict, step=self.global_step) metrics["nll"].update(nll) out.update(nll_dict) return out def compute_reconstruction( self, batch: Data, num_steps: int = 1_000, ) -> dict[str, torch.Tensor | Data]: *_, dims, mask_a_or_f = self.manifold_getter( batch.batch, batch.atom_types, batch.frac_coords, batch.lengths, batch.angles, split_manifold=False, ) if self.cfg.base_distribution_from_data: x0 = self.manifold_getter( batch.batch, batch.atom_types_initial, batch.frac_coords_initial, batch.lengths_initial, batch.angles_initial, split_manifold=True, )[0] else: x0 = None if self.cfg.integrate.get("compute_traj_velo_norms", False): recon, norms_a, norms_f, norms_l = self.sample( batch, num_steps=num_steps, x0=x0 ) norms = {"norms_a": norms_a, "norms_f": norms_f, "norms_l": norms_l} else: recon = self.sample(batch, num_steps=num_steps, x0=x0) norms = {} atom_types, frac_coords, lattices = self.manifold_getter.flatrep_to_crystal( recon, dims, mask_a_or_f ) lengths, angles = lattices_to_params_shape(lattices) out = { "atom_types": atom_types, "frac_coords": frac_coords, "lattices": lattices, "lengths": lengths, "angles": angles, "num_atoms": batch.num_atoms, "input_data_batch": batch, } out.update(norms) return out def compute_recon_trajectory( self, batch: Data, num_steps: int = 1_000, ) -> dict[str, torch.Tensor | Data]: *_, dims, mask_a_or_f = self.manifold_getter( batch.batch, batch.atom_types, batch.frac_coords, batch.lengths, batch.angles, split_manifold=False, ) if self.cfg.integrate.get("compute_traj_velo_norms", False): recon, norms_a, norms_f, norms_l = self.sample( batch, num_steps=num_steps, entire_traj=True ) norms = {"norms_a": norms_a, "norms_f": norms_f, "norms_l": norms_l} else: recon = self.sample(batch, num_steps=num_steps, entire_traj=True) norms = {} atom_types, frac_coords, lattices = [], [], [] lengths, angles = [], [] for recon_step in recon: a, f, l = self.manifold_getter.flatrep_to_crystal( recon_step, dims, mask_a_or_f ) lns, angs = lattices_to_params_shape(l) atom_types.append(a) frac_coords.append(f) lattices.append(l) lengths.append(lns) angles.append(angs) out = { "atom_types": torch.stack(atom_types, dim=0), "frac_coords": torch.stack(frac_coords, dim=0), "lattices": torch.stack(lattices, dim=0), "lengths": torch.stack(lengths, dim=0), "angles": torch.stack(angles, dim=0), "num_atoms": batch.num_atoms, "input_data_batch": batch, } out.update(norms) return out def compute_generation( self, batch: Data, dim_coords: int = 3, num_steps: int = 1_000, ) -> dict[str, torch.Tensor | Data]: *_, dims, mask_a_or_f = self.manifold_getter.from_empty_batch( batch.batch, dim_coords, split_manifold=False ) if self.cfg.base_distribution_from_data: x0 = self.manifold_getter( batch.batch, batch.atom_types_initial, batch.frac_coords_initial, batch.lengths_initial, batch.angles_initial, split_manifold=True, )[0] else: x0 = None if self.cfg.integrate.get("compute_traj_velo_norms", False): recon, norms_a, norms_f, norms_l = self.gen_sample( batch.batch, dim_coords, num_steps=num_steps, x0=x0 ) norms = {"norms_a": norms_a, "norms_f": norms_f, "norms_l": norms_l} else: recon = self.gen_sample( batch.batch, dim_coords, num_steps=num_steps, x0=x0 ) norms = {} atom_types, frac_coords, lattices = self.manifold_getter.flatrep_to_crystal( recon, dims, mask_a_or_f ) lengths, angles = lattices_to_params_shape(lattices) out = { "atom_types": atom_types, "frac_coords": frac_coords, "lattices": lattices, "lengths": lengths, "angles": angles, "num_atoms": batch.num_atoms, "input_data_batch": batch, } out.update(norms) return out def compute_gen_trajectory( self, batch: Data, dim_coords: int = 3, num_steps: int = 1_000, ) -> dict[str, torch.Tensor | Data]: *_, dims, mask_a_or_f = self.manifold_getter.from_empty_batch( batch.batch, dim_coords, split_manifold=False ) if self.cfg.integrate.get("compute_traj_velo_norms", False): recon, norms_a, norms_f, norms_l = self.gen_sample( batch.batch, dim_coords, num_steps=num_steps, entire_traj=True, ) norms = {"norms_a": norms_a, "norms_f": norms_f, "norms_l": norms_l} else: recon = self.gen_sample( batch.batch, dim_coords, num_steps=num_steps, entire_traj=True ) norms = {} atom_types, frac_coords, lattices = [], [], [] lengths, angles = [], [] for recon_step in recon: a, f, l = self.manifold_getter.flatrep_to_crystal( recon_step, dims, mask_a_or_f ) lns, angs = lattices_to_params_shape(l) atom_types.append(a) frac_coords.append(f) lattices.append(l) lengths.append(lns) angles.append(angs) out = { "atom_types": torch.stack(atom_types, dim=0), "frac_coords": torch.stack(frac_coords, dim=0), "lattices": torch.stack(lattices, dim=0), "lengths": torch.stack(lengths, dim=0), "angles": torch.stack(angles, dim=0), "num_atoms": batch.num_atoms, "input_data_batch": batch, } out.update(norms) return out def compute_prediction( self, batch: Data, dim_coords: int = 3, num_steps: int = 1_000, ) -> dict[str, torch.Tensor | Data]: *_, dims, mask_a_or_f = self.manifold_getter.from_only_atom_types( batch.batch, batch.atom_types, dim_coords, split_manifold=False ) if self.cfg.integrate.get("compute_traj_velo_norms", False): recon, norms_a, norms_f, norms_l = self.pred_sample( batch.batch, batch.atom_types, dim_coords, num_steps=num_steps ) norms = {"norms_a": norms_a, "norms_f": norms_f, "norms_l": norms_l} else: recon = self.pred_sample( batch.batch, batch.atom_types, dim_coords, num_steps=num_steps ) norms = {} atom_types, frac_coords, lattices = self.manifold_getter.flatrep_to_crystal( recon, dims, mask_a_or_f ) lengths, angles = lattices_to_params_shape(lattices) out = { "atom_types": atom_types, "frac_coords": frac_coords, "lattices": lattices, "lengths": lengths, "angles": angles, "num_atoms": batch.num_atoms, "input_data_batch": batch, } out.update(norms) return out def validation_step(self, batch: Data, batch_idx: int): return self.shared_eval_step( batch, batch_idx, stage="val", compute_loss=True, compute_nll=self.cfg.val.compute_nll, ) def validation_epoch_end(self, outputs: list[Any]): out = {} for key, val_metric in self.val_metrics.items(): val_metric_value = ( val_metric.compute() ) # get val accuracy from current epoch val_metric_best = self.val_metrics_best[key] val_metric_best.update(val_metric_value) self.log( f"val/best/{key}", val_metric_best.compute(), on_epoch=True, prog_bar=True, sync_dist = True ) val_metric.reset() out[key] = val_metric_value return out def test_step(self, batch: Any, batch_idx: int): return self.shared_eval_step( batch, batch_idx, stage="test", compute_loss=self.cfg.test.get("compute_loss", False), compute_nll=self.cfg.test.get("compute_nll", False), ) def test_epoch_end(self, outputs: list[Any]): for test_metric in self.test_metrics.values(): test_metric.reset() def predict_step(self, batch: Any, batch_idx: int): if not hasattr(batch, "frac_coords"): if "null" in self.cfg.model.manifold_getter.atom_type_manifold: if self.cfg.integrate.get("entire_traj", False): raise NotImplementedError() else: return self.compute_prediction( batch, dim_coords=self.cfg.data.get("dim_coords", 3), num_steps=self.cfg.integrate.get("num_steps", 1_000), ) else: if self.cfg.integrate.get("entire_traj", False): return self.compute_gen_trajectory( batch, dim_coords=self.cfg.data.get("dim_coords", 3), num_steps=self.cfg.integrate.get("num_steps", 1_000), ) else: return self.compute_generation( batch, dim_coords=self.cfg.data.get("dim_coords", 3), num_steps=self.cfg.integrate.get("num_steps", 1_000), ) else: # not generating or predicting new structures if self.cfg.integrate.get("entire_traj", False): return self.compute_recon_trajectory( batch, num_steps=self.cfg.integrate.get("num_steps", 1_000), ) else: return self.compute_reconstruction( batch, num_steps=self.cfg.integrate.get("num_steps", 1_000), ) def configure_optimizers(self): if not self.cfg.optim.optimizer.lr_diff: OmegaConf.set_struct(self.cfg.optim.optimizer, False) # 禁用结构模式 self.cfg.optim.optimizer.pop("lr_diff") self.cfg.optim.optimizer.pop("lr_backbone") self.cfg.optim.optimizer.pop("lr_head") OmegaConf.set_struct(self.cfg.optim.optimizer, True) # 可选：重新启用结构模式 optimizer = hydra.utils.instantiate( self.cfg.optim.optimizer, params=self.parameters(), _convert_="partial", ) else: backbone_param = [] head_param = [] for name, param in self.named_parameters(): if "out" in name: head_param.append(param) else: backbone_param.append(param) param_groups = [ {"params": backbone_param, "lr": self.cfg.optim.optimizer.lr_backbone}, {"params": head_param, "lr": self.cfg.optim.optimizer.lr_head} ] optimizer = torch.optim.AdamW( params = param_groups, weight_decay = self.cfg.optim.optimizer.weight_decay ) if self.cfg.optim.get("lr_scheduler", None) is not None: lr_scheduler = hydra.utils.instantiate( self.cfg.optim.lr_scheduler, optimizer, ) if isinstance(lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): return { "optimizer": optimizer, "lr_scheduler": { "scheduler": lr_scheduler, "interval": "epoch", "monitor": self.cfg.optim.monitor, "frequency": self.cfg.optim.frequency, }, } elif isinstance(lr_scheduler, torch.optim.lr_scheduler.CosineAnnealingLR): return { "optimizer": optimizer, "lr_scheduler": { "scheduler": lr_scheduler, "interval": self.cfg.optim.interval, }, } else: raise NotImplementedError("unsuported lr_scheduler") else: return {"optimizer": optimizer}