Add compression API to BaseModel and AtomicModel #4298
Conversation
for more information, see https://pre-commit.ci
📝 Walkthrough📝 WalkthroughWalkthroughThe pull request introduces a new method called Changes
Possibly related PRs
Suggested labels
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? 🪧 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: 1
🧹 Outside diff range and nitpick comments (3)
deepmd/dpmodel/atomic_model/dp_atomic_model.py (1)
89-119: LGTM! Implementation looks correct.The method properly delegates compression configuration to the descriptor component, maintaining good separation of concerns.
Consider enhancing the docstring with:
- A brief description of what compression does and its benefits
- Parameter value ranges or constraints
- Example usage
def enable_compression( self, min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: - """Call descriptor enable_compression() + """Enable compression for the descriptor to optimize memory usage and performance. + + Compression reduces memory footprint by using lookup tables for atomic interactions. + The accuracy of the compression can be controlled via the stride parameters. Parameters ---------- min_nbor_dist - The nearest distance between atoms + The nearest distance between atoms. Must be positive. table_extrapolate - The scale of model extrapolation + The scale for model extrapolation beyond training data range. Must be >= 1. table_stride_1 - The uniform stride of the first table + The uniform stride of the first lookup table. Smaller values increase accuracy but memory usage. table_stride_2 - The uniform stride of the second table + The uniform stride of the second lookup table. Smaller values increase accuracy but memory usage. check_frequency - The overflow check frequency + How often to check for numerical overflow. Set to -1 to disable checks. + + Example + ------- + >>> model.enable_compression( + ... min_nbor_dist=0.5, + ... table_extrapolate=5, + ... table_stride_1=0.01, + ... table_stride_2=0.1, + ... ) """deepmd/dpmodel/atomic_model/base_atomic_model.py (2)
Line range hint
52-64: Fix mask shape and add documentation.The shape of the "mask" output variable doesn't match its usage in
forward_common_atomic. The mask is used to indicate real vs. virtual atoms and should have shape [nf, nloc] to match the batch size and number of local atoms.Update the shape and add documentation:
return FittingOutputDef( old_list # noqa:RUF005 + [ OutputVariableDef( name="mask", - shape=[1], + shape=["nf", "nloc"], reducible=False, r_differentiable=False, c_differentiable=False, + description="Binary mask indicating real (1) vs virtual (0) atoms", ) ] )
144-167: Consider adding base parameter validation.To ensure consistent validation across all implementations, consider:
- Adding parameter validation in the base class before raising NotImplementedError
- Creating a utility method for compression parameter validation that can be reused by all implementations
Example implementation:
def _validate_compression_params( min_nbor_dist: float, table_extrapolate: float, table_stride_1: float, table_stride_2: float, check_frequency: int, ) -> None: """Validate compression parameters.""" if min_nbor_dist <= 0: raise ValueError("min_nbor_dist must be positive") if table_extrapolate <= 0: raise ValueError("table_extrapolate must be positive") if table_stride_1 <= 0: raise ValueError("table_stride_1 must be positive") if table_stride_2 <= 0: raise ValueError("table_stride_2 must be positive") if check_frequency < -1: raise ValueError("check_frequency must be >= -1")This utility method can be placed in a common module and imported by all model implementations.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (3)
deepmd/dpmodel/atomic_model/base_atomic_model.py(1 hunks)deepmd/dpmodel/atomic_model/dp_atomic_model.py(1 hunks)deepmd/dpmodel/model/make_model.py(1 hunks)
🔇 Additional comments (1)
deepmd/dpmodel/model/make_model.py (1)
189-216: 🛠️ Refactor suggestion
Verify missing method and add parameter validation
A few concerns about the implementation:
- The method calls
get_min_nbor_dist()but this method is not visible in the provided code. - The implementation assumes the atomic model supports compression without checking.
- The parameters could benefit from validation.
Let's verify the existence and implementation of get_min_nbor_dist():
Consider adding parameter validation and interface checking:
def enable_compression(
self,
table_extrapolate: float = 5,
table_stride_1: float = 0.01,
table_stride_2: float = 0.1,
check_frequency: int = -1,
) -> None:
"""Call atomic_model enable_compression()
Parameters
----------
table_extrapolate
The scale of model extrapolation
table_stride_1
The uniform stride of the first table
table_stride_2
The uniform stride of the second table
check_frequency
The overflow check frequency
"""
+ if not hasattr(self.atomic_model, 'enable_compression'):
+ raise NotImplementedError("Compression not supported by the atomic model")
+
+ if table_extrapolate <= 0:
+ raise ValueError("table_extrapolate must be positive")
+ if table_stride_1 <= 0:
+ raise ValueError("table_stride_1 must be positive")
+ if table_stride_2 <= 0:
+ raise ValueError("table_stride_2 must be positive")
+ if check_frequency < -1:
+ raise ValueError("check_frequency must be >= -1")
+
self.atomic_model.enable_compression(
self.get_min_nbor_dist(),
table_extrapolate,
table_stride_1,
table_stride_2,
check_frequency,
)Consider enhancing the docstring to better explain:
- The purpose of compression
- The relationship between the parameters
- The meaning of the default values
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (1)
deepmd/dpmodel/atomic_model/dp_atomic_model.py (1)
89-118: LGTM! Consider enhancing the documentation.The implementation looks good with proper type hints and parameter forwarding to the descriptor.
Consider enhancing the docstring with:
- Return value documentation (even if None)
- Example usage
- A note about the NotImplementedError in the base class
- Explanation of what compression actually does and its impact on performance/accuracy
Example enhancement:
def enable_compression( self, min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: - """Call descriptor enable_compression(). + """Enable compression in the descriptor to optimize memory usage and performance. + + This method configures compression parameters for atomic interactions. Unlike the + base class which raises NotImplementedError, this implementation delegates to + the descriptor's compression mechanism. Parameters ---------- min_nbor_dist The nearest distance between atoms table_extrapolate The scale of model extrapolation table_stride_1 The uniform stride of the first table table_stride_2 The uniform stride of the second table check_frequency The overflow check frequency + + Returns + ------- + None + + Examples + -------- + >>> model = DPAtomicModel(descriptor, fitting, type_map) + >>> model.enable_compression( + ... min_nbor_dist=0.8, + ... table_extrapolate=6, + ... check_frequency=10 + ... ) """
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (3)
deepmd/dpmodel/atomic_model/base_atomic_model.py(1 hunks)deepmd/dpmodel/atomic_model/dp_atomic_model.py(1 hunks)deepmd/dpmodel/model/make_model.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/dpmodel/model/make_model.py
🔇 Additional comments (2)
deepmd/dpmodel/atomic_model/dp_atomic_model.py (1)
89-118: Verify descriptor compression implementation.
Let's verify that the descriptor class properly implements the compression functionality that this method depends on.
✅ Verification successful
Compression functionality is properly implemented across descriptor classes
The verification shows that enable_compression is consistently implemented across the codebase:
- The method is present in all descriptor implementations (se_a, se_r, se_t, se_atten, dpa1)
- The method signature matches the expected parameters (min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2)
- Compression functionality is well-tested as evidenced by numerous test files (test_compressed_descriptor_*.py)
- The implementation follows through from the high-level model down to the specific descriptor classes
The descriptor class properly implements the compression functionality that the dp_atomic_model.py's enable_compression method depends on.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for enable_compression implementation in descriptor classes
echo "Searching for enable_compression implementation in descriptor classes..."
ast-grep --pattern 'class $_ implements BaseDescriptor {
$$$
enable_compression($$$) {
$$$
}
$$$
}'
# Search for compression-related tests
echo "Searching for compression-related tests..."
rg -l "enable_compression.*test"
Length of output: 398
Script:
#!/bin/bash
# Let's try a different approach to find enable_compression implementations
echo "Searching for enable_compression method implementations..."
ast-grep --pattern 'def enable_compression($$$) {
$$$
}'
# Search for descriptor base class and interface
echo "Searching for BaseDescriptor definition..."
rg -A 10 "class BaseDescriptor"
# Search for compression-related code
echo "Searching for compression-related code..."
rg -l "compression" --type python
# Search for descriptor implementations
echo "Searching for descriptor implementations..."
rg -l "class.*Descriptor" --type python
Length of output: 644
Script:
#!/bin/bash
# Let's try without file type filtering since python type is not recognized
echo "Searching for enable_compression method implementations..."
rg -B 2 -A 5 "def enable_compression"
# Search for descriptor base class and interface
echo "Searching for BaseDescriptor definition..."
rg -B 2 -A 10 "class BaseDescriptor"
# Search for compression-related code
echo "Searching for compression-related code..."
rg -l "compression"
# Search for descriptor implementations
echo "Searching for descriptor implementations..."
fd . -e py -x rg -l "class.*Descriptor" {}
Length of output: 16347
deepmd/dpmodel/atomic_model/base_atomic_model.py (1)
144-168: Previous review comment is still applicable.
The issues identified in the previous review regarding parameter validation, documentation, and typo still need to be addressed.
The base class implementation follows good OOP practices.
The abstract implementation in the base class is architecturally sound, forcing derived classes to implement their own compression logic.
The NotImplementedError pattern is correctly used here to establish a contract that derived classes must fulfill, promoting modularity and extensibility.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4298 +/- ##
==========================================
- Coverage 84.35% 84.35% -0.01%
==========================================
Files 559 561 +2
Lines 52509 52622 +113
Branches 3054 3054
==========================================
+ Hits 44292 44387 +95
- Misses 7260 7275 +15
- Partials 957 960 +3 ☔ View full report in Codecov by Sentry. |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (5)
deepmd/dpmodel/atomic_model/make_base_atomic_model.py (2)
174-174: Fix typo in error message.There's a typo in the error message: "atomi" should be "atomic".
- raise NotImplementedError("This atomi model doesn't support compression!") + raise NotImplementedError("This atomic model doesn't support compression!")
161-173: Enhance parameter documentation.The parameter documentation would benefit from additional details about:
- Units for distance-based parameters (
min_nbor_dist,table_stride_1,table_stride_2)- Valid ranges for each parameter
- Explanation of what -1 means for
check_frequencyConsider updating the documentation to include these details. For example:
""" Parameters ---------- min_nbor_dist : float The nearest distance between atoms (in Angstroms). Must be positive. table_extrapolate : float, optional The scale of model extrapolation. Must be positive. Default is 5. table_stride_1 : float, optional The uniform stride of the first table (in Angstroms). Must be positive. Default is 0.01. table_stride_2 : float, optional The uniform stride of the second table (in Angstroms). Must be positive. Default is 0.1. check_frequency : int, optional The overflow check frequency. Use -1 to disable checks. Default is -1. """deepmd/dpmodel/model/base_model.py (1)
194-214: Consider architectural improvements for the compression feature.A few architectural suggestions to enhance the compression feature:
Consider marking this method as
@abstractmethodto enforce implementation in concrete classes, as it's a core feature that derived classes should implement.Consider adding a corresponding
disable_compression()method to allow toggling the feature.Consider storing compression settings as instance variables to:
- Allow querying current compression state
- Support serialization/deserialization of compression settings
- Enable runtime modifications of compression parameters
Example implementation pattern:
@abstractmethod def enable_compression( self, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: """Enable model compression.""" pass @abstractmethod def disable_compression(self) -> None: """Disable model compression.""" pass @property def compression_settings(self) -> Optional[dict]: """Get current compression settings.""" passdeepmd/pt/model/atomic_model/dp_atomic_model.py (2)
Line range hint
58-64: Update docstring to reflect fallback behavior.The method now includes a fallback to
coord_denoise_net.output_def(), but this behavior isn't documented.Consider updating the docstring:
@torch.jit.export def fitting_output_def(self) -> FittingOutputDef: - """Get the output def of the fitting net.""" + """Get the output definition. + + Returns + ------- + FittingOutputDef + The output definition from either the fitting net if available, + or the coordinate denoising net as a fallback. + """
Line range hint
102-115: Add error handling and consider version bump.The serialization now includes the fitting network state, but consider these improvements:
- Add error handling for fitting_net serialization
- Consider bumping the version number since a new field was added
Here's a suggested implementation:
def serialize(self) -> dict: + """Serialize the model to a dictionary. + + Returns + ------- + dict + The serialized model state. + + Raises + ------ + RuntimeError + If serialization of any component fails. + """ dd = BaseAtomicModel.serialize(self) + try: + descriptor_state = self.descriptor.serialize() + fitting_state = self.fitting_net.serialize() + except Exception as e: + raise RuntimeError("Failed to serialize model components") from e dd.update( { "@class": "Model", - "@version": 2, + "@version": 3, # Bump version for new 'fitting' field "type": "standard", "type_map": self.type_map, - "descriptor": self.descriptor.serialize(), - "fitting": self.fitting_net.serialize(), + "descriptor": descriptor_state, + "fitting": fitting_state, } ) return dd
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (4)
deepmd/dpmodel/atomic_model/make_base_atomic_model.py(1 hunks)deepmd/dpmodel/model/base_model.py(1 hunks)deepmd/pt/model/atomic_model/dp_atomic_model.py(1 hunks)deepmd/pt/model/model/make_model.py(1 hunks)
🔇 Additional comments (3)
deepmd/dpmodel/atomic_model/make_base_atomic_model.py (1)
151-175: LGTM! The compression API implementation looks good.
The new enable_compression method:
- Is well-integrated into the existing abstract base class
- Follows the established pattern for abstract methods
- Has comprehensive parameter documentation
- Properly raises NotImplementedError as a placeholder for derived classes
deepmd/pt/model/atomic_model/dp_atomic_model.py (1)
163-193: Verify compression feature usage across codebase.
Let's analyze how this compression feature is used across the codebase to ensure consistent implementation.
✅ Verification successful
Based on the comprehensive search results, I can now provide a final response regarding the compression feature usage across the codebase:
Compression implementation is consistent across frameworks
The compression feature is consistently implemented across TensorFlow, PyTorch and JAX frameworks with matching parameters and behavior:
- All descriptors (se_a, se_t, se_r, se_atten) support compression with identical parameters
- Common parameters across implementations:
min_nbor_dist: Minimum neighbor distancetable_extrapolate: Model extrapolation scale (default: 5)table_stride_1: First table stride (default: 0.01)table_stride_2: Second table stride (default: 0.1)check_frequency: Overflow check frequency (optional)
The implementation in dp_atomic_model.py correctly follows this pattern by delegating to the descriptor's compression functionality with all required parameters.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check usage patterns of compression-related functionality
# Test 1: Find all calls to enable_compression
echo "=== Searching for enable_compression calls ==="
rg "enable_compression" -A 5
# Test 2: Find all compression-related configuration settings
echo "=== Searching for compression configuration ==="
rg -g "*.py" -g "*.json" "table_stride|table_extrapolate|min_nbor_dist"
# Test 3: Find related test files
echo "=== Searching for related test files ==="
fd -g "*test*.py" | xargs rg "compression|enable_compression"
Length of output: 62514
deepmd/pt/model/model/make_model.py (1)
101-128: Improve method documentation and add input validation
The implementation looks good, but there are a few improvements needed:
-
The docstring should be enhanced to:
- Explain the method's purpose and return value
- Document the internally used
get_min_nbor_dist()parameter
-
Consider adding input parameter validation:
- Ensure table_extrapolate is positive
- Verify table_stride_1 and table_stride_2 are positive
- Validate check_frequency is either -1 or positive
Here's a suggested improvement:
def enable_compression(
self,
table_extrapolate: float = 5,
table_stride_1: float = 0.01,
table_stride_2: float = 0.1,
check_frequency: int = -1,
) -> None:
- """Call atomic_model enable_compression().
+ """Enable model compression to reduce memory usage and potentially improve performance.
+
+ This method configures compression settings by calling the atomic_model's enable_compression
+ with the minimum neighbor distance and the provided parameters.
Parameters
----------
table_extrapolate
- The scale of model extrapolation
+ The scale of model extrapolation (must be positive)
table_stride_1
- The uniform stride of the first table
+ The uniform stride of the first table (must be positive)
table_stride_2
- The uniform stride of the second table
+ The uniform stride of the second table (must be positive)
check_frequency
- The overflow check frequency
+ The overflow check frequency (-1 or positive integer)
+
+ Returns
+ -------
+ None
+
+ Raises
+ ------
+ ValueError
+ If any of the input parameters are invalid
"""
+ if table_extrapolate <= 0:
+ raise ValueError("table_extrapolate must be positive")
+ if table_stride_1 <= 0:
+ raise ValueError("table_stride_1 must be positive")
+ if table_stride_2 <= 0:
+ raise ValueError("table_stride_2 must be positive")
+ if check_frequency != -1 and check_frequency < 0:
+ raise ValueError("check_frequency must be -1 or positive")
+
self.atomic_model.enable_compression(
self.get_min_nbor_dist(),
table_extrapolate,
table_stride_1,
table_stride_2,
check_frequency,
)Let's verify the existence of the get_min_nbor_dist method:
…epmd-kit into addCompressionAPI
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (2)
deepmd/dpmodel/atomic_model/linear_atomic_model.py (1)
160-174: Enhance parameter documentation.The docstring should specify:
- Units for distance-based parameters
- Valid ranges for each parameter
- The meaning of special values (e.g., check_frequency=-1)
Consider updating the docstring:
"""Compress model. Parameters ---------- min_nbor_dist - The nearest distance between atoms + The minimum distance between atoms (in Angstroms). Must be positive. table_extrapolate - The scale of model extrapolation + The scale of model extrapolation. Must be positive, typically 5-10. table_stride_1 - The uniform stride of the first table + The uniform stride of the first table (in Angstroms). Must be positive, typically 0.01-0.1. table_stride_2 - The uniform stride of the second table + The uniform stride of the second table (in Angstroms). Must be positive, typically 0.1-1.0. check_frequency - The overflow check frequency + The frequency to check for overflow. Use -1 to disable checking, or a positive integer for check interval. """deepmd/pt/model/atomic_model/linear_atomic_model.py (1)
195-209: Enhance compression method documentationThe documentation should be expanded to include:
- A description of what compression does and its impact on the model
- The Returns section (even if it returns None)
- Example usage
- Any side effects or important notes about the compression process
"""Compress model. + This method enables compression for all sub-models in the linear combination. + Compression helps reduce model size and improve inference performance. Parameters ---------- min_nbor_dist The nearest distance between atoms table_extrapolate The scale of model extrapolation table_stride_1 The uniform stride of the first table table_stride_2 The uniform stride of the second table check_frequency The overflow check frequency + + Returns + ------- + None + + Examples + -------- + >>> model = LinearEnergyAtomicModel(...) + >>> model.enable_compression( + ... min_nbor_dist=0.5, + ... table_extrapolate=5.0, + ... table_stride_1=0.01, + ... table_stride_2=0.1 + ... ) + + Notes + ----- + This operation modifies the internal state of all sub-models. """
📜 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)
Summary by CodeRabbit
New Features
enable_compressionmethod across multiple classes to allow configuration of compression settings for descriptors.DPAtomicModelclass.enable_compressionmethod toLinearEnergyAtomicModelfor improved model compression capabilities.Bug Fixes
fitting_output_defmethod to ensure fallback functionality when the fitting network is unavailable.These updates enhance the functionality and reliability of the model management features.