Merge branch 'master' into feat-other-full-model

Author: Han Wang

.github/skills/add-descriptor/SKILL.md

@@ -5,7 +5,7 @@ license: LGPL-3.0-or-later
 compatibility: Requires Python 3.10+, numpy, pytest. Optional backends for full testing (torch, jax, paddle).
 metadata:
   author: deepmd-kit
-  version: "2.0"
+  version: '2.0'
 ---
 
 # Adding a New Descriptor to deepmd-kit

.github/skills/add-descriptor/references/dpmodel-implementation.md

@@ -93,25 +93,26 @@ Rules:
    # GOOD: xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1)
    ```
 
-2. **`xp.zeros`/`xp.ones` must include `device=`** — omitting device can trigger CUDA init or create tensors on wrong device:
+1. **`xp.zeros`/`xp.ones` must include `device=`** — omitting device can trigger CUDA init or create tensors on wrong device:
 
    ```python
    # BAD:  xp.zeros([2, 1], dtype=nlist.dtype)
    # GOOD: xp.zeros([2, 1], dtype=nlist.dtype, device=array_api_compat.device(nlist))
    ```
 
-3. **`xp.split` with `axis=` keyword doesn't work for torch** — use slicing:
+1. **`xp.split` with `axis=` keyword doesn't work for torch** — use slicing:
 
    ```python
    # BAD:  g2, h2 = xp.split(dmatrix, [1], axis=-1)
    # GOOD: g2, h2 = dmatrix[..., :1], dmatrix[..., 1:]
    ```
 
-4. **`xp_take_along_axis` indices must be int64 for torch**.
+1. **`xp_take_along_axis` indices must be int64 for torch**.
 
-5. **Don't maintain separate ArrayAPI subclasses** — dpmodel classes should be array_api compatible directly.
+1. **Don't maintain separate ArrayAPI subclasses** — dpmodel classes should be array_api compatible directly.
+
+1. **Boolean fancy indexing (`arr[mask]`) is not array-API compatible** — use mask multiplication:
 
-6. **Boolean fancy indexing (`arr[mask]`) is not array-API compatible** — use mask multiplication:
    ```python
    # BAD:  gr[ti_mask] += gr_tmp
    # GOOD: gr += gr_tmp * xp.astype(mask[:, None, None], gr_tmp.dtype)

.pre-commit-config.yaml

@@ -1,5 +1,6 @@
 # See https://pre-commit.com for more information
 # See https://pre-commit.com/hooks.html for more hooks
+exclude: ^source/3rdparty/.+/
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v6.0.0
@@ -53,23 +54,32 @@ repos:
       - id: velin
         args: ["--write"]
         exclude: ^source/3rdparty
-  # Python inside docs
-  - repo: https://github.com/asottile/blacken-docs
-    rev: 1.20.0
+  # markdown
+  - repo: https://github.com/hukkin/mdformat
+    rev: 1.0.0
     hooks:
-      - id: blacken-docs
+      - id: mdformat
+        additional_dependencies:
+          # - mdformat-myst==0.3.0
+          # See https://github.com/executablebooks/mdformat-myst/issues/13
+          - "git+https://github.com/njzjz-bothub/mdformat-myst@d9c414e#egg=mdformat-myst"
+          - mdformat-ruff==0.1.3
+          - mdformat-web==0.2.0
+          - mdformat-config==0.2.1
+          - mdformat-beautysh==1.0.0
+          - mdformat-gfm-alerts==2.0.0
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v21.1.8
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|source/tests/infer/.+\.json$)
-  # markdown, yaml, CSS, javascript
+  # yaml, CSS, javascript
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: v4.0.0-alpha.8
     hooks:
       - id: prettier
-        types_or: [markdown, yaml, css]
+        types_or: [yaml, css]
         # workflow files cannot be modified by pre-commit.ci
         exclude: ^(source/3rdparty|\.clang-format)
   # Shell

AGENTS.md

@@ -23,8 +23,8 @@ DeePMD-kit is a deep learning package for many-body potential energy representat
 ### Lint and Format Code
 
 - Install linter: `uv pip install ruff`
-- Run linting: `ruff check .` -- takes <1 second
-- Format code: `ruff format .` -- takes <1 second
+- Run linting: `ruff check .` -- takes \<1 second
+- Format code: `ruff format .` -- takes \<1 second
 - **Always run `ruff check .` and `ruff format .` before committing changes or the CI will fail.**
 
 ### Training and Validation
@@ -40,19 +40,19 @@ DeePMD-kit is a deep learning package for many-body potential energy representat
 ### Basic Functionality Validation
 
 1. **CLI Interface**: Run `dp --version` and `dp -h` to verify installation
-2. **Python Interface**: Run `python -c "import deepmd; import deepmd.tf; print('Both interfaces work')"`
-3. **Backend Selection**: Test `dp --tf -h`, `dp --pt -h`, `dp --jax -h`, `dp --pd -h`
+1. **Python Interface**: Run `python -c "import deepmd; import deepmd.tf; print('Both interfaces work')"`
+1. **Backend Selection**: Test `dp --tf -h`, `dp --pt -h`, `dp --jax -h`, `dp --pd -h`
 
 ### Training Workflow Validation
 
 1. **TensorFlow Training**: `cd examples/water/se_e2_a && timeout 60 dp train input.json --skip-neighbor-stat` -- should start training and show decreasing loss
-2. **PyTorch Training**: `cd examples/water/se_e2_a && timeout 60 dp --pt train input_torch.json --skip-neighbor-stat` -- should start training and show decreasing loss
-3. **Verify training output**: Look for "batch X: trn: rmse" messages showing decreasing error values
+1. **PyTorch Training**: `cd examples/water/se_e2_a && timeout 60 dp --pt train input_torch.json --skip-neighbor-stat` -- should start training and show decreasing loss
+1. **Verify training output**: Look for "batch X: trn: rmse" messages showing decreasing error values
 
 ### Test-Based Validation
 
 1. **Core Tests**: `pytest source/tests/tf/test_dp_test.py::TestDPTestEner::test_1frame -v` -- should pass in ~10 seconds
-2. **Multi-backend**: Test both TensorFlow and PyTorch components work
+1. **Multi-backend**: Test both TensorFlow and PyTorch components work
 
 ## Common Commands and Timing
 
@@ -127,8 +127,8 @@ source/               # C++ source code and tests
 
 ### Linting and Formatting
 
-- **Ruff check**: <1 second
-- **Ruff format**: <1 second
+- **Ruff check**: \<1 second
+- **Ruff format**: \<1 second
 - **Pre-commit hooks**: May have network issues, use individual tools
 
 ### Commit Messages and PR Titles

CONTRIBUTING.md

@@ -47,7 +47,7 @@ Please perform the following steps to create your Pull Request to this repositor
 ### Step 1: Fork the repository
 
 1. Visit the project: <https://github.com/deepmodeling/deepmd-kit>
-2. Click the **Fork** button on the top right and wait it to finish.
+1. Click the **Fork** button on the top right and wait it to finish.
 
 ### Step 2: Clone the forked repository to local storage and set configurations
 
@@ -58,7 +58,7 @@ Please perform the following steps to create your Pull Request to this repositor
    # Replace `$username` with your GitHub ID
    ```
 
-2. Add deepmodeling's repo as your remote repo, we can name it "upstream". And fetch upstream's latest codes to your workstation.
+1. Add deepmodeling's repo as your remote repo, we can name it "upstream". And fetch upstream's latest codes to your workstation.
 
    ```bash
    git remote add upstream https://github.com/deepmodeling/deepmd-kit.git
@@ -70,9 +70,9 @@ Please perform the following steps to create your Pull Request to this repositor
    git merge upstream/master
    ```
 
-3. Modify your codes and design unit tests.
+1. Modify your codes and design unit tests.
 
-4. Commit your changes to a new branch
+1. Commit your changes to a new branch
 
    ```bash
    git checkout -b branch1
@@ -81,7 +81,8 @@ Please perform the following steps to create your Pull Request to this repositor
    git commit -m "commit-message: update the xx"
    ```
 
-5. Push the changed codes to your original repo on github.
+1. Push the changed codes to your original repo on github.
+
    ```bash
    git push origin branch1
    ```
@@ -97,31 +98,31 @@ Please perform the following steps to create your Pull Request to this repositor
    git rebase upstream/master
    ```
 
-2. Create a new branch based on the master branch.
+1. Create a new branch based on the master branch.
 
    ```bash
    git checkout -b new-branch-name
    ```
 
-3. Modify your codes and design unit tests.
+1. Modify your codes and design unit tests.
 
-4. Commit your changes
+1. Commit your changes
 
    ```bash
    git status # Checks the local status
    git add <file> ... # Adds the file(s) you want to commit. If you want to commit all changes, you can directly use `git add.`
    git commit -m "commit-message: update the xx"
    ```
 
-5. Keep your branch in sync with upstream/master
+1. Keep your branch in sync with upstream/master
 
    ```bash
    # While on your new branch
    git fetch upstream
    git rebase upstream/master
    ```
 
-6. Push your changes to the remote
+1. Push your changes to the remote
 
    ```bash
    git push -u origin new-branch-name # "-u" is used to track the remote branch from origin
@@ -130,7 +131,7 @@ Please perform the following steps to create your Pull Request to this repositor
 ### Step 3: Create a pull request
 
 1. Visit your fork at <https://github.com/$username/deepmd-kit> (replace `$username` with your GitHub ID)
-2. Click `pull requests`, followed by `New pull request` and `Compare & pull request` to create your PR.
+1. Click `pull requests`, followed by `New pull request` and `Compare & pull request` to create your PR.
 
 Now, your PR is successfully submitted! After this PR is merged, you will automatically become a contributor to DeePMD-kit.
 

README.md

@@ -1,6 +1,6 @@
 [<picture><source media="(prefers-color-scheme: dark)" srcset="./doc/_static/logo-dark.svg"><source media="(prefers-color-scheme: light)" srcset="./doc/_static/logo.svg"><img alt="DeePMD-kit logo" src="./doc/_static/logo.svg"></picture>](./doc/logo.md)
 
----
+______________________________________________________________________
 
 # DeePMD-kit
 

deepmd/dpmodel/utils/__init__.py

@@ -36,6 +36,11 @@
     save_dp_model,
     traverse_model_dict,
 )
+from .training_utils import (
+    compute_total_numb_batch,
+    resolve_model_prob,
+    resolve_model_prob_from_epochs,
+)
 
 __all__ = [
     "AtomExcludeMask",
@@ -49,6 +54,7 @@
     "aggregate",
     "build_multiple_neighbor_list",
     "build_neighbor_list",
+    "compute_total_numb_batch",
     "extend_coord_with_ghosts",
     "get_graph_index",
     "get_multiple_nlist_key",
@@ -60,6 +66,8 @@
     "nlist_distinguish_types",
     "normalize_coord",
     "phys2inter",
+    "resolve_model_prob",
+    "resolve_model_prob_from_epochs",
     "save_dp_model",
     "to_face_distance",
     "traverse_model_dict",