feat(dpa-adapt): add grouped embedding for system-level property prediction

[zirenjin]

Summary

This PR adds grouped embedding support to DPA-ADAPT, enabling frame-level atomistic representations to be pooled into group-level embeddings for system-level property prediction.

The main use case is downstream adaptation where one training label corresponds to a collection of structures rather than a single frame, such as polymers, mixtures, reaction assemblies, or catalyst configurations.

Main Changes

• Add grouped property model support with differentiable frame-to-group pooling.
• Add group_id, weight, and pool_mask data fields for grouping frames and controlling atom-level pooling.
• Add grouped data conversion utilities and CLI support.
• Add grouped-aware dataloading and batch sampling so all frames from the same group stay together.
• Support distributed training by ensuring each group is assigned to a single rank.
• Add tests covering grouped data conversion, grouped sampling, grouped loss/model behavior, and DDP-safe batching.

Notes

This feature is intended for DPA-ADAPT tasks where the prediction target is defined at the system/sample level, while the input information is provided by multiple atomistic structures.

Summary by CodeRabbit

• New Features
  • Added grouped-property training support with new group-level model and loss, including group-aware fitting and grouped-data batching.
  • Introduced grouped data tooling (writers, offline grouped dataset loading), plus a new mark-groups CLI command and grouped polymer dataset generation.
• Bug Fixes
  • Hardened grouped pooling/masking for padded/virtual atoms and improved validation for per-group label and fparam consistency; distributed sampling now keeps groups intact.
• Documentation
  • Added a guide for “Assembly Groups for Multi-Component Properties”.
• Tests
  • Added extensive unit and end-to-end coverage for grouped training, batching, and serialization.

[codecov]

Codecov Report

❌ Patch coverage is 83.92857% with 90 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.63%. Comparing base (8de6596) to head (f86712e).
⚠️ Report is 64 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/pt/model/model/group_property_model.py	74.62%	51 Missing ⚠️
deepmd/pt/model/task/group_property.py	87.61%	14 Missing ⚠️
deepmd/pt/loss/group_property.py	83.58%	11 Missing ⚠️
deepmd/pt/utils/grouped.py	90.47%	8 Missing ⚠️
deepmd/pt/utils/dataloader.py	92.30%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5741      +/-   ##
==========================================
- Coverage   81.29%   78.63%   -2.67%     
==========================================
  Files         990     1058      +68     
  Lines      111018   122339   +11321     
  Branches     4234     4411     +177     
==========================================
+ Hits        90250    96197    +5947     
- Misses      19240    24565    +5325     
- Partials     1528     1577      +49     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

dpa_adapt/grouped/_core.py (2)

222-229: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Restrict overrides to ComponentSpec data fields.

hasattr(component, key) also accepts methods like validate, so a typo or bad override can replace behavior instead of being rejected.

Proposed fix

+from dataclasses import (
+    dataclass,
+    field,
+    fields,
+)
...
     def add_component(
         self, component: ComponentSpec, **overrides: Any
     ) -> ComponentSpec:
+        valid_fields = {item.name for item in fields(ComponentSpec)}
         for key, value in overrides.items():
