refact(dpmodel): model output made the same as pt backend#5250
refact(dpmodel): model output made the same as pt backend#5250wanghan-iapcm merged 6 commits intodeepmodeling:masterfrom
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| model_predict["dipole"] = model_ret["dipole"] | ||
| model_predict["global_dipole"] = model_ret["dipole_redu"] | ||
| if self.do_grad_r("dipole") and model_ret["dipole_derv_r"] is not None: | ||
| model_predict["force"] = model_ret["dipole_derv_r"] | ||
| if self.do_grad_c("dipole") and model_ret["dipole_derv_c_redu"] is not None: | ||
| model_predict["virial"] = model_ret["dipole_derv_c_redu"] | ||
| if do_atomic_virial and model_ret["dipole_derv_c"] is not None: | ||
| model_predict["atom_virial"] = model_ret["dipole_derv_c"] | ||
| if "mask" in model_ret: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| coord, | ||
| atype, | ||
| box, | ||
| fparam=fparam, | ||
| aparam=aparam, | ||
| do_atomic_virial=do_atomic_virial, | ||
| ) | ||
| model_predict = {} | ||
| model_predict["atom_dos"] = model_ret["dos"] |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| model_predict["atom_energy"] = model_ret["energy"] | ||
| model_predict["energy"] = model_ret["energy_redu"] | ||
| if self.do_grad_r("energy") and model_ret["energy_derv_r"] is not None: | ||
| model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) | ||
| if self.do_grad_c("energy") and model_ret["energy_derv_c_redu"] is not None: | ||
| model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) | ||
| if do_atomic_virial and model_ret["energy_derv_c"] is not None: | ||
| model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2) | ||
| if "mask" in model_ret: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| fparam=fparam, | ||
| aparam=aparam, | ||
| do_atomic_virial=do_atomic_virial, | ||
| ) | ||
| model_predict = {} | ||
| model_predict["atom_energy"] = model_ret["energy"] | ||
| model_predict["energy"] = model_ret["energy_redu"] | ||
| if self.do_grad_r("energy") and model_ret.get("energy_derv_r") is not None: | ||
| model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| coord, | ||
| atype, | ||
| box, | ||
| fparam=fparam, | ||
| aparam=aparam, | ||
| do_atomic_virial=do_atomic_virial, | ||
| ) | ||
| model_predict = {} | ||
| model_predict["polar"] = model_ret["polarizability"] |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
| ) -> dict[str, Array]: | ||
| model_ret = self.call_common( | ||
| coord, | ||
| atype, | ||
| box, | ||
| fparam=fparam, | ||
| aparam=aparam, | ||
| do_atomic_virial=do_atomic_virial, | ||
| ) |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bea33efe80
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
📝 WalkthroughWalkthroughAdds public high-level call APIs and translated output definitions across many DPModel classes, standardizes backend key resolution in evaluators/serializers, tweaks tensor squeeze axes and mask propagation in PyTorch models, and updates numerous tests to new key/shape conventions. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 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)
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}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
source/tests/jax/test_dp_hessian_model.py (1)
109-114: Minor coverage gap:hessianis not re-checked in thedo_atomic_virial=Trueblock.The second invocation (lines 109-110) exercises
do_atomic_virial=True, but only asserts"atom_virial". Sinceenable_hessian()remains active, the hessian output is still produced. Asserting it here would confirm no interaction betweendo_atomic_virialand hessian computation.💡 Suggested addition
np.testing.assert_allclose( to_numpy_array(ret0["atom_virial"]), to_numpy_array(ret1["atom_virial"]), atol=self.atol, ) + np.testing.assert_allclose( + to_numpy_array(ret0["hessian"]), + to_numpy_array(ret1["hessian"]), + atol=self.atol, + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/jax/test_dp_hessian_model.py` around lines 109 - 114, The do_atomic_virial=True call currently only asserts atom_virial equality; add an assertion that the Hessian outputs from md0.call and md1.call are equal as well to ensure enable_hessian() still produces identical results when do_atomic_virial is enabled—specifically compare ret0["hessian"] to ret1["hessian"] (using the same to_numpy_array wrapper and np.testing.assert_allclose with self.atol) after the existing atom_virial check.source/tests/jax/test_make_hessian_model.py (2)
169-176:test_output_defhas no coverage for the new public"hessian"key.The test still asserts only
model_output_def()["energy_derv_r_derv_r"](the internal key). Now thatcall()returns"hessian"as the public name, consider adding an assertion that the translated/public output def also exposes this key (e.g., viatranslated_output_def()["hessian"]), ensuring the public API contract is verified alongside the internal one.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/jax/test_make_hessian_model.py` around lines 169 - 176, Update the test_output_def to assert the public "hessian" key is present and correctly described by the translated/public output definition: call translated_output_def() on model_hess and add assertions that translated_output_def()["hessian"].category equals OutputVariableCategory.DERV_R_DERV_R and that its r_hessian flag is true (analogous to the existing checks for the internal "energy_derv_r_derv_r" and r_hessian on "energy"); this verifies the public API mapping produced by call() in addition to the internal key.
74-74:nvis now dead code.After the line 130 reshape no longer references
nv, the assignmentnv = self.nvat line 74 is unreferenced throughoutHessianTest.test(). Consider removing both this assignment andself.nv = 1inTestDPModel.setUp().♻️ Proposed cleanup
- nv = self.nv- self.nv = 1🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/jax/test_make_hessian_model.py` at line 74, Remove the now-unused local assignment nv = self.nv in HessianTest.test() and also drop the redundant attribute initialization self.nv = 1 in TestDPModel.setUp(); locate the two occurrences by searching for the exact symbols "nv = self.nv" and "self.nv = 1" and delete them so the tests no longer set or reference the dead nv variable.source/tests/consistent/model/test_dpa1.py (1)
232-239: DP backend skips force/virial cross-validation.Three
SKIP_FLAGs at positions [2], [3], [4] mean force, virial, and atom_virial are never compared against the DP backend. If this is intentional (e.g., DP backend doesn't support force output in this context), a brief inline comment would make the intent clear for future contributors.💡 Suggested comment to document intent
elif backend is self.RefBackend.DP: return ( ret["energy"].ravel(), ret["atom_energy"].ravel(), - SKIP_FLAG, - SKIP_FLAG, - SKIP_FLAG, + SKIP_FLAG, # DP backend does not compute force + SKIP_FLAG, # DP backend does not compute virial + SKIP_FLAG, # DP backend does not compute atom_virial )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_dpa1.py` around lines 232 - 239, The DP backend branch (the elif checking self.RefBackend.DP) currently returns three SKIP_FLAGs for force, virial and atom_virial which hides intent; add a brief inline comment right above that return in test_dpa1.py explaining that the DP reference backend does not provide force/virial/atom_virial outputs in this test context (or note if this is a temporary limitation), so SKIP_FLAG is deliberately used — reference the branch condition self.RefBackend.DP and the SKIP_FLAG placeholders when adding the comment.deepmd/dpmodel/model/dp_zbl_model.py (1)
70-102: Near-duplicate ofEnergyModel.call_lower()— consider extracting shared logic.The
call(),call_lower(), andtranslated_output_def()bodies inDPZBLModelare nearly identical toEnergyModel(minus hessian). A shared helper or mixin would reduce the maintenance surface, but this can be deferred.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/model/dp_zbl_model.py` around lines 70 - 102, DPZBLModel and EnergyModel duplicate the same output-building logic in call(), call_lower(), and translated_output_def(); extract that logic into a shared helper (e.g., EnergyOutputMixin.build_predict_from_ret or BaseEnergyModel._build_model_predict_from_ret) that accepts model_ret, do_atomic_virial and the gradient-query helpers (do_grad_r/do_grad_c) and returns the model_predict dict; then replace the duplicated bodies in DPZBLModel.call_lower / call / translated_output_def and EnergyModel equivalents to call this helper, preserving EnergyModel-only behavior (like hessian handling) by applying any extra steps after the shared helper returns and add/adjust unit tests to cover both classes.
🤖 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/model/dipole_model.py`:
- Around line 87-96: call_lower in dipole_model.py builds model_predict but
never adds the "mask" key, causing inconsistency with call and
translated_output_def and breaking lower-level callers; update the call_lower
implementation to check for a mask (e.g., model_ret.get("mask") or similar
presence check used in call) and, when present, set model_predict["mask"] =
model_ret["mask"] (mirror the behavior in polar_model.py's call_lower) so mask
is propagated out of call_lower.
In `@deepmd/jax/infer/deep_eval.py`:
- Line 391: Validate that odef.name exists in the mapping before indexing
self._OUTDEF_DP2BACKEND: in deepmd/jax/infer/deep_eval.py around the line
assigning dp_name = self._OUTDEF_DP2BACKEND[odef.name], replace the direct
lookup with a defensive check (e.g., if odef.name not in
self._OUTDEF_DP2BACKEND: raise a clear ValueError listing the missing name and
available keys, or use self._OUTDEF_DP2BACKEND.get(odef.name, <fallback>) as
appropriate). Alternatively, perform this validation during initialization of
the mapping (where _OUTDEF_DP2BACKEND is populated or extended, see
deepmd/infer/deep_property.py registration logic) to ensure all model output
definitions are registered before inference; include the symbol
_OUTDEF_DP2BACKEND and odef.name in the error message to aid debugging.
---
Duplicate comments:
In `@deepmd/dpmodel/infer/deep_eval.py`:
- Line 361: The lookup self._OUTDEF_DP2BACKEND[odef.name] can raise
AttributeError/KeyError; update DeepEvalBackend to ensure the mapping attribute
exists and is exhaustive for all names in request_defs: either add a class-level
_OUTDEF_DP2BACKEND with all expected keys or validate at runtime before use
(e.g., check hasattr(self, "_OUTDEF_DP2BACKEND") and that odef.name in
self._OUTDEF_DP2BACKEND) and raise a clear, descriptive error if missing;
reference the symbols _OUTDEF_DP2BACKEND, odef.name, DeepEvalBackend and
request_defs when adding the verification or filling missing entries.
In `@deepmd/dpmodel/model/dp_zbl_model.py`:
- Around line 104-120: In translated_output_def the calls like
output_def["force"].squeeze(-2) on OutputVariableDef objects are no-ops because
the squeezed result isn't used; change them to assign the squeezed value back
(e.g., output_def["force"] = output_def["force"].squeeze(-2)) for each
occurrence (force, virial, atom_virial) so the squeezed OutputVariableDef
replaces the original; update the blocks that call self.do_grad_r("energy") and
self.do_grad_c("energy") to perform these assignments.
In `@deepmd/pt/model/model/dp_zbl_model.py`:
- Around line 44-52: The output_def is being mutated in-place by
OutputVariableDef.squeeze() for keys "force", "virial", and "atom_virial", which
can corrupt the shared cached ModelOutputDef; fix this by deepcopying the
assigned OutputVariableDef before calling squeeze (i.e., when setting
output_def["force"] = out_def_data["energy_derv_r"], output_def["virial"] =
out_def_data["energy_derv_c_redu"], and output_def["atom_virial"] =
out_def_data["energy_derv_c"]) so you call squeeze on the copy; mirror the same
deepcopy approach used in EnergyModel.translated_output_def to avoid
shared-state mutation.
---
Nitpick comments:
In `@deepmd/dpmodel/model/dp_zbl_model.py`:
- Around line 70-102: DPZBLModel and EnergyModel duplicate the same
output-building logic in call(), call_lower(), and translated_output_def();
extract that logic into a shared helper (e.g.,
EnergyOutputMixin.build_predict_from_ret or
BaseEnergyModel._build_model_predict_from_ret) that accepts model_ret,
do_atomic_virial and the gradient-query helpers (do_grad_r/do_grad_c) and
returns the model_predict dict; then replace the duplicated bodies in
DPZBLModel.call_lower / call / translated_output_def and EnergyModel equivalents
to call this helper, preserving EnergyModel-only behavior (like hessian
handling) by applying any extra steps after the shared helper returns and
add/adjust unit tests to cover both classes.
In `@source/tests/consistent/model/test_dpa1.py`:
- Around line 232-239: The DP backend branch (the elif checking
self.RefBackend.DP) currently returns three SKIP_FLAGs for force, virial and
atom_virial which hides intent; add a brief inline comment right above that
return in test_dpa1.py explaining that the DP reference backend does not provide
force/virial/atom_virial outputs in this test context (or note if this is a
temporary limitation), so SKIP_FLAG is deliberately used — reference the branch
condition self.RefBackend.DP and the SKIP_FLAG placeholders when adding the
comment.
In `@source/tests/jax/test_dp_hessian_model.py`:
- Around line 109-114: The do_atomic_virial=True call currently only asserts
atom_virial equality; add an assertion that the Hessian outputs from md0.call
and md1.call are equal as well to ensure enable_hessian() still produces
identical results when do_atomic_virial is enabled—specifically compare
ret0["hessian"] to ret1["hessian"] (using the same to_numpy_array wrapper and
np.testing.assert_allclose with self.atol) after the existing atom_virial check.
In `@source/tests/jax/test_make_hessian_model.py`:
- Around line 169-176: Update the test_output_def to assert the public "hessian"
key is present and correctly described by the translated/public output
definition: call translated_output_def() on model_hess and add assertions that
translated_output_def()["hessian"].category equals
OutputVariableCategory.DERV_R_DERV_R and that its r_hessian flag is true
(analogous to the existing checks for the internal "energy_derv_r_derv_r" and
r_hessian on "energy"); this verifies the public API mapping produced by call()
in addition to the internal key.
- Line 74: Remove the now-unused local assignment nv = self.nv in
HessianTest.test() and also drop the redundant attribute initialization self.nv
= 1 in TestDPModel.setUp(); locate the two occurrences by searching for the
exact symbols "nv = self.nv" and "self.nv = 1" and delete them so the tests no
longer set or reference the dead nv variable.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
deepmd/pt/model/model/polar_model.py (1)
73-75:⚠️ Potential issue | 🟡 MinorRemove the unreachable else branch in
forwardandforward_lower(lines 73–75 and 107–108).The else branch at line 73–75 is unreachable dead code.
DPPolarAtomicModel.__init__enforces that the fitting parameter must be aPolarFittingNetinstance (raisesTypeErrorotherwise), soget_fitting_net()can never returnNone. The same dead code exists inener_model.py,dos_model.py, anddipole_model.py. Additionally, the two branches are asymmetric:forwardline 75 attempts to mutateupdated_coord(which doesn't belong in a polar model), whereasforward_lowerline 108 correctly does not. Remove the entire else block, or if conditional logic is needed for safety, align both methods.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt/model/model/polar_model.py` around lines 73 - 75, The else branch in DPPolarAtomicModel.forward and forward_lower is dead/unreachable because DPPolarAtomicModel.__init__ enforces get_fitting_net() returns a PolarFittingNet (TypeError otherwise); remove the entire else block that mutates model_predict["updated_coord"] (the model_ret-to-model_predict fallback) so both forward and forward_lower only use the primary model_ret path, or if you prefer defensive coding keep a single guarded check that does not mutate fields inappropriate for polar models (do not add updated_coord). Update references in forward, forward_lower, and any similar patterns in ener_model.py, dos_model.py, and dipole_model.py to remove the unreachable else branch and ensure consistency.deepmd/pt/model/model/dipole_model.py (1)
72-87:⚠️ Potential issue | 🔴 CriticalAdd
.squeeze(-2)to force, virial, and atom_virial outputs in DipoleModel.forward().DipoleModel is missing the squeeze operations applied in all other derivative models (EnergyModel, DPZBLModel, DPLinearModel). Lines 77, 79, and 81 should squeeze the last dimension to match output shapes from other models:
Expected pattern (from EnergyModel)
if self.do_grad_r("dipole"): model_predict["force"] = model_ret["dipole_derv_r"].squeeze(-2) if self.do_grad_c("dipole"): model_predict["virial"] = model_ret["dipole_derv_c_redu"].squeeze(-2) if do_atomic_virial: model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(-2)Without this fix, DipoleModel outputs will have shape (N, 1, 3) and (N, 1, 9) instead of (N, 3) and (N, 9), causing downstream shape mismatches in tests and evaluators.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt/model/model/dipole_model.py` around lines 72 - 87, In DipoleModel.forward(), the tensor dims for force, virial and atom_virial are not squeezed and thus keep an extra length-1 axis; update the assignments so that model_predict["force"] uses model_ret["dipole_derv_r"].squeeze(-2), model_predict["virial"] uses model_ret["dipole_derv_c_redu"].squeeze(-2), and model_predict["atom_virial"] uses model_ret["dipole_derv_c"].squeeze(-2) (only when do_atomic_virial is true) so shapes match the other models (see EnergyModel/DPZBLModel pattern).source/tests/universal/dpmodel/model/test_model.py (1)
274-274:⚠️ Potential issue | 🟠 Majordpmodel's
TestSpinEnergyModelDPuses rawmodel_output_def()while PT's equivalent uses translated keys, missing parity.Line 167 (
TestEnergyModelDP) usescls.module.translated_output_def(), which returns human-friendly keys. Line 274 (TestSpinEnergyModelDP) still callscls.module.model_output_def().get_data(), returning raw internal keys. PT'sTestSpinEnergyModelDPusesSpinEnergyModelwithtranslated_output_def(), returning spin-specific keys (atom_energy,energy,mask_mag,force). dpmodel'sSpinModellacks atranslated_output_def()override—it would need one that maps spin-specific output keys (like PT'sSpinEnergyModel.translated_output_def()), or dpmodel needs a newSpinEnergyModelsubclass mirroring PT's design.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/universal/dpmodel/model/test_model.py` at line 274, TestSpinEnergyModelDP is calling cls.module.model_output_def().get_data() which returns raw internal keys while TestEnergyModelDP uses cls.module.translated_output_def() (human-friendly keys); fix by either adding a translated_output_def() override on the dpmodel SpinModel that maps spin-specific outputs (e.g. atom_energy, energy, mask_mag, force) to the internal keys, or create a dpmodel SpinEnergyModel subclass mirroring PT's SpinEnergyModel with translated_output_def(), and update TestSpinEnergyModelDP to use cls.module.translated_output_def() instead of model_output_def().
🧹 Nitpick comments (3)
deepmd/dpmodel/model/dipole_model.py (1)
58-62: Improve consistency: use direct key access or.get()uniformly acrosscallandcall_lower.Lines 58, 60, 62 in
calluse direct dict access (model_ret["dipole_derv_r"], etc.), while lines 90, 92, 94 incall_loweruse.get()for the same keys. Both approaches are safe sincecall_commonandcall_common_loweralways populate these keys (even asNoneviafit_output_to_model_outputandcommunicate_extended_output), but the inconsistency should be reconciled for clarity.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/model/dipole_model.py` around lines 58 - 62, The code mixes direct dict indexing and .get() for model_ret keys between call and call_lower; make them consistent by switching the direct accesses in call (e.g., model_ret["dipole_derv_r"], model_ret["dipole_derv_c_redu"], model_ret["dipole_derv_c"]) to use model_ret.get(...) like call_lower does, and keep the conditional logic that assigns model_predict["force"] and model_predict["virial"] unchanged so behavior remains identical.deepmd/dpmodel/model/make_model.py (1)
368-369: Consider explicit forwarding methods to prevent potential MRO bypass in future subclasses.Lines 368-369 use static class-body aliasing (
call = call_common) that binds to the function object at definition time. If a future subclass were to overridecall_commonwithout also overridingcall, dispatch would silently use the base-class version instead of the override. While the current concrete models (ener_model,property_model,dipole_model,dp_zbl_model,dos_model) all explicitly redefine bothcallandcall_lower—avoiding active breakage—the pattern is a maintenance trap for future subclasses.Replace with explicit forwarding methods to preserve dynamic dispatch through
self:♻️ Suggested refactor
- call = call_common - call_lower = call_common_lower + def call(self, *args: Any, **kwargs: Any) -> dict[str, Array]: + """Alias for call_common; may be overridden in subclasses.""" + return self.call_common(*args, **kwargs) + + def call_lower(self, *args: Any, **kwargs: Any) -> dict[str, Array]: + """Alias for call_common_lower; may be overridden in subclasses.""" + return self.call_common_lower(*args, **kwargs)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/model/make_model.py` around lines 368 - 369, The class currently aliases call = call_common and call_lower = call_common_lower which binds static function objects and can bypass future overrides; replace those aliases with explicit forwarding instance methods (e.g., def call(self, *args, **kwargs): return self.call_common(*args, **kwargs) and similarly for call_lower) so dispatch goes through self and respects MRO and overrides of call_common/call_common_lower; ensure the forwarding method signatures match the original call_common/call_common_lower signatures and preserve docstrings/annotations if present.source/tests/consistent/model/test_ener.py (1)
280-287: DP backend now skips force/virial comparison — ensure this is intentional and documented.The DP backend returns
SKIP_FLAGfor force, virial, and atom_virial in bothTestEner(lines 284–286) andTestEnerLower(lines 522–524). This means these quantities are never cross-validated against DP.If the dpmodel backend is capable of computing force/virial via its
callpath (e.g., through finite-difference or analytical gradient), skipping the comparison could mask regressions. If it genuinely cannot produce these, this is correct. The PR title suggests this is intentional alignment with the PT backend's output contract.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_ener.py` around lines 280 - 287, The DP backend branch in TestEner and TestEnerLower returns SKIP_FLAG for force, virial, and atom_virial (see RefBackend.DP, ret["energy"], ret["atom_energy"]) which prevents cross-validation; confirm intent and either restore comparisons or document/guard the skip: if dpmodel can compute forces/virial via its call path (finite-difference or analytic gradients) remove the SKIP_FLAG returns and return the actual arrays so tests validate them, otherwise add a clear inline comment and/or a test-level assertion explaining that DP does not provide forces/virial (and reference the dpmodel call path) so future reviewers know this is intentional rather than a regression.
🤖 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/universal/dpmodel/model/test_model.py`:
- Line 167: The dpmodel refactor missed implementing translated_output_def on
SpinModel and updating the test to use it; implement a
translated_output_def(self) method on the dpmodel SpinModel class (mirror
EnergyModel.translated_output_def and the PT backend SpinModel implementation)
so it returns the same translated structure currently produced by
model_output_def().get_data() (or delegates to model_output_def() then converts
its data to the translated shape), and change the test to call
module.translated_output_def() (instead of model_output_def().get_data()) so
both EnergyModel and SpinModel use the unified translated_output_def API; ensure
the method signature and return format match the PT backend's
translated_output_def.
---
Outside diff comments:
In `@deepmd/pt/model/model/dipole_model.py`:
- Around line 72-87: In DipoleModel.forward(), the tensor dims for force, virial
and atom_virial are not squeezed and thus keep an extra length-1 axis; update
the assignments so that model_predict["force"] uses
model_ret["dipole_derv_r"].squeeze(-2), model_predict["virial"] uses
model_ret["dipole_derv_c_redu"].squeeze(-2), and model_predict["atom_virial"]
uses model_ret["dipole_derv_c"].squeeze(-2) (only when do_atomic_virial is true)
so shapes match the other models (see EnergyModel/DPZBLModel pattern).
In `@deepmd/pt/model/model/polar_model.py`:
- Around line 73-75: The else branch in DPPolarAtomicModel.forward and
forward_lower is dead/unreachable because DPPolarAtomicModel.__init__ enforces
get_fitting_net() returns a PolarFittingNet (TypeError otherwise); remove the
entire else block that mutates model_predict["updated_coord"] (the
model_ret-to-model_predict fallback) so both forward and forward_lower only use
the primary model_ret path, or if you prefer defensive coding keep a single
guarded check that does not mutate fields inappropriate for polar models (do not
add updated_coord). Update references in forward, forward_lower, and any similar
patterns in ener_model.py, dos_model.py, and dipole_model.py to remove the
unreachable else branch and ensure consistency.
In `@source/tests/universal/dpmodel/model/test_model.py`:
- Line 274: TestSpinEnergyModelDP is calling
cls.module.model_output_def().get_data() which returns raw internal keys while
TestEnergyModelDP uses cls.module.translated_output_def() (human-friendly keys);
fix by either adding a translated_output_def() override on the dpmodel SpinModel
that maps spin-specific outputs (e.g. atom_energy, energy, mask_mag, force) to
the internal keys, or create a dpmodel SpinEnergyModel subclass mirroring PT's
SpinEnergyModel with translated_output_def(), and update TestSpinEnergyModelDP
to use cls.module.translated_output_def() instead of model_output_def().
---
Nitpick comments:
In `@deepmd/dpmodel/model/dipole_model.py`:
- Around line 58-62: The code mixes direct dict indexing and .get() for
model_ret keys between call and call_lower; make them consistent by switching
the direct accesses in call (e.g., model_ret["dipole_derv_r"],
model_ret["dipole_derv_c_redu"], model_ret["dipole_derv_c"]) to use
model_ret.get(...) like call_lower does, and keep the conditional logic that
assigns model_predict["force"] and model_predict["virial"] unchanged so behavior
remains identical.
In `@deepmd/dpmodel/model/make_model.py`:
- Around line 368-369: The class currently aliases call = call_common and
call_lower = call_common_lower which binds static function objects and can
bypass future overrides; replace those aliases with explicit forwarding instance
methods (e.g., def call(self, *args, **kwargs): return self.call_common(*args,
**kwargs) and similarly for call_lower) so dispatch goes through self and
respects MRO and overrides of call_common/call_common_lower; ensure the
forwarding method signatures match the original call_common/call_common_lower
signatures and preserve docstrings/annotations if present.
In `@source/tests/consistent/model/test_ener.py`:
- Around line 280-287: The DP backend branch in TestEner and TestEnerLower
returns SKIP_FLAG for force, virial, and atom_virial (see RefBackend.DP,
ret["energy"], ret["atom_energy"]) which prevents cross-validation; confirm
intent and either restore comparisons or document/guard the skip: if dpmodel can
compute forces/virial via its call path (finite-difference or analytic
gradients) remove the SKIP_FLAG returns and return the actual arrays so tests
validate them, otherwise add a clear inline comment and/or a test-level
assertion explaining that DP does not provide forces/virial (and reference the
dpmodel call path) so future reviewers know this is intentional rather than a
regression.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5250 +/- ##
==========================================
+ Coverage 82.12% 82.16% +0.04%
==========================================
Files 740 745 +5
Lines 74473 74825 +352
Branches 3616 3616
==========================================
+ Hits 61162 61483 +321
- Misses 12149 12180 +31
Partials 1162 1162 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
…ency test for spin model
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
source/tests/consistent/model/test_spin_ener.py (1)
94-112:get_reference_backendandpass_data_to_clsare identical in both classes.Both methods are copy-pasted verbatim. Extracting them into a shared base class (e.g.,
_SpinEnerTestBase) would eliminate the duplication.♻️ Proposed refactor
+class _SpinEnerTestBase: + """Shared helpers for spin energy model tests.""" + + def get_reference_backend(self): + if not self.skip_pt: + return self.RefBackend.PT + if not self.skip_dp: + return self.RefBackend.DP + raise ValueError("No available reference") + + def pass_data_to_cls(self, cls, data) -> Any: + data = copy.deepcopy(data) + if cls is SpinModelDP: + return get_model_dp(data) + elif cls is SpinEnergyModelPT: + return get_model_pt(data) + return cls(**data, **self.additional_data) + -class TestSpinEner(CommonTest, ModelTest, unittest.TestCase): +class TestSpinEner(_SpinEnerTestBase, CommonTest, ModelTest, unittest.TestCase): ... - def get_reference_backend(self): ... - def pass_data_to_cls(self, cls, data): ... -class TestSpinEnerLower(CommonTest, ModelTest, unittest.TestCase): +class TestSpinEnerLower(_SpinEnerTestBase, CommonTest, ModelTest, unittest.TestCase): ... - def get_reference_backend(self): ... - def pass_data_to_cls(self, cls, data): ...Also applies to: 236-254
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_spin_ener.py` around lines 94 - 112, The two methods get_reference_backend and pass_data_to_cls are duplicated across test classes; extract them into a shared base class (e.g., create class _SpinEnerTestBase) that defines get_reference_backend and pass_data_to_cls, move the implementations there, and have the existing test classes (the ones currently defining those methods) inherit from _SpinEnerTestBase and remove their local copies; ensure the base class provides any attributes used (skip_pt, skip_dp, RefBackend, additional_data) or document they must exist on subclasses so tests still run.source/tests/pt/model/test_ener_spin_model.py (1)
310-383: The new publicDPSpinModel.call()andcall_lower()methods are not tested against their PT equivalents.
test_dp_consistencycorrectly verifies internal parity between the DP and PT backends usingcall_common()/call_common_lower()paired with PT'sforward_common()/forward_common_lower(). Both return the same internal keys (energy,energy_redu).However, the public API methods
DPSpinModel.call()andcall_lower()translate these internal keys to user-facing keys (e.g.,atom_energy,energy,force,force_mag), matching what PT's publicforward()/forward_lower()methods expose. Currently, no test verifies thatdp_model.call(...)andself.model.forward(...)return consistent translated outputs.Consider adding a test section within
test_dp_consistency(or as a separate method) that callsdp_model.call()with translated output keys and compares againstself.model.forward(...)results, similar to howtest_output_shapeverifies the PT public interface.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt/model/test_ener_spin_model.py` around lines 310 - 383, Add assertions that DPSpinModel.call() and call_lower() produce the same user-facing outputs as PT's forward() and forward_lower(): after creating dp_model (DPSpinModel.deserialize(...)), call dp_model.call(...) and compare its returned keys (e.g., "atom_energy", "energy", "force", "force_mag") to self.model.forward(...) outputs converted via to_numpy_array using np.testing.assert_allclose with rtol/atol=self.prec; likewise call dp_model.call_lower(...) and compare to self.model.forward_lower(...). Use the same inputs already prepared in test_dp_consistency (coord, atype, spin, cell and extended versions) and mirror the existing assert_allclose pattern for energy/energy_redu so the public-key translation is validated.deepmd/dpmodel/model/spin_model.py (2)
390-393: Pervasive code duplication: extract a_get_var_name()helper.The identical 3-line pattern —
model_output_type(),pop("mask"),var_name = model_output_type[0]— is repeated six times acrosscall_common,call,call_common_lower,call_lower,translated_output_def, andmodel_output_def. Extracting it into a private helper eliminates the duplication and the risk of the implementations diverging.♻️ Suggested helper
+ def _get_var_name(self) -> str: + """Return the primary output variable name, stripping the 'mask' sentinel.""" + model_output_type = self.backbone_model.model_output_type() + if "mask" in model_output_type: + model_output_type.pop(model_output_type.index("mask")) + return model_output_type[0]Then replace every repeated block with
var_name = self._get_var_name().Also applies to: 473-476, 552-555, 642-645, 672-675
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/model/spin_model.py` around lines 390 - 393, Extract the repeated 3-line pattern into a private helper named _get_var_name(self) that calls self.backbone_model.model_output_type(), removes "mask" if present, and returns the first element; then replace the duplicated blocks in the methods call_common, call, call_common_lower, call_lower, translated_output_def, and model_output_def with var_name = self._get_var_name(); ensure the helper is used everywhere the pattern appears (previously at the locations around lines 390, 473, 552, 642, 672) to remove duplication and keep behavior identical.
419-423: Nit: reuse already-computednframes/nlocinstead of re-destructuringatype.shape.
nframes_mandnloc_mare identical tonframesandnlocfrom line 376.✨ Suggested simplification
- if "mask_mag" not in model_ret: - nframes_m, nloc_m = atype.shape[:2] - atomic_mask = self.virtual_scale_mask[atype].reshape([nframes_m, nloc_m, 1]) - model_ret["mask_mag"] = atomic_mask > 0.0 + if "mask_mag" not in model_ret: + atomic_mask = self.virtual_scale_mask[atype].reshape([nframes, nloc, 1]) + model_ret["mask_mag"] = atomic_mask > 0.0🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/model/spin_model.py` around lines 419 - 423, The code recomputes nframes_m and nloc_m from atype.shape to build mask_mag; instead reuse the already-computed nframes and nloc variables (from earlier in this scope) when creating atomic_mask via self.virtual_scale_mask[atype].reshape([nframes, nloc, 1]) and set model_ret["mask_mag"] = atomic_mask > 0.0, avoiding the redundant destructuring of atype.shape.
🤖 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_spin_ener.py`:
- Around line 336-341: Remove the unused variable nall which is assigned from
extended_coord.shape[1] but never used; in the block around the assignment to
nall and the subsequent np.take_along_axis call that sets self.extended_spin,
delete the nall line so only the np.take_along_axis call using mapping remains
(keep extended_coord, mapping and self.extended_spin intact).
---
Nitpick comments:
In `@deepmd/dpmodel/model/spin_model.py`:
- Around line 390-393: Extract the repeated 3-line pattern into a private helper
named _get_var_name(self) that calls self.backbone_model.model_output_type(),
removes "mask" if present, and returns the first element; then replace the
duplicated blocks in the methods call_common, call, call_common_lower,
call_lower, translated_output_def, and model_output_def with var_name =
self._get_var_name(); ensure the helper is used everywhere the pattern appears
(previously at the locations around lines 390, 473, 552, 642, 672) to remove
duplication and keep behavior identical.
- Around line 419-423: The code recomputes nframes_m and nloc_m from atype.shape
to build mask_mag; instead reuse the already-computed nframes and nloc variables
(from earlier in this scope) when creating atomic_mask via
self.virtual_scale_mask[atype].reshape([nframes, nloc, 1]) and set
model_ret["mask_mag"] = atomic_mask > 0.0, avoiding the redundant destructuring
of atype.shape.
In `@source/tests/consistent/model/test_spin_ener.py`:
- Around line 94-112: The two methods get_reference_backend and pass_data_to_cls
are duplicated across test classes; extract them into a shared base class (e.g.,
create class _SpinEnerTestBase) that defines get_reference_backend and
pass_data_to_cls, move the implementations there, and have the existing test
classes (the ones currently defining those methods) inherit from
_SpinEnerTestBase and remove their local copies; ensure the base class provides
any attributes used (skip_pt, skip_dp, RefBackend, additional_data) or document
they must exist on subclasses so tests still run.
In `@source/tests/pt/model/test_ener_spin_model.py`:
- Around line 310-383: Add assertions that DPSpinModel.call() and call_lower()
produce the same user-facing outputs as PT's forward() and forward_lower():
after creating dp_model (DPSpinModel.deserialize(...)), call dp_model.call(...)
and compare its returned keys (e.g., "atom_energy", "energy", "force",
"force_mag") to self.model.forward(...) outputs converted via to_numpy_array
using np.testing.assert_allclose with rtol/atol=self.prec; likewise call
dp_model.call_lower(...) and compare to self.model.forward_lower(...). Use the
same inputs already prepared in test_dp_consistency (coord, atype, spin, cell
and extended versions) and mirror the existing assert_allclose pattern for
energy/energy_redu so the public-key translation is validated.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/consistent/model/test_spin_ener.py (1)
105-112:pass_data_to_clsandget_reference_backendare duplicated verbatim.Both
TestSpinEner(lines 94–112) andTestSpinEnerLower(lines 236–254) define identicalpass_data_to_cls,get_reference_backend, andbuild_tfbodies. Extracting them into a shared mixin reduces maintenance surface:♻️ Suggested refactor
+class _SpinEnerTestBase: + def get_reference_backend(self): + if not self.skip_pt: + return self.RefBackend.PT + if not self.skip_dp: + return self.RefBackend.DP + raise ValueError("No available reference") + + def pass_data_to_cls(self, cls, data): + data = copy.deepcopy(data) + if cls is SpinModelDP: + return get_model_dp(data) + elif cls is SpinEnergyModelPT: + return get_model_pt(data) + return cls(**data, **self.additional_data) + + def build_tf(self, obj, suffix): + raise NotImplementedError("no TF in this test") + class TestSpinEner(_SpinEnerTestBase, CommonTest, ModelTest, unittest.TestCase): ... - def get_reference_backend(self): ... - def pass_data_to_cls(self, cls, data): ... - def build_tf(self, obj, suffix): ... class TestSpinEnerLower(_SpinEnerTestBase, CommonTest, ModelTest, unittest.TestCase): ... - def get_reference_backend(self): ... - def pass_data_to_cls(self, cls, data): ... - def build_tf(self, obj, suffix): ...Also applies to: 247-254
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_spin_ener.py` around lines 105 - 112, Two test classes (TestSpinEner and TestSpinEnerLower) duplicate pass_data_to_cls, get_reference_backend, and build_tf; extract these shared methods into a single mixin (e.g., SpinEnerTestMixin) and have both test classes inherit from it. Move the implementations of pass_data_to_cls, get_reference_backend, and build_tf out of TestSpinEner and TestSpinEnerLower into the mixin and update both classes to remove their local definitions and subclass the new mixin so existing references to get_model_dp/get_model_pt, SpinModelDP, SpinEnergyModelPT, and self.additional_data continue to resolve.
🤖 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_spin_ener.py`:
- Around line 199-215: The DP branch in the test's extract_ret currently returns
SKIP_FLAG for force/force_mag (and extended force fields elsewhere) despite
SpinModelDP.call() and SpinModelDP.call_lower() computing and returning these
when the backbone's do_grad_r(var_name) is True; update both extract_ret
overrides to assert/validate the force outputs instead of skipping them (replace
SKIP_FLAG with the corresponding ret["force"].ravel() and
ret["force_mag"].ravel() or their extended counterparts) so the DP path mirrors
the PT path, or alternatively ensure the test sets the backbone model to return
do_grad_r() == False so forces are legitimately absent (reference extract_ret,
SpinModelDP.call, SpinModelDP.call_lower, and translated_output_def to locate
related logic).
---
Nitpick comments:
In `@source/tests/consistent/model/test_spin_ener.py`:
- Around line 105-112: Two test classes (TestSpinEner and TestSpinEnerLower)
duplicate pass_data_to_cls, get_reference_backend, and build_tf; extract these
shared methods into a single mixin (e.g., SpinEnerTestMixin) and have both test
classes inherit from it. Move the implementations of pass_data_to_cls,
get_reference_backend, and build_tf out of TestSpinEner and TestSpinEnerLower
into the mixin and update both classes to remove their local definitions and
subclass the new mixin so existing references to get_model_dp/get_model_pt,
SpinModelDP, SpinEnergyModelPT, and self.additional_data continue to resolve.
3 participants
Summary by CodeRabbit
New Features
Refactor
Tests