Upload folder using huggingface_hub

README.md

@@ -0,0 +1,143 @@
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MatRIS</span>
+  </strong>
+</p>
+
+
+# Model Introduction
+
+MatRIS is a foundation model for materials representation and interaction simulation, short for Materials Representation and Interaction Simulation. It can predict energy, forces, stress, and magnetic moments for crystal structures, and supports structure relaxation based on ASE and pymatgen structure objects.
+
+
+# Model Description
+
+MatRIS is based on a graph neural network architecture and is trained on materials datasets such as OMat24 and MPTrj. It performs energy, force, stress, and magnetic-moment prediction and structure optimization for crystalline materials.
+
+# Applicable Scenarios
+
+| Scenario | Description |
+| :---: | :---: |
+| Crystal energy prediction | Input a CIF, pymatgen Structure, or ASE Atoms object and predict the system energy |
+| Force and stress prediction | Provide force and stress estimates for structure relaxation, molecular dynamics, or downstream simulation |
+| Magnetic moment prediction | Output structure-related magnetic-moment results under the `efsm` task |
+| Structure relaxation pre-processing | Use `StructOptimizer` to optimize atomic positions and unit cells of candidate crystal structures |
+| Environment connectivity check | Use `cif_file/demo.cif` and a lightweight MatRIS model to check whether the OneScience matchem environment is available |
+
+
+# Usage Instructions
+
+## 1. Using OneCode
+
+You can try out intelligent one-click AI4S programming in the OneCode online environment:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- GPU or DCU is recommended.
+- CPU can be used for module verification and small-scale forward checks, but structure relaxation will be slow. CPU is not recommended for formal batch inference.
+- DCU users need to install DTK in advance. DTK 25.04.2 or above, or the OneScience-recommended version matching the current cluster, is suggested.
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/MatRIS --local_dir ./matris
+cd matris
+```
+
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is also supported
+pip install onescience[matchem-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is also supported
+pip install onescience[matchem-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Weights
+
+Inference testing requires a pre-trained model. MatRIS provides the following model keys:
+
+| Model Key | Description |
+| --- | --- |
+| `matris_10m_omat` | Trained on the OMat24 dataset |
+| `matris_10m_oam` | Trained on OMat24 and fine-tuned on sAlex+MPtrj |
+| `matris_10m_mp` | Trained on the MPTrj dataset |
+
+The `weight/` directory in this repository already contains pre-trained weights. `scripts/test_relaxation.py` will load them from `weight/` when run.
+
+
+### Inference
+
+**Run Modular Verification**
+
+```bash
+python scripts/test_modularization.py
+```
+
+This script instantiates a lightweight MatRIS model (randomly initialized weights) and completes one CPU forward pass to verify the connectivity of model modules.
+
+**Run Structure Relaxation Inference**
+
+```bash
+python scripts/test_relaxation.py
+```
+
+This script reads `cif_file/demo.cif` and performs structure relaxation through `StructOptimizer`.
+
+After inference completes, the log will output the relaxation process, and energy, forces, stress, magnetic moments, and the final structure object will be obtained in memory.
+
+**Other Inference Examples**
+
+Besides structure relaxation, you can also use `MatRISCalculator` to predict energy, forces, stress, and magnetic moments for a single structure:
+
+```python
+import torch
+from ase.build import bulk
+from onescience.utils.matris import MatRISCalculator
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+calc = MatRISCalculator(
+    model="matris_10m_oam",
+    task="efsm",
+    device=device,
+)
+
+atoms = bulk("Cu", a=5.43, cubic=True)
+atoms.calc = calc
+
+energy = atoms.get_potential_energy()   # total energy (eV)
+forces = atoms.get_forces()             # forces (eV/Å)
+stress = atoms.get_stress()             # stress (eV/ų)
+magmoms = atoms.get_magnetic_moments()  # magnetic moments (μB)
+```
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- MatRIS upstream materials are released under the BSD-3-Clause License. This repository retains the source attribution and is organized for OneScience ModelScope automatic execution scenarios.
+
+- If you use MatRIS results in scientific research, we recommend citing the original MatRIS project, the relevant OneScience project information, and adding citations for the materials datasets, structure optimization tools, or downstream analysis tools used according to the actual task.

cif_file/demo.cif

@@ -0,0 +1,36 @@
+# generated using pymatgen
+data_Pu2Si3
+_symmetry_space_group_name_H-M   'P 1'
+_cell_length_a   4.26022438
+_cell_length_b   7.16692875
+_cell_length_c   7.16692875
+_cell_angle_alpha   90.00000000
+_cell_angle_beta   90.00000000
+_cell_angle_gamma   90.00000000
+_symmetry_Int_Tables_number   1
+_chemical_formula_structural   Pu2Si3
+_chemical_formula_sum   'Pu4 Si6'
+_cell_volume   218.82586168
+_cell_formula_units_Z   2
+loop_
+ _symmetry_equiv_pos_site_id
+ _symmetry_equiv_pos_as_xyz
+  1  'x, y, z'
+loop_
+ _atom_site_type_symbol
+ _atom_site_label
+ _atom_site_symmetry_multiplicity
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_occupancy
+  Pu  Pu0  1  0.50000000  0.67597000  0.82403000  1.0
+  Pu  Pu1  1  0.50000000  0.32403000  0.17597000  1.0
+  Pu  Pu2  1  0.50000000  0.17597000  0.67597000  1.0
+  Pu  Pu3  1  0.50000000  0.82403000  0.32403000  1.0
+  Si  Si4  1  0.00000000  0.88257100  0.61742900  1.0
+  Si  Si5  1  0.00000000  0.11742900  0.38257100  1.0
+  Si  Si6  1  0.00000000  0.38257100  0.88257100  1.0
+  Si  Si7  1  0.00000000  0.61742900  0.11742900  1.0
+  Si  Si8  1  0.00000000  0.00000000  0.00000000  1.0
+  Si  Si9  1  0.00000000  0.50000000  0.50000000  1.0

model/__init__.py

@@ -0,0 +1,33 @@
+"""
+MatRIS: Material Representation Learning with Interatomic Structure
+
+A deep learning model for predicting material properties including energy, 
+forces, stress, and magnetic moments.
+"""
+
+import sys
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version(__name__)  # read from pyproject.toml
+except PackageNotFoundError:
+    __version__ = "unknown"
+
+# Import main model
+from model.matris import MatRIS
+
+# Import graph components
+from onescience.datapipes.materials.matris import GraphConverter, RadiusGraph
+
+# 把本仓库 model 注册为 onescience.models.matris，使得 OneScience 的 MatRISCalculator、
+# StructOptimizer 等工具类在加载预训练权重时，直接使用本仓库的 MatRIS.load 方法。
+# 默认权重会下载到本仓库根目录下的 weight/ 中。
+sys.modules["onescience.models.matris"] = sys.modules[__name__]
+
+__all__ = [
+    "MatRIS",
+    "GraphConverter", 
+    "RadiusGraph",
+    "__version__",
+]

model/matris.py

@@ -0,0 +1,319 @@
+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

scripts/test_modularization.py

@@ -0,0 +1,165 @@
+"""
+验证 MatRIS 模块化重构后各模块能否被正确导入与运行。
+
+测试内容：
+1. 分散在各目录的 matris 模块能否正常 import
+2. MatRIS 模型能否正常实例化
+3. 模型能否执行一次简单的前向传播（energy/force/stress/magmom）
+"""
+
+import sys
+from pathlib import Path
+
+# 把本仓库根目录放到 sys.path 最前面，避免 PYTHONPATH 中其他同名 model 包干扰
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+import model  # noqa: E402
+
+import torch
+
+def test_imports():
+    """测试所有分散模块的导入路径"""
+    print("[1/4] 测试模块导入...")
+    
+    # layer
+    from onescience.modules.layer.matris_radial import (
+        PolynomialEnvelope, BesselExpansion, GaussianExpansion, 
+        FourierExpansion, SphericalExpansion
+    )
+    from onescience.modules.layer.matris_interaction import Interaction_Block
+    
+    # embedding
+    from onescience.modules.embedding.matris_embedding import (
+        AtomTypeEmbedding, EdgeBasisEmbedding, ThreebodyEmbedding
+    )
+    
+    # func_utils
+    from onescience.modules.func_utils.matris_func_utils import (
+        MLP, GatedMLP, get_activation, get_normalization, aggregate
+    )
+    from onescience.modules.func_utils.matris_reference import AtomRef
+    from onescience.modules.func_utils.matris_graph import process_graphs
+    
+    # head
+    from onescience.modules.head.matris_head import EnergyHead, MagmomHead, ForceStressHead
+    
+    # utils
+    from onescience.utils.matris import StructOptimizer, MatRISCalculator
+    
+    print("    所有模块导入成功 ✓")
+
+
+def test_model_instantiate():
+    """测试模型能否正常实例化"""
+    print("[2/4] 测试模型实例化...")
+    
+    from model import MatRIS
+    
+    model = MatRIS(
+        num_layers=2,           # 用小层数加速测试
+        node_feat_dim=64,
+        edge_feat_dim=64,
+        three_body_feat_dim=64,
+        num_radial=5,
+        num_angular=5,
+        pairwise_cutoff=5.0,
+        three_body_cutoff=3.0,
+        reference_energy=None,   # 不加载参考能量
+    )
+    
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"    MatRIS 实例化成功，参数量: {num_params:,} ✓")
+    return model
+
+
+def test_forward_cpu():
+    """测试 CPU 前向传播"""
+    print("[3/4] 测试 CPU 前向传播...")
+    
+    from model import MatRIS
+    from onescience.datapipes.materials.matris import GraphConverter
+    from pymatgen.core.structure import Structure
+    from pymatgen.core.lattice import Lattice
+    
+    # 构建一个极简晶体：2原子 Si
+    lattice = Lattice.cubic(5.43)
+    structure = Structure(lattice, ["Si", "Si"], [[0, 0, 0], [0.25, 0.25, 0.25]])
+    
+    model = MatRIS(
+        num_layers=2,
+        node_feat_dim=64,
+        edge_feat_dim=64,
+        three_body_feat_dim=64,
+        num_radial=5,
+        num_angular=5,
+        pairwise_cutoff=5.0,
+        three_body_cutoff=3.0,
+        reference_energy=None,
+    )
+    model.eval()
+    
+    # 构建图
+    graph_converter = GraphConverter(
+        atom_graph_cutoff=5.0,
+        line_graph_cutoff=3.0,
+    )
+    graph = graph_converter(structure)
+    
+    out = model([graph], task="efsm")
+    
+    print(f"    预测能量 (eV/atom): {out['e'].item():.4f} ✓")
+    print(f"    预测力数量: {len(out['f'])} 组 ✓")
+    print(f"    预测应力数量: {len(out['s'])} 组 ✓")
+    print(f"    预测磁矩数量: {len(out['m'])} 组 ✓")
+
+
+def test_cuda_available():
+    """若存在 GPU，测试 CUDA 前向传播"""
+    print("[4/4] 测试 CUDA 可用性...")
+    
+    if not torch.cuda.is_available():
+        print("    无可用 GPU，跳过 CUDA 测试")
+        return
+    
+    from model import MatRIS
+    from onescience.datapipes.materials.matris import GraphConverter
+    from pymatgen.core.structure import Structure
+    from pymatgen.core.lattice import Lattice
+    
+    lattice = Lattice.cubic(5.43)
+    structure = Structure(lattice, ["Si", "Si"], [[0, 0, 0], [0.25, 0.25, 0.25]])
+    
+    model = MatRIS(
+        num_layers=2,
+        node_feat_dim=64,
+        edge_feat_dim=64,
+        three_body_feat_dim=64,
+        num_radial=5,
+        num_angular=5,
+        pairwise_cutoff=5.0,
+        three_body_cutoff=3.0,
+        reference_energy=None,
+    ).cuda()
+    model.eval()
+    
+    graph_converter = GraphConverter(5.0, 3.0)
+    graph = graph_converter(structure).to("cuda")
+    
+    out = model([graph], task="efsm")
+    
+    print(f"    CUDA 前向传播成功，能量: {out['e'].item():.4f} ✓")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("MatRIS 模块化重构验证测试")
+    print("=" * 60)
+    
+    test_imports()
+    test_model_instantiate()
+    test_forward_cpu()
+    test_cuda_available()
+    
+    print("=" * 60)
+    print("全部测试通过 ✓")
+    print("=" * 60)

scripts/test_relaxation.py

@@ -0,0 +1,44 @@
+import sys
+from pathlib import Path
+
+# 把本仓库根目录放到 sys.path 最前面，避免 PYTHONPATH 中其他同名 model 包干扰
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+import model  # noqa: E402
+
+import ase
+from ase.build import bulk
+import torch
+from pymatgen.core.structure import Structure
+from onescience.utils.matris import StructOptimizer
+
+model_name = "matris_10m_oam"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+matris_opt = StructOptimizer(
+    model=model_name, 
+    task = "efsm",
+    optimizer = "FIRE", # FIRE, BFGS ...
+    device=device
+)
+
+atom = Structure.from_file("cif_file/demo.cif")
+
+max_steps = 500
+fmax = 0.05
+opt_result = matris_opt.relax(
+        atoms=atom, # pymatgen.Structure or ase.Atoms
+        verbose=True,
+        steps=max_steps,
+        fmax=fmax,
+        relax_cell=max_steps > 0,
+        ase_filter="FrechetCellFilter",
+    )
+
+trajectory = opt_result['trajectory']
+energy = trajectory.energies[-1]
+force = trajectory.forces[-1]
+stress = trajectory.stresses[-1]
+magmom = trajectory.magmoms[-1]
+
+final_structure = opt_result['final_structure']

weight/.MatRIS_10M_OAM.pth.tar.efileUploading

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85c1091b5a28d582f309e3bd0f663236d57313346dfca8f5aa364541234d2af6
+size 8

weight/MatRIS_10M_MP.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:333b6019ec901a4bce9735ef28136d7da637899272825ff7ef9c84587b22f2d9
+size 42184483

weight/MatRIS_10M_OAM.pth.tar.efilePart

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c033abc53601a74f10d9b7fec0f658220c013c3b11d4d405f1d32136d4c2b067
+size 42273174