-            if not hasattr(component, key):
+            if key not in valid_fields:
                 raise TypeError(f"Unknown ComponentSpec field: {key}")
             setattr(component, key, value)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dpa_adapt/grouped/_core.py` around lines 222 - 229, The add_component method
is allowing overrides for any attribute that exists on ComponentSpec, including
methods like validate, which can overwrite behavior instead of rejecting bad
keys. Update add_component to validate overrides only against ComponentSpec’s
data fields (for example, its declared dataclass fields) and keep the TypeError
for any unknown or non-field key so only real component data can be overridden.

---

350-358: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Honor fparam_schema even when all values come from defaults.

_has_fparam() ignores self.fparam_schema, so calling set_fparam_schema() with defaults but no per-group fparam leaves tensor_fields.fparam as None and skips fparam.npy entirely.

Proposed fix

     def _has_fparam(self) -> bool:
-        return any(group.fparam for group in self.groups)
+        return bool(self.fparam_schema) or any(group.fparam for group in self.groups)

Also applies to: 383-405

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dpa_adapt/grouped/_core.py` around lines 350 - 358, The grouped adaptation
path is dropping fparam data whenever only defaults are provided because
`_has_fparam()` does not account for `self.fparam_schema`. Update the logic in
the grouped core flow (the `tensor_fields` setup and the related `fparam_schema`
handling in the same class, including the later save/load path) so
`set_fparam_schema()` alone is enough to enable `fparam` processing and
`fparam.npy` generation. Use the existing `fparam_schema` and `_has_fparam()`
symbols in `dpa_adapt/grouped/_core.py` to ensure `tensor_fields["fparam"]` is
populated whenever the schema is present, even if no per-group `fparam` values
exist.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deepmd/pt/model/model/group_property_model.py`:
- Around line 242-251: The neighbor-list construction in group_property_model’s
input handling is hard-coding mixed neighbor lists, which breaks descriptors
that need type-distinguished neighbors. Update the
extend_input_and_build_neighbor_list call inside the relevant model method to
use the descriptor’s mixed_types() result from the descriptor/model contract, or
add an explicit initialization-time check in the group_property_model path to
reject unsupported non-mixed descriptors. Keep the fix centered on the
mixed_types() behavior so the selected descriptor and neighbor-list mode stay
consistent.

In `@dpa_adapt/grouped/_core.py`:
- Around line 326-343: Sanitized group system paths can collide when different
group.key values map to the same _safe_name(), causing _write_group_system and
the manifest entries built in the self.groups loop to target the same directory.
Update the grouped/_core.py logic around _safe_name, systems_root,
_write_group_system, and _group_manifest so each group gets a unique system path
even after sanitization, for example by incorporating group_idx or another
stable disambiguator into the directory name before writing and recording the
manifest.

---

Outside diff comments:
In `@dpa_adapt/grouped/_core.py`:
- Around line 222-229: The add_component method is allowing overrides for any
attribute that exists on ComponentSpec, including methods like validate, which
can overwrite behavior instead of rejecting bad keys. Update add_component to
validate overrides only against ComponentSpec’s data fields (for example, its
declared dataclass fields) and keep the TypeError for any unknown or non-field
key so only real component data can be overridden.
- Around line 350-358: The grouped adaptation path is dropping fparam data
whenever only defaults are provided because `_has_fparam()` does not account for
`self.fparam_schema`. Update the logic in the grouped core flow (the
`tensor_fields` setup and the related `fparam_schema` handling in the same
class, including the later save/load path) so `set_fparam_schema()` alone is
enough to enable `fparam` processing and `fparam.npy` generation. Use the
existing `fparam_schema` and `_has_fparam()` symbols in
`dpa_adapt/grouped/_core.py` to ensure `tensor_fields["fparam"]` is populated
whenever the schema is present, even if no per-group `fparam` values exist.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 56fdff0e-02f8-4de0-aadb-a2b1ec2d5869

📥 Commits

Reviewing files that changed from the base of the PR and between 82c1d73 and 8bb0d6b.

📒 Files selected for processing (25)

• deepmd/pt/loss/__init__.py
• deepmd/pt/loss/group_property.py
• deepmd/pt/model/model/group_property_model.py
• deepmd/pt/model/task/group_property.py
• deepmd/pt/train/training.py
• deepmd/pt/utils/dataloader.py
• deepmd/pt/utils/grouped.py
• doc/dpa_adapt/assembly_scenarios.md
• dpa_adapt/finetuner.py
• dpa_adapt/grouped/_aggregation.py
• dpa_adapt/grouped/_convert.py
• dpa_adapt/grouped/_core.py
• dpa_adapt/grouped/_offline.py
• dpa_adapt/grouped/_polymer.py
• dpa_adapt/trainer.py
• source/tests/dpa_adapt/test_assemblies.py
• source/tests/dpa_adapt/test_grouped_convert.py
• source/tests/dpa_adapt/test_grouped_dataset.py
• source/tests/dpa_adapt/test_grouped_end_to_end.py
• source/tests/dpa_adapt/test_grouped_finetuner.py
• source/tests/dpa_adapt/test_grouped_hardening.py
• source/tests/dpa_adapt/test_polymer_builder.py
• source/tests/dpa_adapt/test_pooling.py
• source/tests/pt/test_group_property.py
• source/tests/pt/test_group_property_hardening.py

✅ Files skipped from review due to trivial changes (1)

• doc/dpa_adapt/assembly_scenarios.md

🚧 Files skipped from review as they are similar to previous changes (22)

• deepmd/pt/loss/init.py
• source/tests/dpa_adapt/test_grouped_finetuner.py
• source/tests/pt/test_group_property.py
• deepmd/pt/train/training.py
• source/tests/dpa_adapt/test_pooling.py
• source/tests/dpa_adapt/test_grouped_hardening.py
• source/tests/dpa_adapt/test_grouped_dataset.py
• source/tests/dpa_adapt/test_assemblies.py
• source/tests/dpa_adapt/test_polymer_builder.py
• dpa_adapt/grouped/_offline.py
• deepmd/pt/utils/dataloader.py
• source/tests/dpa_adapt/test_grouped_end_to_end.py
• deepmd/pt/loss/group_property.py
• dpa_adapt/trainer.py
• deepmd/pt/utils/grouped.py
• source/tests/dpa_adapt/test_grouped_convert.py
• dpa_adapt/grouped/_polymer.py
• deepmd/pt/model/task/group_property.py
• source/tests/pt/test_group_property_hardening.py
• dpa_adapt/finetuner.py
• dpa_adapt/grouped/_aggregation.py
• dpa_adapt/grouped/_convert.py

[njzjz-bot]

Thanks for the large addition. I focused on the grouped-training integration and found several correctness issues that should be fixed before merge:

1. Grouped frozen-sklearn extraction ignores real_atom_types.npy.
   The grouped writer intentionally writes type.raw as a uniform placeholder and stores the real per-frame atom types, including virtual -1 padding, in real_atom_types.npy. The frozen-sklearn descriptor extraction path still tiles system.data["atom_types"] for every frame, so heterogeneous grouped systems are described as if every atom had the first type. pool_mask only affects final pooling and cannot fix descriptor typing or neighbor-list construction.

2. Masked offline pooling is not NaN-safe.
   The offline masked mean / sum / std paths multiply descriptors by the mask. If an excluded virtual atom has a non-finite descriptor, 0 * NaN remains NaN, and the later nan_to_num can zero/corrupt the whole pooled feature instead of dropping only the masked atoms. Please use the same torch.where(mask > 0, descriptor, 0)-style masking pattern as the grouped PT model path before reduction.

3. Descriptor cache keys miss grouped inputs.
   The descriptor cache currently fingerprints coordinates/types/cells, but grouped frozen-sklearn features also depend on pool_mask.npy and, after fixing item 1, real_atom_types.npy. Changing group markers or mixed-type real atom types can silently reuse stale cached descriptors.

4. Missing/HDF5 group_id handling is inconsistent between sampler and loss.
   The grouped dataloader falls back to an all-zero group id when filesystem group_id.npy cannot be loaded, treating the whole system as one group. The loss, however, falls back to one group per frame when group_id is absent. This can ignore the configured batch size, break DDP when the implicit single group is fewer than the number of ranks, and incorrectly treats HDF5 systems with valid in-data group_id as missing because the loader only scans set.* files.

5. Grouped-mode detection only checks the first set of the first system.
   dpa_adapt decides whether to switch to group_property from a single set directory. If later systems/sets contain markers but the first one does not, grouped training is not enabled; if the first is grouped but later sets are missing markers, training fails later or becomes inconsistent. Please scan all train/valid set dirs and reject mixed marker presence with a clear error.

6. group_reduce is implemented but not exposed in config validation.
   The fitting head supports group_reduce="mean"|"sum", but the argcheck schema reuses the property schema without adding this option, so users cannot set sum through normal input JSON even though direct Python tests cover it.

7. Grouped frozen-sklearn does not support the CLI's comma-separated multi-target API.
   Non-grouped _load_labels() accepts list[str], and the CLI turns --target-key a,b into a list. The grouped route passes that list into GroupedDataset, which resolves a single string key and will fail on a list target.

8. Several property-fitting options are accepted but ignored.
   The grouped fitting arg schema inherits property options, but GroupPropertyFittingNet swallows many of them via **kwargs (seed, numb_aparam, default_fparam, intensive, distinguish_types, resnet_dt, etc.). trainable list semantics are also collapsed to all(trainable). This makes normal configs look accepted while silently changing behavior.

I would recommend adding integration tests for: trained serialize/deserialize prediction equality, JSON config construction with group_reduce, grouped frozen-sklearn on mixed real_atom_types.npy, cache invalidation when masks/types change, HDF5 or in-data group_id, DDP/missing-group fallback, and grouped multi-target CLI input.

Reviewed by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)

[njzjz]

Possible reviewers based on changed lines, exact file history, and exact-file review history:

• @iProzd — 39 commits on changed files (deepmd/pt/loss/__init__.py, deepmd/pt/model/model/__init__.py, deepmd/pt/model/task/__init__.py, deepmd/pt/train/training.py, deepmd/pt/utils/dataloader.py, deepmd/utils/argcheck.py).

No review request was made automatically.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh