deepmd/pt/model/task/fitting.py (2)

437-440: Use a ternary operator for compactness.
Ruff suggests replacing the if callable(...) block with a ternary operator. This is a minor readability enhancement.

- if callable(merged):
-     sampled = merged()
- else:
-     sampled = merged
+ sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.2)

437-440: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

Replace if-else-block with sampled = merged() if callable(merged) else merged

(SIM108)

---

457-457: Implement aparam statistics.
The TODO suggests you plan to handle aparam similarly to fparam. Let me know if you’d like help implementing that.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/pt/model/task/fitting.py (1)

442-456: ⚠️ Potential issue

Add protection against zero standard deviation.

The division by fparam_std could lead to numerical instability or inf values when the standard deviation is zero or very small.

Apply this diff to add protection:

 fparam_std = torch.std(cat_data, axis=0)
+epsilon = 1e-12  # Small constant to prevent division by zero
+fparam_std = torch.where(fparam_std < epsilon, torch.tensor(epsilon, dtype=fparam_std.dtype, device=fparam_std.device), fparam_std)
 fparam_inv_std = 1.0 / fparam_std

deepmd/pt/model/task/fitting.py (2)

456-457: Implement aparam statistics calculation.

The TODO comment indicates missing implementation for aparam statistics, which is needed for complete functionality as indicated by the PR objectives.

Would you like me to generate the implementation for aparam statistics calculation? It would follow a similar pattern to the fparam implementation but handle the different dimensionality of aparam.

---

437-440: Simplify conditional assignment using ternary operator.

The if-else block can be simplified using a ternary operator.

Apply this diff to improve code style:

-if callable(merged):
-    sampled = merged()
-else:
-    sampled = merged
+sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.2)

437-440: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

Replace if-else-block with sampled = merged() if callable(merged) else merged

(SIM108)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/pt/model/task/fitting.py (1)

99-102: Simplify the code using a ternary operator.

The code can be more concise using a ternary operator.

-        if callable(merged):
-            sampled = merged()
-        else:
-            sampled = merged
+        sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.2)

99-102: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

Replace if-else-block with sampled = merged() if callable(merged) else merged

(SIM108)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/pt/model/task/fitting.py (2)

103-116: ⚠️ Potential issue

Add protection against division by zero.

The standard deviation calculation needs protection against zero or near-zero values.

Apply this diff to handle potential division by zero:

 fparam_std = torch.std(cat_data, dim=0, unbiased=False)
+epsilon = 1e-12
+fparam_std = torch.where(fparam_std < epsilon, torch.tensor(epsilon, dtype=fparam_std.dtype, device=fparam_std.device), fparam_std)
 fparam_inv_std = 1.0 / fparam_std

---

118-140: ⚠️ Potential issue

Add protection against division by zero in aparam calculations.

Similar to fparam, the aparam standard deviation calculation needs protection.

Apply this diff to handle potential division by zero:

 aparam_std = torch.sqrt(sumv2 / sumn - (sumv / sumn) ** 2)
+epsilon = 1e-12
+aparam_std = torch.where(aparam_std < epsilon, torch.tensor(epsilon, dtype=aparam_std.dtype, device=aparam_std.device), aparam_std)
 aparam_inv_std = 1.0 / aparam_std

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

71-93: Add more test cases for edge conditions.

While the current test case validates the basic functionality, consider adding tests for:

• Empty data
• Zero standard deviation
• Single frame/atom scenarios

Would you like me to generate additional test cases to improve coverage?

deepmd/pt/model/task/fitting.py (1)

98-101: Simplify conditional with ternary operator.

The if-else block can be simplified.

-if callable(merged):
-    sampled = merged()
-else:
-    sampled = merged
+sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.2)

98-101: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

Replace if-else-block with sampled = merged() if callable(merged) else merged

(SIM108)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/pt/model/task/fitting.py (1)

101-101: 🛠️ Refactor suggestion

Add input validation for empty data.

The method should validate the input data before processing to ensure robustness.

         sampled = merged
+        if not sampled:
+            raise ValueError("No data samples provided")
+        if self.numb_fparam > 0 and not all("fparam" in frame for frame in sampled):
+            raise ValueError("Missing 'fparam' in some data samples")
+        if self.numb_aparam > 0 and not all("aparam" in frame for frame in sampled):
+            raise ValueError("Missing 'aparam' in some data samples")

🧰 Tools

🪛 Ruff (0.8.2)

98-101: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

Replace if-else-block with sampled = merged() if callable(merged) else merged

(SIM108)

deepmd/pt/model/task/fitting.py (1)

98-101: Simplify the code using a ternary operator.

The code can be more concise by using a ternary operator.

-        if callable(merged):
-            sampled = merged()
-        else:
-            sampled = merged
+        sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.2)

98-101: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

Replace if-else-block with sampled = merged() if callable(merged) else merged

(SIM108)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/pt/model/task/fitting.py (2)

95-98: Simplify the callable check using a ternary operator.

The code can be more concise.

-        if callable(merged):
-            sampled = merged()
-        else:
-            sampled = merged
+        sampled = merged() if callable(merged) else merged

🧰 Tools

🪛 Ruff (0.8.2)

95-98: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block

Replace if-else-block with sampled = merged() if callable(merged) else merged

(SIM108)

---

122-151: Add a comment explaining the aparam statistics computation approach.

The implementation is correct, but it would be helpful to add a comment explaining why the statistics are computed differently for atomic parameters compared to frame parameters.

         # stat aparam
         if self.numb_aparam > 0:
+            # Computing statistics for atomic parameters requires accumulating sums
+            # across all systems due to varying number of atoms per frame
             sys_sumv = []
             sys_sumv2 = []
             sys_sumn = []

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro