from typing import Literal, Union, Sequence import torch from torch import Tensor, nn import os from pathlib import Path from collections.abc import Sequence from onescience.datapipes.materials.matris import RadiusGraph, GraphConverter, datatype from onescience.modules.func_utils.matris_reference import AtomRef from onescience.modules.func_utils.matris_graph import process_graphs from onescience.modules.embedding.matris_embedding import ( ThreebodyFourierExpansion, AtomTypeEmbedding, EdgeBasisEmbedding, ThreebodyEmbedding ) from onescience.modules.func_utils.matris_func_utils import ( MLP, GatedMLP, get_normalization ) from onescience.modules.layer.matris_interaction import Interaction_Block from onescience.modules.head.matris_head import ( EnergyHead, MagmomHead, ForceStressHead, ) class MatRIS(nn.Module): """ Init MatRIS Potential """ def __init__( self, num_layers: int = 6, node_feat_dim: int = 128, edge_feat_dim: int = 128, three_body_feat_dim: int = 128, mlp_hidden_dims: Union[int, Sequence[int]] = (128, 128), dropout: float = 0.0, use_bias: bool = False, distance_expansion: str = "Bessel", three_body_expansion: str = "SH", num_radial: int = 7, num_angular: int = 7, max_l: int = 4, max_n: int = 4, envelope_exponent: int = 8, graph_conv_mlp: str = "GateMLP", activation_type: str = "silu", norm_type: str = "rms", pairwise_cutoff: float = 6, three_body_cutoff: float = 4, use_smoothed_for_delta_edge: bool = False, learnable_basis: bool = True, is_intensive: bool = True, is_conservation: bool = True, reference_energy: str | None = None, ): """ Args: num_layers (int): message passing layers. node_feat_dim (int): atom feature embedding dim. edge_feat_dim (int): edge(pairwise) feature embedding dim. three_body_feat_dim (int): angle(three body) feature embedding dim. mlp_hidden_dims (List or int): hidden dims of MLP. Can be 'int' or 'list'. dropout (float): dropout rate in MLP. use_bias (bool): whether use bias in Interaction block. distance_expansion (str): The function of pairwise basis. Can be "Bessel" or "Gaussian". three_body_expansion (str): The function of three body basis. Can be "Fourier(fourier)" or "Spherical Harmonics(sh)". num_radial (int): number of radial basis used in Bessel and Gaussian basis. num_angular (int): number of three_body basis used in Fourier basis. max_l (int): Maximum l value for Spherical Harmonics basis (SH). max_n (int): Maximum n value for Spherical Harmonics basis (SH). envelope_exponent (int): exponent of 'PolynomialEnvelope'. graph_conv_mlp (str): The type of MLP in mp layers. Can be "MLP", "GatedMLP" and "MoE". See fucntion.py for more informations. activation_type (str): activation function. Can be "SiLU(silu)", "Sigmoid(sigmoid)", "ReLU(relu)"... See fucntion.py for more informations. norm_type (str): normalization function used in MLP. Can be "LayerNorm(layer)", "BatchNorm(batch)", "RMSNorm(rms)"... See fucntion.py for more informations. pairwise_cutoff (float): The cutoff of Atom graph. three_body_cutoff (float): The cutoff of Line graph. use_smoothed_for_delta_edge (bool): Whether to use the smoothed features for edge feature update. learnable_basis (bool): Whether the basis functions are learnable. is_intensive (bool): whether the model outputs energy per atom (True) or total energy (False). is_conservation (bool): whether use conservate force and stress. reference_energy (str): refernece energy of 'str'(eg. MPtrj, OMat..) dataset(Caculated by linear regression). more details can be found at reference_energy.py. """ super().__init__() # model configs self.config = { k: v for k, v in locals().items() if k not in ["self", "__class__"] } self.is_intensive = is_intensive self.reference_energy = None if reference_energy is not None: self.reference_energy = AtomRef( reference_energy=reference_energy, is_intensive=is_intensive ) # Define Graph Converter self.graph_converter = GraphConverter( atom_graph_cutoff=pairwise_cutoff, line_graph_cutoff=three_body_cutoff, ) # ====== embedding layers ======== self.atom_embedding = AtomTypeEmbedding(atom_feat_dim=node_feat_dim) self.edge_embedding = EdgeBasisEmbedding( pairwise_cutoff=pairwise_cutoff, three_body_cutoff=three_body_cutoff, num_radial=num_radial, edge_feat_dim=edge_feat_dim, envelope_exponent=envelope_exponent, learnable=learnable_basis, distance_expansion=distance_expansion, ) self.three_body_embedding = ThreebodyEmbedding( num_angular = num_angular, # Fourier max_n=max_n, max_l=max_l, cutoff=pairwise_cutoff, # Spherical Harmonics three_body_feat_dim = three_body_feat_dim, three_body_expansion = three_body_expansion, learnable = learnable_basis ) # ====== Interaction layers ======== interaction_block = [ Interaction_Block( node_feat_dim=node_feat_dim, edge_feat_dim=edge_feat_dim, three_body_feat_dim=three_body_feat_dim, num_radial=num_radial, num_angular=num_angular, dropout=dropout, use_bias=use_bias, use_smoothed_for_delta_edge=use_smoothed_for_delta_edge, mlp_type=graph_conv_mlp, norm_type=norm_type, activation_type=activation_type, ) for _ in range(num_layers) ] self.interaction_block = nn.ModuleList(interaction_block) # ====== Readout layers ======== self.readout_norm = get_normalization(norm_type, dim=node_feat_dim) self.energy_head = EnergyHead( feat_dim = node_feat_dim, hidden_dim = mlp_hidden_dims, output_dim = 1, mlp_type = "mlp", activation_type = activation_type, ) self.magmom_head = MagmomHead( feat_dim = node_feat_dim, hidden_dim = 2 * node_feat_dim, output_dim = 1, mlp_type = "mlp", activation_type = activation_type, ) self.force_stress_head = ForceStressHead( is_conservation = is_conservation, feat_dim = edge_feat_dim, # is_conservation == False hidden_dim = mlp_hidden_dims, # is_conservation == False output_dim = 3, # is_conservation == False mlp_type = "mlp", # is_conservation == False activation_type = activation_type, # is_conservation == False ) if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: print(f"MatRIS initialized with {self.get_params()} parameters") def forward( self, graphs: Sequence[RadiusGraph], task: str = "ef", is_training: bool = False, ) -> dict[str, Tensor]: """ Args: graphs (List): a list of RadiusGraph. task (str): the prediction task. Can be 'e', 'em', 'ef', 'efs', 'efsm'. """ prediction = {} # ======== Graph processing ======== batch_graph = process_graphs(graphs, compute_stress="s" in task) # ======== Feature embedding ======== node_feat = self.atom_embedding( batch_graph['atomic_numbers'] - 1 ) # atom type feature init (use 0 for 'H') edge_feat, smooth_weight = self.edge_embedding(graphs=batch_graph) # pairwise feature init threebody_feat = None if len(batch_graph['line_graph_dict']['line_graph']) != 0: threebody_feat = self.three_body_embedding(graphs=batch_graph) # three body feature init # ======== Interaction Block ======= for mp_layer in self.interaction_block: node_feat, edge_feat, threebody_feat = mp_layer( batch_graph=batch_graph, node_feat=node_feat, edge_feat=edge_feat, threebody_feat=threebody_feat, smooth_weight=smooth_weight, ) # ======== Readout Block ======= node_feat = self.readout_norm(node_feat) total_energy = self.energy_head(batch_graph = batch_graph, node_feat = node_feat) force_stress_dict = self.force_stress_head( batch_graph = batch_graph, compute_force="f" in task, compute_stress="s" in task, total_energy = total_energy, node_feat = node_feat, edge_feat = edge_feat, is_training = is_training) prediction.update(force_stress_dict) if "m" in task: magmom = self.magmom_head(batch_graph = batch_graph, node_feat = node_feat) prediction["m"] = magmom atoms_per_graph_tensor = torch.tensor(batch_graph['atoms_per_graph'], dtype=torch.int32, device=total_energy.device) if self.is_intensive: energy_per_atom = total_energy / atoms_per_graph_tensor prediction["e"] = energy_per_atom else: prediction["e"] = total_energy prediction["atoms_per_graph"] = atoms_per_graph_tensor ref_energy = ( 0 if self.reference_energy is None else self.reference_energy(graphs) ) prediction["e"] += ref_energy prediction["ref_energy"] = ref_energy return prediction def get_params(self) -> int: """Return the number of parameters in the model.""" return sum(p.numel() for p in self.parameters()) @classmethod def from_dict(cls, dct: dict): matris = MatRIS(**dct["config"]) matris.load_state_dict(dct["state_dict"]) return matris @classmethod def load( cls, model_name: str = "matris_10m_oam", device: str | None = None, ): """Load pretrained model.""" model_name = model_name.lower() supported_models = ["matris_10m_oam", "matris_10m_mp"] if model_name not in supported_models: raise ValueError(f"Unsupported model_name: {model_name}. Supported models are: {supported_models}") if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" # 默认权重目录：本仓库根目录下的 weight/；可通过 ONESCIENCE_MODELS_DIR 覆盖 default_weight_dir = Path(__file__).resolve().parent.parent / "weight" cache_dir = os.environ.get("ONESCIENCE_MODELS_DIR", str(default_weight_dir)) os.makedirs(cache_dir, exist_ok=True) checkpoint_files = { "matris_10m_omat": "MatRIS_10M_OMAT.pth.tar", "matris_10m_oam": "MatRIS_10M_OAM.pth.tar", "matris_10m_mp": "MatRIS_10M_MP.pth.tar", "matris_6m_mp": "MatRIS_6M_MP.pth.tar", } DOWNLOAD_URLS = { "matris_10m_omat": "", # TODO "matris_10m_oam": "https://figshare.com/ndownloader/files/59142728", "matris_10m_mp": "https://figshare.com/ndownloader/files/59143058", "matris_6m_mp": "", # TODO } ckpt_filename = checkpoint_files[model_name] ckpt_path = os.path.join(cache_dir, ckpt_filename) if not os.path.exists(ckpt_path): url = DOWNLOAD_URLS.get(model_name) if not url: raise ValueError(f"No download URL provided for model: {model_name}") print(f"Checkpoint not found, downloading to {ckpt_path} ...") torch.hub.download_url_to_file(url, ckpt_path) ckpt_state = torch.load( ckpt_path, map_location=torch.device("cpu"), weights_only=False ) model = MatRIS.from_dict(ckpt_state) model = model.to(device) print(f"Loading {model_name} successfully, running on {device}.") return model