"""Copyright (c) Meta Platforms, Inc. and affiliates.""" from __future__ import annotations from pathlib import Path from typing import Literal import torch from omegaconf import OmegaConf from tqdm import tqdm class AffineStatsParallelAlgorithm(torch.nn.Module): def __init__( self, shape: tuple | int, stable: bool = True, epsilon: float = 1e-10, ): """Accumulate mean and variance online using the "parallel algorithm" algorithm from [1]. Args: shape: shape of mean, variance, and std array. do not include batch dimension! stable: compute using the stable version of the algorithm [1] epsilon: added to the computation of the standard deviation for numerical stability. References: [1] https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm """ super().__init__() self.register_buffer("n", torch.tensor(0, dtype=torch.long)) self.register_buffer("_mean", torch.zeros(shape)) self.register_buffer("_M2", torch.zeros(shape)) self.register_buffer("epsilon", torch.tensor(epsilon)) if isinstance(shape, int): self.shape = (shape,) else: self.shape = shape self.stable = stable def _parallel_algorithm( self, x: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: assert x.shape[-len(self.shape) :] == self.shape na = self.n.clone() nb = x.shape[0] nab = na + nb xa = self._mean.clone() xb = x.mean(dim=0) delta = xb - xa if self.stable: xab = (na * xa + nb * xb) / nab else: xab = xa + delta * nb / nab m2a = self._M2.clone() m2b = ( x.var(dim=(0,), unbiased=False) * nb ) # do not use bessel's correction then multiply by total number of items in batch. m2ab = m2a + m2b + delta**2 * na * nb / nab return nab, xab, m2ab def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: self.n, self._mean, self._M2 = self._parallel_algorithm(x) return self.mean, self.std @property def mean(self) -> torch.Tensor: return self._mean @property def var(self) -> torch.Tensor: if self.n > 1: return self._M2 / (self.n - 1) else: return torch.zeros_like(self._M2) @property def std(self) -> torch.Tensor: return torch.sqrt(self.var + self.epsilon) collect_stats_on_options = Literal["coord", "lattice"] def compute_affine_stats( dataset: dataset_options, collect_stats_on: collect_stats_on_options, atom_type_manifold: atom_type_manifold_types, coord_manifold: coord_manifold_types, lattice_manifold: lattice_manifold_types, epochs: int = 5, analog_bits_scale: float = 1.0, ): from flowmm.cfg_utils import init_loaders from flowmm.rfm.manifold_getter import ManifoldGetter train_loader, *_ = init_loaders(dataset=dataset, batch_size=8192) manifold_getter = ManifoldGetter( atom_type_manifold=atom_type_manifold, coord_manifold=coord_manifold, lattice_manifold=lattice_manifold, analog_bits_scale=analog_bits_scale, dataset=dataset, ) inited = False for _ in tqdm(range(epochs)): for batch in tqdm(train_loader): ( x1, manifold, a_manifold, f_manifold, l_manifold, dims, mask_a_or_f, ) = manifold_getter( batch.batch, batch.atom_types, batch.frac_coords, batch.lengths, batch.angles, split_manifold=True, ) x0 = manifold.random(*x1.shape, dtype=x1.dtype, device=x1.device) 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) atom_types, frac_coords, lattices = manifold_getter.flatrep_to_georep( x_t, dims=dims, mask_a_or_f=mask_a_or_f, ) vals_x_t = { "atom": atom_types, "coord": frac_coords, "lattice": lattices, } u_atom_types, u_frac_coords, u_lattices = manifold_getter.flatrep_to_georep( u_t, dims=dims, mask_a_or_f=mask_a_or_f, ) vals_u_t = { "atom": u_atom_types, "coord": u_frac_coords, "lattice": u_lattices, } if not inited: aspa_x_t = AffineStatsParallelAlgorithm( vals_x_t[collect_stats_on].shape[-1] ) aspa_u_t = AffineStatsParallelAlgorithm( vals_u_t[collect_stats_on].shape[-1] ) inited = True aspa_x_t(vals_x_t[collect_stats_on]) aspa_u_t(vals_u_t[collect_stats_on]) print( f"{aspa_x_t.mean=}, {aspa_x_t.std=}; {aspa_u_t.mean=}, {aspa_u_t.std=}" ) return { "x_t_mean": aspa_x_t.mean, "x_t_std": aspa_x_t.std, "u_t_mean": aspa_u_t.mean, "u_t_std": aspa_u_t.std, } def get_affine_stats_filename( dataset: dataset_options, manifold_type: atom_type_manifold_types | coord_manifold_types | lattice_manifold_types, ) -> str: return f"stats_{dataset}_{manifold_type}.yaml" def get_affine_stats( dataset: dataset_options, manifold_type: atom_type_manifold_types | coord_manifold_types | lattice_manifold_types, path: str | Path | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if path is None: file = Path(__file__).parent / get_affine_stats_filename(dataset, manifold_type) file = file.resolve() else: file = Path(path) file = file.resolve() if file.exists(): stats = OmegaConf.load(str(file)) else: raise FileNotFoundError( f"{file=} does not exist, have you computed the stats with `compute_affine_stats` already?" ) return {k: torch.tensor(v) for k, v in stats.items()} if __name__ == "__main__": import yaml from flowmm.cfg_utils import dataset_options from flowmm.rfm.manifold_getter import ( atom_type_manifold_types, coord_manifold_types, lattice_manifold_types, ) file = Path(__file__).parent / "spd_stats.yaml" file = file.resolve() # calculate the stats of each dataset for dataset in tqdm(list(dataset_options.__args__)): for collect_stats_on in collect_stats_on_options.__args__: # if collect_stats_on == "atom": # pass # categorical are not standardized if collect_stats_on == "coord": # flat_torus_01 stats = compute_affine_stats( dataset=dataset, collect_stats_on=collect_stats_on, atom_type_manifold="null_manifold", coord_manifold="flat_torus_01", lattice_manifold="non_symmetric", ) # since translation is removed from the geodesic stats["u_t_mean"] = torch.zeros_like(stats["u_t_mean"]) # we do not standardize these since they goes through a sinusoidal embedding stats["x_t_mean"] = torch.zeros_like(stats["x_t_mean"]) stats["x_t_std"] = torch.ones_like(stats["x_t_std"]) file = Path(__file__).parent / get_affine_stats_filename( dataset, "flat_torus_01" ) file = file.resolve() with open(file, "w") as f: yaml.dump({k: v.tolist() for k, v in stats.items()}, f) # flat_torus_01_normal stats = compute_affine_stats( dataset=dataset, collect_stats_on=collect_stats_on, atom_type_manifold="null_manifold", coord_manifold="flat_torus_01_normal", lattice_manifold="non_symmetric", ) # since translation is removed from the geodesic stats["u_t_mean"] = torch.zeros_like(stats["u_t_mean"]) # we do not standardize these since they goes through a sinusoidal embedding stats["x_t_mean"] = torch.zeros_like(stats["x_t_mean"]) stats["x_t_std"] = torch.ones_like(stats["x_t_std"]) file = Path(__file__).parent / get_affine_stats_filename( dataset, "flat_torus_01_normal" ) file = file.resolve() with open(file, "w") as f: yaml.dump({k: v.tolist() for k, v in stats.items()}, f) # flat_torus_01_fixfirst stats = compute_affine_stats( dataset=dataset, collect_stats_on=collect_stats_on, atom_type_manifold="null_manifold", coord_manifold="flat_torus_01_fixfirst", lattice_manifold="non_symmetric", ) # since translation is removed from the geodesic stats["u_t_mean"] = torch.zeros_like(stats["u_t_mean"]) # we do not standardize these since they goes through a sinusoidal embedding stats["x_t_mean"] = torch.zeros_like(stats["x_t_mean"]) stats["x_t_std"] = torch.ones_like(stats["x_t_std"]) file = Path(__file__).parent / get_affine_stats_filename( dataset, "flat_torus_01_fixfirst" ) file = file.resolve() # flat_torus_01_fixfirst_normal stats = compute_affine_stats( dataset=dataset, collect_stats_on=collect_stats_on, atom_type_manifold="null_manifold", coord_manifold="flat_torus_01_fixfirst_normal", lattice_manifold="non_symmetric", ) # since translation is removed from the geodesic stats["u_t_mean"] = torch.zeros_like(stats["u_t_mean"]) # we do not standardize these since they goes through a sinusoidal embedding stats["x_t_mean"] = torch.zeros_like(stats["x_t_mean"]) stats["x_t_std"] = torch.ones_like(stats["x_t_std"]) file = Path(__file__).parent / get_affine_stats_filename( dataset, "flat_torus_01_fixfirst_normal" ) file = file.resolve() with open(file, "w") as f: yaml.dump({k: v.tolist() for k, v in stats.items()}, f) elif collect_stats_on == "lattice": # non_symmetric # this one is difficult because its multiplied by L # so we just leave it untransformed stats = { "x_t_mean": torch.zeros(9), "x_t_std": torch.ones(9), "u_t_mean": torch.zeros(9), "u_t_std": torch.ones(9), } file = Path(__file__).parent / get_affine_stats_filename( dataset, "non_symmetric" ) file = file.resolve() with open(file, "w") as f: yaml.dump({k: v.tolist() for k, v in stats.items()}, f) # spd_euclidean_geo stats = compute_affine_stats( dataset=dataset, collect_stats_on=collect_stats_on, atom_type_manifold="null_manifold", coord_manifold="flat_torus_01_fixfirst", lattice_manifold="spd_euclidean_geo", ) file = Path(__file__).parent / get_affine_stats_filename( dataset, "spd_euclidean_geo" ) file = file.resolve() with open(file, "w") as f: yaml.dump({k: v.tolist() for k, v in stats.items()}, f) # spd_riemanian_geo stats = compute_affine_stats( dataset=dataset, collect_stats_on=collect_stats_on, atom_type_manifold="null_manifold", coord_manifold="flat_torus_01_fixfirst", lattice_manifold="spd_riemanian_geo", ) file = Path(__file__).parent / get_affine_stats_filename( dataset, "spd_riemanian_geo" ) file = file.resolve() with open(file, "w") as f: yaml.dump({k: v.tolist() for k, v in stats.items()}, f) # lattice_params stats = compute_affine_stats( dataset=dataset, collect_stats_on=collect_stats_on, atom_type_manifold="null_manifold", coord_manifold="flat_torus_01_fixfirst", lattice_manifold="lattice_params", ) file = Path(__file__).parent / get_affine_stats_filename( dataset, "lattice_params" ) file = file.resolve() with open(file, "w") as f: yaml.dump({k: v.tolist() for k, v in stats.items()}, f) # lattice_params_normal_base stats = compute_affine_stats( dataset=dataset, collect_stats_on=collect_stats_on, atom_type_manifold="null_manifold", coord_manifold="flat_torus_01_fixfirst", lattice_manifold="lattice_params_normal_base", ) file = Path(__file__).parent / get_affine_stats_filename( dataset, "lattice_params_normal_base" ) file = file.resolve() with open(file, "w") as f: yaml.dump({k: v.tolist() for k, v in stats.items()}, f) else: raise ValueError()