diff --git a/ppmat/datasets/__init__.py b/ppmat/datasets/__init__.py index 05d3de49..e0f8e660 100644 --- a/ppmat/datasets/__init__.py +++ b/ppmat/datasets/__init__.py @@ -51,6 +51,7 @@ from ppmat.datasets.split_mptrj_data import none_to_zero from ppmat.datasets.transform import build_transforms from ppmat.utils import logger +from ppmat.datasets.omatg_dataset import StructureDataset as OMATGStructureDataset __all__ = [ "MP20Dataset", @@ -65,10 +66,11 @@ "HighLevelWaterDataset", "MSDnmrDataset", "MatbenchDataset", - "DensityDataset", + "DensityDataset", "SmallDensityDataset", "SFINDataset", "OMol25Dataset", + "OMATGStructureDataset", ] INFO_CLASS_REGISTRY: Dict[str, type] = { diff --git a/ppmat/datasets/omatg_dataset.py b/ppmat/datasets/omatg_dataset.py new file mode 100644 index 00000000..eccba7ec --- /dev/null +++ b/ppmat/datasets/omatg_dataset.py @@ -0,0 +1,561 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the specific language governing permissions and +# limitations under the License. + +""" +Structure Dataset for OMatG (Open Materials Generation). + +Combines Structure, OMATGData, and StructureDataset into a single module. +""" + +import pickle +from pathlib import Path +from typing import Any, Optional, Sequence + +import lmdb +import numpy as np +import paddle +import pandas as pd +from ase import Atoms +from ase.symbols import Symbols + +from ppmat.datasets.geometric_data_type import Data + + +class Structure: + """Storage for crystalline structure with cell, atomic numbers, coordinates, properties, and metadata. + Supports Cartesian/fractional coordinate conversion and Niggli reduction. + """ + + def __init__( + self, + cell: paddle.Tensor, + atomic_numbers: paddle.Tensor, + pos: paddle.Tensor, + property_dict: Optional[dict[str, Any]] = None, + metadata: Optional[dict[str, Any]] = None, + pos_is_fractional: bool = False, + ) -> None: + if cell.shape != (3, 3): + raise ValueError(f"cell must be 3x3, got {cell.shape}") + if atomic_numbers.dim() != 1: + raise ValueError(f"atomic_numbers must be 1D, got {atomic_numbers.dim()}D") + if pos.shape[0] != len(atomic_numbers) or pos.shape[1] != 3: + raise ValueError(f"pos must be (N, 3), got {pos.shape}") + + self._cell = cell + self._atomic_numbers = atomic_numbers + self._pos = pos + self._property_dict = property_dict if property_dict is not None else {} + self._metadata = metadata if metadata is not None else {} + self._fractional = pos_is_fractional + + @property + def cell(self) -> paddle.Tensor: + return self._cell + + @property + def atomic_numbers(self) -> paddle.Tensor: + return self._atomic_numbers + + @property + def symbols(self) -> list[str]: + return list(Symbols(self._atomic_numbers.numpy())) + + @property + def pos(self) -> paddle.Tensor: + return self._pos + + @property + def pos_is_fractional(self) -> bool: + return self._fractional + + @property + def property_dict(self) -> dict[str, Any]: + return self._property_dict + + @property + def metadata(self) -> dict[str, Any]: + return self._metadata + + def to(self, floating_point_precision: str) -> None: + valid_precisions = ["float32", "float64", "float16", "bfloat16"] + if floating_point_precision not in valid_precisions: + raise ValueError( + f"Unsupported floating point precision: {floating_point_precision}. " + f"Supported precisions are {valid_precisions}." + ) + self._cell = self._cell.cast(floating_point_precision) + self._pos = self._pos.cast(floating_point_precision) + for key, value in self._property_dict.items(): + if paddle.is_tensor(value) and value.dtype in [ + paddle.float32, + paddle.float64, + paddle.float16, + paddle.bfloat16, + ]: + self._property_dict[key] = value.cast(floating_point_precision) + + def get_ase_atoms(self) -> Atoms: + if self._fractional: + return Atoms( + numbers=self.atomic_numbers.tolist(), + scaled_positions=self.pos.numpy(), + cell=self.cell.numpy(), + pbc=True, + info=self.property_dict | self.metadata, + ) + else: + return Atoms( + numbers=self.atomic_numbers.tolist(), + positions=self.pos.numpy(), + cell=self.cell.numpy(), + pbc=True, + info=self.property_dict | self.metadata, + ) + + def niggli_reduce(self) -> None: + from pymatgen.core import Structure as PmgStructure, Element + + species = [Element.from_Z(int(z)).symbol for z in self._atomic_numbers.numpy()] + pmg = PmgStructure( + lattice=self._cell.numpy(), + species=species, + coords=self._pos.numpy(), + coords_are_cartesian=not self._fractional, + ) + reduced = pmg.get_reduced_structure(reduction_algo="niggli") + self._cell = paddle.to_tensor(reduced.lattice.matrix, dtype=self._cell.dtype) + self._pos = paddle.to_tensor(reduced.frac_coords, dtype=self._pos.dtype) + self._fractional = True + + def convert_to_fractional(self) -> None: + if not self._fractional: + with paddle.no_grad(): + self._pos = paddle.remainder( + paddle.linalg.solve(self._cell, self._pos.T).T, + 1.0, + ) + self._fractional = True + + def convert_to_cartesian(self) -> None: + if self._fractional: + with paddle.no_grad(): + self._pos = paddle.matmul(self._pos, self._cell) + self._fractional = False + + def __repr__(self) -> str: + return ( + f"Structure(cell={self._cell.shape}, " + f"atomic_numbers={self._atomic_numbers.shape}, " + f"pos={self._pos.shape}, " + f"fractional={self._fractional})" + ) + + +class OMATGData(Data): + """Representation of single/batch crystal structures. + Batch format: n_atoms(batch_size,), species(total_atoms,), cell(batch,3,3), + pos(total_atoms,3), pos_is_fractional(batch_size,), ptr(batch+1,). + """ + + _FIELD_NAMES = ("n_atoms", "species", "cell", "pos", "pos_is_fractional", "batch", "ptr") + + def __init__(self, structure: Optional[Structure] = None, **kwargs) -> None: + super().__init__(**kwargs) + if structure is None: + self.n_atoms = None + self.species = None + self.cell = None + self.pos = None + self.pos_is_fractional = None + self.property_dict = None + self.batch = None + self.ptr = None + else: + self._from_structure(structure) + + def _from_structure(self, structure: Structure) -> None: + n = len(structure.atomic_numbers) + self.n_atoms = paddle.to_tensor([n], dtype="int64") + self.species = structure.atomic_numbers + self.cell = structure.cell.unsqueeze(0) + self.pos = structure.pos + self.pos_is_fractional = paddle.to_tensor( + [structure.pos_is_fractional], dtype="bool" + ) + self.property_dict = [structure.property_dict] + self.batch = paddle.zeros([n], dtype="int64") + self.ptr = paddle.to_tensor([0, n], dtype="int64") + + @classmethod + def from_batch( + cls, structures: list[Structure], concatenate: bool = True + ) -> "OMATGData": + if not concatenate: + return [cls(s) for s in structures] + + if len(structures) == 0: + return cls() + + if len(structures) == 1: + return cls(structures[0]) + + cells = [] + species_list = [] + pos_list = [] + pos_is_fractional_list = [] + n_atoms_list = [] + batch_indices = [] + properties_list = [] + + for i, struct in enumerate(structures): + n_atoms_list.append(len(struct.atomic_numbers)) + species_list.append(struct.atomic_numbers) + cells.append(struct.cell) + pos_list.append(struct.pos) + pos_is_fractional_list.append( + paddle.to_tensor([struct.pos_is_fractional], dtype="bool") + ) + batch_indices.append( + paddle.full([len(struct.atomic_numbers)], i, dtype="int64") + ) + properties_list.append(struct.property_dict) + + data = cls() + data.n_atoms = paddle.to_tensor(n_atoms_list, dtype="int64") + data.species = paddle.concat(species_list) + data.cell = paddle.stack(cells) + data.pos = paddle.concat(pos_list) + data.pos_is_fractional = paddle.concat(pos_is_fractional_list) + data.batch = paddle.concat(batch_indices) + data.ptr = paddle.concat( + [ + paddle.to_tensor([0], dtype="int64"), + paddle.cumsum(data.n_atoms, axis=0).cast("int64"), + ] + ) + data.property_dict = properties_list + + return data + + @classmethod + def from_collate_dict(cls, data: dict) -> "OMATGData": + num_atoms = data["num_atoms"] + batch = data["node2graph"] + ptr = paddle.concat([ + paddle.to_tensor([0], dtype="int64"), + paddle.cumsum(num_atoms, axis=0).cast("int64"), + ]) + d = cls() + d.n_atoms = num_atoms + d.species = data["atom_types"] + d.cell = data["lattices"] + d.pos = data["frac_coords"] + d.pos_is_fractional = paddle.ones_like(num_atoms, dtype="bool") + d.batch = batch + d.ptr = ptr + d.property_dict = {} + return d + + @property + def num_graphs(self) -> int: + if self.n_atoms is None: + return 0 + return len(self.n_atoms) + + @property + def num_atoms(self) -> int: + if self.species is None: + return 0 + return len(self.species) + + def get_graph(self, idx: int) -> Structure: + if idx < 0 or idx >= self.num_graphs: + raise IndexError(f"Index {idx} out of range for batch of {self.num_graphs}") + + start = int(self.ptr[idx]) + end = int(self.ptr[idx + 1]) + + species = self.species[start:end] + pos = self.pos[start:end] + cell = self.cell[idx] + pos_is_fractional = bool(self.pos_is_fractional[idx]) + + if self.property_dict and idx < len(self.property_dict): + prop = self.property_dict[idx] + else: + prop = {} + + return Structure( + cell=cell, + atomic_numbers=species, + pos=pos, + property_dict=prop, + metadata={}, + pos_is_fractional=pos_is_fractional, + ) + + def slice(self, idx: int) -> slice: + start = int(self.ptr[idx]) + end = int(self.ptr[idx + 1]) + return slice(start, end) + + @classmethod + def _from_dict(cls, data_dict: dict[str, Any]) -> "OMATGData": + data = cls() + for name in cls._FIELD_NAMES: + setattr(data, name, data_dict.get(name)) + data.property_dict = data_dict.get("property_dict") + return data + + def set_field(self, field_name: str, value: paddle.Tensor) -> None: + if field_name not in self._FIELD_NAMES: + raise ValueError(f"Unknown field: {field_name}") + setattr(self, field_name, value) + + def get_field(self, field_name: str) -> paddle.Tensor: + if field_name not in self._FIELD_NAMES: + raise ValueError(f"Unknown field: {field_name}") + return getattr(self, field_name) + + def __repr__(self) -> str: + return ( + f"OMATGData(num_graphs={self.num_graphs}, " + f"num_atoms={self.num_atoms}, " + f"cell={self.cell.shape if self.cell is not None else None})" + ) + + +class StructureDataset(paddle.io.Dataset): + """ + Dataset for reading crystalline structures from several file formats. + + This dataset optionally allows for lazy reading of the structures from LMDB files. + + :param file_path: + Path to the file containing the structures. + Supported formats are .lmdb, .csv, and .parquet. + :type file_path: str + :param property_keys: + An optional sequence of property keys that should be read from the file. + Defaults to None. + :type property_keys: Optional[Sequence[str]] + :param lazy_storage: + Whether to read the structures lazily from a LMDB file when they are requested. + Defaults to True. + :type lazy_storage: bool + :param convert_to_fractional: + Whether to convert the atomic positions to fractional coordinates. + Defaults to True. + :type convert_to_fractional: bool + :param niggli_reduce: + Whether to apply a Niggli reduction to the returned structures. + Defaults to False. + :type niggli_reduce: bool + """ + + def __init__( + self, + file_path: str, + property_keys: Optional[Sequence[str]] = None, + lazy_storage: bool = True, + convert_to_fractional: bool = True, + niggli_reduce: bool = False, + ) -> None: + self.file_path = file_path + self.property_keys = property_keys if property_keys is not None else [] + self.lazy_storage = lazy_storage + self.convert_to_fractional = convert_to_fractional + self.niggli_reduce = niggli_reduce + self._env = None + + file_format = Path(file_path).suffix.lower() + + if file_format == ".lmdb": + self._init_from_lmdb() + elif file_format == ".csv": + self._init_from_csv() + elif file_format == ".parquet": + self._init_from_parquet() + else: + raise ValueError(f"Unsupported file format: {file_format}") + + @property + def env(self): + if self._env is None: + self._env = lmdb.open( + self.file_path, + readonly=True, + lock=False, + readahead=False, + meminit=False, + ) + return self._env + + def _init_from_lmdb(self) -> None: + temp_env = lmdb.open( + self.file_path, + readonly=True, + lock=False, + readahead=False, + meminit=False, + ) + + try: + with temp_env.begin() as txn: + all_keys = [ + k.decode("ascii") if isinstance(k, bytes) else k + for k in txn.cursor().iternext(values=False) + ] + self.keys = [k for k in all_keys if not k.startswith("__")] + + if not self.lazy_storage: + self.structures = [] + with temp_env.begin() as txn: + for key in self.keys: + key_bytes = ( + key.encode("ascii") if isinstance(key, str) else key + ) + data = pickle.loads(txn.get(key_bytes)) + structure = self._create_structure(data) + self.structures.append(structure) + finally: + temp_env.close() + + def _extract_property_dict(self, data) -> dict[str, Any]: + property_dict = {} + for key in self.property_keys: + if key in data: + val = data[key] + if isinstance(val, (int, float, list, np.ndarray)): + property_dict[key] = paddle.to_tensor(val, dtype="float32") + return property_dict + + def _apply_transforms(self, structure: Structure) -> None: + if self.convert_to_fractional: + structure.convert_to_fractional() + if self.niggli_reduce: + structure.niggli_reduce() + + def _init_from_csv(self) -> None: + from ppmat.datasets.build_structure import BuildStructure + + df = pd.read_csv(self.file_path) + if "cif" not in df.columns: + raise KeyError( + f"CSV file does not contain 'cif' column. " + f"Available columns: {list(df.columns)}" + ) + + self.structures = [] + for _, row in df.iterrows(): + pmg = BuildStructure.build_one(row["cif"], "cif_str") + + cell = paddle.to_tensor(pmg.lattice.matrix, dtype="float32") + atomic_numbers = paddle.to_tensor(pmg.atomic_numbers, dtype="int64") + pos = paddle.to_tensor(pmg.cart_coords, dtype="float32") + + property_dict = self._extract_property_dict(row) + structure = Structure( + cell=cell, + atomic_numbers=atomic_numbers, + pos=pos, + property_dict=property_dict if property_dict else None, + pos_is_fractional=False, + ) + self._apply_transforms(structure) + self.structures.append(OMATGData(structure)) + + self.keys = list(range(len(self.structures))) + self.lazy_storage = False + + def _init_from_parquet(self) -> None: + df = pd.read_parquet(self.file_path) + required_cols = ["positions", "cell", "atomic_numbers"] + for col in required_cols: + if col not in df.columns: + raise KeyError( + f"Parquet file missing '{col}'. " + f"Available columns: {list(df.columns)}" + ) + + self.structures = [] + for _, row in df.iterrows(): + atomic_numbers = paddle.to_tensor( + np.asarray(row["atomic_numbers"], dtype=np.int64), dtype="int64" + ) + pos = paddle.to_tensor(np.stack(row["positions"]), dtype="float32") + cell = paddle.to_tensor(np.stack(row["cell"]), dtype="float32") + + property_dict = self._extract_property_dict(row) + structure = Structure( + cell=cell, + atomic_numbers=atomic_numbers, + pos=pos, + property_dict=property_dict if property_dict else None, + pos_is_fractional=False, + ) + self._apply_transforms(structure) + self.structures.append(OMATGData(structure)) + + self.keys = list(range(len(self.structures))) + self.lazy_storage = False + + def _create_structure(self, data: dict[str, Any]) -> OMATGData: + cell = paddle.to_tensor(data["cell"], dtype="float32") + atomic_numbers = paddle.to_tensor(data["atomic_numbers"], dtype="int64") + pos = paddle.to_tensor(data["pos"], dtype="float32") + + property_dict = self._extract_property_dict(data) + structure = Structure( + cell=cell, + atomic_numbers=atomic_numbers, + pos=pos, + property_dict=property_dict if property_dict else None, + pos_is_fractional=False, + ) + self._apply_transforms(structure) + return OMATGData(structure) + + def __len__(self) -> int: + if self.lazy_storage: + return len(self.keys) + else: + return len(self.structures) + + def __getitem__(self, idx: int) -> OMATGData: + if self.lazy_storage: + with self.env.begin() as txn: + key = self.keys[idx] + key_bytes = key.encode("ascii") if isinstance(key, str) else key + + value = txn.get(key_bytes) + if value is None: + raise KeyError(f"Key {key} not found in LMDB") + + data = pickle.loads(value) + + return self._create_structure(data) + else: + return self.structures[idx] + + def __del__(self) -> None: + if self._env is not None: + self._env.close() + self._env = None + + +__all__ = [ + "OMATGData", +] diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py index d258bd2a..9cd658de 100644 --- a/ppmat/models/__init__.py +++ b/ppmat/models/__init__.py @@ -42,6 +42,7 @@ from ppmat.models.megnet.megnet import MEGNetPlus from ppmat.models.infgcn.infgcn import InfGCN from ppmat.models.mateno.mateno import MatENO +from ppmat.models.omatg.model import OMATGCSPNet, OMATGCSPNetFull from ppmat.models.sfin.sfin import SFIN from ppmat.utils import download from ppmat.utils import logger @@ -68,6 +69,8 @@ "DiffNMR", "InfGCN", "MatENO", + "OMATGCSPNet", + "OMATGCSPNetFull", "SFIN", ] diff --git a/ppmat/models/omatg/__init__.py b/ppmat/models/omatg/__init__.py new file mode 100644 index 00000000..fbc68230 --- /dev/null +++ b/ppmat/models/omatg/__init__.py @@ -0,0 +1,232 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OMatG: Open Materials Generation for crystal structure prediction. +Based on Stochastic Interpolants (ICML 2025, NeurIPS 2025). +""" + +import os +import os.path as osp + +from ppmat.models.omatg.model import OMATGCSPNetFull +from ppmat.utils import download +from ppmat.utils import logger +from ppmat.utils import save_load + +# OMATG_WEIGHTS: OMatG has 57+ weight files across 6 datasets and 11+ variants. +# These are NOT registered in ppmat/models/__init__.py MODEL_REGISTRY because +# MODEL_REGISTRY binds one model_name -> one zip URL, while OMatG needs +# dataset x variant cross-product per .pdparams file (not zip bundles with yaml). +# A single zip containing all variants would be >10GB, and MODEL_REGISTRY +# requires each entry to be a self-contained zip with model_name.yaml. +# We keep per-weight URLs here for fine-grained download via build_omatg_model(). +OMATG_WEIGHTS = { + "perov_5_csp": { + "encdec_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/EncDec-ODE-Gamma.pdparams", + "encdec_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/EncDec-SDE-Gamma.pdparams", + "linear_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Linear-ODE-Gamma.pdparams", + "linear_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Linear-ODE.pdparams", + "linear_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Linear-SDE-Gamma.pdparams", + "trig_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Trig-ODE-Gamma.pdparams", + "trig_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Trig-ODE.pdparams", + "trig_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Trig-SDE-Gamma.pdparams", + "vesbd_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/VESBD-ODE.pdparams", + "vpsbd_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/VPSBD-ODE.pdparams", + "vpsbd_sde": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/VPSBD-SDE.pdparams", + }, + "mpts_52_csp": { + "encdec_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/EncDec-ODE-Gamma.pdparams", + "encdec_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/EncDec-SDE-Gamma.pdparams", + "linear_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Linear-ODE-Gamma.pdparams", + "linear_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Linear-ODE.pdparams", + "linear_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Linear-SDE-Gamma.pdparams", + "trig_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Trig-ODE-Gamma.pdparams", + "trig_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Trig-ODE.pdparams", + "trig_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Trig-SDE-Gamma.pdparams", + }, + "mp_20_dng": { + "encdec_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/EncDec-ODE-Gamma.pdparams", + "encdec_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/EncDec-SDE-Gamma.pdparams", + "linear_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Linear-ODE-Gamma.pdparams", + "linear_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Linear-ODE.pdparams", + "linear_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Linear-SDE-Gamma.pdparams", + "trig_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Trig-ODE-Gamma.pdparams", + "trig_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Trig-ODE.pdparams", + "trig_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Trig-SDE-Gamma.pdparams", + "vesbd_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/VESBD-ODE.pdparams", + "vpsbd_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/VPSBD-ODE.pdparams", + "vpsbd_sde": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/VPSBD-SDE.pdparams", + }, + "mp_20_csp": { + "encdec_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/EncDec-ODE-Gamma.pdparams", + "encdec_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/EncDec-SDE-Gamma.pdparams", + "linear_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Linear-ODE-Gamma.pdparams", + "linear_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Linear-ODE.pdparams", + "linear_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Linear-SDE-Gamma.pdparams", + "trig_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Trig-ODE-Gamma.pdparams", + "trig_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Trig-ODE.pdparams", + "trig_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Trig-SDE-Gamma.pdparams", + "vesbd_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/VESBD-ODE.pdparams", + "vpsbd_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/VPSBD-ODE.pdparams", + "vpsbd_sde": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/VPSBD-SDE.pdparams", + }, + "alex_mp_20_csp": { + "encdec_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/EncDec-ODE-Gamma.pdparams", + "encdec_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/EncDec-SDE-Gamma.pdparams", + "linear_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Linear-ODE-Gamma.pdparams", + "linear_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Linear-ODE.pdparams", + "linear_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Linear-SDE-Gamma.pdparams", + "trig_ode_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Trig-ODE-Gamma.pdparams", + "trig_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Trig-ODE.pdparams", + "trig_sde_gamma": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Trig-SDE-Gamma.pdparams", + "vesbd_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/VESBD-ODE.pdparams", + "vpsbd_ode": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/VPSBD-ODE.pdparams", + "vpsbd_sde": "https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/VPSBD-SDE.pdparams", + }, +} + +__all__ = [ + "OMATG_WEIGHTS", + "build_omatg_model", + "get_omatg_model_url", +] + + +def build_omatg_model(dataset: str, variant: str, weights_name: str = None): + """Build OMatG model with automatic weight downloading. + + Uses the project-wide factory pattern (build_model from ppmat.models) to + construct the model from configuration. + + Args: + dataset: Dataset name, e.g., "mp_20_csp", "perov_5_csp", "mpts_52_csp", + "mp_20_dng", "alex_mp_20_csp" + variant: Model variant, e.g., "linear_ode", "trig_ode_gamma", "encdec_sde_gamma" + weights_name: Specific weight file name (optional). If None, uses the default + weight file corresponding to the variant. + + Returns: + Loaded model and configuration dictionary + + Example: + >>> model, config = build_omatg_model("mp_20_csp", "linear_ode") + >>> model, config = build_omatg_model("perov_5_csp", "trig_ode_gamma", "custom.pdparams") + """ + # Validate dataset and variant + if dataset not in OMATG_WEIGHTS: + available_datasets = list(OMATG_WEIGHTS.keys()) + raise ValueError( + f"Unknown dataset: {dataset}. Available datasets: {available_datasets}" + ) + + if variant not in OMATG_WEIGHTS[dataset]: + available_variants = list(OMATG_WEIGHTS[dataset].keys()) + raise ValueError( + f"Unknown variant: {variant} for dataset {dataset}. " + f"Available variants: {available_variants}" + ) + + # Get weight URL + weight_url = OMATG_WEIGHTS[dataset][variant] + + logger.info(f"Building OMatG model: {dataset} / {variant}") + logger.info(f"Weight URL: {weight_url}") + + # Download weight to dataset-specific cache dir to avoid filename collision. + weight_path = osp.join(download.WEIGHTS_HOME, f"omatg_{dataset}", + weight_url.split("/")[-1]) + if not osp.exists(weight_path): + os.makedirs(osp.dirname(weight_path), exist_ok=True) + weight_path = download._download(weight_url, osp.dirname(weight_path)) + else: + logger.message(f"Found {weight_path} exists, skip downloading.") + logger.info(f"Weight saved to: {weight_path}") + + # If custom weights_name is specified, use it; otherwise use the variant name + if weights_name is None: + weight_filename = weight_url.split("/")[-1] + weights_name = weight_filename.replace(".pdparams", "") + logger.info(f"Using default weights: {weights_name}") + + # DNG (Discrete Flow Matching with Mask) variants need pred_type=True + is_dng = "dng" in dataset.lower() + + # Build OMATGCSPNetFull directly with pretrained-compatible default params + try: + model = OMATGCSPNetFull( + hidden_dim=512, + num_layers=6, + max_atoms=100, + act_fn="silu", + dis_emb="sin", + num_freqs=128, + edge_style="fc", + cutoff=7.0, + max_neighbors=20, + ln=True, + ip=True, + smooth=False, + pred_type=is_dng, + pred_scalar=False, + time_embed_dim=256, + ) + if is_dng: + model.enable_masked_species() + + save_load.load_pretrain(model, weight_path, weights_name) + + logger.info(f"Successfully built and loaded OMatG model: {dataset}/{variant}") + + return model, { + "dataset": dataset, + "variant": variant, + "weights_name": weights_name, + } + + except Exception as e: + logger.error(f"Failed to build OMatG model: {e}") + raise + + +def get_omatg_model_url(dataset: str, variant: str) -> str: + """Get the weight URL for an OMatG model variant. + + This is a convenience function for quickly getting the weight URL + without building the full model. + + Args: + dataset: Dataset name + variant: Model variant + + Returns: + Weight URL string + + Example: + >>> url = get_omatg_model_url("mp_20_csp", "linear_ode") + >>> weight_path = download.get_weights_path_from_url(url) + """ + if dataset not in OMATG_WEIGHTS: + available_datasets = list(OMATG_WEIGHTS.keys()) + raise ValueError( + f"Unknown dataset: {dataset}. Available datasets: {available_datasets}" + ) + + if variant not in OMATG_WEIGHTS[dataset]: + available_variants = list(OMATG_WEIGHTS[dataset].keys()) + raise ValueError( + f"Unknown variant: {variant} for dataset {dataset}. " + f"Available variants: {available_variants}" + ) + + return OMATG_WEIGHTS[dataset][variant] diff --git a/ppmat/models/omatg/model.py b/ppmat/models/omatg/model.py new file mode 100644 index 00000000..52cbf224 --- /dev/null +++ b/ppmat/models/omatg/model.py @@ -0,0 +1,679 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OMATG-specific CSPNet with knn graph, internal time embedding, dual outputs. + +Extends diffcsp.CSPNet to add: + - knn graph construction via radius_graph_pbc + - internal time embedding (SinusoidalTimeEmbeddings) + - dual output heads (coord_out_2, lattice_out_2, type_out_2) + - species_shift and enable_masked_species() for DNG mode +""" + +import paddle +import paddle.nn as nn + +from ase.geometry.cell import cellpar_to_cell + +from ppmat.datasets.omatg_dataset import OMATGData, Structure +from ppmat.losses import MSELoss +from ppmat.models.diffcsp.diffcsp import CSPNet +from ppmat.utils.crystal import radius_graph_pbc, frac_to_cart_coords_with_lattice +from ppmat.utils.misc import repeat_blocks + +__all__ = [ + "OMATGCSPNetFull", + "IndependentSampler", +] + + +class OMATGCSPNet(CSPNet): + """OMATG-specific CSPNet extending diffcsp backbone.""" + + def __init__( + self, + hidden_dim=128, + num_layers=4, + max_atoms=100, + act_fn="silu", + dis_emb="sin", + num_freqs=10, + edge_style="fc", + cutoff=6.0, + max_neighbors=20, + ln=False, + ip=True, + smooth=False, + pred_type=False, + pred_scalar=False, + time_embed_dim=None, + ): + super().__init__( + hidden_dim=hidden_dim, + latent_dim=time_embed_dim if time_embed_dim is not None else 1, + num_layers=num_layers, + act_fn=act_fn, + dis_emb=dis_emb, + num_freqs=num_freqs, + edge_style=edge_style, + ln=ln, + ip=ip, + smooth=smooth, + pred_type=pred_type, + prop_dim=hidden_dim, + pred_scalar=pred_scalar, + num_classes=max_atoms, + ) + # OMATG does not use property embedding; clear prop_mlp created by CSPNet + for i in range(num_layers): + getattr(self, "csp_layer_%d" % i).prop_mlp = None + + self.hidden_dim = hidden_dim + self.cutoff = cutoff + self.max_neighbors = max_neighbors + self.species_shift = 1 + + self.time_embed_dim = time_embed_dim + if time_embed_dim is not None: + from ppmat.models.common.time_embedding import SinusoidalTimeEmbeddings + + self.time_embedder = SinusoidalTimeEmbeddings(time_embed_dim) + self.atom_latent_emb = nn.Linear( + hidden_dim + time_embed_dim, hidden_dim + ) + else: + self.time_embedder = None + self.atom_latent_emb = nn.Linear(hidden_dim + 1, hidden_dim) + + self.dis_dim = (num_freqs * 2 * 3) if dis_emb == "sin" else 0 + + self.coord_out_2 = nn.Linear(hidden_dim, 3, bias_attr=False) + self.lattice_out_2 = nn.Linear(hidden_dim, 9, bias_attr=False) + if self.pred_type: + self.type_out_2 = nn.Linear(hidden_dim, max_atoms) + + def enable_masked_species(self): + self.node_embedding = nn.Embedding(self.num_classes + 1, self.hidden_dim) + self.species_shift = 0 + + def reorder_symmetric_edges(self, edge_index, cell_offsets, neighbors, edge_vector): + mask_sep_atoms = edge_index[0] < edge_index[1] + cell_earlier = ( + (cell_offsets[:, 0] < 0) + | ((cell_offsets[:, 0] == 0) & (cell_offsets[:, 1] < 0)) + | ( + (cell_offsets[:, 0] == 0) + & (cell_offsets[:, 1] == 0) + & (cell_offsets[:, 2] < 0) + ) + ) + mask_same_atoms = edge_index[0] == edge_index[1] + mask_same_atoms = mask_same_atoms & cell_earlier + mask = mask_sep_atoms | mask_same_atoms + edge_index_new = edge_index[mask[None, :].expand([2, -1])].reshape([2, -1]) + edge_index_cat = paddle.concat( + [ + edge_index_new, + paddle.stack([edge_index_new[1], edge_index_new[0]], axis=0), + ], + axis=1, + ) + batch_edge = paddle.repeat_interleave( + paddle.arange(neighbors.shape[0]), neighbors + ) + batch_edge = batch_edge[mask] + neighbors_new = 2 * paddle.bincount(batch_edge, minlength=neighbors.shape[0]) + edge_reorder_idx = repeat_blocks( + neighbors_new // 2, + repeats=2, + continuous_indexing=True, + repeat_inc=edge_index_new.shape[1], + ) + edge_index_new = edge_index_cat[:, edge_reorder_idx] + cell_offsets_new = CSPNet.select_symmetric_edges( + self, cell_offsets, mask, edge_reorder_idx, True + ) + edge_vector_new = CSPNet.select_symmetric_edges( + self, edge_vector, mask, edge_reorder_idx, True + ) + return edge_index_new, cell_offsets_new, neighbors_new, edge_vector_new + + def gen_edges(self, num_atoms, frac_coords, lattices, node2graph): + if self.edge_style == "fc": + return CSPNet.gen_edges(self, num_atoms, frac_coords) + + cart_coords = frac_to_cart_coords_with_lattice(frac_coords, num_atoms, lattices) + edge_index, to_jimages, num_bonds = radius_graph_pbc( + cart_coords, + lattices, + num_atoms, + self.cutoff, + self.max_neighbors, + num_atoms.place, + ) + j_index, i_index = edge_index[0], edge_index[1] + distance_vectors = frac_coords[j_index] - frac_coords[i_index] + distance_vectors = distance_vectors + to_jimages + edge_index_new, _, _, edge_vector_new = self.reorder_symmetric_edges( + edge_index, to_jimages, num_bonds, distance_vectors + ) + return edge_index_new, -edge_vector_new + + def forward(self, t, atom_types, frac_coords, lattices, num_atoms, node2graph): + edges, frac_diff = self.gen_edges(num_atoms, frac_coords, lattices, node2graph) + edge2graph = node2graph[edges[0]] + + if self.smooth: + node_features = self.node_embedding(atom_types.cast("float32")) + else: + node_features = self.node_embedding(atom_types - self.species_shift) + + if t.ndim == 0: + t = t.unsqueeze(0) + + if self.time_embed_dim is not None: + t_embed = self.time_embedder(t) + else: + t_embed = t + + if t_embed.ndim == 1: + t_embed = t_embed.unsqueeze(0) + t_per_atom = paddle.repeat_interleave(t_embed, num_atoms, axis=0) + node_features = paddle.concat([node_features, t_per_atom], axis=1) + node_features = self.atom_latent_emb(node_features) + + for i in range(self.num_layers): + layer = getattr(self, "csp_layer_%d" % i) + node_features = layer( + node_features, + frac_coords, + lattices, + edges, + edge2graph, + frac_diff=frac_diff, + ) + + if self.ln: + node_features = self.final_layer_norm(node_features) + + coord_out = self.coord_out(node_features) + coord_out_2 = self.coord_out_2(node_features) + + graph_features = paddle.geometric.segment_mean(node_features, node2graph) + + if self.pred_scalar: + return self.scalar_out(graph_features) + + lattice_out = self.lattice_out(graph_features) + lattice_out = lattice_out.reshape([-1, 3, 3]) + if self.ip: + lattice_out = paddle.matmul(lattice_out, lattices) + + lattice_out_2 = self.lattice_out_2(graph_features) + lattice_out_2 = lattice_out_2.reshape([-1, 3, 3]) + if self.ip: + lattice_out_2 = paddle.matmul(lattice_out_2, lattices) + + if self.pred_type: + type_out = self.type_out(node_features) + type_out_2 = self.type_out_2(node_features) + return ( + lattice_out, + coord_out, + type_out, + lattice_out_2, + coord_out_2, + type_out_2, + ) + + return lattice_out, coord_out, lattice_out_2, coord_out_2 + + def forward_dict(self, t, atom_types, frac_coords, lattices, num_atoms, node2graph): + preds = OMATGCSPNet.forward( + self, t, atom_types, frac_coords, lattices, num_atoms, node2graph + ) + if self.pred_scalar: + return {"scalar": preds} + if self.pred_type: + return { + "pos_b": preds[1], + "pos_eta": preds[4], + "cell_b": preds[0], + "cell_eta": preds[3], + "species_b": preds[2], + "species_eta": preds[5], + } + return { + "pos_b": preds[1], + "pos_eta": preds[3], + "cell_b": preds[0], + "cell_eta": preds[2], + } + +class OMATGCSPNetFull(OMATGCSPNet): + """Full CSPNet with time embedding integrated for OMatG models. + + This is a convenience wrapper that creates CSPNet with time embedding. + Renamed from CSPNetFull to avoid naming conflicts with other CSPNet implementations. + + Args: + hidden_dim: Hidden dimension for embeddings + num_layers: Number of message passing layers + max_atoms: Maximum number of atoms in crystal + act_fn: Activation function + dis_emb: Distance embedding type + num_freqs: Number of frequency bands for sinusoidal embedding + edge_style: Edge construction style + cutoff: Distance cutoff for graph construction + max_neighbors: Maximum number of neighbors + ln: Whether to use layer normalization + ip: Whether to use inner product + smooth: Whether to use smooth embedding + pred_type: Whether to predict atom types + pred_scalar: Whether to predict scalar properties + time_embed_dim: Time embedding dimension + """ + + def __init__( + self, + hidden_dim: int = 128, + num_layers: int = 4, + max_atoms: int = 100, + act_fn: str = "silu", + dis_emb: str = "sin", + num_freqs: int = 10, + edge_style: str = "fc", + cutoff: float = 6.0, + max_neighbors: int = 20, + ln: bool = False, + ip: bool = True, + smooth: bool = False, + pred_type: bool = False, + pred_scalar: bool = False, + time_embed_dim: int = 256, + use_si: bool = False, + si_cfg: dict = None, + sampler_cfg: dict = None, + ): + super().__init__( + hidden_dim=hidden_dim, + num_layers=num_layers, + max_atoms=max_atoms, + act_fn=act_fn, + dis_emb=dis_emb, + num_freqs=num_freqs, + edge_style=edge_style, + cutoff=cutoff, + max_neighbors=max_neighbors, + ln=ln, + ip=ip, + smooth=smooth, + pred_type=pred_type, + pred_scalar=pred_scalar, + time_embed_dim=time_embed_dim, + ) + self.use_si = use_si + self._si = None + self._sampler = None + self._relative_si_costs = None + self.mse_loss = MSELoss() + if use_si: + self._build_si(si_cfg or {}, sampler_cfg or {}) + + def forward(self, data, t=None): + """Forward pass accepting dict or OMATGData/Batch object. + + Args: + data: dict (for standalone use) or OMATGData/Batch (from collator pipeline). + t: Time tensor (optional, will be sampled if None) + + Returns: + Dictionary containing loss_dict with loss components + """ + if isinstance(data, dict): + atom_types = data["atom_types"] + frac_coords = data["frac_coords"] + lattices = data["lattices"] + num_atoms = data["num_atoms"] + node2graph = data["node2graph"] + else: + atom_types = data.species + frac_coords = data.pos + lattices = data.cell + num_atoms = data.n_atoms + node2graph = data.batch + + # SI-based training path: velocity matching loss via StochasticInterpolants + if self.use_si: + return self._si_forward(data) + + # Sample time if not provided + if t is None: + batch_size = lattices.shape[0] + t = paddle.rand([batch_size]) + + # Call forward_dict for b/eta outputs + predictions = self.forward_dict( + t=t, + atom_types=atom_types, + frac_coords=frac_coords, + lattices=lattices, + num_atoms=num_atoms, + node2graph=node2graph, + ) + + lattices_gt = lattices + frac_coords_gt = frac_coords + + if self.pred_type: + loss_lattice = self.mse_loss(predictions["cell_b"], lattices_gt) + loss_coord = self.mse_loss(predictions["pos_b"], frac_coords_gt) + loss_type = paddle.nn.functional.cross_entropy( + input=predictions["species_b"], label=atom_types - self.species_shift + ) + loss = loss_lattice + loss_coord + loss_type + return { + "loss_dict": { + "loss": loss, + "loss_lattice": loss_lattice, + "loss_coord": loss_coord, + "loss_type": loss_type, + } + } + + # CSP mode: use b components from forward_dict + loss_lattice = self.mse_loss(predictions["cell_b"], lattices_gt) + loss_coord = self.mse_loss(predictions["pos_b"], frac_coords_gt) + loss = loss_lattice + loss_coord + return { + "loss_dict": { + "loss": loss, + "loss_lattice": loss_lattice, + "loss_coord": loss_coord, + } + } + + def sample(self, batch_data, num_inference_steps=100, **kwargs): + """Sample crystal structures. + + Based on the original OMatG implementation in omg_trainer.py and generate_csp.py. + Uses reverse-time ODE integration with Euler method. + + Args: + batch_data: Dictionary (for standalone sampling) or OMATGData/Batch object. + Dictionary supports: + - structure_array with num_atoms and optionally atom_types + - or direct fields: atom_types, frac_coords (ignored), lattices (ignored), + num_atoms, node2graph (ignored) + num_inference_steps: Number of integration steps (default 100) + **kwargs: Additional sampling parameters + - step_lr: step learning rate for Euler integration (default 1e-5) + + Returns: + Dictionary containing generated structures + """ + if self.use_si and self._si is not None: + return self._si_sample(batch_data, num_inference_steps, **kwargs) + + # Support dict (standalone sampling) and OMATGData/Batch (collator output) + if isinstance(batch_data, dict): + if "structure_array" in batch_data: + sa = batch_data["structure_array"] + num_atoms_list = sa["num_atoms"].tolist() + atom_types = sa.get("atom_types") + else: + num_atoms_list = batch_data["num_atoms"].tolist() + atom_types = batch_data.get("atom_types") + else: + num_atoms_list = batch_data.n_atoms.tolist() + atom_types = batch_data.species + + batch_size = len(num_atoms_list) + total_atoms = sum(num_atoms_list) + + if atom_types is None: + atom_types = paddle.randint(1, 100, shape=[total_atoms]) + + num_atoms_tensor = paddle.to_tensor(num_atoms_list, dtype="int64") + node2graph = paddle.repeat_interleave( + paddle.arange(batch_size, dtype="int64"), num_atoms_tensor + ) + + # Sample from base distribution (uniform for positions, randn for lattice) + # Reference: omg_lightning.py:predict_step() -> x_0 = self.sampler.sample_p_0(x) + frac_coords = paddle.rand([total_atoms, 3]) + lattices = paddle.randn([batch_size, 3, 3]) + + # Scale initial lattice to reasonable values + lattices = lattices * 2.0 # Scale to avoid extreme values + + # Integration loop using Euler method + step_lr = 1e-5 + + for step in range(num_inference_steps): + t = paddle.full([batch_size], step / num_inference_steps) + lattice_pred, coord_pred = self._predict( + t, atom_types, frac_coords, lattices, + num_atoms_tensor, node2graph, + ) + + if paddle.any(paddle.isnan(lattice_pred)) or paddle.any( + paddle.isnan(coord_pred) + ): + continue + + frac_coords = frac_coords + coord_pred * step_lr + lattices = lattices + lattice_pred * step_lr + + frac_coords = frac_coords % 1.0 + + # Final clipping to ensure valid lattice + lattices = paddle.clip(lattices, -10.0, 10.0) + + # Convert lattice matrices to lengths and angles for BuildStructure + return self._build_sample_result(num_atoms_list, atom_types, frac_coords, lattices) + + def _predict(self, t, atom_types, frac_coords, lattices, num_atoms, node2graph): + predictions = super().forward( + t=t, + atom_types=atom_types, + frac_coords=frac_coords, + lattices=lattices, + num_atoms=num_atoms, + node2graph=node2graph, + ) + lattice_pred, coord_pred = predictions[0], predictions[1] + if paddle.any(paddle.isnan(lattice_pred)): + lattice_pred = paddle.zeros_like(lattice_pred) + if paddle.any(paddle.isnan(coord_pred)): + coord_pred = paddle.zeros_like(coord_pred) + return lattice_pred, coord_pred + + def _make_model_function(self): + def model_function(x_t, time): + return self.forward_dict( + t=time, + atom_types=x_t.species, + frac_coords=x_t.pos, + lattices=x_t.cell, + num_atoms=x_t.n_atoms, + node2graph=x_t.batch, + ) + return model_function + + @staticmethod + def _build_sample_result(num_atoms_list, atom_types, frac_coords, lattices): + from ppmat.utils.crystal import lattices_to_params_shape_paddle + + lengths, angles = lattices_to_params_shape_paddle(lattices) + start_idx = 0 + result = [] + for i, n in enumerate(num_atoms_list): + end_idx = start_idx + n + result.append({ + "num_atoms": n, + "atom_types": atom_types[start_idx:end_idx].tolist(), + "frac_coords": frac_coords[start_idx:end_idx].tolist(), + "lengths": lengths[i].tolist(), + "angles": angles[i].tolist(), + }) + start_idx += n + return {"result": result} + + def _si_sample(self, batch_data, num_inference_steps, **kwargs): + x_1 = self._data_to_omatg(batch_data) + x_0 = self._sampler.sample_p_0(x_1) + + gen = self._si.integrate(x_0, self._make_model_function(), save_intermediate=False) + + return self._build_sample_result( + gen.n_atoms.tolist(), gen.species, gen.pos, gen.cell + ) + + def _build_si(self, si_cfg: dict, sampler_cfg: dict) -> None: + from ppmat.models.omatg.si.core import build_si_from_cfg, build_sampler_from_cfg + from ppmat.models.omatg.si import StochasticInterpolants + + si_list = si_cfg.get("stochastic_interpolants", []) + if not si_list: + self._si = None + self._relative_si_costs = {} + self._sampler = None + return + + use_factory = isinstance(si_list[0], dict) and "__class_name__" in si_list[0] + if use_factory: + self._si = build_si_from_cfg(si_cfg) + else: + data_fields = si_cfg.get("data_fields") + if not data_fields: + raise ValueError("si_cfg must contain 'data_fields' when use_si=True.") + self._si = StochasticInterpolants( + stochastic_interpolants=si_list, + data_fields=data_fields, + integration_time_steps=si_cfg.get("integration_time_steps", 210), + ) + self._relative_si_costs = si_cfg.get("relative_si_costs", {}) + + if sampler_cfg: + if any(isinstance(v, dict) and "__class_name__" in v for v in sampler_cfg.values()): + self._sampler = build_sampler_from_cfg(sampler_cfg) + else: + self._sampler = IndependentSampler( + dataset_name=sampler_cfg.get("dataset_name"), + mirror_species=sampler_cfg.get("mirror_species", True), + mask_species=sampler_cfg.get("mask_species", False), + ) + + def _data_to_omatg(self, data): + from ppmat.datasets.omatg_dataset import OMATGData + + if isinstance(data, OMATGData): + return data + return OMATGData.from_collate_dict(data) + + def _si_forward(self, data: dict) -> dict: + """SI training step: sample x_0, interpolate, compute velocity matching loss.""" + from ppmat.models.omatg.si import SMALL_TIME, BIG_TIME + + x_1 = self._data_to_omatg(data) + x_0 = self._sampler.sample_p_0(x_1) + batch_size = len(x_1.n_atoms) + t = paddle.rand([batch_size]) * (BIG_TIME - SMALL_TIME) + SMALL_TIME + + losses = self._si.losses(self._make_model_function(), t, x_0, x_1) + total_loss = paddle.to_tensor(0.0) + loss_dict = {} + for key, val in losses.items(): + cost = self._relative_si_costs.get(key, 1.0) + weighted = cost * val + loss_dict[key] = val + total_loss = total_loss + weighted + loss_dict["loss"] = total_loss + return {"loss_dict": loss_dict} + + +"""Independent base distribution sampler using Paddle native APIs.""" + +_LATTICE_PARAMS = { + "carbon_24": { + "means": [0.9852757453918457, 1.3865314722061157, 1.7068126201629639], + "stds": [0.14957907795906067, 0.20431114733219147, 0.2403733879327774], + }, + "mp_20": { + "means": [1.575442910194397, 1.7017393112182617, 1.9781638383865356], + "stds": [0.24437622725963593, 0.26526379585266113, 0.3535512685775757], + }, + "mpts_52": { + "means": [1.6565313339233398, 1.8407557010650635, 2.1225264072418213], + "stds": [0.2952289581298828, 0.3340013027191162, 0.41885802149772644], + }, + "perov_5": { + "means": [1.419227957725525, 1.419227957725525, 1.419227957725525], + "stds": [0.07268335670232773, 0.07268335670232773, 0.07268335670232773], + }, + "alex_mp_20": { + "means": [1.5808929163076058, 1.74672046352959, 2.065243388307474], + "stds": [0.27284015410437057, 0.2944785731740152, 0.30899526911753017], + }, +} + + +def _sample_cell(dataset_name): + params = _LATTICE_PARAMS.get(dataset_name) + if params is None: + return paddle.randn([3, 3]).numpy() * 2.0 + lengths = paddle.exp( + paddle.randn([3]) * paddle.to_tensor(params["stds"]) + + paddle.to_tensor(params["means"]) + ) + angles = paddle.rand([3]) * 60.0 + 60.0 + return cellpar_to_cell(paddle.concat((lengths, angles)).numpy()) + + +class IndependentSampler: + """Sample base distributions for SI training using Paddle native APIs. + + Args: + dataset_name: Optional dataset name for informed lattice distribution. + mirror_species: If True, keep input species unchanged (mirror). + mask_species: If True, replace species with zeros (mask token). + """ + + def __init__(self, dataset_name=None, mirror_species=True, mask_species=False): + self._dataset_name = dataset_name + self._mirror_species = mirror_species + self._mask_species = mask_species + + def sample_p_0(self, x_1: OMATGData) -> OMATGData: + batch_size = len(x_1.n_atoms) + structures = [] + for i in range(batch_size): + sl = x_1.slice(i) + pos = paddle.rand(x_1.pos[sl].shape, dtype=x_1.pos.dtype) + cell = paddle.to_tensor(_sample_cell(self._dataset_name), dtype=x_1.cell.dtype) + if self._mask_species: + species = paddle.zeros_like(x_1.species[sl]) + elif self._mirror_species: + species = x_1.species[sl].clone() + else: + species = paddle.randint(1, 100, x_1.species[sl].shape, dtype=x_1.species[sl].dtype) + sampled = Structure( + cell=cell, + atomic_numbers=species, + pos=pos, + pos_is_fractional=True, + ) + structures.append(sampled) + return OMATGData.from_batch(structures, concatenate=True) diff --git a/ppmat/models/omatg/si/__init__.py b/ppmat/models/omatg/si/__init__.py new file mode 100644 index 00000000..8a7b619b --- /dev/null +++ b/ppmat/models/omatg/si/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import ( + BIG_TIME, SMALL_TIME, + StochasticInterpolants, + SingleStochasticInterpolant, SingleStochasticInterpolantIdentity, + SingleStochasticInterpolantOS, DiscreteFlowMatchingMask, + build_si_from_cfg, build_sampler_from_cfg, +) +from .interpolants import LinearInterpolant, PeriodicLinearInterpolant + +__all__ = [ + "StochasticInterpolants", + "SingleStochasticInterpolant", "SingleStochasticInterpolantIdentity", + "SingleStochasticInterpolantOS", "DiscreteFlowMatchingMask", + "LinearInterpolant", "PeriodicLinearInterpolant", + "BIG_TIME", "SMALL_TIME", + "build_si_from_cfg", "build_sampler_from_cfg", +] diff --git a/ppmat/models/omatg/si/core.py b/ppmat/models/omatg/si/core.py new file mode 100644 index 00000000..281f3df2 --- /dev/null +++ b/ppmat/models/omatg/si/core.py @@ -0,0 +1,1323 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Single Stochastic Interpolant implementation. +""" + +import importlib +from enum import Enum +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union + +import paddle + +from ppmat.datasets.omatg_dataset import OMATGData +from ppmat.utils.scatter import scatter_mean + +from .interpolants import ( + Corrector, + Epsilon, + IdentityCorrector, + Interpolant, + LatentGamma, + ScoreBasedDiffusionModelInterpolantVE, + ScoreBasedDiffusionModelInterpolantVP, + StochasticInterpolant, + StochasticInterpolantSpecies, +) + + +class DifferentialEquationType(Enum): + """Enum for differential equation types.""" + + ODE = "ode" + SDE = "sde" + + +def _compute_mean_velocity( + velocity: paddle.Tensor, batch_indices: paddle.Tensor +) -> paddle.Tensor: + mean_vel = scatter_mean(velocity, batch_indices, dim=0) + return mean_vel[batch_indices] + + +class SingleStochasticInterpolant(StochasticInterpolant): + """Stochastic interpolant x_t = I(t, x_0, x_1) + gamma(t) * z. + Supports ODE or SDE during inference. + """ + + def __init__( + self, + interpolant: Interpolant, + gamma: Optional[LatentGamma] = None, + epsilon: Optional[Epsilon] = None, + differential_equation_type: str = "ODE", + integrator_kwargs: Optional[dict[str, Any]] = None, + correct_center_of_mass_motion: bool = False, + velocity_annealing_factor: float = 0.0, + ) -> None: + """Construct stochastic interpolant.""" + super().__init__() + self._interpolant = interpolant + self._gamma = gamma + if self._gamma is not None: + self._use_antithetic = self._gamma.requires_antithetic() + else: + self._use_antithetic = False + self._epsilon = epsilon + self._differential_equation_type = differential_equation_type.upper() + self._corrector = self._interpolant.get_corrector() + + if self._differential_equation_type == "ODE": + self.loss = self._ode_loss + self.integrate = self._ode_integrate + if self._epsilon is not None: + raise ValueError("Epsilon function should not be provided for ODEs.") + elif self._differential_equation_type == "SDE": + self.loss = self._sde_loss + self.integrate = self._sde_integrate + if self._epsilon is None: + raise ValueError("Epsilon function should be provided for SDEs.") + if self._gamma is None: + raise ValueError("Gamma function should be provided for SDEs.") + else: + raise ValueError( + f"Unknown differential equation type: {differential_equation_type}" + ) + + self._correct_center_of_mass_motion = correct_center_of_mass_motion + self._velocity_annealing_factor = velocity_annealing_factor + + def interpolate( + self, + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Tuple[paddle.Tensor, paddle.Tensor]: + """ + Stochastically interpolate between points x_0 and x_1 from two distributions p_0 and p_1 at times t. + + :param t: + Times in [0,1]. + :type t: paddle.Tensor + :param x_0: + Points from p_0. + :type x_0: paddle.Tensor + :param x_1: + Points from p_1. + :type x_1: paddle.Tensor + :param batch_indices: + Tensor containing the configuration index for every atom in the batch. + :type batch_indices: paddle.Tensor + + :return: + Stochastically interpolated points x_t, random variables z used for interpolation. + :rtype: tuple[paddle.Tensor, paddle.Tensor] + """ + assert x_0.shape == x_1.shape + interpolate = self._interpolant.interpolate(t, x_0, x_1) + if self._gamma is not None: + z = paddle.randn(x_0.shape) + gamma_t = self._gamma.gamma(t) + interpolate = self._corrector.correct(interpolate + gamma_t * z) + else: + z = paddle.zeros_like(x_0) + return interpolate, z + + def loss_keys(self) -> Iterable[str]: + """ + Get the keys of the losses returned by the loss function. + + :return: + Keys of the losses. + :rtype: Iterable[str] + """ + if self._differential_equation_type == "ODE": + yield "loss_b" + else: + yield "loss_b" + yield "loss_z" + + def loss(self, *args, **kwargs): + raise NotImplementedError # Overridden in __init__ via self.loss = _ode_loss / _sde_loss + + def _ode_loss( + self, + model_function: Callable[[paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor]], + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + x_t: paddle.Tensor, + z: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Dict[str, paddle.Tensor]: + """ + Compute the losses for the ODE stochastic interpolant. + + :return: + Losses. + :rtype: Dict[str, paddle.Tensor] + """ + assert x_0.shape == x_1.shape + + if self._use_antithetic: + assert self._gamma is not None + x_t_without_gamma = self._interpolant.interpolate(t, x_0, x_1) + gamma = self._gamma.gamma(t) + x_t_p = self._corrector.correct(x_t_without_gamma + gamma * z) + x_t_m = self._corrector.correct(x_t_without_gamma - gamma * z) + + expected_velocity_without_gamma = self._interpolant.interpolate_derivative( + t, x_0, x_1 + ) + gamma_derivative = self._gamma.gamma_derivative(t) + expected_velocity_p = expected_velocity_without_gamma + gamma_derivative * z + expected_velocity_m = expected_velocity_without_gamma - gamma_derivative * z + + if self._correct_center_of_mass_motion: + mean_velocity_p = _compute_mean_velocity( + expected_velocity_p, batch_indices + ) + expected_velocity_p = expected_velocity_p - mean_velocity_p + mean_velocity_m = _compute_mean_velocity( + expected_velocity_m, batch_indices + ) + expected_velocity_m = expected_velocity_m - mean_velocity_m + + pred_b_p = model_function(x_t_p)[0] + pred_b_m = model_function(x_t_m)[0] + + loss = ( + 0.5 * paddle.mean(pred_b_p**2) + + 0.5 * paddle.mean(pred_b_m**2) + - paddle.mean(pred_b_p * expected_velocity_p) + - paddle.mean(pred_b_m * expected_velocity_m) + ) + else: + expected_velocity = self._interpolant.interpolate_derivative(t, x_0, x_1) + if self._gamma is not None: + expected_velocity += self._gamma.gamma_derivative(t) * z + + pred_b = model_function(x_t)[0] + + if self._correct_center_of_mass_motion: + mean_velocity = _compute_mean_velocity( + expected_velocity, batch_indices + ) + expected_velocity = expected_velocity - mean_velocity + + loss = paddle.mean(pred_b**2) - 2.0 * paddle.mean( + pred_b * expected_velocity + ) + + return {"loss_b": loss} + + def _sde_loss( + self, + model_function: Callable[[paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor]], + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + x_t: paddle.Tensor, + z: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Dict[str, paddle.Tensor]: + """ + Compute the losses for the SDE stochastic interpolant. + + :return: + Losses. + :rtype: Dict[str, paddle.Tensor] + """ + assert x_0.shape == x_1.shape + assert self._gamma is not None + + if self._use_antithetic: + x_t_without_gamma = self._interpolant.interpolate(t, x_0, x_1) + gamma = self._gamma.gamma(t) + x_t_p = self._corrector.correct(x_t_without_gamma + gamma * z) + x_t_m = self._corrector.correct(x_t_without_gamma - gamma * z) + + expected_velocity_without_gamma = self._interpolant.interpolate_derivative( + t, x_0, x_1 + ) + gamma_derivative = self._gamma.gamma_derivative(t) + expected_velocity_p = expected_velocity_without_gamma + gamma_derivative * z + expected_velocity_m = expected_velocity_without_gamma - gamma_derivative * z + + if self._correct_center_of_mass_motion: + mean_velocity_p = _compute_mean_velocity( + expected_velocity_p, batch_indices + ) + expected_velocity_p = expected_velocity_p - mean_velocity_p + mean_velocity_m = _compute_mean_velocity( + expected_velocity_m, batch_indices + ) + expected_velocity_m = expected_velocity_m - mean_velocity_m + + pred_b_p, pred_z = model_function(x_t_p) + pred_b_m, _ = model_function(x_t_m) + + loss_b = ( + 0.5 * paddle.mean(pred_b_p**2) + + 0.5 * paddle.mean(pred_b_m**2) + - paddle.mean(pred_b_p * expected_velocity_p) + - paddle.mean(pred_b_m * expected_velocity_m) + ) + else: + expected_velocity = ( + self._interpolant.interpolate_derivative(t, x_0, x_1) + + self._gamma.gamma_derivative(t) * z + ) + pred_b, pred_z = model_function(x_t) + + if self._correct_center_of_mass_motion: + mean_velocity = _compute_mean_velocity( + expected_velocity, batch_indices + ) + expected_velocity = expected_velocity - mean_velocity + + loss_b = paddle.mean(pred_b**2) - paddle.mean(pred_b * expected_velocity) + + loss_z = paddle.mean(pred_z**2) - 2.0 * paddle.mean(pred_z * z) + + return {"loss_b": loss_b, "loss_z": loss_z} + + def _ode_integrate( + self, + model_function: Callable[ + [paddle.Tensor, paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor] + ], + x_t: paddle.Tensor, + time: paddle.Tensor, + time_step: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> paddle.Tensor: + """ + Integrate the ODE for the current positions. + + :param model_function: + Model function returning the velocity fields b and the denoisers eta. + :param x_t: + Current positions. + :param time: + Initial time. + :param time_step: + Time step. + :param batch_indices: + Batch indices. + + :return: + Integrated position. + """ + # Simple Euler integration using paddle operations + dt = time_step.item() if hasattr(time_step, "item") else float(time_step) + t_val = time.item() if hasattr(time, "item") else float(time) + + t_tensor = paddle.to_tensor([t_val] * x_t.shape[0], dtype=x_t.dtype) + model_result = model_function(t_tensor, self._corrector.correct(x_t)) + velocity = model_result[0] + annealing_factor = 1.0 + self._velocity_annealing_factor * t_val + x_new = x_t + dt * annealing_factor * velocity + + return self._corrector.correct(x_new) + + def _sde_integrate( + self, + model_function: Callable[ + [paddle.Tensor, paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor] + ], + x_t: paddle.Tensor, + time: paddle.Tensor, + time_step: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> paddle.Tensor: + """ + Integrate the SDE for the current positions. + + :param model_function: + Model function returning the velocity fields b and the denoisers eta. + :param x_t: + Current positions. + :param time: + Initial time. + :param time_step: + Time step. + :param batch_indices: + Batch indices. + + :return: + Integrated position. + """ + # Euler-Maruyama integration using paddle operations + dt = time_step.item() if hasattr(time_step, "item") else float(time_step) + t_val = time.item() if hasattr(time, "item") else float(time) + + t_tensor = paddle.to_tensor([t_val] * x_t.shape[0], dtype=x_t.dtype) + model_result = model_function(t_tensor, self._corrector.correct(x_t)) + drift = model_result[0] + eta = model_result[1] if len(model_result) > 1 else paddle.zeros_like(drift) + + epsilon_t = ( + self._epsilon.epsilon(paddle.to_tensor([t_val])).item() + if self._epsilon + else 0.0 + ) + gamma_t = ( + self._gamma.gamma(paddle.to_tensor([t_val])).item() if self._gamma else 1.0 + ) + + # Euler-Maruyama update: x_{t+dt} = x_t + drift * dt + sqrt(2 * epsilon) * dW + diffusion = ( + paddle.sqrt(paddle.to_tensor(2.0 * epsilon_t * dt)) + * paddle.randn(x_t.shape) + ) + x_new = x_t + drift * dt - (epsilon_t / gamma_t) * eta * dt + diffusion + + return self._corrector.correct(x_new) + + def get_corrector(self) -> Corrector: + """ + Get the corrector implied by the stochastic interpolant. + + :return: + Corrector. + :rtype: Corrector + """ + return self._corrector + + def integrate(self, *args, **kwargs): + raise NotImplementedError # Overridden in __init__ + +class SingleStochasticInterpolantOS(StochasticInterpolant): + """One-sided stochastic interpolant where x_0 is Gaussian. + + The latent variable z in SingleStochasticInterpolant is merged with x_0. + Supports ODE or SDE during inference via fixed-step Euler integration. + + :param interpolant: Interpolant I(t, x_0, x_1). + :param epsilon: Optional epsilon function for SDE. + :param differential_equation_type: "ODE" or "SDE". + :param integrator_kwargs: Optional kwargs (unused by fixed-step Euler). + :param correct_center_of_mass_motion: Whether to zero COM velocity in loss. + :param predict_velocity: Whether to compute loss for velocity field b. + :param velocity_annealing_factor: Annealing factor for b during inference. + """ + + def __init__( + self, + interpolant: Interpolant, + epsilon: Optional[Epsilon], + differential_equation_type: str, + integrator_kwargs: Optional[dict[str, Any]] = None, + correct_center_of_mass_motion: bool = False, + predict_velocity: bool = True, + velocity_annealing_factor: Optional[float] = None, + ) -> None: + super().__init__() + self._interpolant = interpolant + self._epsilon = epsilon + self._corrector = self._interpolant.get_corrector() + try: + self._differential_equation_type = DifferentialEquationType[ + differential_equation_type + ] + except KeyError: + raise ValueError( + f"Unknown differential equation type: {differential_equation_type}." + ) + if self._differential_equation_type == DifferentialEquationType.ODE: + self.loss = self._ode_loss + self.integrate = self._ode_integrate + if self._epsilon is not None: + raise ValueError("Epsilon function should not be provided for ODEs.") + else: + assert self._differential_equation_type == DifferentialEquationType.SDE + self.loss = self._sde_loss + self.integrate = self._sde_integrate + if self._epsilon is None: + raise ValueError("Epsilon function should be provided for SDEs.") + self._correct_center_of_mass_motion = correct_center_of_mass_motion + self._predict_velocity = predict_velocity + self._use_antithetic = isinstance( + self._interpolant, + ( + ScoreBasedDiffusionModelInterpolantVP, + ScoreBasedDiffusionModelInterpolantVE, + ), + ) + self._velocity_annealing_factor = velocity_annealing_factor + if not self._predict_velocity and self._velocity_annealing_factor is not None: + raise ValueError( + "Velocity annealing factor should only be set if predict_velocity is True." + ) + if self._predict_velocity and self._velocity_annealing_factor is None: + self._velocity_annealing_factor = 0.0 + + def interpolate( + self, + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Tuple[paddle.Tensor, paddle.Tensor]: + assert x_0.shape == x_1.shape + interpolate = self._interpolant.interpolate(t, x_0, x_1) + return interpolate, x_0.clone() + + def loss_keys(self) -> Iterable[str]: + if self._predict_velocity: + yield "loss_b" + if self._differential_equation_type == DifferentialEquationType.SDE: + yield "loss_z" + else: + yield "loss_z" + + def loss(self, *args, **kwargs): + raise NotImplementedError # Overridden in __init__ via self.loss = _ode_loss / _sde_loss + + def _ode_loss( + self, + model_function: Callable[[paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor]], + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + x_t: paddle.Tensor, + z: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Dict[str, paddle.Tensor]: + assert x_0.shape == x_1.shape + if self._predict_velocity: + if self._use_antithetic: + x_t_p = self._interpolant.interpolate(t, x_0, x_1) + x_t_m = self._interpolant.interpolate(t, -x_0, x_1) + expected_velocity_p = self._interpolant.interpolate_derivative( + t, x_0, x_1 + ) + expected_velocity_m = self._interpolant.interpolate_derivative( + t, -x_0, x_1 + ) + if self._correct_center_of_mass_motion: + expected_velocity_p = expected_velocity_p - _compute_mean_velocity( + expected_velocity_p, batch_indices + ) + expected_velocity_m = expected_velocity_m - _compute_mean_velocity( + expected_velocity_m, batch_indices + ) + pred_b_p = model_function(x_t_p)[0] + pred_b_m = model_function(x_t_m)[0] + loss = ( + 0.5 * paddle.mean(pred_b_p**2) + + 0.5 * paddle.mean(pred_b_m**2) + - paddle.mean(pred_b_p * expected_velocity_p) + - paddle.mean(pred_b_m * expected_velocity_m) + ) + else: + expected_velocity = self._interpolant.interpolate_derivative( + t, x_0, x_1 + ) + pred_b = model_function(x_t)[0] + if self._correct_center_of_mass_motion: + expected_velocity = expected_velocity - _compute_mean_velocity( + expected_velocity, batch_indices + ) + loss = paddle.mean(pred_b**2) - 2.0 * paddle.mean( + pred_b * expected_velocity + ) + return {"loss_b": loss} + else: + pred_z = model_function(x_t)[1] + return {"loss_z": paddle.mean(pred_z**2) - 2.0 * paddle.mean(pred_z * z)} + + def _sde_loss( + self, + model_function: Callable[[paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor]], + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + x_t: paddle.Tensor, + z: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Dict[str, paddle.Tensor]: + assert x_0.shape == x_1.shape + pred_b, pred_z = model_function(x_t) + loss_z = paddle.mean(pred_z**2) - 2.0 * paddle.mean(pred_z * z) + if self._predict_velocity: + if self._use_antithetic: + x_t_m = self._interpolant.interpolate(t, -x_0, x_1) + expected_velocity_p = self._interpolant.interpolate_derivative( + t, x_0, x_1 + ) + expected_velocity_m = self._interpolant.interpolate_derivative( + t, -x_0, x_1 + ) + if self._correct_center_of_mass_motion: + expected_velocity_p = expected_velocity_p - _compute_mean_velocity( + expected_velocity_p, batch_indices + ) + expected_velocity_m = expected_velocity_m - _compute_mean_velocity( + expected_velocity_m, batch_indices + ) + pred_b_m = model_function(x_t_m)[0] + loss_b = ( + 0.5 * paddle.mean(pred_b**2) + + 0.5 * paddle.mean(pred_b_m**2) + - paddle.mean(pred_b * expected_velocity_p) + - paddle.mean(pred_b_m * expected_velocity_m) + ) + else: + expected_velocity = self._interpolant.interpolate_derivative( + t, x_0, x_1 + ) + if self._correct_center_of_mass_motion: + expected_velocity = expected_velocity - _compute_mean_velocity( + expected_velocity, batch_indices + ) + loss_b = paddle.mean(pred_b**2) - 2.0 * paddle.mean( + pred_b * expected_velocity + ) + return {"loss_b": loss_b, "loss_z": loss_z} + else: + return {"loss_z": loss_z} + + def _ode_integrate( + self, + model_function: Callable[ + [paddle.Tensor, paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor] + ], + x_t: paddle.Tensor, + time: paddle.Tensor, + time_step: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> paddle.Tensor: + dt = time_step.item() if hasattr(time_step, "item") else float(time_step) + t_val = time.item() if hasattr(time, "item") else float(time) + t_tensor = paddle.to_tensor([t_val] * x_t.shape[0], dtype=x_t.dtype) + if self._predict_velocity: + assert self._velocity_annealing_factor is not None + model_result = model_function(t_tensor, self._corrector.correct(x_t)) + velocity = model_result[0] + annealing_factor = 1.0 + self._velocity_annealing_factor * t_val + x_new = x_t + dt * annealing_factor * velocity + else: + model_result = model_function(t_tensor, self._corrector.correct(x_t)) + z = model_result[1] + t1 = self._interpolant.alpha_dot(t_tensor) * z + corr = self._corrector.correct(x_t) + t2 = ( + self._interpolant.beta_dot(t_tensor) + / self._interpolant.beta(t_tensor) + * (corr - self._interpolant.alpha(t_tensor) * z) + ) + x_new = x_t + dt * (t1 + t2) + return self._corrector.correct(x_new) + + def _sde_integrate( + self, + model_function: Callable[ + [paddle.Tensor, paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor] + ], + x_t: paddle.Tensor, + time: paddle.Tensor, + time_step: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> paddle.Tensor: + dt = time_step.item() if hasattr(time_step, "item") else float(time_step) + t_val = time.item() if hasattr(time, "item") else float(time) + t_tensor = paddle.to_tensor([t_val] * x_t.shape[0], dtype=x_t.dtype) + model_result = model_function(t_tensor, self._corrector.correct(x_t)) + if self._predict_velocity: + assert self._velocity_annealing_factor is not None + drift = (1.0 + self._velocity_annealing_factor * t_val) * model_result[0] + eta = model_result[1] + else: + z = model_result[1] + t1 = self._interpolant.alpha_dot(t_tensor) * z + corr = self._corrector.correct(x_t) + t2 = ( + self._interpolant.beta_dot(t_tensor) + / self._interpolant.beta(t_tensor) + * (corr - self._interpolant.alpha(t_tensor) * z) + ) + drift = t1 + t2 + eta = z + epsilon_t = ( + self._epsilon.epsilon(paddle.to_tensor([t_val])).item() + if self._epsilon + else 0.0 + ) + alpha_t = self._interpolant.alpha(t_tensor) + alpha_val = alpha_t[0].item() if alpha_t.numel() > 0 else 1.0 + diffusion = ( + paddle.sqrt(paddle.to_tensor(2.0 * epsilon_t * dt)) + * paddle.randn(x_t.shape) + ) + x_new = x_t + drift * dt - (epsilon_t / max(alpha_val, 1e-8)) * eta * dt + diffusion + return self._corrector.correct(x_new) + + def get_corrector(self) -> Corrector: + return self._corrector + + def integrate(self, *args, **kwargs): + raise NotImplementedError # Overridden in __init__ + +class SingleStochasticInterpolantIdentity(StochasticInterpolantSpecies): + """Stochastic interpolant x_t = x_0 = x_1 for atom species which must remain constant.""" + + def __init__(self) -> None: + super().__init__() + + def interpolate( + self, + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Tuple[paddle.Tensor, paddle.Tensor]: + """Interpolate between x_0 and x_1 (must be equal).""" + assert bool(paddle.equal_all(x_0, x_1)) + return x_0.clone(), paddle.zeros_like(x_0) + + def loss_keys(self) -> Iterable[str]: + """ + Get the keys of the losses returned by the loss function. + + :return: + Keys of the losses. + :rtype: Iterable[str] + """ + yield "loss" + + def loss( + self, + model_function: Callable[[paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor]], + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + x_t: paddle.Tensor, + z: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Dict[str, paddle.Tensor]: + """ + Compute the losses for the stochastic interpolant. + + This class always returns a zero loss with the key 'loss'. + + :param model_function: + Model function returning the velocity fields b and the denoisers eta. + :param t: + Times in [0,1]. + :param x_0: + Points from p_0. + :param x_1: + Points from p_1. + :param x_t: + Stochastically interpolated points x_t. + :param z: + Random variable z. + :param batch_indices: + Tensor containing the configuration index for every atom in the batch. + + :return: + Losses. + :rtype: Dict[str, paddle.Tensor] + """ + assert bool(paddle.equal_all(x_0, x_1)) + return {"loss": paddle.to_tensor(0.0, place=x_0.place)} + + def integrate( + self, + model_function: Callable[ + [paddle.Tensor, paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor] + ], + x_t: paddle.Tensor, + time: paddle.Tensor, + time_step: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> paddle.Tensor: + """ + Integrate the current positions x_t at the given time for the given time step. + + :param model_function: + Model function (not used for identity). + :param x_t: + Current positions. + :param time: + Initial time. + :param time_step: + Time step. + :param batch_indices: + Batch indices. + + :return: + Integrated position (unchanged). + :rtype: paddle.Tensor + """ + # Always return new object (unchanged for identity interpolant). + return x_t.clone() + + def get_corrector(self) -> Corrector: + """ + Get the corrector implied by the stochastic interpolant. + + :return: + Identity corrector. + :rtype: Corrector + """ + return IdentityCorrector() + + def uses_masked_species(self) -> bool: + """ + Return whether the stochastic interpolant uses masked species. + + :return: + Whether the stochastic interpolant uses masked species. + :rtype: bool + """ + # Dataset does not contain masked species. + return False + +# Global constants +SMALL_TIME: float = 1.0e-3 +BIG_TIME: float = 1.0 - SMALL_TIME + + +class DataField(Enum): + """Enum for data fields in OMATGData relevant for stochastic interpolants.""" + + pos = "pos" + cell = "cell" + species = "species" + + +def reshape_t( + t: paddle.Tensor, n_atoms: paddle.Tensor, data_field: DataField +) -> paddle.Tensor: + """Reshape times tensor for batch configurations for the given data field.""" + assert len(t.shape) == 1 + assert len(n_atoms.shape) == 1 + + # Repeat t for each atom + t_per_atom = paddle.repeat_interleave(t, n_atoms.cast("int64")) + sum_n_atoms = int(n_atoms.sum().item()) + batch_size = len(t) + + if data_field == DataField.pos: + # Reshape to (sum_n_atoms, 3) by repeating 3 times + return paddle.reshape( + paddle.repeat_interleave(t_per_atom, 3), + [sum_n_atoms, 3], + ) + elif data_field == DataField.cell: + # Reshape to (batch_size, 3, 3) + return paddle.reshape( + paddle.repeat_interleave(t, 9), [batch_size, 3, 3] + ) + else: + # species case + return t_per_atom + + +class StochasticInterpolants: + """ + Collection of several stochastic interpolants between points x_0 and x_1 from two distributions + p_0 and p_1 at times t for different coordinate types x (like atom species, fractional coordinates, + and lattice vectors). + + Every stochastic interpolant is associated with a data field and a cost factor. + + :param stochastic_interpolants: + Sequence of stochastic interpolants for the different coordinate types. + :param data_fields: + Sequence of data fields for the different stochastic interpolants. + :param integration_time_steps: + Number of integration time steps for the integration. + + :raises ValueError: + If the number of stochastic interpolants and costs are not equal. + If the number of stochastic interpolants and data fields are not equal. + If the number of integration time steps is not positive. + """ + + def __init__( + self, + stochastic_interpolants: Sequence, + data_fields: Sequence[str], + integration_time_steps: int, + ) -> None: + """Constructor of the StochasticInterpolants class.""" + super().__init__() + + if not len(stochastic_interpolants) == len(data_fields): + raise ValueError( + "The number of stochastic interpolants and data fields must be equal." + ) + + try: + self._data_fields = [DataField[df.lower()] for df in data_fields] + except KeyError: + raise ValueError( + f"All data fields must be in {[d.value for d in DataField]}." + ) + + if not integration_time_steps > 0: + raise ValueError("The number of integration time steps must be positive.") + + self._stochastic_interpolants = stochastic_interpolants + self._integration_time_steps = integration_time_steps + + def __len__(self) -> int: + """ + Return the number of stochastic interpolants handled by this class. + + :return: + Number of stochastic interpolants. + """ + return len(self._stochastic_interpolants) + + def loss_keys(self) -> List[str]: + """ + Return the keys of the losses returned by this class. + + :return: + Keys of the losses. + """ + loss_keys = [] + for df, si in zip(self._data_fields, self._stochastic_interpolants): + for key in si.loss_keys(): + full_key = f"{df.value}_{key}" + if full_key in loss_keys: + raise ValueError(f"Key {full_key} is already used as a loss key.") + loss_keys.append(full_key) + return loss_keys + + def _interpolate( + self, + t: paddle.Tensor, + x_0: OMATGData, + x_1: OMATGData, + ) -> Tuple[OMATGData, OMATGData]: + """ + Stochastically interpolate between the collection of points x_0 and x_1 from the + collection of two distributions p_0 and p_1 at times t. + + :param t: + Times in [0,1]. + :param x_0: + Collection of points from the collection of distributions p_0. + :param x_1: + Collection of points from the collection of distributions p_1. + + :return: + Collection of stochastically interpolated points x_t, and the collection of z values. + :rtype: tuple[OMATGData, OMATGData] + """ + assert bool(paddle.equal_all(x_0.batch, x_1.batch)) + assert bool(paddle.equal_all(x_0.n_atoms, x_1.n_atoms)) + + n_atoms = x_0.n_atoms + x_t = x_0.clone() + z_data = {} + + for stochastic_interpolant, data_field in zip( + self._stochastic_interpolants, self._data_fields + ): + field_name = data_field.value + reshaped_t = reshape_t(t, n_atoms, data_field) + + # Cell data requires different batch indices. + if data_field == DataField.cell: + batch_indices = paddle.arange(len(x_0.n_atoms)) + else: + batch_indices = x_0.batch + + interpolated_x_t, z = stochastic_interpolant.interpolate( + reshaped_t, + getattr(x_0, field_name), + getattr(x_1, field_name), + batch_indices, + ) + + # Update x_t (using internal method) + x_t.set_field(field_name, interpolated_x_t) + z_data[field_name] = z + + return x_t, OMATGData._from_dict(z_data) + + def losses( + self, + model_function: Callable[[OMATGData, paddle.Tensor], OMATGData], + t: paddle.Tensor, + x_0: OMATGData, + x_1: OMATGData, + ) -> dict[str, paddle.Tensor]: + """ + Compute the losses for the collection of stochastic interpolants. + + :param model_function: + Model function returning the velocity fields b and the denoisers eta. + :param t: + Times in [0,1]. + :param x_0: + Collection of points from the distribution p_0. + :param x_1: + Collection of points from the distribution p_1. + + :return: + The losses for the collection of stochastic interpolants. + :rtype: dict[str, paddle.Tensor] + """ + # Interpolate everything first (asserts on batch/n_atoms are done inside _interpolate) + x_t, z = self._interpolate(t, x_0, x_1) + + n_atoms = x_0.n_atoms + + losses = {} + for stochastic_interpolant, data_field in zip( + self._stochastic_interpolants, self._data_fields + ): + field_name = data_field.value + b_data_field = field_name + "_b" + eta_data_field = field_name + "_eta" + + reshaped_t = reshape_t(t, n_atoms, data_field) + + # Cell data requires different batch indices. + if data_field == DataField.cell: + batch_indices = paddle.arange(len(x_0.n_atoms)) + else: + batch_indices = x_0.batch + + def model_prediction_fn(x): + # Create a copy of x_t and update the field + x_t_clone = x_t.clone() + x_t_clone.set_field(field_name, x) + model_result = model_function(x_t_clone, t) + return model_result[b_data_field], model_result[eta_data_field] + + field_losses = stochastic_interpolant.loss( + model_prediction_fn, + reshaped_t, + getattr(x_0, field_name), + getattr(x_1, field_name), + getattr(x_t, field_name), + z.get_field(field_name), + batch_indices, + ) + + for loss_key, loss_value in field_losses.items(): + assert loss_key not in losses + losses[f"{field_name}_{loss_key}"] = loss_value + + return losses + + def integrate( + self, + x_0: OMATGData, + model_function: Callable[[OMATGData, paddle.Tensor], OMATGData], + save_intermediate: bool = False, + ) -> Union[OMATGData, Tuple[OMATGData, List[OMATGData]]]: + """ + Integrate the collection of points x_0 from time 0 to 1 based on the model + that provides the velocity fields b and denoisers eta. + + :param x_0: + Collection of points from the distribution p_0. + :param model_function: + Model function returning the velocity fields b and the denoisers eta. + :param save_intermediate: + If True, the intermediate points of the integration are saved and returned. + + :return: + Collection of integrated points x_1. + If save_intermediate is True, also returns a list of the intermediate points. + """ + times = paddle.linspace(SMALL_TIME, BIG_TIME, self._integration_time_steps) + dt = (BIG_TIME - SMALL_TIME) / (self._integration_time_steps - 1) + + x_t = x_0.clone() + new_x_t = x_0.clone() + + if save_intermediate: + inter_list = [x_t] + else: + inter_list = None + + with paddle.no_grad(): + for t_index in range(1, len(times)): + t = times[t_index - 1] + + for stochastic_interpolant, data_field in zip( + self._stochastic_interpolants, self._data_fields + ): + field_name = data_field.value + b_data_field = field_name + "_b" + eta_data_field = field_name + "_eta" + + def model_prediction_fn(time, x): + # Repeat time for each element in the batch + t_val = time.item() if time.numel() == 1 else float(time[0]) + time = paddle.full_like( + x_t.n_atoms.cast("float32"), + t_val, + ) + x_int = x_t.clone() + x_int.set_field(field_name, x) + model_result = model_function(x_int, time) + return model_result[b_data_field], model_result[eta_data_field] + + # Cell data requires different batch indices. + if data_field == DataField.cell: + batch_indices = paddle.arange(len(x_0.n_atoms)) + else: + batch_indices = x_0.batch + + new_value = stochastic_interpolant.integrate( + model_prediction_fn, + x_t.get_field(field_name), + t, + dt, + batch_indices, + ) + new_x_t.set_field(field_name, new_value) + + x_t = new_x_t.clone() + + if save_intermediate: + inter_list.append(x_t) + + if save_intermediate: + return x_t, inter_list + else: + return x_t + + def get_stochastic_interpolant(self, data_field: str): + """ + Return the stochastic interpolant associated with the data field. + + :param data_field: + Data field for which the stochastic interpolant is requested. + + :return: + Stochastic interpolant associated with the data field. + """ + try: + df = DataField[data_field.lower()] + except KeyError: + raise ValueError(f"Data field must be in {[d.value for d in DataField]}.") + + index = self._data_fields.index(df) + return self._stochastic_interpolants[index] + +_SI_MODULE = "ppmat.models.omatg.si" +_SUB_MODULES = [ + "ppmat.models.omatg.si.interpolants", + "ppmat.models.omatg.si.core", +] + + +def _resolve_class(class_name: str, default_module: str): + """Resolve a class by name using importlib resolution.""" + if "." in class_name: + module_path, cls_name = class_name.rsplit(".", 1) + module = importlib.import_module(module_path) + return getattr(module, cls_name) + + module = importlib.import_module(default_module) + if hasattr(module, class_name): + return getattr(module, class_name) + + for submodule_path in _SUB_MODULES: + try: + sub = importlib.import_module(submodule_path) + if hasattr(sub, class_name): + return getattr(sub, class_name) + except ImportError: + continue + + raise AttributeError( + f"Class '{class_name}' not found in '{default_module}' or its submodules." + ) + + +def _build_object(cfg: dict, default_module: str): + """Build a single object from a {__class_name__, __init_params__} config. + + Recursively builds nested objects when an __init_params__ entry is itself a + config dict (i.e. contains __class_name__). + """ + cls = _resolve_class(cfg["__class_name__"], default_module) + params = cfg.get("__init_params__", {}) + built_params = {} + for key, val in params.items(): + if isinstance(val, dict) and "__class_name__" in val: + built_params[key] = _build_object(val, default_module) + elif isinstance(val, list): + built_params[key] = [ + _build_object(item, default_module) + if isinstance(item, dict) and "__class_name__" in item + else item + for item in val + ] + else: + built_params[key] = val + return cls(**built_params) + + +def build_si_from_cfg(si_cfg: dict) -> StochasticInterpolants: + """Build StochasticInterpolants from a config dict. + + Expected schema: + stochastic_interpolants: list of {__class_name__, __init_params__} + data_fields: list[str] + integration_time_steps: int + relative_si_costs: dict[str, float] # optional + """ + interpolants = [ + _build_object(cfg_item, _SI_MODULE) + for cfg_item in si_cfg["stochastic_interpolants"] + ] + return StochasticInterpolants( + stochastic_interpolants=interpolants, + data_fields=si_cfg["data_fields"], + integration_time_steps=si_cfg.get("integration_time_steps", 210), + ) + + +def build_sampler_from_cfg(sampler_cfg: dict): + """Build IndependentSampler from a config dict. + + Expected schema (all keys optional): + dataset_name: str | None + mirror_species: bool + mask_species: bool + """ + from ppmat.models.omatg.model import IndependentSampler + + return IndependentSampler( + dataset_name=sampler_cfg.get("dataset_name"), + mirror_species=sampler_cfg.get("mirror_species", True), + mask_species=sampler_cfg.get("mask_species", False), + ) + + +MAX_ATOM_NUM: int = 100 + + +class DiscreteFlowMatchingMask(StochasticInterpolantSpecies): + """Discrete flow matching between masked base p_0 and target p_1 for species. + + The base points x_0 are entirely in the masked state (token 0). + The model prediction returns (sum(n_atoms), MAX_ATOM_NUM) logits. + Loss is cross_entropy(pred, x_1 - 1). + + :param noise: noise parameter scaling added during integration. + """ + + def __init__(self, noise: float = 0.0) -> None: + super().__init__() + if noise < 0.0: + raise ValueError("Noise parameter must be greater than or equal to 0.") + self._mask_index = 0 + self._noise = noise + + def interpolate( + self, + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Tuple[paddle.Tensor, paddle.Tensor]: + assert x_0.shape == x_1.shape + assert paddle.all(x_0 == self._mask_index) + assert paddle.all(x_1 != self._mask_index) + x_t = x_0.clone() + mask = paddle.rand(x_0.shape) < t + x_t[mask] = x_1[mask] + return x_t, paddle.zeros_like(x_t) + + def loss_keys(self) -> Iterable[str]: + yield "loss" + + def loss( + self, + model_function: Callable[[paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor]], + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + x_t: paddle.Tensor, + z: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Dict[str, paddle.Tensor]: + assert x_0.shape == x_1.shape + assert paddle.all(x_0 == self._mask_index) + assert paddle.all(x_1 != self._mask_index) + pred = model_function(x_t)[0] + assert pred.shape == (x_0.shape[0], MAX_ATOM_NUM) + return {"loss": paddle.nn.functional.cross_entropy(input=pred, label=x_1 - 1)} + + def integrate( + self, + model_function: Callable[ + [paddle.Tensor, paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor] + ], + x_t: paddle.Tensor, + time: paddle.Tensor, + time_step: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> paddle.Tensor: + eps = paddle.finfo(paddle.float64).eps + x_1_probs = paddle.nn.functional.softmax( + model_function(time, x_t)[0], axis=-1 + ) + x_1_probs = x_1_probs.reshape((-1, MAX_ATOM_NUM)) + shifted_x_1 = paddle.multinomial( + x_1_probs, num_samples=1, replacement=True + ).squeeze(-1) + shifted_x_t = x_t - 1 + assert shifted_x_1.shape == x_t.shape == shifted_x_t.shape + shifted_x_1_hot = paddle.nn.functional.one_hot( + shifted_x_1, num_classes=MAX_ATOM_NUM + ) + dpt = shifted_x_1_hot - 1.0 / MAX_ATOM_NUM + dpt_xt = dpt.gather(-1, shifted_x_t[:, None]).squeeze(-1) + pt = time * shifted_x_1_hot + (1.0 - time) * (1.0 / MAX_ATOM_NUM) + pt_xt = pt.gather(-1, shifted_x_t[:, None]).squeeze(-1) + S = paddle.count_nonzero(x=pt, axis=-1) + rate = paddle.nn.functional.relu(x=dpt - dpt_xt[:, None]) / ( + S * pt_xt + )[:, None] + rate[(pt_xt == 0.0)[:, None].expand([-1, MAX_ATOM_NUM])] = 0.0 + rate[pt == 0.0] = 0.0 + rate_db = paddle.zeros_like(rate) + if self._noise > 0.0: + rate_db[shifted_x_t == shifted_x_1] = 1.0 + rate_db[shifted_x_1 != shifted_x_t] = (MAX_ATOM_NUM * time + 1.0 - time) / ( + 1.0 - time + eps + ) + rate_db *= self._noise + rate = rate + rate_db + step_probs = (rate * time_step).clip(max=1.0) + step_probs = (rate * time_step).clip(max=1.0) + step_probs[paddle.arange(len(shifted_x_t)), shifted_x_t] = 0.0 + step_probs[paddle.arange(len(shifted_x_t)), shifted_x_t] = ( + (1.0 - step_probs.sum(axis=-1, keepdim=True)).clip(min=0.0) + ).squeeze(-1) + x_t = paddle.multinomial( + step_probs, num_samples=1, replacement=True + ).squeeze(-1) + 1 + return x_t + + def uses_masked_species(self) -> bool: + return True diff --git a/ppmat/models/omatg/si/interpolants.py b/ppmat/models/omatg/si/interpolants.py new file mode 100644 index 00000000..c52ff1b9 --- /dev/null +++ b/ppmat/models/omatg/si/interpolants.py @@ -0,0 +1,522 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Abstract classes for Stochastic Interpolants. +""" + +from abc import ABC, abstractmethod +from typing import Callable, Dict, Iterable, Tuple + +import paddle + + +class TimeChecker: + """Check that all times in a tensor are in [0,1].""" + + @staticmethod + def _check_t(t: paddle.Tensor) -> paddle.Tensor: + """Check that all times are in [0,1].""" + return paddle.all((0.0 <= t) & (t <= 1.0)) + + +class Corrector(ABC): + """Abstract corrector function (e.g., for PBC wrapping).""" + + @abstractmethod + def correct(self, x: paddle.Tensor) -> paddle.Tensor: + """Correct the input x.""" + raise NotImplementedError + + @abstractmethod + def unwrap(self, x_0: paddle.Tensor, x_1: paddle.Tensor) -> paddle.Tensor: + """Correct x_1 based on reference x_0.""" + raise NotImplementedError + + +class Epsilon(ABC, TimeChecker): + """Abstract class for defining an epsilon function epsilon(t).""" + + @abstractmethod + def epsilon(self, t: paddle.Tensor) -> paddle.Tensor: + """Evaluate the epsilon function at times t. + + :param t: Times in [0,1]. + :type t: paddle.Tensor + :return: Epsilon function epsilon(t). + :rtype: paddle.Tensor + """ + raise NotImplementedError + + +class Interpolant(ABC, TimeChecker): + """Abstract class for a diffusion scheduler interpolant. + + I(t, x_0, x_1) = alpha(t) * x_0 + beta(t) * x_1 + """ + + def interpolate( + self, t: paddle.Tensor, x_0: paddle.Tensor, x_1: paddle.Tensor + ) -> paddle.Tensor: + """Interpolate between points x_0 and x_1 at times t. + + :param t: Times in [0,1]. + :param x_0: Points from p_0. + :param x_1: Points from p_1. + :return: Interpolated value. + """ + assert bool(self._check_t(t)) + x_0prime = self.get_corrector().correct(x_0) + x_1prime = self.get_corrector().unwrap(x_0prime, x_1) + x_t = self.alpha(t) * x_0prime + self.beta(t) * x_1prime + return self.get_corrector().correct(x_t) + + def interpolate_derivative( + self, t: paddle.Tensor, x_0: paddle.Tensor, x_1: paddle.Tensor + ) -> paddle.Tensor: + """Compute the derivative of the interpolant with respect to time. + + :param t: Times in [0,1]. + :param x_0: Points from p_0. + :param x_1: Points from p_1. + :return: Derivative of the interpolant. + """ + assert bool(self._check_t(t)) + x_0prime = self.get_corrector().correct(x_0) + x_1prime = self.get_corrector().unwrap(x_0prime, x_1) + return self.alpha_dot(t) * x_0prime + self.beta_dot(t) * x_1prime + + @abstractmethod + def get_corrector(self) -> Corrector: + """Get the corrector implied by the interpolant.""" + raise NotImplementedError + + @abstractmethod + def alpha(self, t: paddle.Tensor) -> paddle.Tensor: + """Alpha function in the linear interpolant. + + :param t: Times in [0,1]. + :return: Values of the alpha function at the given times. + """ + raise NotImplementedError + + @abstractmethod + def alpha_dot(self, t: paddle.Tensor) -> paddle.Tensor: + """Time derivative of the alpha function. + + :param t: Times in [0,1]. + :return: Derivatives of the alpha function at the given times. + """ + raise NotImplementedError + + @abstractmethod + def beta(self, t: paddle.Tensor) -> paddle.Tensor: + """Beta function in the linear interpolant. + + :param t: Times in [0,1]. + :return: Values of the beta function at the given times. + """ + raise NotImplementedError + + @abstractmethod + def beta_dot(self, t: paddle.Tensor): + """Time derivative of the beta function. + + :param t: Times in [0,1]. + :return: Derivatives of the beta function at the given times. + """ + raise NotImplementedError + + +class LatentGamma(ABC, TimeChecker): + """Abstract class for the gamma function gamma(t) in gamma(t) * z.""" + + @abstractmethod + def gamma(self, t: paddle.Tensor) -> paddle.Tensor: + """Evaluate the gamma function at times t. + + :param t: Times in [0,1]. + :return: Gamma function gamma(t). + """ + raise NotImplementedError + + @abstractmethod + def gamma_derivative(self, t: paddle.Tensor) -> paddle.Tensor: + """Compute the derivative of the gamma function with respect to time. + + :param t: Times in [0,1]. + :return: Derivative of the gamma function. + """ + raise NotImplementedError + + @abstractmethod + def requires_antithetic(self) -> bool: + """Whether the gamma function requires antithetic sampling. + + :return: Whether antithetic sampling is required. + """ + raise NotImplementedError + + +class StochasticInterpolant(ABC, TimeChecker): + """Abstract class for a stochastic interpolant between x_0 and x_1.""" + + @abstractmethod + def interpolate( + self, + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Tuple[paddle.Tensor, paddle.Tensor]: + """Stochastically interpolate between x_0 and x_1 at times t. + + :param t: Times in [0,1]. + :param x_0: Points from p_0. + :param x_1: Points from p_1. + :param batch_indices: Configuration index for every atom in the batch. + :return: Stochastically interpolated points x_t, random variables z. + """ + raise NotImplementedError + + @abstractmethod + def loss_keys(self) -> Iterable[str]: + """Get the keys of the losses returned by the loss function. + + :return: Keys of the losses. + """ + raise NotImplementedError + + @abstractmethod + def loss( + self, + model_function: Callable[[paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor]], + t: paddle.Tensor, + x_0: paddle.Tensor, + x_1: paddle.Tensor, + x_t: paddle.Tensor, + z: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> Dict[str, paddle.Tensor]: + """Compute the losses for the stochastic interpolant. + + :param model_function: Model returning velocity fields b and denoisers eta. + :param t: Times in [0,1]. + :param x_0: Points from p_0. + :param x_1: Points from p_1. + :param x_t: Stochastically interpolated points x_t. + :param z: Random variable z used for stochastic interpolation. + :param batch_indices: Configuration index for every atom in the batch. + :return: Losses. + """ + raise NotImplementedError + + @abstractmethod + def integrate( + self, + model_function: Callable[ + [paddle.Tensor, paddle.Tensor], Tuple[paddle.Tensor, paddle.Tensor] + ], + x_t: paddle.Tensor, + time: paddle.Tensor, + time_step: paddle.Tensor, + batch_indices: paddle.Tensor, + ) -> paddle.Tensor: + """Integrate the current positions x_t for the given time step. + + :param model_function: Model returning velocity fields b and denoisers eta. + :param x_t: Current positions. + :param time: Initial time (0-dim tensor). + :param time_step: Time step (0-dim tensor). + :param batch_indices: Configuration index for every atom in the batch. + :return: Integrated position. + """ + raise NotImplementedError + + @abstractmethod + def get_corrector(self) -> Corrector: + """Get the corrector implied by the stochastic interpolant.""" + raise NotImplementedError + + +class StochasticInterpolantSpecies(StochasticInterpolant, ABC): + """Abstract class for a stochastic interpolant between species x_0 and x_1.""" + + def get_corrector(self) -> Corrector: + """Stochastic interpolants for atom species should not define a corrector.""" + raise RuntimeError("Corrector not defined for StochasticInterpolantSpecies.") + + @abstractmethod + def uses_masked_species(self) -> bool: + """Whether the stochastic interpolant uses a masked species. + + :return: Whether masked species are used. + """ + raise NotImplementedError + + +class Sigma(ABC, TimeChecker): + """Noise schedule sigma(s) for one-sided variance-exploding interpolant.""" + + @abstractmethod + def sigma(self, s: paddle.Tensor) -> paddle.Tensor: + """Evaluate the sigma function at times s. + + :param s: Times in [0,1]. + :return: Sigma function sigma(s). + """ + raise NotImplementedError + + @abstractmethod + def sigma_dot(self, s: paddle.Tensor) -> paddle.Tensor: + """Derivative of the sigma function with respect to time. + + :param s: Times in [0,1]. + :return: Derivative of the sigma function. + """ + raise NotImplementedError + + +class Tau(ABC, TimeChecker): + """Tau function for one-sided variance-preserving interpolant. + + x_t = sqrt(1 - tau^2(t)) * x_0 + tau(t) * x_1 + """ + + @abstractmethod + def tau(self, t: paddle.Tensor) -> paddle.Tensor: + """Evaluate the tau function at times t. + + :param t: Times in [0,1]. + :return: Tau function tau(t). + """ + raise NotImplementedError + + @abstractmethod + def tau_dot(self, t: paddle.Tensor) -> paddle.Tensor: + """Derivative of the tau function with respect to time. + + :param t: Times in [0,1]. + :return: Derivative of the tau function. + """ + raise NotImplementedError + + +class IdentityCorrector(Corrector): + def correct(self, x: paddle.Tensor) -> paddle.Tensor: + return x + + def unwrap(self, x_0: paddle.Tensor, x_1: paddle.Tensor) -> paddle.Tensor: + return x_1.clone() + + +class PeriodicBoundaryConditionsCorrector(Corrector): + def __init__(self, min_value: float, max_value: float) -> None: + super().__init__() + if min_value >= max_value: + raise ValueError("Minimum value must be less than maximum value.") + self._min_value = min_value + self._max_value = max_value + self._range = max_value - min_value + + def correct(self, x: paddle.Tensor) -> paddle.Tensor: + return paddle.remainder(x - self._min_value, self._range) + self._min_value + + def unwrap(self, x_0: paddle.Tensor, x_1: paddle.Tensor) -> paddle.Tensor: + half = self._range / 2.0 + sep = x_1 - x_0 + return x_0 + paddle.remainder(sep + half, self._range) - half + + +class VanishingEpsilon(Epsilon): + def __init__(self, c: float = 1.0, sigma: float = 0.01, mu: float = 0.075) -> None: + super().__init__() + self._c = c + self._sigma = sigma + self._mu = mu + + def epsilon(self, t: paddle.Tensor) -> paddle.Tensor: + self._check_t(t) + f1 = paddle.sigmoid((t - self._mu) / self._sigma) + f2 = paddle.sigmoid((1 - self._mu - t) / self._sigma) + return self._c * f1 * f2 + + +class ConstantEpsilon(Epsilon): + def __init__(self, c: float) -> None: + super().__init__() + self._c = c + + def epsilon(self, t: paddle.Tensor) -> paddle.Tensor: + self._check_t(t) + return paddle.full_like(t, self._c) + + +class TauConstantSchedule(Tau): + def tau(self, t: paddle.Tensor) -> paddle.Tensor: + self._check_t(t) + return t.clone() + + def tau_dot(self, t: paddle.Tensor) -> paddle.Tensor: + self._check_t(t) + return paddle.ones_like(t) + + +class LinearInterpolant(Interpolant): + def alpha(self, t: paddle.Tensor) -> paddle.Tensor: + return 1.0 - t + + def alpha_dot(self, t: paddle.Tensor) -> paddle.Tensor: + return -paddle.ones_like(t) + + def beta(self, t: paddle.Tensor) -> paddle.Tensor: + return t.clone() + + def beta_dot(self, t: paddle.Tensor) -> paddle.Tensor: + return paddle.ones_like(t) + + def get_corrector(self) -> Corrector: + return IdentityCorrector() + + +class PeriodicLinearInterpolant(LinearInterpolant): + def __init__(self) -> None: + super().__init__() + self._corrector = PeriodicBoundaryConditionsCorrector(0.0, 1.0) + + def get_corrector(self) -> Corrector: + return PeriodicBoundaryConditionsCorrector(0.0, 1.0) + + +class ScoreBasedDiffusionModelInterpolantVP(Interpolant): + """Variance-preserving SBDM interpolant scheduler. + + Uses a Tau schedule for alpha/beta diffusion coefficients. + """ + + def __init__(self, tau: Tau) -> None: + super().__init__() + self._tau = tau + + def alpha(self, t: paddle.Tensor) -> paddle.Tensor: + return paddle.sqrt(1.0 - self._tau.tau(t) ** 2) + + def alpha_dot(self, t: paddle.Tensor) -> paddle.Tensor: + tau = self._tau.tau(t) + return -tau * self._tau.tau_dot(t) / paddle.sqrt(1.0 - tau**2) + + def beta(self, t: paddle.Tensor) -> paddle.Tensor: + return self._tau.tau(t) + + def beta_dot(self, t: paddle.Tensor) -> paddle.Tensor: + return self._tau.tau_dot(t) + + def get_corrector(self) -> Corrector: + return IdentityCorrector() + + +class ScoreBasedDiffusionModelInterpolantVE(Interpolant): + """Variance-exploding SBDM interpolant scheduler. + + Uses a Sigma schedule for alpha/beta diffusion coefficients. + """ + + def __init__(self, sigma: Sigma) -> None: + super().__init__() + self._sigma = sigma + + def alpha(self, t: paddle.Tensor) -> paddle.Tensor: + return paddle.sqrt( + self._sigma.sigma(1.0 - t) ** 2 + - self._sigma.sigma(paddle.zeros_like(t)) ** 2 + ) + + def alpha_dot(self, t: paddle.Tensor) -> paddle.Tensor: + sigma = self._sigma.sigma(1.0 - t) + alpha = paddle.sqrt(sigma**2 - self._sigma.sigma(paddle.zeros_like(t)) ** 2) + return -sigma * self._sigma.sigma_dot(1.0 - t) / alpha + + def beta(self, t: paddle.Tensor) -> paddle.Tensor: + return paddle.ones_like(t) + + def beta_dot(self, t: paddle.Tensor) -> paddle.Tensor: + return paddle.zeros_like(t) + + def get_corrector(self) -> Corrector: + return IdentityCorrector() + + +class GeometricSigma(Sigma): + """Geometric noise schedule for the VESBD interpolant.""" + + def __init__(self, sigma_min: float, sigma_max: float) -> None: + super().__init__() + if sigma_min <= 0.0 or sigma_max <= 0.0 or sigma_max <= sigma_min: + raise ValueError("sigma_min and sigma_max must be positive, sigma_max > sigma_min.") + self._sigma_min = sigma_min + self._sigma_max = sigma_max + from math import log + self._log_ratio = log(sigma_max / sigma_min) + + def sigma(self, s: paddle.Tensor) -> paddle.Tensor: + self._check_t(s) + return self._sigma_min * (self._sigma_max / self._sigma_min) ** s + + def sigma_dot(self, s: paddle.Tensor) -> paddle.Tensor: + self._check_t(s) + return self._sigma_min * self._log_ratio * (self._sigma_max / self._sigma_min) ** s + + +class LatentGammaSqrt(LatentGamma): + def __init__(self, a: float) -> None: + super().__init__() + if a <= 0.0: + raise ValueError("Constant a must be positive.") + self._a = a + + def gamma(self, t: paddle.Tensor) -> paddle.Tensor: + self._check_t(t) + return paddle.sqrt(self._a * t * (1.0 - t)) + + def gamma_derivative(self, t: paddle.Tensor) -> paddle.Tensor: + self._check_t(t) + return self._a * (1.0 - 2.0 * t) / (2.0 * paddle.sqrt(self._a * t * (1.0 - t))) + + def requires_antithetic(self) -> bool: + return True + + +class LatentGammaEncoderDecoder(LatentGamma): + def __init__(self, a: float = 1.0, switch_time: float = 0.5, power: float = 1.0) -> None: + super().__init__() + if a <= 0.0 or switch_time <= 0.0 or switch_time >= 1.0 or power < 0.5: + raise ValueError("Invalid parameters.") + self._sqrt_a = a**0.5 + self._switch_time = switch_time + self._power = power + + def gamma(self, t: paddle.Tensor) -> paddle.Tensor: + self._check_t(t) + a = (t - self._switch_time * t) ** self._power + b = (self._switch_time - self._switch_time * t) ** self._power + a + return self._sqrt_a * paddle.sin(paddle.pi * a / b) ** 2 + + def gamma_derivative(self, t: paddle.Tensor) -> paddle.Tensor: + self._check_t(t) + a = (t - self._switch_time * t) ** self._power + b = (self._switch_time - self._switch_time * t) ** self._power + c = paddle.sin(2.0 * paddle.pi * a / (a + b)) + return -self._sqrt_a * self._power * paddle.pi * a * b * c / (t * (t - 1.0) * (a + b) ** 2) + + def requires_antithetic(self) -> bool: + return False diff --git a/structure_generation/configs/omatg/README.md b/structure_generation/configs/omatg/README.md new file mode 100644 index 00000000..717a684a --- /dev/null +++ b/structure_generation/configs/omatg/README.md @@ -0,0 +1,306 @@ +# OMatG + +[Open Materials Generation with Stochastic Interpolants](https://openreview.net/forum?id=gHGrzxFujU) (ICML 2025) / +[All that structure matches does not glitter](https://openreview.net/forum?id=ig9ujp50D4) (NeurIPS 2025) + +## Abstract + +A state-of-the-art generative model for crystal structure prediction and *de novo* generation of inorganic crystals. +OMatG implements the [stochastic interpolants (SIs) framework](https://arxiv.org/abs/2303.08797) that bridges samples +from a base distribution to the target data distribution. A stochastic interpolant +$x_t = \alpha(t)\,x_0 + \beta(t)\,x_1 + \gamma(t)\,z$ evolves base samples $x_0$ into data samples $x_1$ over time +$t\in[0,1]$. The time-dependent density is realized via deterministic (ODE) or stochastic (SDE) sampling, requiring +only a learned velocity field $b^\theta(t, x)$ (and optionally a denoiser $z^\theta(t, x)$ for SDE). + +OMatG defines a crystalline material by its unit cell ($\mathbf{L}\in\mathbb{R}^{3\times3}$), fractional coordinates +($\mathbf{X}\in[0,1)^{3\times N}$ with periodic boundary conditions), and discrete atomic species +($\mathbf{A}\in\mathbb{Z}^N_{>0}$). The SI framework handles the continuous variables $\{\mathbf{X}, \mathbf{L}\}$ +while discrete species $\mathbf{A}$ are treated with [discrete flow matching](https://arxiv.org/abs/2402.04997). + +Two crystal generation modes are supported: +1. **CSP** (crystal structure prediction): atomic species are fixed; only coordinates and lattice vectors evolve. +2. **DNG** (*de novo* generation): all species are masked at start and evolve together with structure. + +## Model Architecture + +OMatG consists of the following core components: + +- **CSPNet**: Crystal structure prediction network using message-passing GNN architecture, supporting fully-connected (fc) or k-NN (knn) edge construction. +- **Stochastic Interpolants**: SI framework supporting ODE/SDE integration, with multiple interpolant types (Linear/Trigonometric/EncDec/VESBD/VPSBD). +- **IndependentSampler**: Independent distribution sampler for position/lattice/species base distributions. + +## Datasets + +### Included Datasets + +Several standard material datasets are included as LMDB files: + +| Dataset | Structures | Max Atoms | Description | +|---------|-----------:|:---------:|-------------| +| MP-20 | 45,229 | 20 | [Materials Project](https://pubs.aip.org/aip/apm/article/1/1/011002/119685) structures | +| MPTS-52 | 40,476 | 52 | [Chronological MP split](https://joss.theoj.org/papers/10.21105/joss.05618) | +| Perov-5 | 18,928 | 5 | [Perovskite dataset](https://pubs.rsc.org/en/content/articlelanding/2012/ee/c2ee22341d) | +| Alex-MP-20 | 675,204 | — | Consolidated [Alexandria](https://arxiv.org/abs/2210.00579) + MP-20 | + + +### Supported Datasets + +Pretrained weights are available for 5 dataset × mode combinations. +Training configs are provided for `mp_20` (CSP + DNG); other datasets +(`perov_5_csp`, `mpts_52_csp`, `alex_mp_20_csp`) can be used by changing +the `file_path` in the dataset section of the config. + +| Dataset | Mode | Variants | Weight Index | +|---------|:----:|:--------:|--------------| +| mp_20_csp | CSP | 11 | `build_omatg_model("mp_20_csp", variant)` | +| mp_20_dng | DNG | 11 | `build_omatg_model("mp_20_dng", variant)` | +| perov_5_csp | CSP | 11 | `build_omatg_model("perov_5_csp", variant)` | +| mpts_52_csp | CSP | 8 | `build_omatg_model("mpts_52_csp", variant)` | +| alex_mp_20_csp | CSP | 11 | `build_omatg_model("alex_mp_20_csp", variant)` | + + +## Configuration Files + +| File | Mode | Dataset | Description | +|------|:----:|---------|-------------| +| `omatg_mp20_csp.yaml` | CSP | MP-20 | Simplified MSE (fixed lr) | +| `omatg_mp20_csp_linear_ode.yaml` | CSP | MP-20 | Linear-ODE (Cosine LR) | +| `omatg_mp20_dng_linear_sde.yaml` | DNG | MP-20 | Linear-SDE (species prediction + WeightDecay) | +| `omatg_mp20_csp_sample.yaml` | CSP | MP-20 | Sampling by number of atoms | +| `omatg_mp20_dng_sample.yaml` | DNG | MP-20 | DNG sampling (species prediction) | + +## Results + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
ModelDatasetModeInterpolantConfig (Train / Sample)Weight
mp_20_cspMP-20CSPLinear-ODEtrain / sampleweight
mp_20_dngMP-20DNGLinear-SDEtrain / sampleweight
perov_5_cspPerov-5CSPEncDec-ODE-Gammause mp_20 CSP config, change file_path to ./data/perov_5/weight
mpts_52_cspMPTS-52CSPEncDec-ODE-Gammause mp_20 CSP config, change file_path to ./data/mpts_52/weight
alex_mp_20_cspAlex-MP-20CSPEncDec-ODE-Gammause mp_20 CSP config, change file_path to ./data/alex_mp_20/weight
+ +## Pretrained Weights + +Pre-trained weights hosted on Baidu BOS (50 files). Use `build_omatg_model(dataset, variant)` +for automatic download. Cached to `~/.paddlemat/weights/omatg_{dataset}/` after first download. + +```python +from ppmat.models.omatg import build_omatg_model +model, meta = build_omatg_model("mp_20_csp", "encdec_ode_gamma") # CSP +model, meta = build_omatg_model("mp_20_dng", "encdec_ode_gamma") # DNG +``` + +### Perov-5 + +| Variant | Weight | +|---------|--------| +| encdec_ode_gamma | [EncDec-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/EncDec-ODE-Gamma.pdparams) | +| encdec_sde_gamma | [EncDec-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/EncDec-SDE-Gamma.pdparams) | +| linear_ode | [Linear-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Linear-ODE.pdparams) | +| linear_ode_gamma | [Linear-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Linear-ODE-Gamma.pdparams) | +| linear_sde_gamma | [Linear-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Linear-SDE-Gamma.pdparams) | +| trig_ode | [Trig-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Trig-ODE.pdparams) | +| trig_ode_gamma | [Trig-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Trig-ODE-Gamma.pdparams) | +| trig_sde_gamma | [Trig-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/Trig-SDE-Gamma.pdparams) | +| vesbd_ode | [VESBD-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/VESBD-ODE.pdparams) | +| vpsbd_ode | [VPSBD-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/VPSBD-ODE.pdparams) | +| vpsbd_sde | [VPSBD-SDE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_perov_5_csp/VPSBD-SDE.pdparams) | + +### MPTS-52 + +| Variant | Weight | +|---------|--------| +| encdec_ode_gamma | [EncDec-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/EncDec-ODE-Gamma.pdparams) | +| encdec_sde_gamma | [EncDec-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/EncDec-SDE-Gamma.pdparams) | +| linear_ode | [Linear-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Linear-ODE.pdparams) | +| linear_ode_gamma | [Linear-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Linear-ODE-Gamma.pdparams) | +| linear_sde_gamma | [Linear-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Linear-SDE-Gamma.pdparams) | +| trig_ode | [Trig-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Trig-ODE.pdparams) | +| trig_ode_gamma | [Trig-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Trig-ODE-Gamma.pdparams) | +| trig_sde_gamma | [Trig-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mpts_52_csp/Trig-SDE-Gamma.pdparams) | + +### MP-20 (DNG) + +| Variant | Weight | +|---------|--------| +| encdec_ode_gamma | [EncDec-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/EncDec-ODE-Gamma.pdparams) | +| encdec_sde_gamma | [EncDec-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/EncDec-SDE-Gamma.pdparams) | +| linear_ode | [Linear-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Linear-ODE.pdparams) | +| linear_ode_gamma | [Linear-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Linear-ODE-Gamma.pdparams) | +| linear_sde_gamma | [Linear-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Linear-SDE-Gamma.pdparams) | +| trig_ode | [Trig-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Trig-ODE.pdparams) | +| trig_ode_gamma | [Trig-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Trig-ODE-Gamma.pdparams) | +| trig_sde_gamma | [Trig-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/Trig-SDE-Gamma.pdparams) | +| vesbd_ode | [VESBD-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/VESBD-ODE.pdparams) | +| vpsbd_ode | [VPSBD-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/VPSBD-ODE.pdparams) | +| vpsbd_sde | [VPSBD-SDE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_dng/VPSBD-SDE.pdparams) | + +### MP-20 (CSP) + +| Variant | Weight | +|---------|--------| +| encdec_ode_gamma | [EncDec-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/EncDec-ODE-Gamma.pdparams) | +| encdec_sde_gamma | [EncDec-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/EncDec-SDE-Gamma.pdparams) | +| linear_ode | [Linear-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Linear-ODE.pdparams) | +| linear_ode_gamma | [Linear-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Linear-ODE-Gamma.pdparams) | +| linear_sde_gamma | [Linear-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Linear-SDE-Gamma.pdparams) | +| trig_ode | [Trig-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Trig-ODE.pdparams) | +| trig_ode_gamma | [Trig-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Trig-ODE-Gamma.pdparams) | +| trig_sde_gamma | [Trig-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/Trig-SDE-Gamma.pdparams) | +| vesbd_ode | [VESBD-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/VESBD-ODE.pdparams) | +| vpsbd_ode | [VPSBD-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/VPSBD-ODE.pdparams) | +| vpsbd_sde | [VPSBD-SDE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_mp_20_csp/VPSBD-SDE.pdparams) | + +### Alex-MP-20 + +| Variant | Weight | +|---------|--------| +| encdec_ode_gamma | [EncDec-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/EncDec-ODE-Gamma.pdparams) | +| encdec_sde_gamma | [EncDec-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/EncDec-SDE-Gamma.pdparams) | +| linear_ode | [Linear-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Linear-ODE.pdparams) | +| linear_ode_gamma | [Linear-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Linear-ODE-Gamma.pdparams) | +| linear_sde_gamma | [Linear-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Linear-SDE-Gamma.pdparams) | +| trig_ode | [Trig-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Trig-ODE.pdparams) | +| trig_ode_gamma | [Trig-ODE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Trig-ODE-Gamma.pdparams) | +| trig_sde_gamma | [Trig-SDE-Gamma.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/Trig-SDE-Gamma.pdparams) | +| vesbd_ode | [VESBD-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/VESBD-ODE.pdparams) | +| vpsbd_ode | [VPSBD-ODE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/VPSBD-ODE.pdparams) | +| vpsbd_sde | [VPSBD-SDE.pdparams](https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/structure_generation/OMatG/omatg_alex_mp_20_csp/VPSBD-SDE.pdparams) | + +## Commands + +### Training + +```bash +# Smoke training (1 epoch, small batch, no eval) +python structure_generation/train.py \ + -c structure_generation/configs/omatg/omatg_mp20_csp.yaml \ + Trainer.max_epochs=1 \ + Global.do_eval=False \ + Dataset.train.sampler.__init_params__.batch_size=32 \ + Dataset.val.sampler.__init_params__.batch_size=32 + +# Full training (requires data/mp_20/ dataset) +python structure_generation/train.py \ + -c structure_generation/configs/omatg/omatg_mp20_csp_linear_ode.yaml \ + Trainer.max_epochs=500 \ + Trainer.output_dir=./output/omatg_mp20_csp +``` + +### Validation + +```bash +# Evaluate on the validation split using a saved checkpoint +python structure_generation/train.py \ + -c structure_generation/configs/omatg/omatg_mp20_csp_linear_ode.yaml \ + Global.do_train=False \ + Global.do_eval=True \ + Trainer.pretrained_model_path=./output/omatg_mp20_csp/checkpoints +``` + +### Testing + +```bash +# Evaluate on the test split +python structure_generation/train.py \ + -c structure_generation/configs/omatg/omatg_mp20_csp_linear_ode.yaml \ + Global.do_train=False \ + Global.do_test=True \ + Global.do_eval=False \ + Trainer.pretrained_model_path=./output/omatg_mp20_csp/checkpoints +``` + +### Sampling + +```bash +# Sample by number of atoms (with local checkpoint) +python structure_generation/sample.py \ + --config_path structure_generation/configs/omatg/omatg_mp20_csp_sample.yaml \ + --checkpoint_path ./output/omatg_mp20_csp/checkpoints/best.pdparams \ + --mode by_num_atoms \ + --num_atoms 8 \ + --save_path ./results/omatg_samples + +# Sample by chemical formula (with local checkpoint) +python structure_generation/sample.py \ + --config_path structure_generation/configs/omatg/omatg_mp20_csp_sample.yaml \ + --checkpoint_path ./output/omatg_mp20_csp/checkpoints/best.pdparams \ + --mode by_chemical_formula \ + --chemical_formula LiMnO2 \ + --save_path ./results/omatg_samples + +# Sample using pre-trained weights (HTTP auto-download) +# Use build_omatg_model() in a Python script, then pass the model to the sampler. + +# Batch sample by dataloader (with local checkpoint) +python structure_generation/sample.py \ + --config_path structure_generation/configs/omatg/omatg_mp20_csp_sample.yaml \ + --checkpoint_path ./output/omatg_mp20_csp/checkpoints/best.pdparams \ + --mode by_dataloader \ + --save_path ./results/omatg_samples +``` + + + +## Citation + +```bibtex +@inproceedings{ + hoellmer2025, + title={Open Materials Generation with Stochastic Interpolants}, + author={Philipp H{\"o}llmer and Thomas Egg and Maya Martirossyan and Eric + Fuemmeler and Zeren Shui and Amit Gupta and Pawan Prakash and Adrian + Roitberg and Mingjie Liu and George Karypis and Mark Transtrum and Richard + Hennig and Ellad B. Tadmor and Stefano Martiniani}, + booktitle={Forty-second International Conference on Machine Learning}, + year={2025}, + url={https://openreview.net/forum?id=gHGrzxFujU}, + archivePrefix={arXiv}, + eprint={2502.02582}, + primaryClass={cs.LG}, +} +``` + +```bibtex +@inproceedings{ + martirossyan2025, + title={All that structure matches does not glitter}, + author={Maya Martirossyan and Thomas Egg and Philipp H{\"o}llmer + and George Karypis and Mark Transtrum and Adrian Roitberg + and Mingjie Liu and Richard Hennig and Ellad B. Tadmor and Stefano Martiniani}, + booktitle={Thirty-Ninth Annual Conference on Neural Information Processing Systems}, + year={2025}, + url={https://openreview.net/forum?id=ig9ujp50D4}, + archivePrefix={arXiv}, + eprint={2509.12178}, + primaryClass={cs.LG}, +} +``` diff --git a/structure_generation/configs/omatg/omatg_mp20_csp.yaml b/structure_generation/configs/omatg/omatg_mp20_csp.yaml new file mode 100644 index 00000000..cf345708 --- /dev/null +++ b/structure_generation/configs/omatg/omatg_mp20_csp.yaml @@ -0,0 +1,134 @@ +Global: + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 500 + seed: 42 + output_dir: ./output/omatg_mp20_csp + save_freq: 50 + log_freq: 10 + start_eval_epoch: 1 + eval_freq: 1 + pretrained_model_path: null + resume_from_checkpoint: null + use_amp: False + amp_level: 'O1' + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: "loss" + greater_is_better: False + compute_metric_during_train: False + metric_strategy_during_eval: 'step' + use_visualdl: False + use_wandb: False + use_tensorboard: False + +Model: + __class_name__: OMATGCSPNetFull + __init_params__: + hidden_dim: 512 + num_layers: 6 + max_atoms: 100 + act_fn: silu + dis_emb: sin + num_freqs: 128 + edge_style: fc + cutoff: 7.0 + max_neighbors: 20 + ln: True + ip: True + smooth: False + pred_type: False + pred_scalar: False + time_embed_dim: 256 + +Optimizer: + __class_name__: Adam + __init_params__: + beta1: 0.9 + beta2: 0.999 + lr: 0.0005 + +Dataset: + train: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/train.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: True + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 256 + val: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/val.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: True + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 128 + test: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/test.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: True + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: DistributedBatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 128 + +Sample: + data: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/test.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: True + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 128 + build_structure_cfg: + format: array + niggli: False + model_sample_params: + num_inference_steps: 100 diff --git a/structure_generation/configs/omatg/omatg_mp20_csp_linear_ode.yaml b/structure_generation/configs/omatg/omatg_mp20_csp_linear_ode.yaml new file mode 100644 index 00000000..beaae40c --- /dev/null +++ b/structure_generation/configs/omatg/omatg_mp20_csp_linear_ode.yaml @@ -0,0 +1,139 @@ +Global: + do_train: True + do_eval: True + do_test: False + +Trainer: + max_epochs: 2000 + seed: 42 + output_dir: ./output/omatg_mp20_csp_linear_ode + save_freq: 100 + log_freq: 10 + start_eval_epoch: 1 + eval_freq: 1 + pretrained_model_path: null + resume_from_checkpoint: null + use_amp: False + amp_level: 'O1' + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: "loss" + greater_is_better: False + compute_metric_during_train: False + metric_strategy_during_eval: 'step' + use_visualdl: False + use_wandb: False + use_tensorboard: False + +Model: + __class_name__: OMATGCSPNetFull + __init_params__: + hidden_dim: 512 + num_layers: 6 + max_atoms: 100 + act_fn: silu + dis_emb: sin + num_freqs: 128 + edge_style: fc + cutoff: 7.0 + max_neighbors: 20 + ln: True + ip: True + smooth: False + pred_type: False + pred_scalar: False + time_embed_dim: 256 + +Optimizer: + __class_name__: Adam + __init_params__: + beta1: 0.9 + beta2: 0.999 + weight_decay: 0.0 + lr: + __class_name__: Cosine + __init_params__: + learning_rate: 0.0006689636445843722 + by_epoch: True + +Dataset: + train: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/train.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: False + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 1024 + val: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/val.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: False + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 1024 + test: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/test.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: False + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: DistributedBatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 1024 + +Sample: + data: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/test.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 128 + build_structure_cfg: + format: array + niggli: False + model_sample_params: + num_inference_steps: 100 diff --git a/structure_generation/configs/omatg/omatg_mp20_csp_sample.yaml b/structure_generation/configs/omatg/omatg_mp20_csp_sample.yaml new file mode 100644 index 00000000..35d7959c --- /dev/null +++ b/structure_generation/configs/omatg/omatg_mp20_csp_sample.yaml @@ -0,0 +1,53 @@ +Global: + do_train: False + do_eval: False + do_test: False + +Model: + __class_name__: OMATGCSPNetFull + __init_params__: + hidden_dim: 512 + num_layers: 6 + max_atoms: 100 + act_fn: silu + dis_emb: sin + num_freqs: 128 + edge_style: fc + cutoff: 7.0 + max_neighbors: 20 + ln: True + ip: True + smooth: False + pred_type: False + pred_scalar: False + time_embed_dim: 256 + +Sample: + data: + dataset: + __class_name__: NumAtomsCrystalDataset + __init_params__: + total_num: 10 + prop_names: [] + prop_values: [] + loader: + num_workers: 0 + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 10 + + build_structure_cfg: + format: array + niggli: False + + model_sample_params: + num_inference_steps: 100 + + metrics: + __class_name__: CSPMetric + __init_params__: + gt_file_path: "./data/mp_20/test.csv" diff --git a/structure_generation/configs/omatg/omatg_mp20_dng_linear_sde.yaml b/structure_generation/configs/omatg/omatg_mp20_dng_linear_sde.yaml new file mode 100644 index 00000000..42847256 --- /dev/null +++ b/structure_generation/configs/omatg/omatg_mp20_dng_linear_sde.yaml @@ -0,0 +1,139 @@ +Global: + do_train: True + do_eval: True + do_test: False + +Trainer: + max_epochs: 2000 + seed: 42 + output_dir: ./output/omatg_mp20_dng_linear_sde + save_freq: 100 + log_freq: 10 + start_eval_epoch: 1 + eval_freq: 1 + pretrained_model_path: null + resume_from_checkpoint: null + use_amp: False + amp_level: 'O1' + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: 'eval_loss' + name_for_best_metric: "loss" + greater_is_better: False + compute_metric_during_train: False + metric_strategy_during_eval: 'step' + use_visualdl: False + use_wandb: False + use_tensorboard: False + +Model: + __class_name__: OMATGCSPNetFull + __init_params__: + hidden_dim: 512 + num_layers: 6 + max_atoms: 100 + act_fn: silu + dis_emb: sin + num_freqs: 128 + edge_style: fc + cutoff: 7.0 + max_neighbors: 20 + ln: True + ip: True + smooth: False + pred_type: True + pred_scalar: False + time_embed_dim: 256 + +Optimizer: + __class_name__: Adam + __init_params__: + beta1: 0.9 + beta2: 0.999 + weight_decay: 0.01 + lr: + __class_name__: Cosine + __init_params__: + learning_rate: 0.001 + by_epoch: True + +Dataset: + train: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/train.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: False + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 512 + val: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/val.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: False + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 512 + test: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/test.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: False + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: DistributedBatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 512 + +Sample: + data: + dataset: + __class_name__: OMATGStructureDataset + __init_params__: + file_path: "./data/mp_20/test.lmdb" + property_keys: [] + lazy_storage: True + convert_to_fractional: True + niggli_reduce: False + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 128 + build_structure_cfg: + format: array + niggli: False + model_sample_params: + num_inference_steps: 100 diff --git a/structure_generation/configs/omatg/omatg_mp20_dng_sample.yaml b/structure_generation/configs/omatg/omatg_mp20_dng_sample.yaml new file mode 100644 index 00000000..83bc9a9d --- /dev/null +++ b/structure_generation/configs/omatg/omatg_mp20_dng_sample.yaml @@ -0,0 +1,53 @@ +Global: + do_train: False + do_eval: False + do_test: False + +Model: + __class_name__: OMATGCSPNetFull + __init_params__: + hidden_dim: 512 + num_layers: 6 + max_atoms: 100 + act_fn: silu + dis_emb: sin + num_freqs: 128 + edge_style: fc + cutoff: 7.0 + max_neighbors: 20 + ln: True + ip: True + smooth: False + pred_type: True + pred_scalar: False + time_embed_dim: 256 + +Sample: + data: + dataset: + __class_name__: NumAtomsCrystalDataset + __init_params__: + total_num: 10 + prop_names: [] + prop_values: [] + loader: + num_workers: 0 + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 10 + + build_structure_cfg: + format: array + niggli: False + + model_sample_params: + num_inference_steps: 100 + + metrics: + __class_name__: CSPMetric + __init_params__: + gt_file_path: "./data/mp_20/test.csv" diff --git a/test/omatg/__init__.py b/test/omatg/__init__.py new file mode 100644 index 00000000..6809e580 --- /dev/null +++ b/test/omatg/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/test/omatg/test_omatg.py b/test/omatg/test_omatg.py new file mode 100644 index 00000000..8f2a6d60 --- /dev/null +++ b/test/omatg/test_omatg.py @@ -0,0 +1,751 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import paddle + +from ppmat.models.omatg.model import OMATGCSPNet as CSPNet, OMATGCSPNetFull +from ppmat.utils.crystal import ( + cart_to_frac_coords, + frac_to_cart_coords, + lattice_params_to_matrix_paddle, +) +from ppmat.datasets.omatg_dataset import Structure +from ppmat.datasets.omatg_dataset import OMATGData + + +def _make_batch(batch_size=2, atoms_per_struct=(3, 4), max_z=10): + """Build a minimal OMATGData batch for forward/loss tests.""" + num_list = list(atoms_per_struct[:batch_size]) + total = sum(num_list) + atom_types = paddle.randint(1, max_z, [total], dtype="int64") + frac_coords = paddle.rand([total, 3]) + lattices = paddle.rand([batch_size, 3, 3]) * 2.0 + paddle.eye(3).unsqueeze(0) + num_atoms = paddle.to_tensor(num_list, dtype="int64") + node2graph = paddle.repeat_interleave( + paddle.arange(batch_size, dtype="int64"), num_atoms + ) + d = { + "atom_types": atom_types, + "frac_coords": frac_coords, + "lattices": lattices, + "num_atoms": num_atoms, + "node2graph": node2graph, + } + return OMATGData.from_collate_dict(d) + + +class TestCSPNetForward(unittest.TestCase): + """CSPNet and OMATGCSPNetFull minimal forward / loss smoke tests.""" + + def test_cspnet_forward_returns_b_eta(self): + """CSP mode forward returns 4-tuple (cell_b, pos_b, cell_eta, pos_eta).""" + batch_size = 2 + num_atoms = paddle.to_tensor([3, 4], dtype="int64") + total_atoms = int(num_atoms.sum()) + node2graph = paddle.to_tensor([0, 0, 0, 1, 1, 1, 1], dtype="int64") + + model = CSPNet( + hidden_dim=64, num_layers=2, max_atoms=100, + time_embed_dim=32, edge_style="fc", + ) + t = paddle.rand([batch_size]) + atom_types = paddle.randint(1, 10, [total_atoms]) + frac_coords = paddle.rand([total_atoms, 3]) + lattices = paddle.rand([batch_size, 3, 3]) * 2.0 + + output = model(t, atom_types, frac_coords, lattices, num_atoms, node2graph) + self.assertEqual(len(output), 4) + self.assertEqual(output[0].shape, [batch_size, 3, 3]) + self.assertEqual(output[1].shape, [total_atoms, 3]) + self.assertEqual(output[2].shape, [batch_size, 3, 3]) + self.assertEqual(output[3].shape, [total_atoms, 3]) + + def test_cspnet_forward_dict_keys(self): + """forward_dict returns b/eta dict with correct keys.""" + model = CSPNet( + hidden_dim=64, num_layers=2, max_atoms=100, + time_embed_dim=32, edge_style="fc", + ) + data = _make_batch() + t = paddle.rand([2]) + out = model.forward_dict( + t, data.species, data.pos, data.cell, + data.n_atoms, data.batch, + ) + self.assertIn("pos_b", out) + self.assertIn("pos_eta", out) + self.assertIn("cell_b", out) + self.assertIn("cell_eta", out) + + def test_omgcspnetfull_forward_csp_loss(self): + """CSP mode forward returns loss_dict with loss/loss_lattice/loss_coord.""" + model = OMATGCSPNetFull( + hidden_dim=64, num_layers=2, max_atoms=100, + time_embed_dim=32, edge_style="fc", pred_type=False, + ) + data = _make_batch() + output = model(data) + self.assertIn("loss_dict", output) + self.assertIn("loss", output["loss_dict"]) + self.assertIn("loss_lattice", output["loss_dict"]) + self.assertIn("loss_coord", output["loss_dict"]) + self.assertNotIn("loss_type", output["loss_dict"]) + loss_val = float(output["loss_dict"]["loss"]) + self.assertFalse(loss_val != loss_val, "loss is NaN") + + def test_omgcspnetfull_forward_dng_loss(self): + """DNG mode forward returns loss_dict including loss_type (cross-entropy).""" + model = OMATGCSPNetFull( + hidden_dim=64, num_layers=2, max_atoms=100, + time_embed_dim=32, edge_style="fc", pred_type=True, + ) + model.enable_masked_species() + data = _make_batch() + output = model(data) + self.assertIn("loss_type", output["loss_dict"]) + self.assertIn("loss", output["loss_dict"]) + loss_val = float(output["loss_dict"]["loss"]) + self.assertFalse(loss_val != loss_val, "DNG loss is NaN") + + def test_cspnet_backward_gradient(self): + """Loss backward produces non-None gradients on model parameters.""" + model = OMATGCSPNetFull( + hidden_dim=64, num_layers=2, max_atoms=100, + time_embed_dim=32, edge_style="fc", + ) + data = _make_batch() + output = model(data) + loss = output["loss_dict"]["loss"] + loss.backward() + has_grad = any( + p.grad is not None and float(p.grad.abs().sum()) > 0 + for p in model.parameters() + ) + self.assertTrue(has_grad, "No non-zero gradients after backward") + + +class TestStructureData(unittest.TestCase): + """Structure and OMATGData smoke tests.""" + + def test_structure_creation(self): + cell = paddle.eye(3) * 5.0 + atomic_numbers = paddle.to_tensor([1, 1, 8], dtype="int64") + pos = paddle.rand([3, 3]) + struct = Structure(cell, atomic_numbers, pos, pos_is_fractional=True) + self.assertEqual(struct.cell.shape, [3, 3]) + self.assertTrue(struct.pos_is_fractional) + self.assertEqual(len(struct.atomic_numbers), 3) + + def test_structure_convert_coords(self): + cell = paddle.to_tensor([[5.0, 0, 0], [0, 5.0, 0], [0, 0, 5.0]]) + atomic_numbers = paddle.to_tensor([1], dtype="int64") + pos = paddle.to_tensor([[2.5, 2.5, 2.5]]) + struct = Structure(cell, atomic_numbers, pos, pos_is_fractional=False) + struct.convert_to_fractional() + self.assertTrue(struct.pos_is_fractional) + self.assertAlmostEqual(float(struct.pos[0, 0]), 0.5, places=4) + struct.convert_to_cartesian() + self.assertFalse(struct.pos_is_fractional) + self.assertAlmostEqual(float(struct.pos[0, 0]), 2.5, places=4) + + def test_structure_get_ase_atoms(self): + cell = paddle.eye(3) * 5.0 + atomic_numbers = paddle.to_tensor([1, 8], dtype="int64") + pos = paddle.to_tensor([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]]) + struct = Structure(cell, atomic_numbers, pos, pos_is_fractional=True) + atoms = struct.get_ase_atoms() + self.assertEqual(len(atoms), 2) + + def test_omgdata_from_structure(self): + cell = paddle.eye(3) * 5.0 + atomic_numbers = paddle.to_tensor([1, 1, 8], dtype="int64") + pos = paddle.rand([3, 3]) + struct = Structure(cell, atomic_numbers, pos, pos_is_fractional=True) + data = OMATGData(struct) + self.assertEqual(data.num_graphs, 1) + self.assertEqual(data.num_atoms, 3) + self.assertIsNotNone(data.species) + self.assertIsInstance(data.property_dict, list) + + def test_omgdata_batch(self): + structs = [] + for _ in range(3): + cell = paddle.eye(3) * 5.0 + atomic_numbers = paddle.to_tensor([1, 8], dtype="int64") + pos = paddle.rand([2, 3]) + structs.append(Structure(cell, atomic_numbers, pos, pos_is_fractional=True)) + batch = OMATGData.from_batch(structs, concatenate=True) + self.assertEqual(batch.num_graphs, 3) + self.assertEqual(batch.num_atoms, 6) + self.assertEqual(batch.batch.shape, [6]) + self.assertEqual(batch.ptr.shape, [4]) + + def test_omgdata_get_graph(self): + structs = [] + for _ in range(2): + cell = paddle.eye(3) * 5.0 + atomic_numbers = paddle.to_tensor([1, 8], dtype="int64") + pos = paddle.rand([2, 3]) + structs.append(Structure(cell, atomic_numbers, pos, pos_is_fractional=True)) + batch = OMATGData.from_batch(structs, concatenate=True) + g0 = batch.get_graph(0) + self.assertEqual(len(g0.atomic_numbers), 2) + g1 = batch.get_graph(1) + self.assertEqual(len(g1.atomic_numbers), 2) + + def test_omgdata_clone(self): + cell = paddle.eye(3) * 5.0 + atomic_numbers = paddle.to_tensor([1, 8], dtype="int64") + pos = paddle.rand([2, 3]) + struct = Structure(cell, atomic_numbers, pos, pos_is_fractional=True) + data = OMATGData(struct) + cloned = data.clone() + self.assertEqual(cloned.num_atoms, data.num_atoms) + self.assertTrue(paddle.equal_all(cloned.species, data.species)) + + def test_omgdata_set_get_field(self): + cell = paddle.eye(3) * 5.0 + atomic_numbers = paddle.to_tensor([1, 8], dtype="int64") + pos = paddle.rand([2, 3]) + struct = Structure(cell, atomic_numbers, pos, pos_is_fractional=True) + data = OMATGData(struct) + self.assertTrue(paddle.equal_all(data.get_field("species"), data.species)) + new_pos = paddle.rand([2, 3]) + data.set_field("pos", new_pos) + self.assertTrue(paddle.equal_all(data.get_field("pos"), new_pos)) + + +class TestLatticeUtils(unittest.TestCase): + """Lattice utility smoke tests.""" + + def test_lattice_params_to_matrix(self): + lengths = paddle.to_tensor([[3.0, 3.0, 5.0], [4.0, 4.0, 4.0]]) + angles = paddle.to_tensor([[90.0, 90.0, 90.0], [90.0, 90.0, 120.0]]) + matrices = lattice_params_to_matrix_paddle(lengths, angles) + self.assertEqual(matrices.shape, [2, 3, 3]) + + def test_frac_cart_conversion(self): + lengths = paddle.to_tensor([[3.0, 3.0, 5.0]]) + angles = paddle.to_tensor([[90.0, 90.0, 90.0]]) + num_atoms = paddle.to_tensor([3], dtype="int64") + frac_coords = paddle.rand([3, 3]) + cart = frac_to_cart_coords(frac_coords, num_atoms, lengths=lengths, angles=angles) + self.assertEqual(cart.shape, [3, 3]) + frac_back = cart_to_frac_coords(cart, num_atoms, lengths=lengths, angles=angles) + self.assertEqual(frac_back.shape, [3, 3]) + + +class TestSIComponents(unittest.TestCase): + """SI framework component smoke tests.""" + + def test_gamma_sqrt(self): + from ppmat.models.omatg.si.interpolants import LatentGammaSqrt + t = paddle.to_tensor([0.3, 0.5, 0.7]) + g = LatentGammaSqrt(a=1.0) + self.assertEqual(g.gamma(t).shape, [3]) + self.assertTrue(g.requires_antithetic()) + self.assertEqual(g.gamma_derivative(t).shape, [3]) + + def test_gamma_encdec(self): + from ppmat.models.omatg.si.interpolants import LatentGammaEncoderDecoder + t = paddle.to_tensor([0.3, 0.5, 0.7]) + g = LatentGammaEncoderDecoder() + self.assertEqual(g.gamma(t).shape, [3]) + self.assertFalse(g.requires_antithetic()) + + def test_epsilon_vanishing(self): + from ppmat.models.omatg.si.interpolants import VanishingEpsilon + t = paddle.to_tensor([0.3, 0.5, 0.7]) + eps = VanishingEpsilon(c=1.0) + self.assertEqual(eps.epsilon(t).shape, [3]) + + def test_sigma_geometric(self): + from ppmat.models.omatg.si.interpolants import GeometricSigma + s = paddle.to_tensor([0.0, 0.5, 1.0]) + sig = GeometricSigma(sigma_min=0.1, sigma_max=10.0) + self.assertEqual(sig.sigma(s).shape, [3]) + self.assertEqual(sig.sigma_dot(s).shape, [3]) + + def test_tau_constant(self): + from ppmat.models.omatg.si.interpolants import TauConstantSchedule + t = paddle.to_tensor([0.3, 0.7]) + tau = TauConstantSchedule() + self.assertEqual(tau.tau(t).shape, [2]) + self.assertEqual(tau.tau_dot(t).shape, [2]) + + def test_interpolant_linear(self): + from ppmat.models.omatg.si.interpolants import LinearInterpolant + t = paddle.to_tensor([0.3]) + interp = LinearInterpolant() + self.assertEqual(interp.alpha(t).shape, [1]) + self.assertEqual(interp.beta(t).shape, [1]) + + def test_interpolant_periodic_linear(self): + from ppmat.models.omatg.si.interpolants import PeriodicLinearInterpolant + interp = PeriodicLinearInterpolant() + corr = interp.get_corrector() + self.assertIsNotNone(corr) + + def test_interpolant_vp(self): + from ppmat.models.omatg.si.interpolants import ( + ScoreBasedDiffusionModelInterpolantVP, + ) + from ppmat.models.omatg.si.interpolants import TauConstantSchedule + t = paddle.to_tensor([0.3]) + interp = ScoreBasedDiffusionModelInterpolantVP(TauConstantSchedule()) + self.assertEqual(interp.alpha(t).shape, [1]) + + def test_interpolant_ve(self): + from ppmat.models.omatg.si.interpolants import ( + ScoreBasedDiffusionModelInterpolantVE, + GeometricSigma, + ) + t = paddle.to_tensor([0.3]) + interp = ScoreBasedDiffusionModelInterpolantVE(GeometricSigma(0.1, 10.0)) + self.assertEqual(interp.alpha(t).shape, [1]) + + def test_dfm_mask_loss(self): + """DiscreteFlowMatchingMask cross-entropy loss.""" + from ppmat.models.omatg.si.core import ( + DiscreteFlowMatchingMask, + ) + dfm = DiscreteFlowMatchingMask(noise=0.1) + x_0 = paddle.zeros([5], dtype="int64") + x_1 = paddle.randint(1, 10, [5], dtype="int64") + t = paddle.to_tensor([0.5]) + x_t, z = dfm.interpolate(t, x_0, x_1, paddle.to_tensor([0, 0, 0, 0, 0])) + self.assertEqual(x_t.shape, [5]) + + def model_fn(x_t): + pred = paddle.rand([5, 100]) + return pred, paddle.zeros_like(pred) + losses = dfm.loss(model_fn, t, x_0, x_1, x_t, z, + paddle.to_tensor([0, 0, 0, 0, 0])) + self.assertIn("loss", losses) + self.assertTrue(dfm.uses_masked_species()) + + def test_identity_interpolant(self): + from ppmat.models.omatg.si.core import ( + SingleStochasticInterpolantIdentity, + ) + ident = SingleStochasticInterpolantIdentity() + x = paddle.to_tensor([1, 2, 3], dtype="int64") + t = paddle.to_tensor([0.5]) + x_t, z = ident.interpolate(t, x, x, paddle.to_tensor([0, 0, 0])) + self.assertTrue(paddle.equal_all(x_t, x)) + self.assertFalse(ident.uses_masked_species()) + + +class TestSITrainingPath(unittest.TestCase): + """SI velocity-matching training path smoke tests.""" + + def _build_csp_si_model(self): + """Build a small CSP model with SI (Linear-ODE).""" + from ppmat.models.omatg.si import ( + StochasticInterpolants, SingleStochasticInterpolant, + SingleStochasticInterpolantIdentity, + PeriodicLinearInterpolant, LinearInterpolant, + ) + from ppmat.models.omatg.model import IndependentSampler + si = StochasticInterpolants( + stochastic_interpolants=[ + SingleStochasticInterpolantIdentity(), + SingleStochasticInterpolant( + interpolant=PeriodicLinearInterpolant(), gamma=None, + epsilon=None, differential_equation_type="ODE", + velocity_annealing_factor=10.18, + correct_center_of_mass_motion=True, + ), + SingleStochasticInterpolant( + interpolant=LinearInterpolant(), gamma=None, epsilon=None, + differential_equation_type="ODE", + velocity_annealing_factor=1.82, + ), + ], + data_fields=["species", "pos", "cell"], + integration_time_steps=210, + ) + sampler = IndependentSampler(dataset_name="mp_20", mirror_species=True) + model = OMATGCSPNetFull( + hidden_dim=32, num_layers=1, max_atoms=100, + time_embed_dim=16, pred_type=False, use_si=False, + ) + model._si = si + model._sampler = sampler + model._relative_si_costs = { + "species_loss": 0.0, "pos_loss_b": 0.9994, "cell_loss_b": 0.0006, + } + model.use_si = True + return model + + def test_csp_si_forward_loss(self): + """CSP SI training path produces velocity-matching loss.""" + model = self._build_csp_si_model() + data = _make_batch() + out = model(data) + self.assertIn("loss_dict", out) + self.assertIn("loss", out["loss_dict"]) + self.assertIn("pos_loss_b", out["loss_dict"]) + self.assertIn("cell_loss_b", out["loss_dict"]) + self.assertIn("species_loss", out["loss_dict"]) + loss_val = float(out["loss_dict"]["loss"]) + self.assertFalse(loss_val != loss_val, "SI loss is NaN") + + def test_csp_si_backward(self): + """SI loss backward produces gradients.""" + model = self._build_csp_si_model() + data = _make_batch() + out = model(data) + loss = out["loss_dict"]["loss"] + loss.backward() + has_grad = any( + p.grad is not None and float(p.grad.abs().sum()) > 0 + for p in model.parameters() + ) + self.assertTrue(has_grad, "No gradients after SI backward") + + def test_dng_si_forward_loss(self): + """DNG SI training path (SDE + gamma + DFM mask).""" + from ppmat.models.omatg.si import ( + StochasticInterpolants, SingleStochasticInterpolant, + PeriodicLinearInterpolant, LinearInterpolant, + ) + from ppmat.models.omatg.si.core import DiscreteFlowMatchingMask + from ppmat.models.omatg.si.interpolants import LatentGammaSqrt, VanishingEpsilon + from ppmat.models.omatg.model import IndependentSampler + si = StochasticInterpolants( + stochastic_interpolants=[ + DiscreteFlowMatchingMask(noise=0.189), + SingleStochasticInterpolant( + interpolant=PeriodicLinearInterpolant(), + gamma=LatentGammaSqrt(a=0.018), + epsilon=VanishingEpsilon(c=9.7, mu=0.17, sigma=0.029), + differential_equation_type="SDE", + velocity_annealing_factor=6.33, + correct_center_of_mass_motion=True, + ), + SingleStochasticInterpolant( + interpolant=LinearInterpolant(), gamma=None, epsilon=None, + differential_equation_type="ODE", + velocity_annealing_factor=1.07, + ), + ], + data_fields=["species", "pos", "cell"], + integration_time_steps=710, + ) + sampler = IndependentSampler(dataset_name="mp_20", mask_species=True, mirror_species=False) + model = OMATGCSPNetFull( + hidden_dim=32, num_layers=1, max_atoms=100, + time_embed_dim=16, pred_type=True, use_si=False, + ) + model.enable_masked_species() + model._si = si + model._sampler = sampler + model._relative_si_costs = { + "species_loss": 0.5918, "pos_loss_b": 0.1309, + "pos_loss_z": 0.2708, "cell_loss_b": 0.0065, + } + model.use_si = True + data = _make_batch() + out = model(data) + self.assertIn("loss", out["loss_dict"]) + self.assertIn("pos_loss_z", out["loss_dict"]) + self.assertIn("species_loss", out["loss_dict"]) + + +class TestOSInterpolants(unittest.TestCase): + """One-sided interpolant smoke tests for VESBD/VPSBD variants.""" + + def test_vesbd_ode_forward(self): + """VESBD-ODE SI training via SingleStochasticInterpolantOS.""" + from ppmat.models.omatg.si import ( + StochasticInterpolants, SingleStochasticInterpolant, + SingleStochasticInterpolantIdentity, + SingleStochasticInterpolantOS, + LinearInterpolant, + ) + from ppmat.models.omatg.si.interpolants import ScoreBasedDiffusionModelInterpolantVE + from ppmat.models.omatg.si.interpolants import GeometricSigma + from ppmat.models.omatg.model import IndependentSampler + os_interp = SingleStochasticInterpolantOS( + interpolant=ScoreBasedDiffusionModelInterpolantVE( + GeometricSigma(0.1, 10.0) + ), + epsilon=None, differential_equation_type="ODE", + velocity_annealing_factor=1.0, + ) + si = StochasticInterpolants( + stochastic_interpolants=[ + SingleStochasticInterpolantIdentity(), + os_interp, + SingleStochasticInterpolant( + interpolant=LinearInterpolant(), gamma=None, epsilon=None, + differential_equation_type="ODE", + velocity_annealing_factor=1.82, + ), + ], + data_fields=["species", "pos", "cell"], + integration_time_steps=210, + ) + sampler = IndependentSampler(dataset_name="mp_20", mirror_species=True) + model = OMATGCSPNetFull( + hidden_dim=32, num_layers=1, max_atoms=100, + time_embed_dim=16, pred_type=False, use_si=False, + ) + model._si = si + model._sampler = sampler + model._relative_si_costs = { + "species_loss": 0.0, "pos_loss_b": 0.9994, "cell_loss_b": 0.0006, + } + model.use_si = True + data = _make_batch() + out = model(data) + self.assertIn("loss", out["loss_dict"]) + + def test_vpsbd_sde_forward(self): + """VPSBD-SDE SI training via SingleStochasticInterpolantOS.""" + from ppmat.models.omatg.si import ( + StochasticInterpolants, SingleStochasticInterpolant, + SingleStochasticInterpolantIdentity, + SingleStochasticInterpolantOS, + LinearInterpolant, + ) + from ppmat.models.omatg.si.interpolants import ScoreBasedDiffusionModelInterpolantVP + from ppmat.models.omatg.si.interpolants import TauConstantSchedule, ConstantEpsilon + from ppmat.models.omatg.model import IndependentSampler + os_interp = SingleStochasticInterpolantOS( + interpolant=ScoreBasedDiffusionModelInterpolantVP( + TauConstantSchedule() + ), + epsilon=ConstantEpsilon(1.0), + differential_equation_type="SDE", + velocity_annealing_factor=1.0, + ) + si = StochasticInterpolants( + stochastic_interpolants=[ + SingleStochasticInterpolantIdentity(), + os_interp, + SingleStochasticInterpolant( + interpolant=LinearInterpolant(), gamma=None, epsilon=None, + differential_equation_type="ODE", + velocity_annealing_factor=1.07, + ), + ], + data_fields=["species", "pos", "cell"], + integration_time_steps=710, + ) + sampler = IndependentSampler(dataset_name="mp_20", mirror_species=True) + model = OMATGCSPNetFull( + hidden_dim=32, num_layers=1, max_atoms=100, + time_embed_dim=16, pred_type=False, use_si=False, + ) + model._si = si + model._sampler = sampler + model._relative_si_costs = { + "species_loss": 0.0, "pos_loss_b": 0.5, + "pos_loss_z": 0.3, "cell_loss_b": 0.2, + } + model.use_si = True + data = _make_batch() + out = model(data) + self.assertIn("loss", out["loss_dict"]) + self.assertIn("pos_loss_z", out["loss_dict"]) + + +class TestSIFactory(unittest.TestCase): + """Config-driven SI/sampler factory smoke tests.""" + + def test_build_si_from_cfg_csp(self): + """build_si_from_cfg builds CSP StochasticInterpolants from config.""" + from ppmat.models.omatg.si import build_si_from_cfg, build_sampler_from_cfg + si_cfg = { + "stochastic_interpolants": [ + {"__class_name__": "SingleStochasticInterpolantIdentity"}, + {"__class_name__": "SingleStochasticInterpolant", + "__init_params__": { + "interpolant": {"__class_name__": "PeriodicLinearInterpolant"}, + "gamma": None, "epsilon": None, + "differential_equation_type": "ODE", + "velocity_annealing_factor": 10.18, + "correct_center_of_mass_motion": True, + }}, + {"__class_name__": "SingleStochasticInterpolant", + "__init_params__": { + "interpolant": {"__class_name__": "LinearInterpolant"}, + "gamma": None, "epsilon": None, + "differential_equation_type": "ODE", + "velocity_annealing_factor": 1.82, + }}, + ], + "data_fields": ["species", "pos", "cell"], + "integration_time_steps": 210, + } + sampler_cfg = { + "position_distribution": {"__class_name__": "UniformPositionDistribution"}, + "cell_distribution": { + "__class_name__": "InformedLatticeDistribution", + "__init_params__": {"dataset_name": "mp_20"}, + }, + "species_distribution": {"__class_name__": "MirrorSpecies"}, + } + si = build_si_from_cfg(si_cfg) + self.assertEqual(len(si), 3) + sampler = build_sampler_from_cfg(sampler_cfg) + self.assertIsNotNone(sampler) + + def test_build_si_from_cfg_dng(self): + """build_si_from_cfg builds DNG config with nested gamma/epsilon/DFM.""" + from ppmat.models.omatg.si import build_si_from_cfg + si_cfg = { + "stochastic_interpolants": [ + {"__class_name__": "DiscreteFlowMatchingMask", + "__init_params__": {"noise": 0.189}}, + {"__class_name__": "SingleStochasticInterpolant", + "__init_params__": { + "interpolant": {"__class_name__": "PeriodicLinearInterpolant"}, + "gamma": {"__class_name__": "LatentGammaSqrt", + "__init_params__": {"a": 0.018}}, + "epsilon": {"__class_name__": "VanishingEpsilon", + "__init_params__": {"c": 9.7}}, + "differential_equation_type": "SDE", + "velocity_annealing_factor": 6.33, + }}, + {"__class_name__": "SingleStochasticInterpolant", + "__init_params__": { + "interpolant": {"__class_name__": "LinearInterpolant"}, + "differential_equation_type": "ODE", + }}, + ], + "data_fields": ["species", "pos", "cell"], + "integration_time_steps": 710, + } + si = build_si_from_cfg(si_cfg) + self.assertEqual(len(si), 3) + + def test_config_driven_si_forward(self): + """Full config-driven SI forward via OMATGCSPNetFull constructor.""" + si_cfg = { + "stochastic_interpolants": [ + {"__class_name__": "SingleStochasticInterpolantIdentity"}, + {"__class_name__": "SingleStochasticInterpolant", + "__init_params__": { + "interpolant": {"__class_name__": "PeriodicLinearInterpolant"}, + "gamma": None, "epsilon": None, + "differential_equation_type": "ODE", + "velocity_annealing_factor": 10.18, + }}, + {"__class_name__": "SingleStochasticInterpolant", + "__init_params__": { + "interpolant": {"__class_name__": "LinearInterpolant"}, + "differential_equation_type": "ODE", + "velocity_annealing_factor": 1.82, + }}, + ], + "data_fields": ["species", "pos", "cell"], + "integration_time_steps": 210, + "relative_si_costs": { + "species_loss": 0.0, "pos_loss_b": 0.9994, "cell_loss_b": 0.0006, + }, + } + sampler_cfg = { + "dataset_name": "mp_20", + "mirror_species": True, + "mask_species": False, + } + model = OMATGCSPNetFull( + hidden_dim=32, num_layers=1, max_atoms=100, + time_embed_dim=16, pred_type=False, + use_si=True, si_cfg=si_cfg, sampler_cfg=sampler_cfg, + ) + data = _make_batch() + out = model(data) + self.assertIn("loss", out["loss_dict"]) + + +class TestSISampling(unittest.TestCase): + """SI integrate-based sampling smoke tests.""" + + def test_csp_si_sample(self): + """CSP SI sampling produces structures via si.integrate.""" + from ppmat.models.omatg.si import ( + StochasticInterpolants, SingleStochasticInterpolant, + SingleStochasticInterpolantIdentity, + PeriodicLinearInterpolant, LinearInterpolant, + ) + from ppmat.models.omatg.model import IndependentSampler + si = StochasticInterpolants( + stochastic_interpolants=[ + SingleStochasticInterpolantIdentity(), + SingleStochasticInterpolant( + interpolant=PeriodicLinearInterpolant(), gamma=None, + epsilon=None, differential_equation_type="ODE", + velocity_annealing_factor=10.18, + ), + SingleStochasticInterpolant( + interpolant=LinearInterpolant(), gamma=None, epsilon=None, + differential_equation_type="ODE", + velocity_annealing_factor=1.82, + ), + ], + data_fields=["species", "pos", "cell"], + integration_time_steps=10, + ) + sampler = IndependentSampler(dataset_name="mp_20", mirror_species=True) + model = OMATGCSPNetFull( + hidden_dim=32, num_layers=1, max_atoms=100, + time_embed_dim=16, pred_type=False, use_si=False, + ) + model._si = si + model._sampler = sampler + model._relative_si_costs = { + "species_loss": 0.0, "pos_loss_b": 0.9994, "cell_loss_b": 0.0006, + } + model.use_si = True + data = _make_batch() + result = model.sample(data, num_inference_steps=10) + self.assertIn("result", result) + self.assertEqual(len(result["result"]), 2) + self.assertIn("num_atoms", result["result"][0]) + self.assertIn("frac_coords", result["result"][0]) + + +class TestSamplers(unittest.TestCase): + """Sampler smoke tests using Paddle native APIs.""" + + def test_independent_sampler_sample_p_0(self): + """IndependentSampler.sample_p_0 returns OMATGData with correct batch.""" + from ppmat.models.omatg.model import IndependentSampler + sampler = IndependentSampler(dataset_name="mp_20", mirror_species=True) + data = _make_batch() + x_1 = OMATGData() + x_1.n_atoms = data.n_atoms + x_1.species = data.species + x_1.cell = data.cell + x_1.pos = data.pos + x_1.pos_is_fractional = paddle.ones_like(data.n_atoms, dtype="bool") + x_1.batch = data.batch + x_1.ptr = paddle.concat([ + paddle.to_tensor([0], dtype="int64"), + paddle.cumsum(data.n_atoms, axis=0).cast("int64"), + ]) + x_1.property_dict = [{} for _ in range(len(data.n_atoms))] + x_0 = sampler.sample_p_0(x_1) + self.assertEqual(x_0.num_graphs, 2) + self.assertEqual(x_0.num_atoms, x_1.num_atoms) + + +if __name__ == "__main__": + unittest.main()