refactor(data): replace fragile regex in __inc__ and _get_cat_dim with explicit key sets

Open
#58 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
3/5
Estimated time
1-2 days
Newbie friendliness
72/100
Issue type
Refactor
Clarity
Clearly specified
Activity status
Quiet
Tech stack
python
Domain
backend, data

Research direction

Start in nvalchemi/data/data.py at DataMixin.inc() and compare its current regex behavior with the explicit key sets used by _get_field_level(). Then inspect nvalchemi/data/datapipes/backends/zarr.py at _get_cat_dim(); done means both sites use the agreed explicit index-key set and no longer match unrelated field names by substring.

Written by the indexing model from the issue text.

Description

enhancement

Context

During the attribute rename work in #56, we added neighbor_list to the regex pattern in DataMixin.__inc__() (nvalchemi/data/data.py:155):

return self.num_nodes if bool(re.search("(index|face|neighbor_list)", key)) else 0

This method determines which fields receive cumulative node-index offsets when batching. The regex approach is implicit and fragile — it matches any field whose name contains the substring, so a field like neighbor_list_shifts would incorrectly match (flagged by @zubatyuk in PR #56 review).

A similar pattern exists in the Zarr backend's _get_cat_dim() (nvalchemi/data/datapipes/backends/zarr.py:231):

if bool(re.search("(index|face)", key)):
    return -1

In contrast, _get_field_level() in the same file already does this correctly — it checks membership in AtomicData._default_node_keys, _default_edge_keys, and _default_system_keys.

Proposed fix

Replace both regex sites with explicit set membership, as @zubatyuk suggested:

  1. Add an _index_keys ClassVar on AtomicData (or DataMixin) containing the canonical set of fields that require node-index offsets (currently {"neighbor_list"}).
  2. Update __inc__ to use key in self._index_keys instead of regex.
  3. Update _get_cat_dim to use AtomicData._index_keys (or an equivalent edge-key set) instead of regex.

This makes the behavior explicit, self-documenting, and safe against substring collisions.

Affected files

  • nvalchemi/data/data.py — DataMixin.__inc__()
  • nvalchemi/data/datapipes/backends/zarr.py — _get_cat_dim()

Origin

Identified during review of PR #56 (attribute rename). Deferred as out of scope for that PR.

Dominant language
Python
Stars
167
Forks
38
Avg merge
8d 15h
Merged PRs (30d)
7

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/nvalchemi-toolkit

All issues in NVIDIA/nvalchemi-toolkit

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.