feat(pt/test): add unit test for the compression of se_e3_tebd

[OutisLi]

Summary by CodeRabbit

• Tests
  • Added comprehensive test suite validating model compression functionality, including consistency checks between original and compressed models across multiple configurations, support for various periodic boundary conditions, and integrated validation of physical properties.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Note

Other AI code review bot(s) detected

CodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review.

📝 Walkthrough

Walkthrough

A new comprehensive test module is added to validate model compression and TEBD-based neural descriptors in DeepMD. The module includes setup helpers, five test case classes, and cleanup routines for testing model attributes, evaluations across multiple scenarios, consistency checks, and ASE-based integration testing.

Changes

Cohort / File(s)

Change Summary

Test Module for SE-e3 TEBD Model Compression source/tests/pt/test_model_compression_se_e3_tebd.py

New test file introducing TestDeepPotAPBC, TestDeepPotANoPBC, TestDeepPotALargeBoxNoPBC, TestDeepPotAPBCExcludeTypes, and TestSkipNeighborStat classes. Includes setup helpers for model generation and training, attribute validation tests, 1-frame and 2-frame evaluation tests for PBC/non-PBC scenarios, per-atom output verification, consistency checks between original and compressed models, ASE-based integration testing, and module-level cleanup routines.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

• Verify that all setup/teardown logic properly manages test-generated artifacts and model cleanup
• Validate the comprehensiveness of test assertions across diverse scenarios (PBC, non-PBC, exclude_types, skip-neighbor-stat)
• Ensure the ASE integration test correctly compares DP(Frozen) and DP(Compressed) model outputs
• Check that per-atom attribute tests (ae, av, ff, ee, vv) cover expected dimensions and values

Possibly related PRs

• feat(pt): add compression support for se_e3_tebd #4992: Directly exercises the compression APIs and TEBD descriptor functionality introduced by this test module's target implementation.
• feat(dp/pt): refactor se_e3 descriptor #3813: Validates SE-e3 descriptor refactoring and exclude_types handling that this test module now covers comprehensively.

Suggested reviewers

• njzjz
• iProzd
• wanghan-iapcm

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title clearly and specifically describes the main change: adding unit tests for SE_E3_TEBD model compression in the PyTorch module.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[Copilot]

Pull request overview

This PR adds comprehensive unit tests for model compression functionality of the se_e3_tebd descriptor in PyTorch. The tests verify that compressed models produce results consistent with the original frozen models.

Key Changes

• Added unit tests covering compression of se_e3_tebd descriptor models
• Tests include scenarios with PBC, no PBC, exclude_types, and skip-neighbor-stat options
• Verifies consistency between original and compressed models for energy, forces, and virial predictions

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.32%. Comparing base (da452d7) to head (6d4bf1a).
⚠️ Report is 10 commits behind head on devel.

Additional details and impacted files

@@            Coverage Diff             @@
##            devel    #5060      +/-   ##
==========================================
+ Coverage   84.24%   84.32%   +0.07%     
==========================================
  Files         709      709              
  Lines       70236    70278      +42     
  Branches     3623     3619       -4     
==========================================
+ Hits        59169    59259      +90     
+ Misses       9900     9851      -49     
- Partials     1167     1168       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (3)

source/tests/pt/test_model_compression_se_e3_tebd.py (3)

35-175: Shared input.json and duplicated cleanup paths are a bit surprising

All three _init_models* helpers write to the same tests_path / "input.json" and setUpModule binds both INPUT and INPUT_ET to that path. tearDownModule then deletes both, which works because _file_delete is tolerant of missing files, but the implicit sharing is non-obvious and could become fragile if other tests also rely on input.json in the same directory or if CWD differs.

Similarly, the cleanup of artifacts like "checkpoint" and "model-compression" uses bare relative paths, which assumes a particular working directory layout.

Consider, as a low-priority cleanup:

• Giving each scenario its own input filename (e.g. input-se-e3-tebd.json, input-se-e3-tebd-exclude-types.json, input-se-e3-tebd-skip-neighbor-stat.json) and deleting those explicitly.
• Optionally prefixing the artifact paths with tests_path (if appropriate for this test suite) to avoid depending on the process CWD.

