Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,170 @@
+---
+license: mit
+tasks:
+  - materials-simulation
+  - molecular-dynamics
+  - energy-prediction
+  - force-prediction
+  - structure-relaxation
+  - fine-tuning
+frameworks:
+  - pytorch
+language:
+  - en
+tags:
+  - OneScience
+  - MatterSim
+  - materials-science
+  - molecular-simulation
+  - machine-learning-potential
+  - graph-neural-network
+  - equivariant-neural-network
+  - inference
+  - fine-tuning
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MatterSim</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+MatterSim is a deep-learning interatomic potential developed by Microsoft Research for a broad range of elements, temperatures, and pressures. It predicts energies and forces for inorganic materials, molecules, and periodic systems.
+
+Paper: *MatterSim: A deep-learning atomistic model across elements, temperatures, and pressures*  
+Reference implementation: https://github.com/microsoft/mattersim
+
+# Model Description
+
+MatterSim uses a deep-learning architecture trained on multiple materials and molecular datasets. It supports energy and force prediction, structure relaxation, molecular dynamics, and fine-tuning on custom datasets for inorganic materials, molecules, and periodic systems.
+
+# Use Cases
+
+| Use case | Description |
+| :---: | :--- |
+| Single-point energy/force prediction | Quickly predict the energy and atomic forces of a given atomic structure |
+| Batch structure inference | Predict energies and forces for multiple structures in a batch |
+| Structure relaxation | Optimize atomic positions and cell shape with FIRE/BFGS |
+| Molecular dynamics | Run short MD sampling in the NVT ensemble |
+| Fine-tuning on custom data | Fine-tune a pretrained MatterSim model on your own dataset |
+| Environment connectivity check | Use the single-point and relaxation scripts to verify the OneScience MatChem environment, model loading, and CUDA/DCU availability |
+
+# 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 import checks and small-configuration connectivity tests, but full training and inference will be slow.
+- 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/Mattersim --local-dir ./mattersim
+cd mattersim
+```
+
+### 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
+```
+
+### Training Data
+
+By default, this repository only includes the `high_level_water.xyz` sample data for quickly validating model loading, single-point inference, structure relaxation, molecular dynamics, and fine-tuning workflows. Download any additional training data separately and place it in `data/`.
+
+### Trained Weights
+
+The repository includes `weight/mattersim-v1.0.0-1M.pth`. All scripts also accept a custom model weight through `--checkpoint`.
+
+### Inference
+
+```bash
+cd scripts
+python single_point.py --checkpoint ../weight/mattersim-v1.0.0-1M.pth
+```
+
+```bash
+cd scripts
+python batch_inference.py --checkpoint ../weight/mattersim-v1.0.0-1M.pth
+```
+
+**Structure relaxation**
+
+```bash
+cd scripts
+python relax.py --checkpoint ../weight/mattersim-v1.0.0-1M.pth --device cuda
+```
+
+> The default checkpoint is `../weight/mattersim-v1.0.0-1M.pth`.
+
+**Molecular dynamics**
+
+```bash
+cd scripts
+python md.py --checkpoint ../weight/mattersim-v1.0.0-1M.pth --device cuda
+```
+
+> This script also uses `../weight/mattersim-v1.0.0-1M.pth` by default.
+
+### Fine-Tuning
+
+Edit the paths and parameters in `scripts/finetune_config.yaml`, including `train_data_path` and `checkpoint`:
+
+```bash
+cd scripts
+# Edit train_data_path, checkpoint, and other fields in finetune_config.yaml
+```
+
+Single GPU:
+
+```bash
+python finetune.py --config finetune_config.yaml
+```
+
+Multi-GPU DDP:
+
+```bash
+torchrun --nproc_per_node=4 finetune.py --config finetune_config.yaml
+```
+
+# 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 MatterSim-related code comes from the MatChem examples in the OneScience project and refers to the upstream MatterSim project (https://github.com/microsoft/mattersim). The upstream MatterSim code is released under the [MIT License](https://github.com/microsoft/mattersim/blob/main/LICENSE).
+- If you use MatterSim training or inference results in research, please cite the original MatterSim paper, the relevant OneScience projects, and the datasets used.

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

