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,146 @@
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MatterSim</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+MatterSim is a deep-learning interatomic potential model across elements, temperatures, and pressures proposed by Microsoft Research. It predicts energy 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 is based on a deep-learning architecture and is trained on multiple materials and molecular datasets. It performs energy and force prediction, structure relaxation, molecular dynamics, and custom-dataset fine-tuning for inorganic materials, molecules, and periodic systems.
+
+# Applicable Scenarios
+
+| Scenario | Description |
+| :---: | :--- |
+| Single-point energy/force prediction | Quickly predict energy and atomic forces for a given atomic structure |
+| Batch structure inference | Perform batch energy/force prediction for multiple structures |
+| Structure relaxation | Optimize atomic positions and unit-cell shape using FIRE/BFGS |
+| Molecular dynamics | Run short-range MD sampling under the NVT ensemble |
+| Custom data fine-tuning | Fine-tune a pre-trained MatterSim model on your own dataset |
+| Environment connectivity check | Use the single-point/relaxation scripts to check the OneScience matchem environment, model loading, and CUDA/DCU availability |
+
+# 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 import and small-configuration connectivity checks; full training and inference will be slow.
+- 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/Mattersim --local_dir ./mattersim
+cd mattersim
+```
+
+### 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 Data Description
+
+By default, this repository only includes the example data file `high_level_water.xyz`, used for quickly verifying model loading, single-point inference, structure relaxation, molecular dynamics, and fine-tuning workflows. For other training data, please download it yourself and place it in the `data/` directory.
+
+### Training Weights
+
+The repository includes `weight/mattersim-v1.0.0-1M.pth`. All scripts also support specifying model weights via `--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 weight 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
+```
+
+> The default weight is also `../weight/mattersim-v1.0.0-1M.pth`.
+
+### Fine-tuning
+
+Edit the paths and parameters in `scripts/finetune_config.yaml` directly (e.g., `train_data_path`, `checkpoint`, etc.):
+
+```bash
+cd scripts
+# Edit fields such as train_data_path and checkpoint 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
+```
+
+# 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
+
+- The MatterSim-related code comes from the matchem example implementation in the OneScience project and references 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 scientific research, we recommend citing the original MatterSim paper, the relevant OneScience project information, and the sources of the datasets actually 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