fix(pt/dp): share params of repinit_three_body

[iProzd]

Fix #4137.

Summary by CodeRabbit

• New Features

  • Enhanced descriptor handling to support three-body interactions, improving model flexibility.
  • Introduced a new model configuration, model_dpa2tebd, with advanced parameters for better performance.

• Tests

  • Added a new test class, TestMultiTaskDPA2Tebd, to expand testing coverage for the new model configuration.

[coderabbitai]

Walkthrough

Walkthrough

The changes in this pull request enhance the dpa2.py file by adding support for a three-body descriptor in the methods set_stat_mean_and_stddev, get_stat_mean_and_stddev, share_params, and compute_input_stats. Additionally, a new model configuration model_dpa2tebd is introduced in the test files, along with a dedicated test class TestMultiTaskDPA2Tebd to ensure proper testing of the new model configuration.

Changes

Files	Change Summary
deepmd/dpmodel/descriptor/dpa2.py	Enhanced methods to support three-body descriptors, modifying set_stat_mean_and_stddev, get_stat_mean_and_stddev, share_params, and compute_input_stats.
source/tests/pt/model/test_permutation.py	Added new model configuration model_dpa2tebd with detailed descriptor and fitting network parameters.
source/tests/pt/test_multitask.py	Introduced TestMultiTaskDPA2Tebd class for testing the model_dpa2tebd configuration, including setup and teardown methods.

Possibly related PRs

• fix: correct exclude_types in descriptors #3841: Involves making exclude_types consistent across descriptors, related to descriptor management in the main PR.
• feat(dp/pt): refactor se_e3 descriptor #3813: Refactors the se_e3 descriptor, sharing functionalities with enhancements made to descriptor handling in the main PR.

Suggested labels

Python, Docs, Examples

Suggested reviewers

---

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share

• X
• Mastodon
• Reddit
• LinkedIn

Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  -- I pushed a fix in commit <commit_id>, please review it.
  -- Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  -- @coderabbitai generate unit testing code for this file.
  -- @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  -- @coderabbitai generate interesting stats about this repository and render them as a table.
  -- @coderabbitai read src/utils.ts and generate unit testing code.
  -- @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  -- @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 0

Outside diff range and nitpick comments (1)

deepmd/pt/model/descriptor/dpa2.py (1)

513-513: Rename unused loop variable.

As per the Ruff suggestion, please rename the unused loop variable ii to _ii to indicate that it is intentionally unused.

-for ii, descrpt in enumerate(descrpt_list):
+for _ii, descrpt in enumerate(descrpt_list):

Tools

Ruff

513-513: Loop control variable ii not used within loop body

Rename unused ii to _ii

(B007)

Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between 2c9be6f and 6d3e452.

Files selected for processing (4)

• deepmd/dpmodel/descriptor/dpa2.py (1 hunks)
• deepmd/pt/model/descriptor/dpa2.py (4 hunks)
• source/tests/pt/model/test_permutation.py (1 hunks)
• source/tests/pt/test_multitask.py (2 hunks)

Additional context used

Ruff

deepmd/pt/model/descriptor/dpa2.py

513-513: Loop control variable ii not used within loop body

Rename unused ii to _ii

(B007)

Additional comments not posted (6)

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

304-341: LGTM!

The new test class TestMultiTaskDPA2Tebd is well-structured and correctly sets up the necessary configuration and paths for testing the model_dpa2tebd configuration. The setUp and tearDown methods are properly implemented, and there are no apparent issues with the logic or syntax.

This addition enhances the testing framework by providing a dedicated test case for the model_dpa2tebd configuration, thereby expanding the coverage of the multitask training functionality.

source/tests/pt/model/test_permutation.py (1)

167-217: Verify the impact of the key changes in model_dpa2tebd.

The model_dpa2tebd configuration introduces several changes compared to model_dpa2:

• Addition of 3-body interactions in the descriptor, which could improve accuracy for certain systems but increases computational cost.
• Changes in the "repformer" section, such as reducing the number of layers, disabling certain attention mechanisms, and introducing new configuration parameters.
• A simpler fitting network with one less hidden layer.

Please verify the model's performance with these changes and ensure that the trade-offs align with the desired objectives.

To verify the impact of the changes, you can run the following script:

The script searches for evidence that:

1. Both models are being tested and compared in the test files.
2. Both models are being trained and evaluated on the same dataset.
3. There is some documentation or comments discussing the changes and their expected impact.

If the script finds relevant results, it would indicate that the impact of the changes has been considered and verified. If not, it suggests that more testing and documentation may be needed.

deepmd/pt/model/descriptor/dpa2.py (2)

392-395: LGTM!

The changes to conditionally share parameters of repinit_three_body based on use_three_body look good. This ensures consistency when the three-body interaction is enabled.

Also applies to: 405-408

---

510-513: Looks good!

The changes to include repinit_three_body in compute_input_stats, set_stat_mean_and_stddev, get_stat_mean_and_stddev and serialize when use_three_body is enabled are correct and consistent.

Also applies to: 522-525, 531-539

Tools

Ruff

513-513: Loop control variable ii not used within loop body

Rename unused ii to _ii

(B007)

deepmd/dpmodel/descriptor/dpa2.py (2)

736-739: LGTM!

The change correctly updates the mean and stddev for the self.repinit_three_body descriptor when self.use_three_body is True. This ensures that the statistics are properly maintained for all relevant descriptors.

---

745-753: LGTM!

The change correctly includes the mean and stddev of the self.repinit_three_body descriptor in the returned lists when self.use_three_body is True. This ensures that the retrieved statistics are comprehensive and consistent with the enabled descriptors.