Feat: add pairtab compression #4432
Conversation
for more information, see https://pre-commit.ci
📝 WalkthroughWalkthroughThe changes introduce a new method called Changes
Possibly related PRs
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (2)
✅ Files skipped from review due to trivial changes (2)
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? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
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)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (5)
deepmd/dpmodel/atomic_model/linear_atomic_model.py (4)
Line range hint
442-445: Improve assertion message clarityThe assertion message could be more descriptive to help developers understand the configuration requirements better.
- assert ( - self.sw_rmax > self.sw_rmin - ), "The upper boundary `sw_rmax` must be greater than the lower boundary `sw_rmin`." + assert ( + self.sw_rmax > self.sw_rmin + ), f"Invalid boundary configuration: sw_rmax ({self.sw_rmax}) must be greater than sw_rmin ({self.sw_rmin})"
Line range hint
474-481: Handle division by zero more explicitlyThe current implementation uses numpy's error state management to handle division by zero, but this could be more explicit and safer.
Consider this approach:
- with np.errstate(divide="ignore", invalid="ignore"): - sigma = numerator / denominator + # Avoid division by zero + mask = denominator != 0 + sigma = np.zeros_like(numerator) + sigma[mask] = numerator[mask] / denominator[mask]
Line range hint
482-492: Consider extracting smoothing functionThe smoothing calculation is complex and could benefit from being extracted into a separate method for better readability and maintainability.
Consider extracting the smoothing logic:
+ def _compute_smooth_coefficient(self, u: np.ndarray) -> np.ndarray: + """Compute smooth interpolation coefficient. + + Parameters + ---------- + u : np.ndarray + Normalized distance values + + Returns + ------- + np.ndarray + Smooth interpolation coefficient + """ + return -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 def _compute_weight(self, ...): # ... existing code ... with np.errstate(invalid="ignore"): - smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 + smooth = self._compute_smooth_coefficient(u)
Line range hint
493-496: Document the weight calculation logicThe weight calculation logic is complex but lacks documentation explaining the mathematical reasoning behind the coefficients.
Add a docstring section explaining the mathematical formula:
def _compute_weight(self, ...): """ZBL weight. + + The weight is calculated using a smooth step function: + 1. For distances < sw_rmin: weight = 1 + 2. For distances > sw_rmax: weight = 0 + 3. For distances in between: weight = smooth polynomial interpolation + using a 5th-degree polynomial: -6x^5 + 15x^4 - 10x^3 + 1 + + This ensures continuous first and second derivatives at the boundaries. Returns ------- list[np.ndarray] the atomic ZBL weight for interpolation. (nframes, nloc, 1) """deepmd/pt/model/atomic_model/linear_atomic_model.py (1)
Line range hint
588-652: Improve documentation for better maintainabilityThe weight computation logic is complex but lacks detailed documentation. Consider:
- Adding mathematical formulas in docstring to explain the interpolation
- Documenting the meaning of variables (u, sigma, coef)
- Explaining the choice of the polynomial coefficients (-6, 15, -10)
Add this documentation:
def _compute_weight( self, extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlists_: list[torch.Tensor], ) -> list[torch.Tensor]: - """ZBL weight. + """Compute interpolation weights between DP and ZBL models. + + The interpolation uses a smooth polynomial function: + - For r < sw_rmin: weight_zbl = 1 + - For sw_rmin <= r < sw_rmax: weight_zbl = -6u⁵ + 15u⁴ - 10u³ + 1 + where u = (r - sw_rmin)/(sw_rmax - sw_rmin) + - For r >= sw_rmax: weight_zbl = 0 + + The distance r is computed as a softmin over all neighbor distances: + r = Σ(d_i * exp(-d_i/α)) / Σ(exp(-d_i/α)) + where d_i are pairwise distances and α is smin_alpha. Returns ------- list[torch.Tensor] the atomic ZBL weight for interpolation. (nframes, nloc, 1) """
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (2)
deepmd/dpmodel/atomic_model/linear_atomic_model.py(1 hunks)deepmd/pt/model/atomic_model/linear_atomic_model.py(1 hunks)
🔇 Additional comments (1)
deepmd/pt/model/atomic_model/linear_atomic_model.py (1)
577-587: Verify compression method usage across codebase
Let's ensure that no code relies on the compression functionality for PairTab models.
✅ Verification successful
Compression is correctly handled for PairTab models
The verification shows that the implementation is consistent across the codebase:
- The
enable_compressionmethod is correctly implemented as a no-op with a docstring explicitly stating "Pairtab model does not support compression" in both PyTorch and TensorFlow implementations - No code attempts to use compression functionality with PairTab models
- The compression-related configuration and documentation properly handle this limitation
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for calls to enable_compression on PairTab instances
rg -l "PairTabAtomicModel" | xargs rg "enable_compression"
# Search for compression-related configuration
rg -l "pairtab" | xargs rg -i "compress"
Length of output: 3014
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4432 +/- ##
==========================================
- Coverage 84.64% 84.64% -0.01%
==========================================
Files 614 614
Lines 57138 57141 +3
Branches 3487 3486 -1
==========================================
Hits 48367 48367
- Misses 7646 7647 +1
- Partials 1125 1127 +2 ☔ View full report in Codecov by Sentry. |
Summary by CodeRabbit
New Features
enable_compressionin the PairTabAtomicModel class, indicating that the model does not support compression settings.Documentation
enable_compressionmethod to clarify its purpose.