(breaking change)feat: Support element type-wise bias in property fitting#5322
Open
Chengqian-Zhang wants to merge 10 commits intodeepmodeling:masterfrom
Open
(breaking change)feat: Support element type-wise bias in property fitting#5322Chengqian-Zhang wants to merge 10 commits intodeepmodeling:masterfrom
Chengqian-Zhang wants to merge 10 commits intodeepmodeling:masterfrom
Conversation
Chengqian-Zhang
changed the title
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a boolean distinguish_types flag to PropertyFittingNet, threads it through atomic models and stat utilities, bumps serialization version, changes per-key bias/std indexing to be per-type when enabled, and introduces an intensive mode affecting natoms normalization in output-stat computations. Changes
Sequence DiagramsequenceDiagram
participant Config as Config\n(ARGs)
participant FittingNet as PropertyFittingNet
participant Atomic as AtomicModel
participant Stat as StatUtils
Config->>FittingNet: Initialize(distinguish_types)
Note over FittingNet: store & serialize distinguish_types (`@version`=6)
Atomic->>FittingNet: get_distinguish_types()
FittingNet-->>Atomic: return bool
alt distinguish_types == true
Atomic->>Stat: apply_out_stat per-type bias (out_bias[kk][atype]) + out_std
else
Atomic->>Stat: apply_out_stat global bias (out_bias[kk][0]) + out_std
end
Stat->>Stat: compute/unbias/reconstruct (intensive toggles natoms normalization)
Stat-->>Atomic: adjusted outputs (bias/std applied)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/dpmodel/utils/stat.py (1)
481-487:⚠️ Potential issue | 🟠 MajorUse the intensive reduction for
model_predtoo.When
intensive=True, the regression now works on type fractions, butmodel_pred_gis still formed a few lines above with a rawnp.sum(..., axis=1). Inchange-by-statisticmode that subtracts/adds an extensive prediction against an intensive label, so the fitted delta bias is off by the frame atom count. Apply the same intensive reduction tomodel_predbefore bothstats_inputconstruction and this reconstruction path.Also applies to: 497-506
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/stat.py` around lines 481 - 487, The model currently sums per-atom predictions into model_pred_g with a raw np.sum(..., axis=1), but when intensive=True the rest of the stats pipeline works on per-type fractions; therefore apply the same intensive reduction used for stats_input to model_pred (model_pred_g) so it becomes intensive (e.g., divide/type-normalize or use the same per-type reduction function) before constructing stats_input and likewise apply the intensive reduction in the reconstruction path where bias_atom_e/std_atom_e are computed (the block that calls compute_stats_from_redu and the earlier block that forms model_pred_g and stats_input). Ensure you reuse the same reduction logic/utility used for stats_input so both the stats construction and the bias/std reconstruction use identical intensive aggregation when intensive=True.deepmd/pt/utils/stat.py (1)
517-523:⚠️ Potential issue | 🟠 MajorNormalize
model_predthe same way as the intensive label.This path now regresses on type fractions, but the global
model_predis still built from a raw atomic sum. Inchange-by-statisticmode that means the PT backend subtracts/adds an extensive prediction against an intensive label, so the learned delta bias is scaled incorrectly. Reducemodel_predwith the same intensive rule before both the bias solve and this reconstruction.Also applies to: 531-540
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt/utils/stat.py` around lines 517 - 523, The model_pred used in the bias solve and subsequent reconstruction is still an extensive raw atomic sum while the regression now targets intensive (type-fraction) labels; before calling compute_stats_from_redu for both the bias solve (the call that sets bias_atom_e/std_atom_e) and the reconstruction block (also at ~531-540), reduce/normalize model_pred with the same intensive rule used for labels (respecting the intensive flag and merged_natoms/assigned_atom_ener logic) so that the inputs to compute_stats_from_redu and the later reconstruction use the intensive-scaled model_pred; update the code paths that build model_pred (and any variables passed as stats_input) to apply that reduction whenever intensive is True so the change-by-statistic delta is scaled correctly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/atomic_model/property_atomic_model.py`:
- Around line 54-59: The condition incorrectly uses the bound method
self.get_compute_stats_distinguish_types without calling it and the per-type std
is not used; change the branch to call
self.get_compute_stats_distinguish_types() and, inside the loop over
self.bias_keys in the distinguished branch (where atype is available), index
out_std by atype (use out_std[kk][atype]) to match out_bias[kk][atype], so each
ret[kk] uses the per-type std and bias.
In `@deepmd/dpmodel/fitting/property_fitting.py`:
- Around line 93-99: Change the constructor parameter distinguish_types to be
optional (bool | None = None) in PropertyFitting (the __init__ where
task_dim/intensive/distinguish_types are set and the related block at lines
~136-147) so we can detect "missing" vs explicit True/False; set
self.distinguish_types to the explicit boolean when provided, but during
deserialization of older payloads (the load/from_dict/from_json routine that
handles payload versions) explicitly backfill distinguish_types=False for pre-v6
payloads; also update deepmd/utils/argcheck.py to use the same optional
default/behavior so both creation and arg checking are aligned.
In `@deepmd/pt/model/atomic_model/property_atomic_model.py`:
- Around line 52-57: The condition currently checks the bound method object
instead of calling it and uses a hard-coded out_std[kk][0] even when applying
per-type bias; change the predicate check to call
self.get_compute_stats_distinguish_types() and make the std scaling type-aware
in the true branch by using out_std[kk][atype] (while keeping the false branch
using out_std[kk][0] and out_bias[kk][0]). Update the loop over self.bias_keys
where ret[kk] is computed to reference atype for both out_std and out_bias when
the predicate returns True, and keep the existing [0] indexing for both when the
predicate returns False.
In `@deepmd/utils/out_stat.py`:
- Around line 50-51: Before performing the intensive normalization when
intensive is True, compute totals = np.sum(natoms, axis=1) and check for any
zero entries; if any zero-total frames exist, raise a ValueError listing the
offending frame indices (e.g., using np.where(totals == 0)[0]) instead of
dividing by zero, otherwise proceed with natoms = natoms / totals[:, None] so
broadcasting is explicit; reference natoms and the intensive branch in
out_stat.py.
In `@source/tests/common/test_out_stat.py`:
- Around line 145-169: The test test_compute_stats_from_redu_intensive is
feeding an extensive label (self.output_redu) into the intensive path; change
the label passed to compute_stats_from_redu to an actual intensive target (e.g.,
self.output.mean(axis=1) or self.output_redu / self.natoms.sum(axis=1,
keepdims=True)) when calling compute_stats_from_redu(..., intensive=True), and
update the assertions to compare the recovered bias against self.mean (or the
corresponding intensive reference) rather than the huge coefficients currently
asserted.
---
Outside diff comments:
In `@deepmd/dpmodel/utils/stat.py`:
- Around line 481-487: The model currently sums per-atom predictions into
model_pred_g with a raw np.sum(..., axis=1), but when intensive=True the rest of
the stats pipeline works on per-type fractions; therefore apply the same
intensive reduction used for stats_input to model_pred (model_pred_g) so it
becomes intensive (e.g., divide/type-normalize or use the same per-type
reduction function) before constructing stats_input and likewise apply the
intensive reduction in the reconstruction path where bias_atom_e/std_atom_e are
computed (the block that calls compute_stats_from_redu and the earlier block
that forms model_pred_g and stats_input). Ensure you reuse the same reduction
logic/utility used for stats_input so both the stats construction and the
bias/std reconstruction use identical intensive aggregation when intensive=True.
In `@deepmd/pt/utils/stat.py`:
- Around line 517-523: The model_pred used in the bias solve and subsequent
reconstruction is still an extensive raw atomic sum while the regression now
targets intensive (type-fraction) labels; before calling compute_stats_from_redu
for both the bias solve (the call that sets bias_atom_e/std_atom_e) and the
reconstruction block (also at ~531-540), reduce/normalize model_pred with the
same intensive rule used for labels (respecting the intensive flag and
merged_natoms/assigned_atom_ener logic) so that the inputs to
compute_stats_from_redu and the later reconstruction use the intensive-scaled
model_pred; update the code paths that build model_pred (and any variables
passed as stats_input) to apply that reduction whenever intensive is True so the
change-by-statistic delta is scaled correctly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 67e04055-7dc9-448e-aece-69b8d660d36c
📒 Files selected for processing (9)
deepmd/dpmodel/atomic_model/property_atomic_model.pydeepmd/dpmodel/fitting/property_fitting.pydeepmd/dpmodel/utils/stat.pydeepmd/pt/model/atomic_model/property_atomic_model.pydeepmd/pt/model/task/property.pydeepmd/pt/utils/stat.pydeepmd/utils/argcheck.pydeepmd/utils/out_stat.pysource/tests/common/test_out_stat.py
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
deepmd/pt/model/atomic_model/property_atomic_model.py (1)
52-57:⚠️ Potential issue | 🟠 MajorUse type-aware
out_stdin the distinguish-types branch.In the
Truebranch, bias is type-indexed (out_bias[kk][atype]) but std is still fixed to type 0 (out_std[kk][0]). This can silently ignore per-type std scaling.Suggested fix
if self.get_compute_stats_distinguish_types(): for kk in self.bias_keys: - ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][atype] + ret[kk] = ret[kk] * out_std[kk][atype] + out_bias[kk][atype] else: for kk in self.bias_keys: ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt/model/atomic_model/property_atomic_model.py` around lines 52 - 57, The branch inside get_compute_stats_distinguish_types() incorrectly uses out_std[kk][0] while indexing out_bias by atype, causing per-type std to be ignored; update the distinguish-types branch so each kk uses out_std[kk][atype] (matching out_bias[kk][atype]) when computing ret[kk] = ret[kk] * out_std[kk][atype] + out_bias[kk][atype] so per-type scaling is applied consistently for bias_keys.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@deepmd/pt/model/atomic_model/property_atomic_model.py`:
- Around line 52-57: The branch inside get_compute_stats_distinguish_types()
incorrectly uses out_std[kk][0] while indexing out_bias by atype, causing
per-type std to be ignored; update the distinguish-types branch so each kk uses
out_std[kk][atype] (matching out_bias[kk][atype]) when computing ret[kk] =
ret[kk] * out_std[kk][atype] + out_bias[kk][atype] so per-type scaling is
applied consistently for bias_keys.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 3e24f042-92e8-4de4-9ba6-5818891b4d97
📒 Files selected for processing (2)
deepmd/dpmodel/atomic_model/property_atomic_model.pydeepmd/pt/model/atomic_model/property_atomic_model.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/dpmodel/atomic_model/property_atomic_model.py
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
deepmd/pd/utils/stat.py (2)
156-159: Consider adding a shape assertion to guard against unexpected std shapes.The tiling logic assumes
out_std[kk]has shape(1, ...)or is otherwise broadcastable. Ifout_std[kk]unexpectedly has a partial type dimension (e.g., shape(k, odim)where1 < k < ntypes),np.tilewould produce an incorrect shape(k * ntypes, odim)instead of(ntypes, odim).🛡️ Proposed defensive assertion
else: ntypes = vv.shape[0] + # Expect out_std to be global (no type dim) or have leading dim of 1 + assert out_std[kk].shape[0] == 1 or out_std[kk].ndim < vv.ndim, ( + f"Unexpected std shape {out_std[kk].shape} for bias shape {vv.shape}" + ) reps = [ntypes] + [1] * (vv.ndim - 1) new_std[kk] = np.tile(out_std[kk], reps)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/utils/stat.py` around lines 156 - 159, Add a defensive shape check before tiling in the block that computes new_std[kk]: verify the leading dimension of out_std[kk] is either 1 or equals ntypes (where ntypes = vv.shape[0]) and raise a clear AssertionError or ValueError if it is not; if it is 1, proceed to construct reps and call np.tile as currently done, and if it already equals ntypes skip tiling and assign directly to new_std[kk] (use the symbols out_std, new_std, vv, ntypes, reps, np.tile to locate the code).
140-143: Type hints do not match actual usage.The function signature declares
paddle.Tensorparameters and return types, but the function actually operates ondict[str, np.ndarray](as seen from the call at line 513 and the use ofnp.tileat line 159).📝 Suggested type hint correction
def _post_process_stat( - out_bias: paddle.Tensor, - out_std: paddle.Tensor, -) -> tuple[paddle.Tensor, paddle.Tensor]: + out_bias: dict[str, np.ndarray], + out_std: dict[str, np.ndarray], +) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/utils/stat.py` around lines 140 - 143, The _post_process_stat function currently types its parameters and return as paddle.Tensor but actually accepts and returns dict[str, np.ndarray] (it uses np.tile and is called with dicts), so change the signature of _post_process_stat(out_bias, out_std) to use dict[str, np.ndarray] for both parameters and the returned tuple; update any imports or typing annotations accordingly and ensure callers (the code that constructs/passes those dicts) and internal usages (e.g., np.tile) remain consistent with the new dict[str, np.ndarray] types.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@deepmd/pd/utils/stat.py`:
- Around line 156-159: Add a defensive shape check before tiling in the block
that computes new_std[kk]: verify the leading dimension of out_std[kk] is either
1 or equals ntypes (where ntypes = vv.shape[0]) and raise a clear AssertionError
or ValueError if it is not; if it is 1, proceed to construct reps and call
np.tile as currently done, and if it already equals ntypes skip tiling and
assign directly to new_std[kk] (use the symbols out_std, new_std, vv, ntypes,
reps, np.tile to locate the code).
- Around line 140-143: The _post_process_stat function currently types its
parameters and return as paddle.Tensor but actually accepts and returns
dict[str, np.ndarray] (it uses np.tile and is called with dicts), so change the
signature of _post_process_stat(out_bias, out_std) to use dict[str, np.ndarray]
for both parameters and the returned tuple; update any imports or typing
annotations accordingly and ensure callers (the code that constructs/passes
those dicts) and internal usages (e.g., np.tile) remain consistent with the new
dict[str, np.ndarray] types.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 46e28ca9-002d-4a7c-88c7-734d7d066335
📒 Files selected for processing (1)
deepmd/pd/utils/stat.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5322 +/- ##
==========================================
- Coverage 82.33% 82.33% -0.01%
==========================================
Files 777 777
Lines 77814 77854 +40
Branches 3680 3678 -2
==========================================
+ Hits 64068 64099 +31
- Misses 12574 12581 +7
- Partials 1172 1174 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/pt/model/test_atomic_model_global_stat.py (1)
233-233: Consider adding spaces after commas for consistency with PEP 8.The array literal lacks spaces after commas, which may be flagged by
ruff format.🔧 Suggested formatting
- expected_std = np.array([[[0,1],[0,1]],[[1,1],[1,1]],[[0,0],[0,0]]]) # 3 keys, 2 atypes, 2 max dims. + expected_std = np.array([[[0, 1], [0, 1]], [[1, 1], [1, 1]], [[0, 0], [0, 0]]]) # 3 keys, 2 atypes, 2 max dims.As per coding guidelines: "Always run
ruff check .andruff format .before committing changes or CI will fail."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt/model/test_atomic_model_global_stat.py` at line 233, The array literal assigned to expected_std is missing spaces after commas and should be formatted to comply with PEP 8/ruff (e.g., add spaces after each comma in the np.array([...]) literal); update the expected_std definition to include spaces after commas and then run `ruff format .` (or `ruff check .`) to ensure the file passes style checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@source/tests/consistent/model/test_property.py`:
- Around line 1208-1210: Update the misleading comment above the assertion that
checks dp_observed: change the comment referencing "both types appear observed"
to reflect that only type 0 is present so dp_observed should be ["O"]
(mentioning stats_distinguish_types=False and the dp_observed variable so you
can find the block in the test for the property model).
---
Nitpick comments:
In `@source/tests/pt/model/test_atomic_model_global_stat.py`:
- Line 233: The array literal assigned to expected_std is missing spaces after
commas and should be formatted to comply with PEP 8/ruff (e.g., add spaces after
each comma in the np.array([...]) literal); update the expected_std definition
to include spaces after commas and then run `ruff format .` (or `ruff check .`)
to ensure the file passes style checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: b5919558-d0ac-4257-af68-ad702ea15fb3
📒 Files selected for processing (10)
deepmd/dpmodel/atomic_model/property_atomic_model.pydeepmd/pt/model/atomic_model/property_atomic_model.pysource/tests/common/dpmodel/test_atomic_model_atomic_stat.pysource/tests/common/dpmodel/test_atomic_model_global_stat.pysource/tests/common/test_out_stat.pysource/tests/consistent/model/test_property.pysource/tests/pd/model/test_atomic_model_atomic_stat.pysource/tests/pd/model/test_atomic_model_global_stat.pysource/tests/pt/model/test_atomic_model_atomic_stat.pysource/tests/pt/model/test_atomic_model_global_stat.py
Member
|
pre-commit.ci autofix |
Chengqian-Zhang
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Introduced the
distinguish_typesoption for property fitting to support more flexible atom contribution calculations.distinguish_types=True(Default) has consistent behavior with energy fitting(if the property is extensive), calculating individual atom contributions based on each specific atom type.distinguish_types=Falsetreats atom contributions as element-agnostic. This is particularly useful for element extrapolation scenarios where the model needs to generalize across different chemical species.Key changes
deepmd/pt/model/task/property.py:distinguish_typesparameter.deepmd/pt/utils/stat.py:output_std: Replaced from the constant "ones" padding with actual standard deviation values(required by functionapply_out_statindeepmd/pt/model/atomic_model/property_atomic_model.py)deepmd/utils/out_stat.py:compute_stats_from_reduto include theintensiveparameter.Summary by CodeRabbit
New Features
Bug Fixes
Tests