feat(pt_expt): add dipole, polar, dos, property and dp-zbl models with cross-backend consistency tests#5260
Open
wanghan-iapcm wants to merge 46 commits intodeepmodeling:masterfrom
Open
Conversation
added 14 commits
February 22, 2026 16:10
…ts to dpmodel's folder.
…om the base models of the corresponding backend
Add TestEnerComputeOrLoadStat to the consistency test framework, comparing
dp, pt, and pt_expt backends after compute_or_load_stat. Tests cover
descriptor stats, fparam/aparam fitting stats, output bias, and forward
consistency, parameterized over exclusion types and fparam source
(default injection vs explicit data). Both compute and load-from-file
paths are tested.
Three dpmodel bugs found and fixed:
- repflows.py: compute_input_stats now respects set_stddev_constant,
matching the pt backend behavior
- stat.py: compute_output_stats_global now applies atom_exclude_types
mask to natoms before computing output bias
- general_fitting.py: compute_input_stats now supports save/load of
fparam/aparam stats via stat_file_path, matching the pt backend
…-other-full-model
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
CodeQL found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.
![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: b2028a80f7
ℹ️ 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".
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5260 +/- ##
==========================================
+ Coverage 82.00% 82.19% +0.19%
==========================================
Files 750 753 +3
Lines 75082 75461 +379
Branches 3615 3615
==========================================
+ Hits 61572 62027 +455
+ Misses 12347 12271 -76
Partials 1163 1163 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
added 5 commits
February 24, 2026 12:12
Move get_observed_type_list from a PT-only method to a backend-independent abstract API on BaseBaseModel, with a concrete implementation in dpmodel's make_model CM using array_api_compat for torch compatibility. Add a cross-backend consistency test that verifies dp, pt, and pt_expt return identical results when only a subset of types is observed.
…bare np ops dpmodel's model-level change_type_map was not forwarding model_with_new_type_stat to the atomic model, so fine-tuning with new atom types would silently lose reference statistics. Align with the pt backend by unwrapping .atomic_model and passing it through. Also fix array API violations in fitting change_type_map methods: np.zeros/np.ones/np.concatenate fail when arrays are torch tensors (pt_expt backend). Replace with xp.zeros/xp.ones/xp.concat using proper array namespace and device. Add cross-backend test (test_change_type_map_extend_stat) that exercises the model-level change_type_map with model_with_new_type_stat across dp, pt, and pt_expt.
Add get_out_bias() and set_out_bias() methods to dpmodel's base_atomic_model, and update make_model to call them instead of accessing the attribute directly. For PT, add get_out_bias() to base_atomic_model and remove the redundant implementations from dp_atomic_model, pairtab_atomic_model, and linear_atomic_model.
…et-by-statistic The PT backend calls atomic_model.compute_fitting_input_stat(merged) in change_out_bias when mode is set-by-statistic, but dpmodel/pt_expt did not. This meant fparam/aparam statistics (avg, inv_std) were never updated during bias adjustment in these backends. Add compute_fitting_input_stat to dpmodel's DPAtomicModel and call it from make_model.change_out_bias. Enhance test_change_out_bias with fparam/aparam data, pt_expt coverage, and verification that fitting input stats are updated after set-by-statistic but unchanged after change-by-statistic.
added 17 commits
February 24, 2026 15:43
…-other-full-model
…-other-full-model
…bd to model tests
…the logic of change_out_bias clean
wanghan-iapcm
changed the title
Contributor
📝 WalkthroughWalkthroughAdds out_bias accessors and data_stat_protect to atomic models, wraps samplers with a cached injector for exclusions/default fparam, renames internal fitting → fitting_net and adds compute_or_load_stat/stat-file persistence, extends make_model bases, introduces PT-EXPT model classes, and large cross-backend test additions. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant AtomicModel
participant WrappedSampler
participant Sampler
participant Filesystem
participant FittingNet
rect rgba(200,200,255,0.5)
User->>AtomicModel: compute_or_load_stat(sampled_func, stat_file_path?, compute_or_load_out_stat)
end
AtomicModel->>WrappedSampler: _make_wrapped_sampler(sampled_func)
WrappedSampler->>Sampler: call sampled_func (cached via lru_cache)
Sampler-->>WrappedSampler: samples (injected pair/atom exclusions + default fparam)
WrappedSampler-->>AtomicModel: samples
alt stat_file_path exists
AtomicModel->>Filesystem: load stat items
Filesystem-->>AtomicModel: return persisted StatItem lists
else compute stats
AtomicModel->>FittingNet: compute_input_stats(merged_samples, protection)
FittingNet-->>AtomicModel: computed avg/std (and out_stat if requested)
AtomicModel->>Filesystem: optionally save StatItem lists
Filesystem-->>AtomicModel: ack
end
AtomicModel-->>User: stats available / out_stat applied
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 14
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/dpmodel/fitting/general_fitting.py (1)
258-316:⚠️ Potential issue | 🟠 MajorAvoid calling
merged()twice in one stats pass.Line 268 and Line 315 independently call
merged()on cache-miss paths. Ifmergedis stochastic or expensive, this can yield inconsistentfparam/aparamstats and doubles sampling cost.Proposed fix (cache sampled data once per invocation)
@@ - # stat fparam + sampled_cache: list[dict] | None = None + + def _get_sampled() -> list[dict]: + nonlocal sampled_cache + if sampled_cache is None: + sampled_cache = merged() if callable(merged) else merged + return sampled_cache + + # stat fparam @@ - sampled = merged() if callable(merged) else merged + sampled = _get_sampled() @@ - sampled = merged() if callable(merged) else merged + sampled = _get_sampled()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/fitting/general_fitting.py` around lines 258 - 316, The code calls merged() separately in the fparam and aparam branches, causing double sampling; compute sampled once and reuse it. Move the line sampled = merged() if callable(merged) else merged to a single point before the fparam/aparam stat blocks (or memoize merged result in a local variable used by both branches), then replace the duplicate sampled assignments inside the fparam and aparam else-paths so both StatItem computations use the same sampled data (references: merged, sampled, fparam_stats/aparam_stats, self.numb_fparam, self.numb_aparam).
♻️ Duplicate comments (1)
source/tests/consistent/model/test_ener.py (1)
1535-1560:_compare_variables_recursiveis duplicated — see earlier comment ontest_zbl_ener.py.Same utility function as in
test_zbl_ener.py. Consider extracting to a shared location.🤖 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 1535 - 1560, The _compare_variables_recursive function is duplicated across tests (also present in test_zbl_ener.py); extract it into a shared test utility module (e.g., tests/utils/compare.py or tests/conftest utilities) and replace the duplicate definitions with an import and call to that single function; update references in source/tests/consistent/model/test_ener.py (keep the function name _compare_variables_recursive) and in test_zbl_ener.py to import from the new shared module, run tests to ensure no behavioral changes, and delete the duplicated local definitions.
🧹 Nitpick comments (12)
deepmd/dpmodel/fitting/polarizability_fitting.py (1)
241-255: Allocate extension rows only for truly new types.At Line 241 and Line 248, using
len(type_map)over-allocatesscaleandconstant_matrixwhen only some entries are newly introduced. This is functionally fine but less efficient and less explicit.Suggested diff
if has_new_type: xp = array_api_compat.array_namespace(self.scale) - extend_shape = [len(type_map), *list(self.scale.shape[1:])] + old_ntypes = self.scale.shape[0] + n_new_types = sum(idx >= old_ntypes for idx in remap_index) + extend_shape = [n_new_types, *list(self.scale.shape[1:])] extend_scale = xp.ones( extend_shape, dtype=self.scale.dtype, device=array_api_compat.device(self.scale), ) self.scale = xp.concat([self.scale, extend_scale], axis=0) - extend_shape = [len(type_map), *list(self.constant_matrix.shape[1:])] + extend_shape = [n_new_types, *list(self.constant_matrix.shape[1:])] extend_constant_matrix = xp.zeros( extend_shape, dtype=self.constant_matrix.dtype, device=array_api_compat.device(self.constant_matrix), )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/fitting/polarizability_fitting.py` around lines 241 - 255, The code over-allocates rows by using len(type_map) when extending self.scale and self.constant_matrix; change the logic to compute how many new type rows are actually needed (e.g., needed = max(0, len(type_map) - current_rows) where current_rows = self.scale.shape[0]) and only create extend_scale and extend_constant_matrix with shape [needed, ...] and concat when needed>0, preserving dtype and device via array_api_compat.device(self.scale) and dtype=self.scale.dtype for scale and analogous for constant_matrix; update the concat calls on self.scale and self.constant_matrix to use these smaller extensions and skip concatenation when needed == 0.deepmd/dpmodel/atomic_model/linear_atomic_model.py (1)
356-366: Consider sharing one cached sampler across sub-model and output-stat passes.Right now
sampled_funccan be invoked independently by each sub-model path and again by the linear output-stat path. Caching once at this level improves reproducibility and reduces repeated sampling overhead.♻️ Suggested refactor
+import functools ... def compute_or_load_stat( self, sampled_func: Callable[[], list[dict]], stat_file_path: DPPath | None = None, compute_or_load_out_stat: bool = True, ) -> None: + `@functools.lru_cache` + def cached_sampled_func() -> list[dict]: + return sampled_func() + for md in self.models: md.compute_or_load_stat( - sampled_func, stat_file_path, compute_or_load_out_stat=False + cached_sampled_func, stat_file_path, compute_or_load_out_stat=False ) ... if compute_or_load_out_stat: - wrapped_sampler = self._make_wrapped_sampler(sampled_func) + wrapped_sampler = self._make_wrapped_sampler(cached_sampled_func) self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/atomic_model/linear_atomic_model.py` around lines 356 - 366, The code repeatedly calls sampled_func separately for each sub-model and again for the linear output-stat pass; instead create a single cached sampler and reuse it for all calls. Implement a cached sampler (e.g., a lightweight wrapper that memoizes sampled_func outputs) and assign it to a local variable (call it cached_sampler); pass cached_sampler into each md.compute_or_load_stat(...) instead of sampled_func and also pass the same cached_sampler into self.compute_or_load_out_stat(...) (use self._make_wrapped_sampler(cached_sampler) only if necessary for current wrapping behavior). Ensure this reuse respects the existing stat_file_path/type_map handling and the compute_or_load_out_stat flag.source/tests/consistent/model/test_frozen.py (1)
64-65: Scope teardown cleanup to test-owned temp files only.Line 64 currently matches any
tmp*.pbin the working directory, and Line 65 deletes unconditionally. This can remove unrelated local artifacts. Prefer tracking exact temp paths created by this test (or a dedicated, unique prefix) and deleting only those.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_frozen.py` around lines 64 - 65, The cleanup loop currently removes any file matching tempfile.gettempprefix() + "*.pb" which can delete unrelated files; change the test to only remove temp files it created by either (a) using a unique test-specific prefix (e.g., pass a prefix to tempfile functions) or (b) collecting created temp paths in a list and iterating over that list for deletion (replace the glob loop and variable tmp_pb with deletion of the recorded paths). Ensure the creation sites (where temp files are opened/NamedTemporaryFile/TemporaryDirectory) are updated to return or append the exact paths so the teardown only deletes those.source/tests/pt_expt/model/test_property_model.py (1)
29-93: Consider extracting shared setup/helpers into a common test mixin.
setUp,_make_dp_model, and_prepare_lower_inputsare very similar to the new dipole/polar/dos PT-EXPT tests; centralizing them would reduce drift and maintenance cost.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt_expt/model/test_property_model.py` around lines 29 - 93, Extract the duplicated test setup into a reusable test mixin: create a TestPropertyMixin that contains setUp, _make_dp_model, and _prepare_lower_inputs (moving the DPDescrptSeA/DPPropertyFittingNet/DPPropertyModel construction and the coordinate/atype/cell generation and normalization/neighbor-list code into the mixin) then have TestPropertyModel inherit from TestPropertyMixin and unittest.TestCase (and update other dipole/polar/dos PT-EXPT test classes to reuse the mixin); ensure unique symbols remain the same (setUp, _make_dp_model, _prepare_lower_inputs, DPDescrptSeA, DPPropertyFittingNet, DPPropertyModel) and adjust any imports/attributes references so existing tests continue to call the helper methods without changing their signatures.source/tests/consistent/model/test_property.py (1)
234-1579: Consider extracting shared DP/PT/PT_EXPT API-parity test helpers.This class duplicates a large amount of setup/assertion logic that is repeated in sibling model consistency files, which will increase maintenance cost as APIs evolve.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_property.py` around lines 234 - 1579, Tests duplicate heavy setup and repeated cross-backend assertions across TestPropertyModelAPIs and TestPropertyComputeOrLoadStat; extract shared helpers to reduce duplication. Create a small test helper module (e.g., test_consistent_helpers) that provides functions to build paired models from model_args and serialization (used by PropertyModelPT.deserialize and PropertyModelPTExpt.deserialize), construct coords/atype/box and extended neighbor structures (used in setUp logic), wrap evaluators _eval_dp/_eval_pt/_eval_pt_expt, and assertion helpers for comparing outputs and serialized `@variables` (replacing uses of _compare_variables_recursive and repeating np.testing.assert_allclose calls). Replace in-class duplicated code by calling these helpers from TestPropertyModelAPIs.setUp, TestPropertyComputeOrLoadStat.setUp and the various tests like test_change_out_bias, test_compute_stat, test_load_stat_from_file, test_set_case_embd, test_change_type_map to call assert_backend_outputs_equal / compare_serializations to keep test logic the same but centralized.source/tests/common/dpmodel/test_atomic_model_global_stat.py (1)
500-502: Unused variableret0flagged by static analysis.On line 502, the result of
forward_common_atomicis assigned toret0but never used. If this is intentional (smoke test), use_to silence the linter. As per coding guidelines,ruff check .will flag this.Suggested fix
- ret0 = md0.forward_common_atomic(*args) + _ = md0.forward_common_atomic(*args)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/common/dpmodel/test_atomic_model_global_stat.py` around lines 500 - 502, The variable ret0 is assigned from forward_common_atomic but never used; replace the unused assignment with a throwaway variable or use the returned value in an assertion. Specifically, in the test where you call md0.forward_common_atomic(*args) (currently assigned to ret0), change the assignment to use "_" (e.g., _ = md0.forward_common_atomic(*args)) if the call is only a smoke test, or assert expected properties of the return value from forward_common_atomic instead of leaving ret0 unused.deepmd/pt_expt/model/dipole_model.py (1)
55-98: Inconsistent key access pattern betweenforwardandforward_lower.In
forward(lines 58-63), derivative keys are accessed directly viamodel_ret["dipole_derv_r"], which raisesKeyErrorif the key is absent. Inforward_lower(lines 90-95), the safermodel_ret.get(...)pattern is used. Consider aligning the access patterns for consistency.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt_expt/model/dipole_model.py` around lines 55 - 98, The forward method uses direct dict indexing (model_ret["dipole_derv_r"], model_ret["dipole_derv_c_redu"], model_ret["dipole_derv_c"]) which can raise KeyError and is inconsistent with forward_lower; change forward to use model_ret.get("dipole_derv_r"), model_ret.get("dipole_derv_c_redu") and model_ret.get("dipole_derv_c") (and adjust the conditional checks accordingly) so the presence checks mirror forward_lower and avoid exceptions when keys are missing.source/tests/consistent/model/test_zbl_ener.py (1)
1020-1045:_compare_variables_recursiveis duplicated across test files.This utility function is identically defined in both
test_zbl_ener.py(line 1020) andtest_ener.py(line 1535). Consider extracting it to a shared test helper module (e.g.,source/tests/consistent/model/common.py) to avoid duplication.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_zbl_ener.py` around lines 1020 - 1045, The helper function _compare_variables_recursive is duplicated in test_zbl_ener.py and test_ener.py; extract it into a shared test helper module (e.g., create source/tests/consistent/model/common.py) and import it in both test files: move the _compare_variables_recursive definition to that new module, update both test_zbl_ener.py and test_ener.py to remove their local definitions and add from source.tests.consistent.model.common import _compare_variables_recursive (or equivalent relative import), and run tests to ensure the function name and signature (args: d1, d2, path, rtol, atol) remain identical so all callers continue to work.deepmd/dpmodel/model/make_model.py (2)
381-398:get_observed_type_listrelies on a hard-coded1e-6threshold to identify observed types.Types whose computed bias happens to be very small (below
1e-6) but non-zero would be incorrectly classified as "unobserved." This threshold is not configurable and could produce surprising results for models with near-zero biases for certain types. Consider documenting this assumption or making the threshold a class-level constant.🤖 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 381 - 398, get_observed_type_list currently uses a hard-coded 1e-6 cutoff to decide whether a type is "observed"; make this threshold configurable by introducing a class-level constant (e.g., OBSERVED_TYPE_BIAS_THRESHOLD = 1e-6) on the model class and replace the literal 1e-6 in get_observed_type_list with this constant; ensure the constant is documented in the class docstring and used when computing bias_mask (bias_mask = xp.any(xp.abs(out_bias) > self.OBSERVED_TYPE_BIAS_THRESHOLD, axis=-1)) so callers can adjust the sensitivity without changing the method body.
618-629:change_type_mapunconditionally accesses.atomic_modelonmodel_with_new_type_stat.Line 626 accesses
model_with_new_type_stat.atomic_modelwithout ahasattrguard. This is fine for the expected model-level usage but will produce an unclearAttributeErrorif called with an unexpected type. Given theAnytype annotation, a defensive check or more specific type annotation would improve robustness.🤖 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 618 - 629, The method change_type_map currently assumes model_with_new_type_stat has an atomic_model attribute; update change_type_map to defensively obtain the atomic_model: if model_with_new_type_stat is None pass None, else if hasattr(model_with_new_type_stat, "atomic_model") use model_with_new_type_stat.atomic_model, otherwise raise a clear TypeError (or convert/unwrap if you expect a different type) before calling self.atomic_model.change_type_map; reference change_type_map, model_with_new_type_stat, and atomic_model in your change.source/tests/consistent/model/test_polar.py (2)
672-672: UseassertGreaterfor a more informative failure message.
assertTrue(len(common_keys) > 0)won't show the actual length on failure.Proposed fix
- self.assertTrue(len(common_keys) > 0) + self.assertGreater(len(common_keys), 0)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_polar.py` at line 672, Replace the assertion that checks common_keys with a more informative unittest assertion: change the use of assertTrue(len(common_keys) > 0) to assertGreater(len(common_keys), 0) so failures will report the actual length; locate the assertion referencing common_keys in the test_polar test (the line with assertTrue(len(common_keys) > 0)) and make the substitution.
1230-1255:_compare_variables_recursivesilently skips keys present in only one dict.This is a design choice, but it means if one backend serializes an extra
@variableskey that the other omits (or vice-versa), the mismatch won't be caught. If strict parity is desired, consider adding an assertion thatd1.keys() == d2.keys()at the@variableslevel, or at least logging skipped keys.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_polar.py` around lines 1230 - 1255, In _compare_variables_recursive, ensure missing or extra keys in the "@variables" section are detected by asserting the key sets before comparing values: when key == "@variables" and both v1 and v2 are dicts, compare set(v1.keys()) and set(v2.keys()) and raise/assert with a clear message including child_path (e.g. f"@variables keys mismatch at {child_path}: {sorted(...) } vs {sorted(...)}") if they differ, then proceed to the existing per-key numeric comparisons; this will catch parity issues where one backend omits or adds variable entries.
🤖 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/base_atomic_model.py`:
- Around line 56-64: The new parameter data_stat_protect on BaseAtomicModel must
be persisted across save/load: add data_stat_protect to the model's
serialization and deserialization paths (e.g., include "data_stat_protect" in
whatever dict/state object is returned by the class's
serialize/to_dict/state_dict method and read back in the corresponding
deserialize/from_dict/load_state_dict method), and use that stored value to set
self.data_stat_protect during object reconstruction; update both serialization
and deserialization routines in base_atomic_model.py so non-default values are
preserved.
- Around line 336-338: The code indexes sampled[0] without checking for an empty
sampler result; update the logic around the sampled variable so you first guard
against empty batches (e.g., if not sampled or len(sampled)==0) and handle that
case (raise a clear error or return early) before evaluating checks like
`"find_fparam" not in sampled[0] and "fparam" not in sampled[0] and
self.has_default_fparam()`; modify the surrounding method in
base_atomic_model.py (the block that references sampled, sampled[0], and
self.has_default_fparam()) to perform this emptiness check and then proceed with
the existing conditions.
In `@deepmd/dpmodel/descriptor/repflows.py`:
- Around line 441-443: The early return when both self.set_stddev_constant and
self.set_davg_zero prevents self.stats from being set and can leave get_stats()
consumers with stale state; fix by ensuring EnvMatStatSe(self) is constructed
and assigned to self.stats before any return. Concretely, in the method
containing the lines referencing self.set_stddev_constant, self.set_davg_zero
and EnvMatStatSe, move or add the assignment self.stats = EnvMatStatSe(self) so
it always runs (even if you still return immediately afterward), or restructure
the condition to check flags after assigning self.stats; this ensures
compute_input_stats()/get_stats() see a valid stats object.
In `@deepmd/dpmodel/utils/stat.py`:
- Around line 366-371: The loop that applies AtomExcludeMask to entries in
sampled currently mutates the caller-owned array in place (system[natoms_key][:,
2:] *= ...), so change it to create a masked copy instead: compute type_mask via
AtomExcludeMask(ntypes, system["atom_exclude_types"]).get_type_mask(), then
assign a new array to a local/input variable (e.g., input_natoms =
system[natoms_key].copy(); input_natoms[:, 2:] = input_natoms[:, 2:] *
type_mask.reshape(1, -1)) and use that copy downstream instead of altering
system[natoms_key]; ensure any downstream code that expects the masked values
uses the copied variable rather than mutating sampled.
In `@deepmd/jax/model/hlo.py`:
- Around line 268-270: The HLO model's get_observed_type_list currently raises
NotImplementedError but must implement the BaseModel contract; change
get_observed_type_list in the HLO class to return list(self.type_map) as a
fallback (using the existing self.type_map that holds declared element types) so
generic tooling (e.g., deep_eval) can call it without error.
In `@source/tests/common/dpmodel/test_atomic_model_atomic_stat.py`:
- Around line 158-159: The context manager usages "with h5py.File(h5file, 'w')
as f:" bind an unused variable f (F841); remove the unused binding by changing
those lines to "with h5py.File(h5file, 'w'):" (or "with h5py.File(h5file, 'w')
as _:") at each occurrence so the file is still opened/closed but no unused name
is created; update both places where "as f" appears in this test file.
- Line 164: The unpacking assigns an unused variable `nnei` (e.g., in the
statement `nf, nloc, nnei = self.nlist.shape`) which should use a throwaway
binding per RUF059; replace `nnei` with `_` (or `_nnei`) in this unpack and the
other occurrence (the similar unpack at the later site) so the value is
intentionally ignored and then run `ruff check .` and `ruff format .` before
committing.
In `@source/tests/consistent/model/test_dipole.py`:
- Around line 780-797: Remove the two dead local assignments dp_bias_init and
dp_bias_before to satisfy Ruff F841: delete the unused variable dp_bias_init
(assignment from to_numpy_array(self.dp_model.get_out_bias()).copy()) and delete
dp_bias_before (assignment from dp_bias.copy()); if those values were intended
for later assertions, instead use or assert them where needed (e.g., compare
pre/post bias directly) but do not leave unused temporaries in test_dipole.py
around the change_out_bias / get_out_bias calls.
In `@source/tests/consistent/model/test_dos.py`:
- Around line 764-781: Remove the dead/unused assignments causing Ruff F841:
delete the unused variables dp_bias_init (assigned from
self.dp_model.get_out_bias()) and dp_bias_before (assigned from dp_bias.copy()),
or if they were intended for debugging/verification, replace their assignments
with an explicit use (e.g., assert or logging) so they are no longer unused;
locate the assignments in the test function around the calls to
self.dp_model.get_out_bias() and the "change-by-statistic" section and remove or
use dp_bias_init and dp_bias_before accordingly.
In `@source/tests/consistent/model/test_ener.py`:
- Around line 1768-1771: The test mutates the sampled natoms in-place via
compute_or_load_stat when atom_exclude_types is set, so fix by passing deep
copies of the sampled data to each backend call: wrap self.np_sampled and
self.pt_sampled in deep copies when calling dp_model.compute_or_load_stat,
pt_model.compute_or_load_stat, and pt_expt_model.compute_or_load_stat so each
backend gets an unmodified copy; reference the compute_or_load_stat calls on
dp_model, pt_model and pt_expt_model and copy the np_sampled/pt_sampled objects
before passing them.
In `@source/tests/consistent/model/test_polar.py`:
- Around line 790-791: The variable dp_bias_before is assigned but never used in
the test_polar.py snippet (assigned from dp_bias) causing a ruff F841; either
remove the unused assignment to dp_bias_before, or if the intent was to verify
the "change-by-statistic" mode mutated dp_bias, capture the pre-call state into
dp_bias_before and add an assertion comparing dp_bias_before to dp_bias after
the operation (e.g., assert dp_bias != dp_bias_before) so the value is read;
update the code around the "change-by-statistic" test block where dp_bias_before
is defined and used accordingly.
- Around line 773-774: The test assigns dp_bias_init from
dp_model.get_out_bias() but never uses it (unused variable dp_bias_init), so
remove the assignment or use it in an assertion; locate the line creating
dp_bias_init (dp_bias_init =
to_numpy_array(self.dp_model.get_out_bias()).copy()) in test_polar and either
delete that statement if the initial bias isn't needed, or add an assertion
comparing dp_bias_init to the post-update bias (via
get_out_bias()/to_numpy_array) to validate that the bias changed.
- Around line 698-714: The helper method _get_fitting_stats defined in
test_polar.py (and duplicated across other test files) is unused; either delete
this dead helper from all test files or add a test that calls it (e.g., enhance
test_change_out_bias or another relevant test to assert fparam/aparam stats) to
validate the fitting net outputs. To fix: remove the _get_fitting_stats function
definitions across test_polar.py, test_property.py, test_dos.py, test_dipole.py,
and test_ener.py if no tests need it, or update an existing test (reference
test_change_out_bias) to call model.get_fitting_net() and use _get_fitting_stats
(or inline its logic) to assert fparam_avg, fparam_inv_std, aparam_avg, and
aparam_inv_std match expected numpy values.
In `@source/tests/consistent/model/test_property.py`:
- Around line 761-778: Remove the two unused local variable assignments in the
test: drop the dp_bias_init and dp_bias_before assignments in
source/tests/consistent/model/test_property.py (they are assigned but never
used), so delete the lines that set dp_bias_init =
to_numpy_array(self.dp_model.get_out_bias()).copy() and dp_bias_before =
dp_bias.copy(); ensure no other logic depends on those variables and run ruff
check/format after the change.
---
Outside diff comments:
In `@deepmd/dpmodel/fitting/general_fitting.py`:
- Around line 258-316: The code calls merged() separately in the fparam and
aparam branches, causing double sampling; compute sampled once and reuse it.
Move the line sampled = merged() if callable(merged) else merged to a single
point before the fparam/aparam stat blocks (or memoize merged result in a local
variable used by both branches), then replace the duplicate sampled assignments
inside the fparam and aparam else-paths so both StatItem computations use the
same sampled data (references: merged, sampled, fparam_stats/aparam_stats,
self.numb_fparam, self.numb_aparam).
---
Duplicate comments:
In `@source/tests/consistent/model/test_ener.py`:
- Around line 1535-1560: The _compare_variables_recursive function is duplicated
across tests (also present in test_zbl_ener.py); extract it into a shared test
utility module (e.g., tests/utils/compare.py or tests/conftest utilities) and
replace the duplicate definitions with an import and call to that single
function; update references in source/tests/consistent/model/test_ener.py (keep
the function name _compare_variables_recursive) and in test_zbl_ener.py to
import from the new shared module, run tests to ensure no behavioral changes,
and delete the duplicated local definitions.
---
Nitpick comments:
In `@deepmd/dpmodel/atomic_model/linear_atomic_model.py`:
- Around line 356-366: The code repeatedly calls sampled_func separately for
each sub-model and again for the linear output-stat pass; instead create a
single cached sampler and reuse it for all calls. Implement a cached sampler
(e.g., a lightweight wrapper that memoizes sampled_func outputs) and assign it
to a local variable (call it cached_sampler); pass cached_sampler into each
md.compute_or_load_stat(...) instead of sampled_func and also pass the same
cached_sampler into self.compute_or_load_out_stat(...) (use
self._make_wrapped_sampler(cached_sampler) only if necessary for current
wrapping behavior). Ensure this reuse respects the existing
stat_file_path/type_map handling and the compute_or_load_out_stat flag.
In `@deepmd/dpmodel/fitting/polarizability_fitting.py`:
- Around line 241-255: The code over-allocates rows by using len(type_map) when
extending self.scale and self.constant_matrix; change the logic to compute how
many new type rows are actually needed (e.g., needed = max(0, len(type_map) -
current_rows) where current_rows = self.scale.shape[0]) and only create
extend_scale and extend_constant_matrix with shape [needed, ...] and concat when
needed>0, preserving dtype and device via array_api_compat.device(self.scale)
and dtype=self.scale.dtype for scale and analogous for constant_matrix; update
the concat calls on self.scale and self.constant_matrix to use these smaller
extensions and skip concatenation when needed == 0.
In `@deepmd/dpmodel/model/make_model.py`:
- Around line 381-398: get_observed_type_list currently uses a hard-coded 1e-6
cutoff to decide whether a type is "observed"; make this threshold configurable
by introducing a class-level constant (e.g., OBSERVED_TYPE_BIAS_THRESHOLD =
1e-6) on the model class and replace the literal 1e-6 in get_observed_type_list
with this constant; ensure the constant is documented in the class docstring and
used when computing bias_mask (bias_mask = xp.any(xp.abs(out_bias) >
self.OBSERVED_TYPE_BIAS_THRESHOLD, axis=-1)) so callers can adjust the
sensitivity without changing the method body.
- Around line 618-629: The method change_type_map currently assumes
model_with_new_type_stat has an atomic_model attribute; update change_type_map
to defensively obtain the atomic_model: if model_with_new_type_stat is None pass
None, else if hasattr(model_with_new_type_stat, "atomic_model") use
model_with_new_type_stat.atomic_model, otherwise raise a clear TypeError (or
convert/unwrap if you expect a different type) before calling
self.atomic_model.change_type_map; reference change_type_map,
model_with_new_type_stat, and atomic_model in your change.
In `@deepmd/pt_expt/model/dipole_model.py`:
- Around line 55-98: The forward method uses direct dict indexing
(model_ret["dipole_derv_r"], model_ret["dipole_derv_c_redu"],
model_ret["dipole_derv_c"]) which can raise KeyError and is inconsistent with
forward_lower; change forward to use model_ret.get("dipole_derv_r"),
model_ret.get("dipole_derv_c_redu") and model_ret.get("dipole_derv_c") (and
adjust the conditional checks accordingly) so the presence checks mirror
forward_lower and avoid exceptions when keys are missing.
In `@source/tests/common/dpmodel/test_atomic_model_global_stat.py`:
- Around line 500-502: The variable ret0 is assigned from forward_common_atomic
but never used; replace the unused assignment with a throwaway variable or use
the returned value in an assertion. Specifically, in the test where you call
md0.forward_common_atomic(*args) (currently assigned to ret0), change the
assignment to use "_" (e.g., _ = md0.forward_common_atomic(*args)) if the call
is only a smoke test, or assert expected properties of the return value from
forward_common_atomic instead of leaving ret0 unused.
In `@source/tests/consistent/model/test_frozen.py`:
- Around line 64-65: The cleanup loop currently removes any file matching
tempfile.gettempprefix() + "*.pb" which can delete unrelated files; change the
test to only remove temp files it created by either (a) using a unique
test-specific prefix (e.g., pass a prefix to tempfile functions) or (b)
collecting created temp paths in a list and iterating over that list for
deletion (replace the glob loop and variable tmp_pb with deletion of the
recorded paths). Ensure the creation sites (where temp files are
opened/NamedTemporaryFile/TemporaryDirectory) are updated to return or append
the exact paths so the teardown only deletes those.
In `@source/tests/consistent/model/test_polar.py`:
- Line 672: Replace the assertion that checks common_keys with a more
informative unittest assertion: change the use of assertTrue(len(common_keys) >
0) to assertGreater(len(common_keys), 0) so failures will report the actual
length; locate the assertion referencing common_keys in the test_polar test (the
line with assertTrue(len(common_keys) > 0)) and make the substitution.
- Around line 1230-1255: In _compare_variables_recursive, ensure missing or
extra keys in the "@variables" section are detected by asserting the key sets
before comparing values: when key == "@variables" and both v1 and v2 are dicts,
compare set(v1.keys()) and set(v2.keys()) and raise/assert with a clear message
including child_path (e.g. f"@variables keys mismatch at {child_path}:
{sorted(...) } vs {sorted(...)}") if they differ, then proceed to the existing
per-key numeric comparisons; this will catch parity issues where one backend
omits or adds variable entries.
In `@source/tests/consistent/model/test_property.py`:
- Around line 234-1579: Tests duplicate heavy setup and repeated cross-backend
assertions across TestPropertyModelAPIs and TestPropertyComputeOrLoadStat;
extract shared helpers to reduce duplication. Create a small test helper module
(e.g., test_consistent_helpers) that provides functions to build paired models
from model_args and serialization (used by PropertyModelPT.deserialize and
PropertyModelPTExpt.deserialize), construct coords/atype/box and extended
neighbor structures (used in setUp logic), wrap evaluators
_eval_dp/_eval_pt/_eval_pt_expt, and assertion helpers for comparing outputs and
serialized `@variables` (replacing uses of _compare_variables_recursive and
repeating np.testing.assert_allclose calls). Replace in-class duplicated code by
calling these helpers from TestPropertyModelAPIs.setUp,
TestPropertyComputeOrLoadStat.setUp and the various tests like
test_change_out_bias, test_compute_stat, test_load_stat_from_file,
test_set_case_embd, test_change_type_map to call assert_backend_outputs_equal /
compare_serializations to keep test logic the same but centralized.
In `@source/tests/consistent/model/test_zbl_ener.py`:
- Around line 1020-1045: The helper function _compare_variables_recursive is
duplicated in test_zbl_ener.py and test_ener.py; extract it into a shared test
helper module (e.g., create source/tests/consistent/model/common.py) and import
it in both test files: move the _compare_variables_recursive definition to that
new module, update both test_zbl_ener.py and test_ener.py to remove their local
definitions and add from source.tests.consistent.model.common import
_compare_variables_recursive (or equivalent relative import), and run tests to
ensure the function name and signature (args: d1, d2, path, rtol, atol) remain
identical so all callers continue to work.
In `@source/tests/pt_expt/model/test_property_model.py`:
- Around line 29-93: Extract the duplicated test setup into a reusable test
mixin: create a TestPropertyMixin that contains setUp, _make_dp_model, and
_prepare_lower_inputs (moving the
DPDescrptSeA/DPPropertyFittingNet/DPPropertyModel construction and the
coordinate/atype/cell generation and normalization/neighbor-list code into the
mixin) then have TestPropertyModel inherit from TestPropertyMixin and
unittest.TestCase (and update other dipole/polar/dos PT-EXPT test classes to
reuse the mixin); ensure unique symbols remain the same (setUp, _make_dp_model,
_prepare_lower_inputs, DPDescrptSeA, DPPropertyFittingNet, DPPropertyModel) and
adjust any imports/attributes references so existing tests continue to call the
helper methods without changing their signatures.
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (68)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/atomic_model/dp_atomic_model.pydeepmd/dpmodel/atomic_model/linear_atomic_model.pydeepmd/dpmodel/atomic_model/pairtab_atomic_model.pydeepmd/dpmodel/atomic_model/polar_atomic_model.pydeepmd/dpmodel/atomic_model/property_atomic_model.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/repflows.pydeepmd/dpmodel/fitting/dos_fitting.pydeepmd/dpmodel/fitting/general_fitting.pydeepmd/dpmodel/fitting/polarizability_fitting.pydeepmd/dpmodel/model/base_model.pydeepmd/dpmodel/model/dipole_model.pydeepmd/dpmodel/model/dos_model.pydeepmd/dpmodel/model/dp_model.pydeepmd/dpmodel/model/dp_zbl_model.pydeepmd/dpmodel/model/ener_model.pydeepmd/dpmodel/model/make_model.pydeepmd/dpmodel/model/model.pydeepmd/dpmodel/model/polar_model.pydeepmd/dpmodel/model/property_model.pydeepmd/dpmodel/model/spin_model.pydeepmd/dpmodel/utils/stat.pydeepmd/jax/model/hlo.pydeepmd/pd/model/model/frozen.pydeepmd/pd/model/model/make_model.pydeepmd/pd/train/training.pydeepmd/pt/model/atomic_model/base_atomic_model.pydeepmd/pt/model/atomic_model/dp_atomic_model.pydeepmd/pt/model/atomic_model/linear_atomic_model.pydeepmd/pt/model/atomic_model/pairtab_atomic_model.pydeepmd/pt/model/model/ener_model.pydeepmd/pt/model/model/frozen.pydeepmd/pt/model/model/make_model.pydeepmd/pt/train/training.pydeepmd/pt_expt/atomic_model/__init__.pydeepmd/pt_expt/atomic_model/dp_atomic_model.pydeepmd/pt_expt/atomic_model/energy_atomic_model.pydeepmd/pt_expt/common.pydeepmd/pt_expt/model/__init__.pydeepmd/pt_expt/model/dipole_model.pydeepmd/pt_expt/model/dos_model.pydeepmd/pt_expt/model/dp_zbl_model.pydeepmd/pt_expt/model/ener_model.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/model/model.pydeepmd/pt_expt/model/polar_model.pydeepmd/pt_expt/model/property_model.pysource/tests/common/dpmodel/test_atomic_model_atomic_stat.pysource/tests/common/dpmodel/test_atomic_model_global_stat.pysource/tests/common/dpmodel/test_dp_atomic_model.pysource/tests/consistent/model/test_dipole.pysource/tests/consistent/model/test_dos.pysource/tests/consistent/model/test_dpa1.pysource/tests/consistent/model/test_ener.pysource/tests/consistent/model/test_frozen.pysource/tests/consistent/model/test_polar.pysource/tests/consistent/model/test_property.pysource/tests/consistent/model/test_zbl_ener.pysource/tests/pd/test_training.pysource/tests/pt/test_training.pysource/tests/pt_expt/atomic_model/test_dp_atomic_model.pysource/tests/pt_expt/model/test_dipole_model.pysource/tests/pt_expt/model/test_dos_model.pysource/tests/pt_expt/model/test_dp_zbl_model.pysource/tests/pt_expt/model/test_ener_model.pysource/tests/pt_expt/model/test_polar_model.pysource/tests/pt_expt/model/test_property_model.py
💤 Files with no reviewable changes (8)
- deepmd/pt/model/atomic_model/dp_atomic_model.py
- deepmd/pt/model/model/ener_model.py
- deepmd/pt_expt/atomic_model/init.py
- deepmd/pt/model/atomic_model/linear_atomic_model.py
- deepmd/pt_expt/atomic_model/energy_atomic_model.py
- source/tests/pt_expt/atomic_model/test_dp_atomic_model.py
- deepmd/pt/model/atomic_model/pairtab_atomic_model.py
- deepmd/pt_expt/atomic_model/dp_atomic_model.py
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0ec574876a
ℹ️ 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".
added 2 commits
February 25, 2026 16:23
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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_dipole.py`:
- Around line 1248-1256: The comparator loop that iterates over keys in d1
currently skips keys missing from d2; change it to fail fast by asserting
presence instead of continuing: when iterating the top-level keys (variables d1,
d2 and path) assert that each key from d1 is present in d2 (e.g., raise an
AssertionError or use assert) and likewise inside the "@variables" special-case
loop assert that each vk in v1 exists in v2 instead of continuing; update the
blocks referencing d1, d2, path, and the "@variables" branch so key-set
mismatches immediately cause test failures.
In `@source/tests/consistent/model/test_dos.py`:
- Around line 1237-1245: The loop in _compare_variables_recursive silently
continues when a key present in d1 (or v1) is missing in d2, making comparisons
non-strict; change the behavior so missing keys raise/assert test failures
instead of continuing: in the outer loop over d1 keys (inside
_compare_variables_recursive) replace the early "continue" for key not in d2
with a failure/assertion that includes child_path and the missing key, and
similarly in the nested loop over v1 keys (the "@variables" branch) replace the
"continue" for vk not in v2 with a failure/assertion reporting the missing
variable key and its path. Ensure messages reference child_path and vk so test
output identifies the mismatch.
In `@source/tests/consistent/model/test_polar.py`:
- Around line 1242-1250: In _compare_variables_recursive, stop skipping missing
keys and instead assert that the key sets match before comparing values: for the
outer loop over d1/d2 (vars d1, d2, key, child_path, v1, v2) replace the "if key
not in d2: continue" with an assertion that key in d2 (or assert set(d1.keys())
== set(d2.keys())) so missing keys fail the test; likewise, inside the
"@variables" branch (vk, v1, v2) assert vk exists in v2 (or assert
set(v1.keys()) == set(v2.keys())) before value comparisons so absent keys do not
get silently ignored.
In `@source/tests/consistent/model/test_property.py`:
- Around line 1234-1242: In _compare_variables_recursive, do not silently skip
keys that are missing between v1 and v2; replace the current "if vk not in v2:
continue" behavior with an explicit assertion or test that raises/fails when
keys differ so the test enforces exact key parity for "@variables" entries;
locate the block inside _compare_variables_recursive where key == "@variables"
and the nested loop iterates over vk, and change it to compare the sets of keys
(v1.keys() vs v2.keys()) or assert presence of each vk in v2 with a failure
message that includes child_path and the missing key to surface missing
serialized stats/variables between backends.
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
source/tests/consistent/model/test_dipole.pysource/tests/consistent/model/test_dos.pysource/tests/consistent/model/test_polar.pysource/tests/consistent/model/test_property.py
added 2 commits
February 25, 2026 16:30
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (5)
source/tests/consistent/model/test_property.py (2)
738-780: Test logic forchange_out_biasis correct; both locals are used.The previous review flagged
dp_bias_init(Line 743) anddp_bias_before(Line 764) as unused, but in the current code they are consumed by theassertFalse(np.allclose(...))checks on Lines 759 and 778 respectively. No issue here.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_property.py` around lines 738 - 780, This is a duplicate/false-positive review: the local variables dp_bias_init and dp_bias_before are used in assertions, so mark the comment as resolved and do not change the test; leave the change_out_bias calls and the subsequent np.allclose assertions (and assertFalse checks referencing dp_bias_init and dp_bias_before) intact.
1212-1237:_compare_variables_recursivesilently skips mismatched keys between dicts.Lines 1217-1218 and 1223-1224 both
continuewhen a key is absent from the other dict, meaning extra or missing keys in either direction are never flagged. This could mask genuine serialization divergences between backends.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_property.py` around lines 1212 - 1237, The helper _compare_variables_recursive currently ignores missing keys by using continue when a key in d1 is not in d2 (and likewise inside the `@variables` loop), which hides asymmetric differences; update _compare_variables_recursive to explicitly check key set equality (or assert presence) before descending: when iterating over d1 keys, assert the key exists in d2 (include child_path in the assertion message) or compute missing = set(d1) - set(d2) and extra = set(d2) - set(d1) and raise/assert with clear messages; similarly inside the "@variables" branch for vk ensure vk exists in both dicts (or compare sets) so missing/extra variable names cause test failures rather than being silently skipped.source/tests/consistent/model/test_dos.py (1)
1215-1241:_compare_variables_recursivesilently skips missing keys, weakening parity enforcement.Keys present in
d1but absent fromd2are silently skipped, so serialization regressions (missing variables, structural drift) will not be caught.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_dos.py` around lines 1215 - 1241, The helper _compare_variables_recursive currently ignores keys present in d1 but missing from d2 which hides regressions; change the logic to assert/fail when a key from d1 is not found in d2 (use child_path to give context), and similarly assert when a variable key vk under an "@variables" dict is missing in v2; update the two places that now "continue" (the top-level if key not in d2 and the inner if vk not in v2) to raise an AssertionError or call pytest.fail with a descriptive message like "@variables missing at {child_path}/{vk}" so missing keys fail the test rather than being silently skipped.source/tests/consistent/model/test_dipole.py (1)
1226-1251:_compare_variables_recursiveis duplicated from the other two files and still silently skips missing keys.See the consolidated comments on
test_dos.pylines 1215–1241 for the fix.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_dipole.py` around lines 1226 - 1251, The _compare_variables_recursive function currently silently skips keys missing in d2; update it to assert when expected keys are absent instead of continuing silently: in the top-level for loop replace the "if key not in d2: continue" with an assertion (e.g., raise AssertionError or use a test assert) that includes the missing key and current path, and likewise inside the "@variables" block replace "if vk not in v2: continue" with an assertion that vk exists in v2 (including child_path/vk in the message); keep the existing numeric comparisons and recursive call behavior otherwise so mismatches still surface.source/tests/consistent/model/test_polar.py (1)
1220-1245:_compare_variables_recursiveis duplicated fromtest_dos.pyand still silently skips missing keys.Both concerns are tracked across all three files — see comment on
test_dos.pylines 1215–1241 for the consolidation suggestion and the fix for key-set strictness.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_polar.py` around lines 1220 - 1245, The helper _compare_variables_recursive is duplicated and currently silently skips missing keys; extract this function into a shared test utility (e.g., tests.utils) and update its logic to enforce strict key-set equality: when iterating keys in d1 and d2, compute the union and assert (or raise AssertionError) on any key present in one dict but not the other (including inside "@variables"), then compare corresponding values as before (using np.asarray and np.testing.assert_allclose); update callers in test_polar (and test_dos/test_* duplicates) to import the consolidated function.
🧹 Nitpick comments (10)
source/tests/consistent/model/test_property.py (2)
1395-1424: Inconsistentdo_atomic_virialbetween_eval_dpand_eval_pt/_eval_pt_expt.
_eval_dpomitsdo_atomic_virialwhile_eval_ptand_eval_pt_exptpassdo_atomic_virial=True. For a property model this is unlikely to affect the compared keys ("foo","atom_foo"), but the asymmetry is surprising. If it's intentional (e.g., the dpmodel__call__doesn't accept the kwarg), a brief comment would help.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_property.py` around lines 1395 - 1424, The three helpers _eval_dp, _eval_pt and _eval_pt_expt are inconsistent: _eval_pt and _eval_pt_expt pass do_atomic_virial=True while _eval_dp does not; either modify the dp_model invocation in _eval_dp to pass do_atomic_virial=True (if dp_model.__call__/dp_model accepts that kwarg) or, if dp_model does not accept that kwarg, add a concise comment above _eval_dp explaining why do_atomic_virial is intentionally omitted (and mention dp_model.__call__ inability), so the asymmetry is explicit and reviewers aren’t surprised.
359-391: Consider adding pt_expt assertions intest_get_out_biasandtest_set_out_bias.These tests validate dp vs pt but skip pt_expt. Since
self.pt_expt_modelis built insetUpand other tests (e.g.,test_translated_output_def,test_get_model_def_script) already cover pt_expt, the omission here may leave a gap in bias-accessor parity testing for the pt_expt backend.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_property.py` around lines 359 - 391, Add parallel assertions for the pt_expt backend in both test_get_out_bias and test_set_out_bias: call self.pt_expt_model.get_out_bias, convert to numpy via torch_to_numpy (or to_numpy_array if needed), and assert_allclose against dp_bias/expected new_bias with the same rtol/atol; in test_set_out_bias also call self.pt_expt_model.set_out_bias(numpy_to_torch(new_bias)) before asserting. Reference the existing test functions test_get_out_bias and test_set_out_bias and the methods get_out_bias and set_out_bias on self.pt_expt_model to locate where to add these checks.source/tests/consistent/model/test_ener.py (2)
1055-1056:pe_merged = dp_mergedis a shallow alias — fragile if exclusion types are later added.Currently safe because this model config has no
atom_exclude_types, but if someone extends this test,dp_model.change_out_biasmutatingnatomsin-place would corrupt the data seen bypt_expt_model.change_out_bias. Consider using a separate construction or acopy.deepcopyfor defense.Suggested defensive fix
- # pt_expt stat data (numpy, same as dp) - pe_merged = dp_merged + # pt_expt stat data (numpy, same structure as dp — use deepcopy for safety) + import copy + pe_merged = copy.deepcopy(dp_merged)🤖 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 1055 - 1056, pe_merged is currently a shallow alias to dp_merged which is fragile if exclusion types are added because in-place mutations (e.g. dp_model.change_out_bias) can corrupt the pt_expt data; fix by constructing an independent copy for pe_merged instead of assignment — use a deep copy of dp_merged (or rebuild pe_merged from the same source) so later calls like dp_model.change_out_bias or pt_expt_model.change_out_bias no longer share the same underlying arrays (ensure natoms and any mutable arrays are duplicated).
1517-1542: MixedNonecase in_compare_variables_recursivewould produce a confusing error.Lines 1532-1533 handle the case where both values are
None, but if only one isNone,np.testing.assert_allclose(array, None)will raise aTypeErrorrather than a clear assertion message. Consider adding an explicit check.Suggested improvement
a1 = np.asarray(v1[vk]) if v1[vk] is not None else None a2 = np.asarray(v2[vk]) if v2[vk] is not None else None if a1 is None and a2 is None: continue + assert a1 is not None and a2 is not None, ( + f"None mismatch at {child_path}/{vk}: " + f"a1 is {'None' if a1 is None else 'not None'}, " + f"a2 is {'None' if a2 is None else 'not None'}" + ) np.testing.assert_allclose(🤖 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 1517 - 1542, The helper _compare_variables_recursive currently calls np.testing.assert_allclose(a1, a2) even when one of a1 or a2 is None which raises a TypeError; before calling np.testing.assert_allclose in the "@variables" vk loop, add an explicit check for the mixed-None case (if (a1 is None) != (a2 is None)) and raise a clear AssertionError (or use np.testing.fail) with a message like "@variables mismatch at {child_path}/{vk}: one value is None and the other is not"; only call np.testing.assert_allclose when both a1 and a2 are not None, passing through rtol and atol as before.source/tests/consistent/model/test_dos.py (2)
516-519: Redundant localmodel_argsimport — already imported at module level (Line 59).🧹 Proposed fix
- from deepmd.utils.argcheck import ( - model_args, - ) - # Build a model with dim_case_embd > 0 data = model_args().normalize_value(🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_dos.py` around lines 516 - 519, The local import of model_args inside the test at test_dos.py is redundant because model_args is already imported at module level; remove the inner import statement that references model_args so the test uses the module-level symbol (ensure no local shadowing or redefinition occurs and run tests to confirm no other references rely on the local import).
352-625:pt_expt_modelis under-exercised across mostTestDOSModelAPIstest methods.The class initializes
self.pt_expt_modelinsetUp, but the majority of individual test methods (e.g.,test_get_descriptor,test_get_fitting_net,test_set_out_bias,test_model_output_def,test_model_output_type,test_do_grad_r,test_do_grad_c,test_get_rcut,test_get_type_map, etc.) only comparedp_modelvspt_model. Onlytest_get_model_def_script,test_get_min_nbor_dist, andtest_set_case_embdincludept_exptassertions.Consider extending each getter/scalar test to also assert the pt_expt variant, or add a comment explaining the intentional omission.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_dos.py` around lines 352 - 625, Many tests in TestDOSModelAPIs compare dp_model vs pt_model but omit pt_expt_model; update the getter/scalar tests (e.g., test_get_descriptor, test_get_fitting_net, test_set_out_bias, test_model_output_def, test_model_output_type, test_do_grad_r, test_do_grad_c, test_get_rcut, test_get_type_map, test_get_sel, test_get_nsel, test_get_nnei, test_mixed_types, test_has_message_passing, test_need_sorted_nlist_for_lower, test_get_dim_fparam, test_get_dim_aparam, test_get_sel_type, test_is_aparam_nall, test_atomic_output_def, etc.) to also assert equality (or appropriate numeric closeness) with self.pt_expt_model, using the same helper conversions (torch_to_numpy/numpy_to_torch/to_numpy_array) and the same tolerances, or alternatively add a short comment above each test explaining why pt_expt_model is intentionally excluded; locate assertions around self.dp_model and self.pt_model in each test and replicate/adjust them to include self.pt_expt_model (e.g., compare dp_val vs pt_expt_val or call self.pt_expt_model.method() and assert equality).source/tests/consistent/model/test_dipole.py (2)
533-536: Redundant localmodel_argsimport — already imported at module level (Line 59).🧹 Proposed fix
- from deepmd.utils.argcheck import ( - model_args, - ) - # Build a model with dim_case_embd > 0 data = model_args().normalize_value(🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_dipole.py` around lines 533 - 536, Remove the redundant local import of model_args in test_dipole.py (the local "from deepmd.utils.argcheck import model_args" shown in the diff) since model_args is already imported at module scope; delete that local import block to avoid shadowing/redundant imports and rely on the existing module-level model_args symbol.
358-632:pt_expt_modelis under-exercised in mostTestDipoleModelAPIstest methods — same coverage gap as the DOS and Polar API test classes.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_dipole.py` around lines 358 - 632, Many tests only compare dp_model vs pt_model and omit pt_expt_model; update each API test (e.g., test_translated_output_def, test_get_descriptor, test_get_fitting_net, test_get_out_bias, test_set_out_bias, test_model_output_def, test_model_output_type, test_do_grad_r, test_do_grad_c, test_get_rcut, test_get_type_map, test_get_sel, test_get_nsel, test_get_nnei, test_mixed_types, test_has_message_passing, test_need_sorted_nlist_for_lower, test_get_dim_fparam, test_get_dim_aparam, test_get_sel_type, test_is_aparam_nall) to also check pt_expt_model for parity with dp_model by calling the same methods (e.g., translated_output_def(), get_descriptor(), get_fitting_net(), get_out_bias(), model_output_def(), model_output_type(), do_grad_r("dipole"), do_grad_c("dipole"), get_rcut(), get_type_map(), get_sel(), get_nsel(), get_nnei(), mixed_types(), has_message_passing(), need_sorted_nlist_for_lower(), get_dim_fparam(), get_dim_aparam(), get_sel_type(), is_aparam_nall()) and asserting equality of keys, shapes, numerical arrays (use torch_to_numpy/pt_expt_numpy_to_torch conversions as needed) or non-None where appropriate so pt_expt_model is exercised exactly like pt_model in each test.source/tests/consistent/model/test_polar.py (2)
527-530: Redundant localmodel_argsimport — already imported at module level (Line 59).🧹 Proposed fix
- from deepmd.utils.argcheck import ( - model_args, - ) - # Build a model with dim_case_embd > 0 data = model_args().normalize_value(🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_polar.py` around lines 527 - 530, Remove the redundant local import of model_args in test_polar.py: the symbol model_args is already imported at module scope, so delete the duplicate "from deepmd.utils.argcheck import model_args" inside the test to avoid shadowing/redundant imports and keep the module-level import as the single source of truth; after removing the line, run the tests to ensure no unused-import warnings or linter errors remain.
352-626:pt_expt_modelis under-exercised in mostTestPolarModelAPIstest methods — same coverage gap asTestDOSModelAPIs. Most getter/boolean tests only comparedp_modelvspt_model, omittingpt_expt_modelassertions.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/consistent/model/test_polar.py` around lines 352 - 626, Many tests only compare dp_model vs pt_model and omit pt_expt_model; update each API test (e.g., translated_output_def(), get_descriptor(), get_fitting_net(), get_out_bias()/set_out_bias(), model_output_def(), model_output_type(), do_grad_r(), do_grad_c(), get_rcut(), get_type_map(), get_sel(), get_nsel(), get_nnei(), mixed_types(), has_message_passing(), need_sorted_nlist_for_lower(), get_dim_fparam(), get_dim_aparam(), get_sel_type(), is_aparam_nall(), get_model_def_script(), get_min_nbor_dist()) to also assert the same properties against pt_expt_model (keys, shapes, numerical equality or booleans as appropriate) by calling the same methods on pt_expt_model and adding the same equality/shape/allclose checks used for dp_model vs pt_model; ensure set_out_bias() and get_out_bias() tests also call pt_expt_model where relevant and mirror the numpy/torch conversion 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_dos.py`:
- Around line 1215-1241: Move the duplicated helper _compare_variables_recursive
into a single shared module (e.g. create source/tests/consistent/model/common.py
or source/tests/consistent/model/test_helpers.py), keep the exact function
signature and ensure it imports numpy as np; then remove the copies from
test_dos.py, test_polar.py and test_dipole.py and replace them with a single
import such as "from .common import _compare_variables_recursive" (or
appropriate relative path) in each test file so all three tests reuse the shared
implementation.
In `@source/tests/consistent/model/test_ener.py`:
- Around line 1831-1840: In test_load_stat_from_file, guard against in-place
mutation by wrapping the sampled inputs in deepcopy before passing to
compute_or_load_stat: replace direct uses of self.np_sampled / self.pt_sampled
when calling dp_model.compute_or_load_stat, pt_model.compute_or_load_stat and
pt_expt_model.compute_or_load_stat with deepcopy(self.np_sampled) or
deepcopy(self.pt_sampled) as appropriate (same pattern used in
test_compute_stat) so that compute_or_load_stat cannot mutate the original
self.* sample objects when atom_exclude_types is non-empty.
---
Duplicate comments:
In `@source/tests/consistent/model/test_dipole.py`:
- Around line 1226-1251: The _compare_variables_recursive function currently
silently skips keys missing in d2; update it to assert when expected keys are
absent instead of continuing silently: in the top-level for loop replace the "if
key not in d2: continue" with an assertion (e.g., raise AssertionError or use a
test assert) that includes the missing key and current path, and likewise inside
the "@variables" block replace "if vk not in v2: continue" with an assertion
that vk exists in v2 (including child_path/vk in the message); keep the existing
numeric comparisons and recursive call behavior otherwise so mismatches still
surface.
In `@source/tests/consistent/model/test_dos.py`:
- Around line 1215-1241: The helper _compare_variables_recursive currently
ignores keys present in d1 but missing from d2 which hides regressions; change
the logic to assert/fail when a key from d1 is not found in d2 (use child_path
to give context), and similarly assert when a variable key vk under an
"@variables" dict is missing in v2; update the two places that now "continue"
(the top-level if key not in d2 and the inner if vk not in v2) to raise an
AssertionError or call pytest.fail with a descriptive message like "@variables
missing at {child_path}/{vk}" so missing keys fail the test rather than being
silently skipped.
In `@source/tests/consistent/model/test_polar.py`:
- Around line 1220-1245: The helper _compare_variables_recursive is duplicated
and currently silently skips missing keys; extract this function into a shared
test utility (e.g., tests.utils) and update its logic to enforce strict key-set
equality: when iterating keys in d1 and d2, compute the union and assert (or
raise AssertionError) on any key present in one dict but not the other
(including inside "@variables"), then compare corresponding values as before
(using np.asarray and np.testing.assert_allclose); update callers in test_polar
(and test_dos/test_* duplicates) to import the consolidated function.
In `@source/tests/consistent/model/test_property.py`:
- Around line 738-780: This is a duplicate/false-positive review: the local
variables dp_bias_init and dp_bias_before are used in assertions, so mark the
comment as resolved and do not change the test; leave the change_out_bias calls
and the subsequent np.allclose assertions (and assertFalse checks referencing
dp_bias_init and dp_bias_before) intact.
- Around line 1212-1237: The helper _compare_variables_recursive currently
ignores missing keys by using continue when a key in d1 is not in d2 (and
likewise inside the `@variables` loop), which hides asymmetric differences; update
_compare_variables_recursive to explicitly check key set equality (or assert
presence) before descending: when iterating over d1 keys, assert the key exists
in d2 (include child_path in the assertion message) or compute missing = set(d1)
- set(d2) and extra = set(d2) - set(d1) and raise/assert with clear messages;
similarly inside the "@variables" branch for vk ensure vk exists in both dicts
(or compare sets) so missing/extra variable names cause test failures rather
than being silently skipped.
---
Nitpick comments:
In `@source/tests/consistent/model/test_dipole.py`:
- Around line 533-536: Remove the redundant local import of model_args in
test_dipole.py (the local "from deepmd.utils.argcheck import model_args" shown
in the diff) since model_args is already imported at module scope; delete that
local import block to avoid shadowing/redundant imports and rely on the existing
module-level model_args symbol.
- Around line 358-632: Many tests only compare dp_model vs pt_model and omit
pt_expt_model; update each API test (e.g., test_translated_output_def,
test_get_descriptor, test_get_fitting_net, test_get_out_bias, test_set_out_bias,
test_model_output_def, test_model_output_type, test_do_grad_r, test_do_grad_c,
test_get_rcut, test_get_type_map, test_get_sel, test_get_nsel, test_get_nnei,
test_mixed_types, test_has_message_passing, test_need_sorted_nlist_for_lower,
test_get_dim_fparam, test_get_dim_aparam, test_get_sel_type,
test_is_aparam_nall) to also check pt_expt_model for parity with dp_model by
calling the same methods (e.g., translated_output_def(), get_descriptor(),
get_fitting_net(), get_out_bias(), model_output_def(), model_output_type(),
do_grad_r("dipole"), do_grad_c("dipole"), get_rcut(), get_type_map(), get_sel(),
get_nsel(), get_nnei(), mixed_types(), has_message_passing(),
need_sorted_nlist_for_lower(), get_dim_fparam(), get_dim_aparam(),
get_sel_type(), is_aparam_nall()) and asserting equality of keys, shapes,
numerical arrays (use torch_to_numpy/pt_expt_numpy_to_torch conversions as
needed) or non-None where appropriate so pt_expt_model is exercised exactly like
pt_model in each test.
In `@source/tests/consistent/model/test_dos.py`:
- Around line 516-519: The local import of model_args inside the test at
test_dos.py is redundant because model_args is already imported at module level;
remove the inner import statement that references model_args so the test uses
the module-level symbol (ensure no local shadowing or redefinition occurs and
run tests to confirm no other references rely on the local import).
- Around line 352-625: Many tests in TestDOSModelAPIs compare dp_model vs
pt_model but omit pt_expt_model; update the getter/scalar tests (e.g.,
test_get_descriptor, test_get_fitting_net, test_set_out_bias,
test_model_output_def, test_model_output_type, test_do_grad_r, test_do_grad_c,
test_get_rcut, test_get_type_map, test_get_sel, test_get_nsel, test_get_nnei,
test_mixed_types, test_has_message_passing, test_need_sorted_nlist_for_lower,
test_get_dim_fparam, test_get_dim_aparam, test_get_sel_type,
test_is_aparam_nall, test_atomic_output_def, etc.) to also assert equality (or
appropriate numeric closeness) with self.pt_expt_model, using the same helper
conversions (torch_to_numpy/numpy_to_torch/to_numpy_array) and the same
tolerances, or alternatively add a short comment above each test explaining why
pt_expt_model is intentionally excluded; locate assertions around self.dp_model
and self.pt_model in each test and replicate/adjust them to include
self.pt_expt_model (e.g., compare dp_val vs pt_expt_val or call
self.pt_expt_model.method() and assert equality).
In `@source/tests/consistent/model/test_ener.py`:
- Around line 1055-1056: pe_merged is currently a shallow alias to dp_merged
which is fragile if exclusion types are added because in-place mutations (e.g.
dp_model.change_out_bias) can corrupt the pt_expt data; fix by constructing an
independent copy for pe_merged instead of assignment — use a deep copy of
dp_merged (or rebuild pe_merged from the same source) so later calls like
dp_model.change_out_bias or pt_expt_model.change_out_bias no longer share the
same underlying arrays (ensure natoms and any mutable arrays are duplicated).
- Around line 1517-1542: The helper _compare_variables_recursive currently calls
np.testing.assert_allclose(a1, a2) even when one of a1 or a2 is None which
raises a TypeError; before calling np.testing.assert_allclose in the
"@variables" vk loop, add an explicit check for the mixed-None case (if (a1 is
None) != (a2 is None)) and raise a clear AssertionError (or use np.testing.fail)
with a message like "@variables mismatch at {child_path}/{vk}: one value is None
and the other is not"; only call np.testing.assert_allclose when both a1 and a2
are not None, passing through rtol and atol as before.
In `@source/tests/consistent/model/test_polar.py`:
- Around line 527-530: Remove the redundant local import of model_args in
test_polar.py: the symbol model_args is already imported at module scope, so
delete the duplicate "from deepmd.utils.argcheck import model_args" inside the
test to avoid shadowing/redundant imports and keep the module-level import as
the single source of truth; after removing the line, run the tests to ensure no
unused-import warnings or linter errors remain.
- Around line 352-626: Many tests only compare dp_model vs pt_model and omit
pt_expt_model; update each API test (e.g., translated_output_def(),
get_descriptor(), get_fitting_net(), get_out_bias()/set_out_bias(),
model_output_def(), model_output_type(), do_grad_r(), do_grad_c(), get_rcut(),
get_type_map(), get_sel(), get_nsel(), get_nnei(), mixed_types(),
has_message_passing(), need_sorted_nlist_for_lower(), get_dim_fparam(),
get_dim_aparam(), get_sel_type(), is_aparam_nall(), get_model_def_script(),
get_min_nbor_dist()) to also assert the same properties against pt_expt_model
(keys, shapes, numerical equality or booleans as appropriate) by calling the
same methods on pt_expt_model and adding the same equality/shape/allclose checks
used for dp_model vs pt_model; ensure set_out_bias() and get_out_bias() tests
also call pt_expt_model where relevant and mirror the numpy/torch conversion
checks.
In `@source/tests/consistent/model/test_property.py`:
- Around line 1395-1424: The three helpers _eval_dp, _eval_pt and _eval_pt_expt
are inconsistent: _eval_pt and _eval_pt_expt pass do_atomic_virial=True while
_eval_dp does not; either modify the dp_model invocation in _eval_dp to pass
do_atomic_virial=True (if dp_model.__call__/dp_model accepts that kwarg) or, if
dp_model does not accept that kwarg, add a concise comment above _eval_dp
explaining why do_atomic_virial is intentionally omitted (and mention
dp_model.__call__ inability), so the asymmetry is explicit and reviewers aren’t
surprised.
- Around line 359-391: Add parallel assertions for the pt_expt backend in both
test_get_out_bias and test_set_out_bias: call self.pt_expt_model.get_out_bias,
convert to numpy via torch_to_numpy (or to_numpy_array if needed), and
assert_allclose against dp_bias/expected new_bias with the same rtol/atol; in
test_set_out_bias also call
self.pt_expt_model.set_out_bias(numpy_to_torch(new_bias)) before asserting.
Reference the existing test functions test_get_out_bias and test_set_out_bias
and the methods get_out_bias and set_out_bias on self.pt_expt_model to locate
where to add these checks.
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
source/tests/consistent/model/test_dipole.pysource/tests/consistent/model/test_dos.pysource/tests/consistent/model/test_ener.pysource/tests/consistent/model/test_polar.pysource/tests/consistent/model/test_property.py
Comment on lines
+1215
to
+1241
| def _compare_variables_recursive( | ||
| d1: dict, d2: dict, path: str = "", rtol: float = 1e-10, atol: float = 1e-10 | ||
| ) -> None: | ||
| """Recursively compare ``@variables`` sections in two serialized dicts.""" | ||
| for key in d1: | ||
| if key not in d2: | ||
| continue | ||
| child_path = f"{path}/{key}" if path else key | ||
| v1, v2 = d1[key], d2[key] | ||
| if key == "@variables" and isinstance(v1, dict) and isinstance(v2, dict): | ||
| for vk in v1: | ||
| if vk not in v2: | ||
| continue | ||
| a1 = np.asarray(v1[vk]) if v1[vk] is not None else None | ||
| a2 = np.asarray(v2[vk]) if v2[vk] is not None else None | ||
| if a1 is None and a2 is None: | ||
| continue | ||
| np.testing.assert_allclose( | ||
| a1, | ||
| a2, | ||
| rtol=rtol, | ||
| atol=atol, | ||
| err_msg=f"@variables mismatch at {child_path}/{vk}", | ||
| ) | ||
| elif isinstance(v1, dict) and isinstance(v2, dict): | ||
| _compare_variables_recursive(v1, v2, child_path, rtol, atol) | ||
|
|
Contributor
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
_compare_variables_recursive is copy-pasted identically across test_dos.py, test_polar.py, and test_dipole.py.
Extract it to source/tests/consistent/model/common.py (or a shared test_helpers.py) and import it in each test file to eliminate the triplication.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@source/tests/consistent/model/test_dos.py` around lines 1215 - 1241, Move the
duplicated helper _compare_variables_recursive into a single shared module (e.g.
create source/tests/consistent/model/common.py or
source/tests/consistent/model/test_helpers.py), keep the exact function
signature and ensure it imports numpy as np; then remove the copies from
test_dos.py, test_polar.py and test_dipole.py and replace them with a single
import such as "from .common import _compare_variables_recursive" (or
appropriate relative path) in each test file so all three tests reuse the shared
implementation.
Comment on lines
+1831
to
+1840
| # 1. Compute stats and save to file | ||
| self.dp_model.compute_or_load_stat( | ||
| lambda: self.np_sampled, stat_file_path=DPPath(dp_h5, "a") | ||
| ) | ||
| self.pt_model.compute_or_load_stat( | ||
| lambda: self.pt_sampled, stat_file_path=DPPath(pt_h5, "a") | ||
| ) | ||
| self.pt_expt_model.compute_or_load_stat( | ||
| lambda: self.np_sampled, stat_file_path=DPPath(pe_h5, "a") | ||
| ) |
Contributor
There was a problem hiding this comment.
Missing deepcopy in test_load_stat_from_file — same in-place mutation bug that was fixed in test_compute_stat.
test_compute_stat (line 1757) correctly wraps the sampled data in deepcopy to guard against stat.py mutating natoms in-place when atom_exclude_types is non-empty. However, test_load_stat_from_file passes the raw self.np_sampled / self.pt_sampled references. When parameterized with atom_exclude_types=[1], the first compute_or_load_stat call (dp, line 1832) will mutate self.np_sampled, and the pt_expt call (line 1838) will operate on corrupted data.
Proposed fix
+ from copy import deepcopy
+
# 1. Compute stats and save to file
self.dp_model.compute_or_load_stat(
- lambda: self.np_sampled, stat_file_path=DPPath(dp_h5, "a")
+ lambda: deepcopy(self.np_sampled), stat_file_path=DPPath(dp_h5, "a")
)
self.pt_model.compute_or_load_stat(
- lambda: self.pt_sampled, stat_file_path=DPPath(pt_h5, "a")
+ lambda: deepcopy(self.pt_sampled), stat_file_path=DPPath(pt_h5, "a")
)
self.pt_expt_model.compute_or_load_stat(
- lambda: self.np_sampled, stat_file_path=DPPath(pe_h5, "a")
+ lambda: deepcopy(self.np_sampled), stat_file_path=DPPath(pe_h5, "a")
)🤖 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 1831 - 1840, In
test_load_stat_from_file, guard against in-place mutation by wrapping the
sampled inputs in deepcopy before passing to compute_or_load_stat: replace
direct uses of self.np_sampled / self.pt_sampled when calling
dp_model.compute_or_load_stat, pt_model.compute_or_load_stat and
pt_expt_model.compute_or_load_stat with deepcopy(self.np_sampled) or
deepcopy(self.pt_sampled) as appropriate (same pattern used in
test_compute_stat) so that compute_or_load_stat cannot mutate the original
self.* sample objects when atom_exclude_types is non-empty.
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
1 participant
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
make_model(), removing the intermediate pt_expt atomic model layerget_out_bias,set_out_bias,get_observed_type_list,compute_or_load_stat) into shared dpmodel base classescompute_fitting_input_statforset-by-statisticmode from model-levelchange_out_biasto training-levelmodel_change_out_bias(pt and pd backends), keeping thechange_out_biaslogic focused on bias only (copied from chore(pt): mv the input stat update to model_change_out_bias #5266)general_fitting.change_type_map(barenp.zeros/np.ones/np.concatenate→xpequivalents with device)change_type_mapnot forwardingmodel_with_new_type_statthrough the call chainchange_out_bias,change_type_map,compute_or_load_stat, model API methods.Changes
New pt_expt models
deepmd/pt_expt/model/dipole_model.pydeepmd/pt_expt/model/polar_model.pydeepmd/pt_expt/model/dos_model.pydeepmd/pt_expt/model/property_model.pydeepmd/pt_expt/model/dp_zbl_model.pyArchitecture refactoring
deepmd/pt_expt/atomic_model/layer — models now wrap dpmodel atomic models directlyBaseModel: remove concrete methods/data, add plugin registrymake_modelso backends (dp, pt_expt) inherit shared model logic from dpmodelget_out_bias/set_out_biasintobase_atomic_model.pyget_observed_type_listto abstract API and implement in dpmodel, pt, pdmodel_change_out_biasin pt/pd training code (chore(pt): mv the input stat update to model_change_out_bias #5266)Bug fixes
general_fitting.change_type_map: use array-api-compat ops instead of bare numpy (breaks pt_expt)make_model.change_type_map: properly forwardmodel_with_new_type_statto atomic modelstat.py: fix in-place mutation issueTests
test_dipole.py,test_polar.py,test_dos.py,test_property.pyandtest_zbl_ener.py(~1400 lines each)test_ener.pywith pt_expt and full model API coveragetest_dipole_model.py,test_polar_model.py,test_dos_model.py,test_property_model.py,test_dp_zbl_model.pytest_get_model_def_script,test_get_min_nbor_dist,test_set_case_embdacross all 6 model test filesmodel_change_out_biastests in pt/pd training tests (chore(pt): mv the input stat update to model_change_out_bias #5266)Summary by CodeRabbit
New Features
Refactor
Tests