model/__init__.py

@@ -0,0 +1,31 @@
+"""MatterSim model and integration entry points."""
+
+MATTERSIM_SOURCE_VERSION = "1.2.3"
+MATTERSIM_INTEGRATION_VERSION = "dcu2"
+
+from .adapter import (
+    DEFAULT_CHECKPOINT,
+    load_calculator,
+    load_potential,
+    predict_structures,
+    resolve_checkpoint,
+)
+__all__ = [
+    "DEFAULT_CHECKPOINT",
+    "M3Gnet",
+    "MATTERSIM_INTEGRATION_VERSION",
+    "MATTERSIM_SOURCE_VERSION",
+    "MatterSim",
+    "load_calculator",
+    "load_potential",
+    "predict_structures",
+    "resolve_checkpoint",
+]
+
+
+def __getattr__(name):
+    if name in {"M3Gnet", "MatterSim"}:
+        from .mattersim import M3Gnet, MatterSim
+
+        return {"M3Gnet": M3Gnet, "MatterSim": MatterSim}[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

model/adapter.py

@@ -0,0 +1,94 @@
+"""Checkpoint and model factories for the OneScience MatterSim integration."""
+
+import os
+from pathlib import Path
+
+import torch
+
+DEFAULT_CHECKPOINT = "mattersim-v1.0.0-1M.pth"
+
+
+def resolve_checkpoint(checkpoint: str | os.PathLike | None = None) -> str:
+    """Resolve a MatterSim checkpoint without changing its native format.
+
+    Explicit checkpoint values are passed through so MatterSim model aliases keep
+    working. When no value is supplied, the shared OneScience model store is
+    checked before falling back to MatterSim's native alias and download logic.
+    """
+    if checkpoint is not None:
+        return str(Path(checkpoint).expanduser())
+
+    models_dir = os.environ.get("ONESCIENCE_MODELS_DIR")
+    if models_dir:
+        shared_checkpoint = Path(models_dir).expanduser() / "mattersim" / DEFAULT_CHECKPOINT
+        if shared_checkpoint.is_file():
+            return str(shared_checkpoint)
+
+    return DEFAULT_CHECKPOINT
+
+
+def _device(device: str | None) -> str:
+    return device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def load_potential(
+    checkpoint: str | os.PathLike | None = None,
+    device: str | None = None,
+    load_training_state: bool = False,
+    **kwargs,
+):
+    """Load a MatterSim ``Potential`` from a resolved checkpoint."""
+    from onescience.utils.mattersim.potential import Potential
+
+    return Potential.from_checkpoint(
+        load_path=resolve_checkpoint(checkpoint),
+        device=_device(device),
+        load_training_state=load_training_state,
+        **kwargs,
+    )
+
+
+def load_calculator(
+    checkpoint: str | os.PathLike | None = None,
+    device: str | None = None,
+    **kwargs,
+):
+    """Create an ASE-compatible ``MatterSimCalculator``."""
+    from onescience.utils.mattersim.calculator import MatterSimCalculator
+
+    return MatterSimCalculator.from_checkpoint(
+        resolve_checkpoint(checkpoint), device=_device(device), **kwargs
+    )
+
+
+def predict_structures(
+    atoms,
+    checkpoint: str | os.PathLike | None = None,
+    device: str | None = None,
+    batch_size: int = 16,
+    include_forces: bool = True,
+    include_stresses: bool = False,
+    cutoff: float = 5.0,
+    threebody_cutoff: float = 4.0,
+):
+    """Predict ASE structures with the MatterSim potential."""
+    from onescience.datapipes.materials.mattersim import build_dataloader
+
+    potential = load_potential(checkpoint=checkpoint, device=device)
+    dataloader = build_dataloader(
+        atoms=list(atoms),
+        batch_size=batch_size,
+        cutoff=cutoff,
+        threebody_cutoff=threebody_cutoff,
+        only_inference=True,
+    )
+    energies, forces, stresses = potential.predict_properties(
+        dataloader,
+        include_forces=include_forces,
+        include_stresses=include_stresses,
+    )
+    return {
+        "energies": energies,
+        "forces": forces,
+        "stresses": stresses,
+    }

model/mattersim.py

@@ -0,0 +1,201 @@
+# -*- coding: utf-8 -*-
+from typing import Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_runstats.scatter import scatter
+
+from onescience.modules.block.mattersim_block import MainBlock
+from onescience.modules.embedding.mattersim_embedding import (
+    SmoothBesselBasis,
+    SphericalBasisLayer,
+)
+from onescience.modules.func_utils.mattersim_jit import compile_mode
+from onescience.modules.func_utils.mattersim_scaling import AtomScaling
+from onescience.modules.layer.mattersim_layer import GatedMLP, MLP
+
+
+@compile_mode("script")
+class M3Gnet(nn.Module):
+    """
+    M3Gnet
+    """
+
+    def __init__(
+        self,
+        num_blocks: int = 4,
+        units: int = 128,
+        max_l: int = 4,
+        max_n: int = 4,
+        cutoff: float = 5.0,
+        device: str = "cuda" if torch.cuda.is_available() else "cpu",
+        max_z: int = 94,
+        threebody_cutoff: float = 4.0,
+        **kwargs,
+    ):
+        super().__init__()
+        self.rbf = SmoothBesselBasis(r_max=cutoff, max_n=max_n)
+        self.sbf = SphericalBasisLayer(max_n=max_n, max_l=max_l, cutoff=cutoff)
+        self.edge_encoder = MLP(
+            in_dim=max_n, out_dims=[units], activation="swish", use_bias=False
+        )
+        module_list = [
+            MainBlock(max_n, max_l, cutoff, units, max_n, threebody_cutoff)
+            for i in range(num_blocks)
+        ]
+        self.graph_conv = nn.ModuleList(module_list)
+        self.final = GatedMLP(
+            in_dim=units,
+            out_dims=[units, units, 1],
+            activation=["swish", "swish", None],
+        )
+        self.apply(self.init_weights)
+        self.atom_embedding = MLP(
+            in_dim=max_z + 1, out_dims=[units], activation=None, use_bias=False
+        )
+        self.atom_embedding.apply(self.init_weights_uniform)
+        self.normalizer = AtomScaling(verbose=False, max_z=max_z, device=device)
+        self.max_z = max_z
+        self.device = device
+        self.model_args = {
+            "num_blocks": num_blocks,
+            "units": units,
+            "max_l": max_l,
+            "max_n": max_n,
+            "cutoff": cutoff,
+            "max_z": max_z,
+            "threebody_cutoff": threebody_cutoff,
+        }
+
+    def forward(
+        self,
+        input: Dict[str, torch.Tensor],
+        dataset_idx: int = -1,
+    ) -> torch.Tensor:
+        # Exact data from input_dictionary
+        pos = input["atom_pos"]
+        cell = input["cell"]
+        pbc_offsets = input["pbc_offsets"].float()
+        atom_attr = input["atom_attr"]
+        edge_index = input["edge_index"].long()
+        three_body_indices = input["three_body_indices"].long()
+        num_bonds = input["num_bonds"]
+        num_triple_ij = input["num_triple_ij"]
+        num_atoms = input["num_atoms"]
+        num_graphs = input["num_graphs"]
+        batch = input["batch"]
+
+        # Use precomputed values if available, otherwise compute on the fly
+        # (backward-compat for callers not using batch_to_dict)
+        total_num_atoms = input.get("total_num_atoms", int(num_atoms.sum()))
+        total_num_bonds = input.get("total_num_bonds", int(num_bonds.sum()))
+
+        bond_index_bias = input.get("bond_index_bias", None)
+        if bond_index_bias is None:
+            cumsum = torch.cumsum(num_bonds, dim=0) - num_bonds
+            bond_index_bias = torch.repeat_interleave(
+                cumsum, input["num_three_body"], dim=0
+            ).unsqueeze(-1)
+
+        three_body_edge_map = input.get("three_body_edge_map", None)
+
+        # -------------------------------------------------------------#
+        three_body_indices = three_body_indices + bond_index_bias
+
+        # === Refer to the implementation of M3GNet,        ===
+        # === we should re-compute the following attributes ===
+        # edge_length, edge_vector(optional), triple_edge_length, theta_jik
+        edge_batch = batch[edge_index[0]]
+        edge_vector = pos[edge_index[0]] - (
+            pos[edge_index[1]]
+            + torch.einsum("bi, bij->bj", pbc_offsets, cell[edge_batch])
+        )
+        edge_length = torch.linalg.norm(edge_vector, dim=1)
+        vij = edge_vector[three_body_indices[:, 0].clone()]
+        vik = edge_vector[three_body_indices[:, 1].clone()]
+        rij = edge_length[three_body_indices[:, 0].clone()]
+        rik = edge_length[three_body_indices[:, 1].clone()]
+        cos_jik = torch.sum(vij * vik, dim=1) / (rij * rik)
+        # eps = 1e-7 avoid nan in torch.acos function
+        cos_jik = torch.clamp(cos_jik, min=-1.0 + 1e-7, max=1.0 - 1e-7)
+        triple_edge_length = rik.view(-1)
+        edge_length = edge_length.unsqueeze(-1)
+        atomic_numbers = atom_attr.squeeze(1).long()
+
+        # featurize
+        atom_attr = self.atom_embedding(self.one_hot_atoms(atomic_numbers))
+        edge_attr = self.rbf(edge_length.view(-1))
+        edge_attr_zero = edge_attr  # e_ij^0
+        edge_attr = self.edge_encoder(edge_attr)
+        three_basis = self.sbf(triple_edge_length, torch.acos(cos_jik))
+
+        # Main Loop
+        for idx, conv in enumerate(self.graph_conv):
+            atom_attr, edge_attr = conv(
+                atom_attr,
+                edge_attr,
+                edge_attr_zero,
+                edge_index,
+                three_basis,
+                three_body_indices,
+                edge_length,
+                num_bonds,
+                num_triple_ij,
+                num_atoms,
+                total_num_atoms=total_num_atoms,
+                total_num_bonds=total_num_bonds,
+                three_body_edge_map=three_body_edge_map,
+            )
+
+        energies_i = self.final(atom_attr).view(-1)  # [batch_size*num_atoms]
+        energies_i = self.normalizer(energies_i, atomic_numbers)
+        energies = scatter(energies_i, batch, dim=0, dim_size=num_graphs)
+
+        return energies  # [batch_size]
+
+    def init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+
+    def init_weights_uniform(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.uniform_(m.weight, a=-0.05, b=0.05)
+
+    @torch.jit.export
+    def one_hot_atoms(self, species):
+        # one_hots = []
+        # for i in range(species.shape[0]):
+        #     one_hots.append(
+        #         F.one_hot(
+        #             species[i],
+        #             num_classes=self.max_z+1).float().to(species.device)
+        #     )
+        # return torch.cat(one_hots, dim=0)
+        return F.one_hot(species, num_classes=self.max_z + 1).float()
+
+    def print(self):
+        from prettytable import PrettyTable
+
+        table = PrettyTable(["Modules", "Parameters"])
+        total_params = 0
+        for name, parameter in self.named_parameters():
+            if not parameter.requires_grad:
+                continue
+            params = parameter.numel()
+            table.add_row([name, params])
+            total_params += params
+        print(table)
+        print(f"Total Trainable Params: {total_params}")
+
+    @torch.jit.export
+    def set_normalizer(self, normalizer: AtomScaling):
+        self.normalizer = normalizer
+
+    def get_model_args(self):
+        return self.model_args
+
+
+MatterSim = M3Gnet
+
+__all__ = ["M3Gnet", "MatterSim"]

scripts/batch_inference.py

@@ -0,0 +1,32 @@
+import argparse
+
+from ase.build import bulk
+
+from onescience.models.mattersim import predict_structures
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="MatterSim batch inference")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
+    parser.add_argument("--batch-size", type=int, default=16)
+    args = parser.parse_args()
+
+    structures = [bulk("Si", "diamond", a=a) for a in (5.40, 5.43, 5.46)]
+    result = predict_structures(
+        structures,
+        checkpoint=args.checkpoint,
+        device=args.device,
+        batch_size=args.batch_size,
+    )
+    for index, (energy, forces) in enumerate(
+        zip(result["energies"], result["forces"], strict=True)
+    ):
+        print(
+            f"Structure {index}: energy={energy:.6f} eV, "
+            f"max_force={abs(forces).max():.6e} eV/angstrom"
+        )
+
+
+if __name__ == "__main__":
+    main()

scripts/finetune.py

@@ -0,0 +1,74 @@
+import argparse
+from pathlib import Path
+
+import yaml
+
+from onescience.utils.mattersim import FineTuneConfig, MatterSimTrainer
+
+
+def _load_yaml_config(path: str) -> dict:
+    with open(path, "r", encoding="utf-8") as stream:
+        return yaml.safe_load(stream) or {}
+
+
+def _build_parser(base_config: dict) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Fine-tune MatterSim with OneScience")
+    parser.add_argument("--config", help="Path to YAML config file")
+    parser.add_argument(
+        "--train-data-path", default=base_config.get("train_data_path")
+    )
+    parser.add_argument(
+        "--valid-data-path", default=base_config.get("valid_data_path")
+    )
+    parser.add_argument("--checkpoint", default=base_config.get("checkpoint"))
+    parser.add_argument("--save-path", default=base_config.get("save_path", "./results/mattersim"))
+    parser.add_argument("--run-name", default=base_config.get("run_name", "onescience-mattersim"))
+    parser.add_argument("--epochs", type=int, default=base_config.get("epochs", 1000))
+    parser.add_argument("--batch-size", type=int, default=base_config.get("batch_size", 16))
+    parser.add_argument("--lr", type=float, default=base_config.get("lr", 2e-4))
+    parser.add_argument(
+        "--device", choices=("cpu", "cuda"), default=base_config.get("device", "cuda")
+    )
+    parser.add_argument("--seed", type=int, default=base_config.get("seed", 42))
+    parser.add_argument(
+        "--include-stresses",
+        action="store_true",
+        default=base_config.get("include_stresses", False),
+    )
+    parser.add_argument(
+        "--no-include-forces",
+        action="store_false",
+        dest="include_forces",
+        default=base_config.get("include_forces", True),
+    )
+    parser.add_argument(
+        "--re-normalize",
+        action="store_true",
+        default=base_config.get("re_normalize", False),
+    )
+    parser.add_argument(
+        "--no-save-checkpoint",
+        action="store_false",
+        dest="save_checkpoint",
+        default=base_config.get("save_checkpoint", True),
+    )
+    return parser
+
+
+def main() -> None:
+    # Two-phase parsing: first get --config, then use YAML defaults for the rest.
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument("--config")
+    pre_args, remaining = pre_parser.parse_known_args()
+
+    base_config = _load_yaml_config(pre_args.config) if pre_args.config else {}
+    parser = _build_parser(base_config)
+    args = parser.parse_args(remaining)
+
+    # Drop None values and the config key itself.
+    kwargs = {k: v for k, v in vars(args).items() if v is not None and k != "config"}
+    MatterSimTrainer(FineTuneConfig(**kwargs)).fit()
+
+
+if __name__ == "__main__":
+    main()

scripts/finetune_config.yaml

@@ -0,0 +1,44 @@
+# MatterSim 微调基础配置
+# 请复制此文件后，修改下方带 "TODO:" 标记的路径为你本地实际路径。
+
+# 数据路径（TODO: 修改为你的训练数据路径）
+# 示例使用仓库自带的 data/high_level_water.xyz
+train_data_path: ../data/high_level_water.xyz
+
+# 验证数据路径，不需要可保持 null（TODO: 按需修改）
+valid_data_path: null
+
+# checkpoint 路径（TODO: 修改为你的预训练模型路径）
+# 示例使用仓库自带的 weight/mattersim-v1.0.0-1M.pth
+checkpoint: ../weight/mattersim-v1.0.0-1M.pth
+
+# 输出设置
+save_path: ./results/mattersim
+run_name: onescience-mattersim
+
+# 训练超参数
+epochs: 1
+batch_size: 16
+lr: 2.0e-4
+seed: 42
+
+# 设备：cuda 或 cpu，留空则自动检测
+device: cuda
+
+# 训练目标
+cutoff: 5.0
+threebody_cutoff: 4.0
+include_forces: true
+include_stresses: false
+
+# 损失权重
+force_loss_ratio: 1.0
+stress_loss_ratio: 0.1
+
+# 保存与早停
+save_checkpoint: true
+ckpt_interval: 10
+early_stop_patience: 10
+
+# 重新归一化（可选）
+re_normalize: false

scripts/md.py

@@ -0,0 +1,27 @@
+import argparse
+
+from ase.build import bulk
+
+from onescience.utils.mattersim import molecular_dynamics
+
+
+parser = argparse.ArgumentParser(description="MatterSim molecular dynamics")
+parser.add_argument(
+    "--checkpoint",
+    default="../weight/mattersim-v1.0.0-1M.pth",
+    help="Path to MatterSim checkpoint",
+)
+parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
+args = parser.parse_args()
+
+atoms = bulk("Si", "diamond", a=5.43).repeat((2, 2, 2))
+dynamics = molecular_dynamics(
+    atoms,
+    checkpoint=args.checkpoint,
+    ensemble="nvt_berendsen",
+    temperature=300,
+    timestep=1.0,
+    logfile="-",
+    device=args.device,
+)
+dynamics.run(10)

scripts/relax.py

@@ -0,0 +1,26 @@
+import argparse
+
+from ase.build import bulk
+
+from onescience.utils.mattersim import relax
+
+
+parser = argparse.ArgumentParser(description="MatterSim structure relaxation")
+parser.add_argument(
+    "--checkpoint",
+    default="../weight/mattersim-v1.0.0-1M.pth",
+    help="Path to MatterSim checkpoint",
+)
+parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
+args = parser.parse_args()
+
+converged, atoms = relax(
+    bulk("Si", "diamond", a=5.43),
+    checkpoint=args.checkpoint,
+    filter="FrechetCellFilter",
+    fmax=0.01,
+    steps=20,
+    device=args.device,
+)
+print(f"Converged: {converged}")
+print(f"Energy: {atoms.get_potential_energy():.6f} eV")

scripts/single_point.py

@@ -0,0 +1,30 @@
+import argparse
+
+import torch
+from ase.build import bulk
+from ase.units import GPa
+
+from onescience.models.mattersim import load_calculator
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="MatterSim single-point inference")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda")
+    args = parser.parse_args()
+
+    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
+    atoms = bulk("Si", "diamond", a=5.43)
+    atoms.calc = load_calculator(checkpoint=args.checkpoint, device=device)
+
+    energy = atoms.get_potential_energy()
+    print(f"Energy (eV)                 = {energy}")
+    print(f"Energy per atom (eV/atom)   = {energy / len(atoms)}")
+    print(f"Forces of first atom (eV/A) = {atoms.get_forces()[0]}")
+    stress = atoms.get_stress(voigt=False)
+    print(f"Stress[0][0] (eV/A^3)       = {stress[0][0]}")
+    print(f"Stress[0][0] (GPa)          = {stress[0][0] / GPa}")
+
+
+if __name__ == "__main__":
+    main()

weight/mattersim-v1.0.0-1M.pth

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