Upload folder using huggingface_hub

README.md

@@ -0,0 +1,161 @@
+---
+license: bsd-3-clause
+tasks:
+  - materials-simulation
+  - structure-relaxation
+  - property-prediction
+frameworks:
+  - pytorch
+language:
+  - en
+tags:
+  - OneScience
+  - materials-science
+  - materials-representation
+  - crystal-structure
+  - energy-prediction
+  - force-prediction
+  - stress-prediction
+  - magnetic-moment-prediction
+  - structure-relaxation
+datasets:
+  - OneScience-Sugon/pbe
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MatRIS</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+MatRIS, short for Materials Representation and Interaction Simulation, is a foundation model for material representation and interaction simulation. It predicts the energy, forces, stress, and magnetic moments of crystal structures and supports structure relaxation using ASE Atoms and pymatgen Structure objects.
+
+# Model Description
+
+MatRIS uses a graph neural network architecture trained on materials datasets such as OMat24 and MPTrj. It supports energy, force, stress, and magnetic-moment prediction and structure optimization for crystalline materials.
+
+# Use Cases
+
+| Use case | Description |
+| :---: | :--- |
+| Crystal energy prediction | Predict system energy from a CIF file, pymatgen Structure, or ASE Atoms object |
+| Force and stress prediction | Estimate forces and stress for structure relaxation, molecular dynamics, or downstream simulations |
+| Magnetic-moment prediction | Return structure-dependent magnetic moments for the `efsm` task |
+| Structure-relaxation preprocessing | Optimize atomic positions and unit cells of candidate crystals with `StructOptimizer` |
+| Environment connectivity check | Use `cif_file/demo.cif` and a lightweight MatRIS model to verify the OneScience MatChem environment |
+
+# Usage
+
+## 1. Using OneCode
+
+Try 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**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for module checks and small forward-pass tests, but structure relaxation will be slow and is not recommended for production batch inference.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Sugon/MatRIS --local-dir ./matris
+cd matris
+```
+
+### Install the Runtime Environment
+
+**DCU environment**
+
+```bash
+# 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
+# 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
+```
+
+### Trained Weights
+
+Inference requires a pretrained 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 pretrained weights. `scripts/test_relaxation.py` loads them from that directory.
+
+### Inference
+
+**Run the modularity check**
+
+```bash
+python scripts/test_modularization.py
+```
+
+This script instantiates a lightweight MatRIS model with randomly initialized weights and runs one CPU forward pass to verify that the model modules work together.
+
+**Run structure-relaxation inference**
+
+```bash
+python scripts/test_relaxation.py
+```
+
+This script reads `cif_file/demo.cif` and performs structure relaxation with `StructOptimizer`.
+
+After inference, the log shows the relaxation process, while the energy, forces, stress, magnetic moments, and final structure object are available in memory.
+
+**Other inference examples**
+
+In addition to structure relaxation, you can use `MatRISCalculator` to predict the energy, forces, stress, and magnetic moments of 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)
+```
+
+# Official OneScience Resources
+
+| 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
+
+- The upstream MatRIS materials use the BSD 3-Clause License. This repository retains source attribution and has been organized for automated execution in the OneScience Hugging Face environment.
+- If you use MatRIS results in research, please cite the original MatRIS project and the relevant OneScience projects, as well as the materials datasets, structure-optimization tools, or downstream analysis tools used for your 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_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

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