feat(dpa-adapt): add grouped embedding for system-level property prediction#5741
Open
zirenjin wants to merge 275 commits into
Open
feat(dpa-adapt): add grouped embedding for system-level property prediction#5741zirenjin wants to merge 275 commits into
zirenjin wants to merge 275 commits into
Conversation
The old _load_descriptor_model, _validate_type_map, _remap_atom_types, _extract_features_cached, and _extract_features method bodies were left in place alongside the new thin delegators, causing CodeQL 'variable defined multiple times' warnings. Removed the old bodies; kept _extract_features_cached on DPAFineTuner directly so that test patches on DPAFineTuner._extract_features are honoured through the cache wrapper.
… method - Replace try/except ImportError in _unwrap_multioutput with direct import (sklearn is always available when dpa_tools is loaded) - Remove _FrozenSklearnPipeline.extract_features_cached (dead code; the caching wrapper lives on DPAFineTuner so test patches work)
The workflow still referenced the deleted deepmd_property_tools/ directory. Updated paths trigger to deepmd/dpa_tools/** and test command to source/tests/dpa_tools/. Added torch to lightweight dependencies.
numpy 2.3+ requires Python>=3.11, but the property_tools_tests workflow runs on Python 3.10. Pin numpy>=1.21,<2.2 to keep the lightweight dependency install working on older Python.
refactor: unify dpa_tools CLI/API and merge deepmd_property_tools
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…flow-cleanup Update dpa_tools
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…flow-cleanup Fix CI dependency installation
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Fix dpa tools CI test stability
Docs/dpa tools readme
dpa_tools quickstart demo bundle, convert() glob multi-match, docs cleanup
…dding # Conflicts: # README.md # doc/dpa_adapt/index.rst # doc/dpa_adapt/input_formats.md # doc/dpa_adapt/overview.md # dpa_adapt/__init__.py # dpa_adapt/cli.py # dpa_adapt/data/__init__.py # dpa_adapt/data/convert.py # dpa_adapt/data/dataset.py # dpa_adapt/data/desc_cache.py # dpa_adapt/data/loader.py # dpa_adapt/data/smiles.py # dpa_adapt/data/type_map.py # dpa_adapt/finetuner.py # dpa_adapt/mft.py # dpa_adapt/trainer.py # source/tests/dpa_adapt/test_cli_smoke.py # source/tests/dpa_adapt/test_conditions.py # source/tests/dpa_adapt/test_convert.py # source/tests/dpa_adapt/test_type_map.py
test_pooling.py imported POOLING_PRIMITIVES, parse_pooling and
_pool_descriptor from dpa_adapt.finetuner, but those symbols never
existed -- the module only had a fixed _VALID_POOLING allowlist. The
failed import broke pytest collection, which under pytest-split failed
every Test Python split.
Add the module-level pooling API the tests expect:
- POOLING_PRIMITIVES = ("mean","sum","std","max","min") canonical order
- parse_pooling(): split on "+" / accept a sequence, dedup, canonical
order, reject unknown and empty specs
- _pool_descriptor(): reduce over the atom axis and concat in primitive
order, byte-identical to the previous per-string pooling
Route frozen-sklearn feature pooling through _pool_descriptor and replace
DPAFineTuner's allowlist check with parse_pooling plus an intensive-pooling
strategy guard (composite pooling requires strategy='frozen_sklearn').
Three grouped tests (added in 2cd5d60) exercised a grouped adaptation API that was never implemented, so they surfaced once the pooling collection error was cleared: - DPATrainer auto-detects grouped systems (set.*/{group_id,weight, pool_mask}.npy) and emits the group_property fitting_net + loss, and auto-sizes numb_fparam from set.*/fparam.npy; a `grouped` override is accepted. - DPAFineTuner accepts target= (alias for property_name); fit() accepts train=/valid= aliases. - frozen_sklearn fit()/predict() route grouped input through GroupedDataset (weighted per-group pooling + shared labels), so target_key/labels are optional for grouped data.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds grouped-property training in the PyTorch backend and grouped-data tooling in DPA-ADAPT. It introduces new grouped loss, model, batching, conversion, finetuning, and dataset-writing paths, plus documentation and tests. ChangesPyTorch backend grouped-property training
DPA-ADAPT grouped data toolkit
Estimated code review effort: 5 (Critical) | ~120 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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: 13
🧹 Nitpick comments (4)
deepmd/pt/model/model/group_property_model.py (1)
298-316: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winVectorize the per-group fparam consistency check.
This loop is O(n_groups) Python iterations, each re-scanning all
nframesrows viainverse == group_index. The scatter+gather+allclosepattern already used inGroupPropertyLoss._group_labels(deepmd/pt/loss/group_property.py) can validate per-group consistency without a Python loop, which matters for training with large numbers of groups per batch.🤖 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 `@deepmd/pt/model/model/group_property_model.py` around lines 298 - 316, The per-group fparam consistency check in GroupPropertyModel is still using a Python loop with repeated inverse == group_index scans, which is too slow for large batches. Refactor the fparam handling in the GroupPropertyModel path to use the same vectorized scatter/gather-style validation pattern as GroupPropertyLoss._group_labels, so consistency is checked without iterating over group_order in Python. Keep the existing reshape/to/cat flow, but replace the looped first/allclose check with a batched group-wise comparison that preserves the same ValueError behavior for inconsistent group rows.deepmd/pt/model/task/group_property.py (1)
29-38: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
ActivationFnhere._activation()only acceptstanh/relu/gelu/linear/none, while the shared PT utility already handlesrelu6,softplus,sigmoid,silu,gelu_tf, and custom SiLU variants. Using the common helper keeps this task aligned with the rest of the PT stack and avoids a second activation whitelist.🤖 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 `@deepmd/pt/model/task/group_property.py` around lines 29 - 38, The _activation helper in group_property.py is using a separate hardcoded whitelist, so update it to reuse the shared ActivationFn utility instead of manually mapping only tanh/relu/gelu/linear/none. Locate the _activation function and replace the custom branching with the common PT activation helper so it supports the broader set of activations already handled elsewhere, while preserving the same return type and fallback error behavior for unsupported names.source/tests/pt/test_group_property.py (2)
163-163: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winQuadratic list-flattening pattern (
sum(..., [])).Ruff flags all three occurrences as RUF017 (quadratic list summation). Per coding guidelines,
ruff check .must pass before committing or CI will fail.♻️ Proposed fix
+import itertools + rank_frames = [set(sum(rank0, [])), set(sum(rank1, []))] + rank_frames = [set(itertools.chain.from_iterable(rank0)), set(itertools.chain.from_iterable(rank1))]Apply the same
itertools.chain.from_iterable(...)substitution at lines 188 and 212.Also applies to: 188-188, 212-212
🤖 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 `@source/tests/pt/test_group_property.py` at line 163, The test helpers in test_group_property.py use quadratic list flattening via sum(..., []) in the rank frame construction, which triggers Ruff RUF017. Update the affected rank_frames and related flattening sites to use itertools.chain.from_iterable for flattening instead, and apply the same substitution to the other two matching occurrences in this test module so ruff check . passes cleanly.Sources: Coding guidelines, Linters/SAST tools
192-192: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUnused loop variable
rank(B007).
rankis never referenced in the loop body; ruff flags this and, per coding guidelines,ruff check .must pass before committing.♻️ Proposed fix
- for rank, frames in enumerate(frames_by_rank): + for _rank, frames in enumerate(frames_by_rank):🤖 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 `@source/tests/pt/test_group_property.py` at line 192, The loop in test_group_property.py introduces an unused variable `rank`, which triggers ruff B007. Update the `for` loop that iterates over `frames_by_rank` to remove the unused `rank` binding, since only `frames` is needed in the loop body. Keep the rest of the test logic unchanged and ensure the `frames_by_rank` iteration still works with the new loop variable pattern.Sources: Coding guidelines, Linters/SAST tools
🤖 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 329-336: `GroupPropertyModel` currently has `serialize()` but no
matching `deserialize()` classmethod, so serialized checkpoints cannot be
restored; add the inverse loader alongside `serialize()` and align it with the
patterns used by other `BaseModel` subclasses in this module. Reconstruct the
model from the serialized dict by restoring the descriptor, fitting network, and
type map, and make sure the new `deserialize()` is discoverable on
`GroupPropertyModel` so saved state can round-trip cleanly.
- Around line 182-201: The `charge_spin` argument is being accepted by
`GroupPropertyModel` but is not passed through to the descriptor call, so
charge-spin conditioning is silently ignored. Update the model’s forward/path
that invokes `self.descriptor(...)` to forward `charge_spin` whenever
`has_chg_spin_ebd()`, `has_default_chg_spin()`, or `get_default_chg_spin()`
indicate support, and keep the behavior consistent with the descriptor’s
expected inputs. Use the `GroupPropertyModel` and `self.descriptor` call site to
locate the fix.
- Around line 248-266: The fallback pool mask in group_property_model.py
currently uses all ones, which treats padded/virtual atoms as real atoms when
pool_mask is omitted. Update the logic in the frame-embedding path around
normalize_pool_mask_tensor and frame_embedding so the default mask excludes
virtual atoms by deriving it from atype >= 0 (or make pool_mask mandatory),
while keeping the existing zero-mask validation and weighted sum behavior
intact.
In `@deepmd/pt/utils/dataloader.py`:
- Around line 194-209: The distributed grouped batching path in
distributed_grouped_frame_batches() can create empty rank-local samplers when
the number of groups is smaller than the world size, which leads to
cycle_iterator() never making progress on some ranks. Add an early validation in
the dataloader path around the GroupDistributedBatchSampler construction to
reject this configuration, or adjust the grouping logic so each rank gets at
least one batch before DDP starts. Use the distributed branch in
deepmd/pt/utils/dataloader.py and the GroupDistributedBatchSampler /
cycle_iterator flow to locate the fix.
- Around line 354-369: `GroupDistributedBatchSampler.__iter__()` currently
increments a local `_epoch`, which can cause different ranks to generate
different shuffle seeds and drift out of sync; make the epoch advancement
synchronized across ranks so `_make_batches()` uses the same `base_seed +
_epoch` everywhere. Also add a guard before calling `cycle_iterator()`/batch
generation for the case where `len(unique_groups) < num_replicas`, and handle
empty-rank cases explicitly so ranks with no assigned groups do not loop
forever.
In `@dpa_adapt/finetuner.py`:
- Around line 1681-1695: The grouped evaluate path in finetuner.py is still
loading frame-level labels via _load_labels and reshaping them to match grouped
predictions, which can mismatch because predict() on GroupedDataset returns one
row per group. Update evaluate() to mirror the grouped training flow by using
the existing _grouped branch and GroupedDataset.get_labels() when getattr(self,
"_grouped", False) is true, so the labels align with the predictions before any
metric computation.
In `@dpa_adapt/grouped/_convert.py`:
- Around line 205-245: The loop in _convert.py is unpacking natoms from per_set
but never uses it, which triggers Ruff B007. Update the iteration over per_set
in the group-id writing block to stop binding the unused natoms value, or
replace it with a throwaway variable, and keep the rest of the logic in the same
GroupMarkerResult path unchanged.
- Around line 229-234: The pool mask generation in mark_groups/_convert.py is
still gated on real_types containing negatives, so overwrite=True can leave an
old pool_mask.npy behind when the current data no longer needs a mask. Update
the pool mask handling near the real_types check to also regenerate or remove
the file deterministically when overwrite is requested, even if no negative
entries are present. Use the existing _should_write(...), POOL_MASK_KEY, and
result.wrote_pool_mask flow in the same block so overwrite=True always
reconciles the on-disk pool_mask.npy with the current derivation.
In `@dpa_adapt/grouped/_core.py`:
- Around line 316-322: The manifest and per-group write logic for fparam are
inconsistent: _has_fparam() marks tensor_fields.fparam as present for the whole
dataset, but _write_group_system only writes fparam.npy when a given
group.fparam is truthy. Update the grouped writer in _core.py so the behavior is
consistent across all groups, either by validating that every group in the
Assembly has fparam when tensor_fields.fparam is set, or by applying
fparam_schema defaults before writing. Use _has_fparam, _write_group_system, and
the tensor_fields construction to locate and align the manifest and file-writing
paths.
- Around line 280-292: In `write`, make overwrite fully reset the output by
removing any existing `systems/*` contents before recreating the directory
structure so stale grouped systems are not rediscovered later. Also update
`_write_group_system()` to treat `manifest.tensor_fields.fparam` as
authoritative and always write `fparam.npy` when the manifest expects it, even
if an individual group has no local `fparam`, so the later load/validation path
stays consistent.
In `@dpa_adapt/grouped/_offline.py`:
- Around line 131-163: The system-discovery logic in _candidate_paths and
_paths_from_directory only checks immediate subdirectories, which can miss
nested grouped systems that mark_groups already finds recursively. Update
_paths_from_directory to mirror the recursive traversal used by _walk_systems in
dpa_adapt/grouped/_convert.py so GroupedDataset and has_grouped_markers discover
all marked systems at any depth. Keep the existing filtering for valid system
directories, but recurse through child directories instead of stopping at
path.iterdir() level. Add the needed os-based walk/import in the module if you
follow the same approach as _walk_systems.
- Around line 82-89: The grouped offline descriptor extraction in
load_or_extract/extract_features is ignoring the per-system pool mask, so padded
or virtual atoms still affect the averaged embedding. Update the offline sklearn
path in grouped/_offline.py to thread pool_mask.npy through the descriptor
loading/extraction flow, and make extract_features apply that mask before
averaging so only real atoms contribute to the final descriptors. Keep the
change localized around load_or_extract(), extract_features(), and the grouped
system handling that builds descriptors from self.systems.
In `@dpa_adapt/trainer.py`:
- Around line 406-412: Grouped fparam inference is happening too late in fit(),
so _validate_fparam(...) sees self.fparam_dim as 0 and skips validation for
grouped fparam.npy inputs. Move the grouped detection and fparam_dim inference
in Trainer.fit (using _systems_are_grouped(train_sys) and
_detect_fparam_dim(train_sys)) before the preflight validation, or rerun
_validate_fparam after those values are set, so grouped inputs are validated
with the correct shape/frame expectations.
---
Nitpick comments:
In `@deepmd/pt/model/model/group_property_model.py`:
- Around line 298-316: The per-group fparam consistency check in
GroupPropertyModel is still using a Python loop with repeated inverse ==
group_index scans, which is too slow for large batches. Refactor the fparam
handling in the GroupPropertyModel path to use the same vectorized
scatter/gather-style validation pattern as GroupPropertyLoss._group_labels, so
consistency is checked without iterating over group_order in Python. Keep the
existing reshape/to/cat flow, but replace the looped first/allclose check with a
batched group-wise comparison that preserves the same ValueError behavior for
inconsistent group rows.
In `@deepmd/pt/model/task/group_property.py`:
- Around line 29-38: The _activation helper in group_property.py is using a
separate hardcoded whitelist, so update it to reuse the shared ActivationFn
utility instead of manually mapping only tanh/relu/gelu/linear/none. Locate the
_activation function and replace the custom branching with the common PT
activation helper so it supports the broader set of activations already handled
elsewhere, while preserving the same return type and fallback error behavior for
unsupported names.
In `@source/tests/pt/test_group_property.py`:
- Line 163: The test helpers in test_group_property.py use quadratic list
flattening via sum(..., []) in the rank frame construction, which triggers Ruff
RUF017. Update the affected rank_frames and related flattening sites to use
itertools.chain.from_iterable for flattening instead, and apply the same
substitution to the other two matching occurrences in this test module so ruff
check . passes cleanly.
- Line 192: The loop in test_group_property.py introduces an unused variable
`rank`, which triggers ruff B007. Update the `for` loop that iterates over
`frames_by_rank` to remove the unused `rank` binding, since only `frames` is
needed in the loop body. Keep the rest of the test logic unchanged and ensure
the `frames_by_rank` iteration still works with the new loop variable pattern.
🪄 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: b314c46e-0047-4939-81cc-82941207452c
📒 Files selected for processing (36)
deepmd/pt/loss/__init__.pydeepmd/pt/loss/group_property.pydeepmd/pt/model/model/__init__.pydeepmd/pt/model/model/group_property_model.pydeepmd/pt/model/task/__init__.pydeepmd/pt/model/task/group_property.pydeepmd/pt/train/training.pydeepmd/pt/utils/dataloader.pydeepmd/pt/utils/grouped.pydeepmd/utils/argcheck.pydoc/dpa_adapt/assembly_scenarios.mddpa_adapt/__init__.pydpa_adapt/cli.pydpa_adapt/data/aggregation.pydpa_adapt/data/assemblies.pydpa_adapt/data/grouped_convert.pydpa_adapt/data/grouped_dataset.pydpa_adapt/data/polymer.pydpa_adapt/finetuner.pydpa_adapt/grouped/__init__.pydpa_adapt/grouped/_aggregation.pydpa_adapt/grouped/_convert.pydpa_adapt/grouped/_core.pydpa_adapt/grouped/_offline.pydpa_adapt/grouped/_polymer.pydpa_adapt/trainer.pysource/tests/dpa_adapt/test_assemblies.pysource/tests/dpa_adapt/test_grouped_convert.pysource/tests/dpa_adapt/test_grouped_dataset.pysource/tests/dpa_adapt/test_grouped_end_to_end.pysource/tests/dpa_adapt/test_grouped_finetuner.pysource/tests/dpa_adapt/test_grouped_hardening.pysource/tests/dpa_adapt/test_polymer_builder.pysource/tests/dpa_adapt/test_pooling.pysource/tests/pt/test_group_property.pysource/tests/pt/test_group_property_hardening.py
Comment on lines
+354
to
+369
| def _make_batches(self) -> list[list[int]]: | ||
| base_seed = 0 if self.seed is None else self.seed | ||
| rng = np.random.default_rng(base_seed + self._epoch) | ||
| return distributed_grouped_frame_batches( | ||
| self.group_ids, | ||
| self.max_frames, | ||
| num_replicas=self.num_replicas, | ||
| rank=self.rank, | ||
| shuffle=self.shuffle, | ||
| rng=rng, | ||
| ) | ||
|
|
||
| def __iter__(self): | ||
| self._batches = self._make_batches() | ||
| self._epoch += 1 | ||
| yield from self._batches |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the target file and related symbols first.
ast-grep outline deepmd/pt/utils/dataloader.py --view expanded || true
printf '\n=== Search symbols ===\n'
rg -n "GroupDistributedBatchSampler|distributed_grouped_frame_batches|cycle_iterator|set_epoch|_epoch|group_ids|max_frames" deepmd -S
printf '\n=== Candidate tests ===\n'
rg -n "GroupDistributedBatchSampler|distributed_grouped_frame_batches|cycle_iterator|group.*replica|num_replicas" tests deepmd/tests -S || true
printf '\n=== File size ===\n'
wc -l deepmd/pt/utils/dataloader.py
printf '\n=== Relevant slice around lines 320-420 ===\n'
sed -n '320,420p' deepmd/pt/utils/dataloader.pyRepository: deepmodeling/deepmd-kit
Length of output: 18952
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search for tests / usages in the repo.
rg -n "GroupDistributedBatchSampler|distributed_grouped_frame_batches|cycle_iterator|set_epoch|_epoch|group_ids|max_frames" . -S --glob '!**/.git/**'
# Print the relevant section of the target file with line numbers.
cat -n deepmd/pt/utils/dataloader.py | sed -n '320,420p'Repository: deepmodeling/deepmd-kit
Length of output: 30779
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '=== grouped.py outline ===\n'
ast-grep outline deepmd/pt/utils/grouped.py --view expanded || true
printf '\n=== grouped.py relevant slice ===\n'
wc -l deepmd/pt/utils/grouped.py
cat -n deepmd/pt/utils/grouped.py | sed -n '145,210p'
printf '\n=== training.py cycle_iterator slice ===\n'
cat -n deepmd/pt/train/training.py | sed -n '255,380p'
printf '\n=== group property tests slice ===\n'
cat -n source/tests/pt/test_group_property.py | sed -n '130,230p'Repository: deepmodeling/deepmd-kit
Length of output: 13832
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '=== group_property_hardening slice ===\n'
cat -n source/tests/pt/test_group_property_hardening.py | sed -n '220,280p'
printf '\n=== lmdb dataloader set_epoch test slice ===\n'
cat -n source/tests/pt/test_lmdb_dataloader.py | sed -n '490,530p'Repository: deepmodeling/deepmd-kit
Length of output: 3498
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '=== grouped.py full relevant region ===\n'
cat -n deepmd/pt/utils/grouped.py | sed -n '100,210p'
printf '\n=== search for set_epoch support on batch samplers ===\n'
rg -n "set_epoch" deepmd/pt source/tests/pt -SRepository: deepmodeling/deepmd-kit
Length of output: 4912
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '=== DpLoaderSet _build_dataloaders group branch ===\n'
cat -n deepmd/pt/utils/dataloader.py | sed -n '177,214p'
printf '\n=== any guards for too-few groups / empty batch samplers ===\n'
rg -n "num_replicas|group_ids|empty|zero batches|len\\(.*\\) == 0|no batches" deepmd/pt source/tests/pt -SRepository: deepmodeling/deepmd-kit
Length of output: 30294
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import numpy as np
def grouped_frame_batches(group_ids, max_frames, shuffle=True, rng=None):
groups = {}
for frame_idx, group_id in enumerate(np.asarray(group_ids).astype(np.int64, copy=False)):
groups.setdefault(int(group_id), []).append(frame_idx)
group_items = list(groups.values())
if shuffle:
rng = rng or np.random.default_rng()
order = rng.permutation(len(group_items))
group_items = [group_items[int(ii)] for ii in order]
batches = []
current = []
limit = max(int(max_frames), 1)
for indices in group_items:
if current and len(current) + len(indices) > limit:
batches.append(current)
current = []
current.extend(indices)
if len(current) >= limit:
batches.append(current)
current = []
if current:
batches.append(current)
return batches
def distributed_grouped_frame_batches(group_ids, max_frames, num_replicas, rank, shuffle=True, rng=None):
if shuffle and rng is None:
rng = np.random.default_rng(0)
group_items = grouped_frame_batches(group_ids, max_frames, shuffle=shuffle, rng=rng) # wrong helper? no
PYRepository: deepmodeling/deepmd-kit
Length of output: 161
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import numpy as np
def group_frame_indices(group_ids):
groups = {}
for frame_idx, group_id in enumerate(np.asarray(group_ids).astype(np.int64, copy=False)):
groups.setdefault(int(group_id), []).append(frame_idx)
return list(groups.values())
def shuffle_group_items(group_items, shuffle, rng):
if not shuffle:
return list(group_items)
rng = rng or np.random.default_rng()
order = rng.permutation(len(group_items))
return [group_items[int(ii)] for ii in order]
def pack_group_items(group_items, max_frames):
batches = []
current = []
limit = max(int(max_frames), 1)
for indices in group_items:
if current and len(current) + len(indices) > limit:
batches.append(current)
current = []
current.extend(indices)
if len(current) >= limit:
batches.append(current)
current = []
if current:
batches.append(current)
return batches
def distributed_grouped_frame_batches(group_ids, max_frames, num_replicas, rank, shuffle=True, rng=None):
num_replicas = int(num_replicas)
rank = int(rank)
if shuffle and rng is None:
rng = np.random.default_rng(0)
group_items = shuffle_group_items(group_frame_indices(group_ids), shuffle=shuffle, rng=rng)
rank_items = group_items[rank::num_replicas]
return pack_group_items(rank_items, max_frames)
group_ids = np.array([0,0,1,1,2,3,3,4], dtype=np.int64)
for epoch0, epoch1 in [(0,0), (0,1), (1,0), (1,1)]:
out = []
for rank, epoch in [(0, epoch0), (1, epoch1)]:
rng = np.random.default_rng(11 + epoch)
out.append(distributed_grouped_frame_batches(group_ids, 3, 2, rank, True, rng))
frames = [set(sum(batches, [])) for batches in out]
print(f"epochs rank0={epoch0}, rank1={epoch1}")
print(" rank0", out[0], "frames", sorted(frames[0]))
print(" rank1", out[1], "frames", sorted(frames[1]))
print(" disjoint", frames[0].isdisjoint(frames[1]), "cover", frames[0] | frames[1] == set(range(len(group_ids))))
print()
# Empty-rank case: fewer groups than replicas
group_ids = np.array([0,0,1,1], dtype=np.int64) # 2 groups, 4 replicas
for rank in range(4):
batches = distributed_grouped_frame_batches(group_ids, 3, 4, rank, False, np.random.default_rng(0))
print("rank", rank, "batches", batches, "len", len(batches))
PYRepository: deepmodeling/deepmd-kit
Length of output: 837
Sync GroupDistributedBatchSampler epochs across ranks
__iter__() advances a local _epoch, but each rank restarts cycle_iterator() independently. Once ranks wrap at different times, base_seed + _epoch diverges, the shuffled group order can differ before group_items[rank::num_replicas], and groups can be duplicated or dropped across ranks.
Also guard len(unique_groups) < num_replicas: some ranks get an empty batch list, and cycle_iterator() will spin forever there.
🤖 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 `@deepmd/pt/utils/dataloader.py` around lines 354 - 369,
`GroupDistributedBatchSampler.__iter__()` currently increments a local `_epoch`,
which can cause different ranks to generate different shuffle seeds and drift
out of sync; make the epoch advancement synchronized across ranks so
`_make_batches()` uses the same `base_seed + _epoch` everywhere. Also add a
guard before calling `cycle_iterator()`/batch generation for the case where
`len(unique_groups) < num_replicas`, and handle empty-rank cases explicitly so
ranks with no assigned groups do not loop forever.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5741 +/- ##
==========================================
- Coverage 81.29% 79.51% -1.79%
==========================================
Files 990 1018 +28
Lines 111018 116074 +5056
Branches 4234 4274 +40
==========================================
+ Hits 90250 92292 +2042
- Misses 19240 22240 +3000
- Partials 1528 1542 +14 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
CodeRabbit (correctness / data integrity): - group_property_model: forward charge_spin into the descriptor; derive the default pool_mask from atype>=0 so padded/virtual atoms are excluded from frame embeddings; add deserialize() to match serialize(). - dataloader: partition grouped batches across ranks with an epoch-independent seed and reshuffle only each rank's own batch order, so ranks that wrap epochs at different times can no longer duplicate/drop groups; grouped.py rejects setups with fewer groups than ranks. - grouped/_core: clear stale system dirs on overwrite; write fparam.npy for every group (schema order, defaults filled) when the manifest advertises a uniform fparam field. - grouped/_convert: regenerate/remove stale pool_mask.npy on overwrite; drop unused loop var (Ruff B007). - grouped/_offline: discover systems recursively (match mark_groups); thread the pool mask into offline descriptor pooling. - finetuner: exclude virtual atoms from offline pooling; grouped evaluate() uses group-level labels. - trainer: infer grouped numb_fparam before the fparam preflight so grouped fparam.npy files are validated. CodeQL / pre-commit: - remove deprecated dpa_adapt.data grouped shims and unused imports; break the finetuner<->_offline import cycle; drop unnecessary del and redundant annotation quotes; guard possibly-uninitialized test vars; isort/ruff/ ruff-format/mdformat.
Contributor
There was a problem hiding this comment.
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 winRestrict overrides to
ComponentSpecdata fields.
hasattr(component, key)also accepts methods likevalidate, 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 winHonor
fparam_schemaeven when all values come from defaults.
_has_fparam()ignoresself.fparam_schema, so callingset_fparam_schema()with defaults but no per-groupfparamleavestensor_fields.fparamasNoneand skipsfparam.npyentirely.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
📒 Files selected for processing (25)
deepmd/pt/loss/__init__.pydeepmd/pt/loss/group_property.pydeepmd/pt/model/model/group_property_model.pydeepmd/pt/model/task/group_property.pydeepmd/pt/train/training.pydeepmd/pt/utils/dataloader.pydeepmd/pt/utils/grouped.pydoc/dpa_adapt/assembly_scenarios.mddpa_adapt/finetuner.pydpa_adapt/grouped/_aggregation.pydpa_adapt/grouped/_convert.pydpa_adapt/grouped/_core.pydpa_adapt/grouped/_offline.pydpa_adapt/grouped/_polymer.pydpa_adapt/trainer.pysource/tests/dpa_adapt/test_assemblies.pysource/tests/dpa_adapt/test_grouped_convert.pysource/tests/dpa_adapt/test_grouped_dataset.pysource/tests/dpa_adapt/test_grouped_end_to_end.pysource/tests/dpa_adapt/test_grouped_finetuner.pysource/tests/dpa_adapt/test_grouped_hardening.pysource/tests/dpa_adapt/test_polymer_builder.pysource/tests/dpa_adapt/test_pooling.pysource/tests/pt/test_group_property.pysource/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
- isort: reorder imports in group_property_model.py and finetuner.py - disallow-caps: spell the product name "DeePMD" (not "DeepMD") in the grouped module/doc/test comments and docstrings
- test_grouped_end_to_end: use pytest.importorskip so the backend names are always bound (resolves py/uninitialized-local-variable; try/except + skip isn't modelled as terminating by CodeQL). - finetuner._fit_sklearn: return without a value in the grouped branch so the None-returning method no longer mixes implicit/explicit returns (py/mixed-returns).
The grouped property head pools an un-normalized frame/group embedding concatenated with fparam (no descriptor-style input statistics). At that input scale tanh saturates, so the head learns only its bias and collapses predictions to the global label mean. Direct-overfit evidence: on the same few samples, tanh plateaus at ~30 RMSE (constant output) while GELU reaches <0.5 and fits distinct labels. - deepmd/pt/model/task/group_property.py: default activation_function -> "gelu" - deepmd/utils/argcheck.py: fitting_group_property defaults activation to gelu (property head keeps tanh) - dpa_adapt/trainer.py: grouped config emits activation_function="gelu" (a user fitting_net_params override still wins) - test: assert the generated grouped config uses gelu
njzjz-bot
suggested changes
Jul 8, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
Thanks for the large addition. I focused on the grouped-training integration and found several correctness issues that should be fixed before merge:
-
Grouped frozen-sklearn extraction ignores
real_atom_types.npy.
The grouped writer intentionally writestype.rawas a uniform placeholder and stores the real per-frame atom types, including virtual-1padding, inreal_atom_types.npy. The frozen-sklearn descriptor extraction path still tilessystem.data["atom_types"]for every frame, so heterogeneous grouped systems are described as if every atom had the first type.pool_maskonly affects final pooling and cannot fix descriptor typing or neighbor-list construction. -
Masked offline pooling is not NaN-safe.
The offline maskedmean/sum/stdpaths multiply descriptors by the mask. If an excluded virtual atom has a non-finite descriptor,0 * NaNremainsNaN, and the laternan_to_numcan zero/corrupt the whole pooled feature instead of dropping only the masked atoms. Please use the sametorch.where(mask > 0, descriptor, 0)-style masking pattern as the grouped PT model path before reduction. -
Descriptor cache keys miss grouped inputs.
The descriptor cache currently fingerprints coordinates/types/cells, but grouped frozen-sklearn features also depend onpool_mask.npyand, after fixing item 1,real_atom_types.npy. Changing group markers or mixed-type real atom types can silently reuse stale cached descriptors. -
Missing/HDF5
group_idhandling is inconsistent between sampler and loss.
The grouped dataloader falls back to an all-zero group id when filesystemgroup_id.npycannot be loaded, treating the whole system as one group. The loss, however, falls back to one group per frame whengroup_idis 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-datagroup_idas missing because the loader only scansset.*files. -
Grouped-mode detection only checks the first set of the first system.
dpa_adaptdecides whether to switch togroup_propertyfrom 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. -
group_reduceis implemented but not exposed in config validation.
The fitting head supportsgroup_reduce="mean"|"sum", but the argcheck schema reuses the property schema without adding this option, so users cannot setsumthrough normal input JSON even though direct Python tests cover it. -
Grouped frozen-sklearn does not support the CLI's comma-separated multi-target API.
Non-grouped_load_labels()acceptslist[str], and the CLI turns--target-key a,binto a list. The grouped route passes that list intoGroupedDataset, which resolves a single string key and will fail on a list target. -
Several property-fitting options are accepted but ignored.
The grouped fitting arg schema inherits property options, butGroupPropertyFittingNetswallows many of them via**kwargs(seed,numb_aparam,default_fparam,intensive,distinguish_types,resnet_dt, etc.).trainablelist semantics are also collapsed toall(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)
FinetuneRuleBuilder already auto-selects a multitask pretrained checkpoints only branch (and re-initializes the fitting net) when a single-task target like GroupPropertyModel finetunes with no --model-branch/finetune_head set, but this exact path -- direct DPA foundation-model checkpoint to GroupPropertyModel -- had no test coverage. Add two tests: the auto-pick + random-fitting case, and documenting the deterministic first-branch pick when several branches are available with no explicit selection.
…e group system names - GroupPropertyModel.forward hard-coded mixed_types=True when building the neighbor list instead of calling self.mixed_types(), so descriptors that require type-distinguished neighbor lists silently got a mixed one. - Assembly.write() sanitized group keys into system directory names without checking for collisions; two keys that sanitize to the same name (e.g. differing only in stripped punctuation) silently overwrote one another and the manifest recorded a duplicate path. Addresses the two still-open threads from the PR deepmodeling#5741 review.
…, group_id fallback Four correctness issues flagged in review, all specific to heterogeneous grouped systems (padded/virtual atoms, uniform type.raw placeholder): - frozen-sklearn extract_features() tiled the single, uniform atom_types array for every frame instead of reading the real, per-frame local types (with -1 padding) grouped systems store in real_atom_types.npy, so every atom in every frame was described as if it had the placeholder type. Added _real_atom_types_for_system() + remap_atom_types_preserving_padding() (which, unlike remap_atom_types, does not wrap -1 to the last checkpoint type via numpy negative-index fancy indexing). - _pool_descriptor() multiplied the raw descriptor by the pool mask before reducing (mean/sum/std); a non-finite value on a masked/virtual atom kept 0 * NaN == NaN and poisoned the whole frames pooled feature. Sanitize masked rows via torch.where(mask > 0, descrpt, 0) before any reduction, matching the pattern already used in GroupPropertyModel.forward. - The descriptor cache fingerprinted coords/types/cells but not pool_mask.npy or real_atom_types.npy, so changing group markers or real atom types on otherwise-identical padded geometry silently reused a stale cached descriptor. - GroupCompleteBatchSampler/GroupDistributedBatchSampler fell back to one giant group (all-zero group_id) when group_id.npy is missing, while GroupPropertyLoss falls back to one group per frame (torch.arange). The mismatch let the sampler ignore batch_size and could break DDP by handing an implicit single group to more ranks than it has frames for. Also rewrote load_group_ids_for_system() to read set.* through the systems own DPPath dirs (already backend-resolved by DeepmdData) instead of a raw pathlib glob, so HDF5-backed systems are checked the same way as on-disk ones instead of an in-data group_id being silently reported as missing.
…st the first _systems_are_grouped() only looked at the first set.* directory of the first train system to decide whether to switch the fitting/loss to group_property. If a later system had markers while the first did not, grouped mode stayed off and that systems group labels were silently dropped; if the first had markers but a later one did not, grouped mode turned on and then failed (or trained inconsistently) once that system was reached. valid_systems was not checked at all, so a grouped/ungrouped mismatch between train and valid went unnoticed too. Scan every set.* directory across both train_systems and valid_systems, and raise a clear DPADataError on any mix of grouped/ungrouped sets or a partially-marked set (some but not all of group_id/weight/pool_mask) instead of guessing.
GroupPropertyFittingNet supports group_reduce="mean"|"sum" (the frame-embedding -> group-embedding reduction), but fitting_group_property() reused the property schema verbatim aside from the activation-function default, with no group_reduce field. dargs strict-mode input.json validation (the same check dp --pt train runs) rejected it as an unknown key, so group_reduce="sum" was only reachable by constructing the model in Python directly, never through normal config files.
…rouped fit Non-grouped _load_labels() already accepts target_key as a list (the CLI turns --target-key a,b into ["a", "b"]), stacking each keys column into a (n_frames, n_keys) label array. GroupedDataset only accepted a single str and resolved/loaded one property.npy-style file, so passing the same list through the grouped route (_fit_sklearn_grouped already forwards whatever target_key it was given) failed. GroupedDataset.target_keys is now always a list; _read_system_group_rows reads one label file per key from each set.* dir and stacks them per frame, same convention as _load_labels (single key keeps its original shape, multiple keys are column-stacked) -- aggregate_weighted_groups already supported (n_items, task_dim) labels, so this was the only missing piece.
… options GroupPropertyFittingNet took **kwargs and accepted the full property-schema field list (numb_aparam, default_fparam, dim_case_embd, resnet_dt, intensive, distinguish_types, seed, ...) since it is a standalone MLP that never went through GeneralFitting, the base class that actually implements those. A config setting any of them passed validation and construction, then silently had no effect. trainable as a per-layer list was also collapsed to all(trainable), losing the per-layer freeze the property schema documents. - Replace **kwargs with two explicit, named, no-op parameters (type, mixed_types) for the two fields the generic model-building path always injects; any other unrecognized field now fails construction immediately with a normal TypeError instead of vanishing into kwargs. - fitting_group_property() drops numb_aparam/default_fparam/dim_case_embd/ resnet_dt/intensive/distinguish_types from the schema entirely, so setting one is a dargs strict-mode validation error, not a silently accepted no-op. - seed is now used: layer init is deterministic when given (scoped via fork_rng so it does not perturb the caller's global RNG state), and unseeded construction still draws from the global stream as before. - trainable=[...] now freezes each Linear layer individually instead of collapsing to all(trainable), matching the documented per-layer semantics and raising if the list length does not match len(neuron)+1.
| # coords/atom_types/cells alone are identical across these two | ||
| # systems -- only set.000/real_atom_types.npy (per-frame, with -1 | ||
| # padding) differs. The fingerprint must still change. | ||
| s1 = _make_system(tmp_path, "s1", natoms=3, nframes=2) |
| # systems -- only set.000/real_atom_types.npy (per-frame, with -1 | ||
| # padding) differs. The fingerprint must still change. | ||
| s1 = _make_system(tmp_path, "s1", natoms=3, nframes=2) | ||
| s2 = _make_system(tmp_path, "s2", natoms=3, nframes=2) |
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
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
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
group_id,weight, andpool_maskdata fields for grouping frames and controlling atom-level pooling.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
mark-groupsCLI command and grouped polymer dataset generation.fparamconsistency; distributed sampling now keeps groups intact.