This would make the setup/teardown behavior more explicit and easier to maintain.

Also applies to: 177-198

---

205-225: Deduplicate repeated coordinate and atype literals across classes

The same coords and atype arrays are hard-coded in each test class. This works, but it does increase maintenance overhead and the risk of accidental drift if one copy is edited without updating the others.

If you want to reduce duplication, consider factoring these into a shared helper or base mixin, e.g. a small function that returns (coords, atype) or a common test base class that defines them as class attributes reused by all the subclasses.

Based on learnings, similar refactors between closely related PBC/no-PBC tests have been done elsewhere in this repo, so this would be consistent with that direction.

Also applies to: 329-349, 438-457, 545-565, 670-689

---

510-538: Add pbc=True to Atoms objects to exercise PBC code path

The test_ase test currently exercises only the non-PBC branch of DP.calculate(). Although cell is provided to both Atoms objects, the pbc parameter defaults to False in ASE. The DP.calculate() method checks if sum(self.atoms.get_pbc()) > 0: (lines 131-134) and only uses the cell when at least one PBC dimension is enabled; otherwise, it sets cell = None before evaluation.

To validate PBC behavior for the ASE-calculator integration, explicitly set pbc=True on both Atoms objects:

         water0 = Atoms(
             "OHHOHH",
             positions=self.coords.reshape((-1, 3)),
             cell=self.box.reshape((3, 3)),
+            pbc=True,
             calculator=DP(FROZEN_MODEL),
         )
         water1 = Atoms(
             "OHHOHH",
             positions=self.coords.reshape((-1, 3)),
             cell=self.box.reshape((3, 3)),
+            pbc=True,
             calculator=DP(COMPRESSED_MODEL),
         )

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e98dc5a and 6d4bf1a.

📒 Files selected for processing (1)

• source/tests/pt/test_model_compression_se_e3_tebd.py (1 hunks)

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2024-10-08T15:32:11.479Z

Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-10-08T15:32:11.479Z
Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905.

Applied to files:

• source/tests/pt/test_model_compression_se_e3_tebd.py

🧬 Code graph analysis (1)

source/tests/pt/test_model_compression_se_e3_tebd.py (1)

deepmd/calculator.py (1)

• DP (36-159)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)

• GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
• GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
• GitHub Check: Test Python (3, 3.9)
• GitHub Check: Test Python (5, 3.12)
• GitHub Check: Test Python (4, 3.12)
• GitHub Check: Test Python (6, 3.12)
• GitHub Check: Test Python (5, 3.9)
• GitHub Check: Test Python (2, 3.12)
• GitHub Check: Test Python (6, 3.9)
• GitHub Check: Test Python (4, 3.9)
• GitHub Check: Test Python (2, 3.9)
• GitHub Check: Test Python (3, 3.12)
• GitHub Check: Test Python (1, 3.12)
• GitHub Check: Test Python (1, 3.9)
• GitHub Check: Test C++ (true)
• GitHub Check: Test C++ (false)
• GitHub Check: Build wheels for cp311-win_amd64
• GitHub Check: Analyze (python)
• GitHub Check: Build wheels for cp310-manylinux_aarch64
• GitHub Check: Analyze (c-cpp)
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build wheels for cp311-manylinux_x86_64
• GitHub Check: Build wheels for cp311-macosx_x86_64
• GitHub Check: Build wheels for cp311-macosx_arm64
• GitHub Check: Build C++ (cuda120, cuda)
• GitHub Check: Build C++ (rocm, rocm)
• GitHub Check: Build C++ (cpu, cpu)
• GitHub Check: Build C++ (clang, clang)
• GitHub Check: Build C++ (cuda, cuda)

🔇 Additional comments (1)

source/tests/pt/test_model_compression_se_e3_tebd.py (1)

28-33: Utility _file_delete is simple and robust

The helper correctly handles both files and directories and is safely idempotent for missing paths. No changes needed here.