Skip to content

feat: support tensor custom properties#36

Merged
Ramlaoui merged 2 commits into
mainfrom
feat/tensor-property-core
Jun 5, 2026
Merged

feat: support tensor custom properties#36
Ramlaoui merged 2 commits into
mainfrom
feat/tensor-property-core

Conversation

@Ramlaoui
Copy link
Copy Markdown
Collaborator

@Ramlaoui Ramlaoui commented Jun 5, 2026

Summary

  • add tensor custom property storage tags with per-value shape metadata
  • allow molecule and atom custom tensor properties through set_property(..., scope=...)
  • make get_property/property_keys/has_property/delete_property work across molecule and atom scopes with unique keys
  • keep tensor shape out of schema equality so same-key tensor values can vary by record until a concatenating API asks for uniform shapes

Validation

  • cargo fmt --all --check
  • cargo check -p atompack -p atompack-py
  • cargo clippy --workspace --all-targets -- -D warnings
  • cargo test --workspace
  • uv run --no-sync --extra dev --locked pytest tests/test_atom_molecule.py tests/test_database.py -q
  • uv run --no-sync --extra dev --locked ruff format --check python
  • uv run --no-sync --extra dev --locked ruff check python

@Ramlaoui
Copy link
Copy Markdown
Collaborator Author

Ramlaoui commented Jun 5, 2026

Follow-up performance diagnosis for add_molecules:

The apparent slowdown was not the Rust storage path globally. The Rust smoke with builtin energy/forces stayed flat or faster, and an isolated Python split showed builtins-only object writes were also fine. The slowdown appeared when Python-owned molecules carried custom properties.

Root cause: the Python Database.add_molecules binding cloned every owned PyMolecule before writing. Adding tensor support increases PropertyValue from 32 bytes to 56 bytes, so even legacy scalar/vector custom properties became more expensive to clone. The canonical Python smoke prebuilds molecules outside the timed region, but the timed add_molecules call still paid this clone cost.

Fix added in 792b1bb: pass borrowed &Molecule references for owned Python molecules instead of cloning them; raw SOA view handling is unchanged.

Validation after the fix:

  • cargo fmt --all --check
  • cargo check -p atompack-py
  • uv run --no-sync --extra dev --locked pytest tests/test_database.py tests/test_atom_molecule.py -q -> 70 passed
  • ATOMPACK_PERF_COLOR=never make perf-smoke-py -> passed, write add_molecules ~611,725 mol/s

Ad hoc split on the patched branch:

  • builtins-only object writes: median ~760k mol/s
  • custom-property object writes: median ~622k mol/s

So the tensor feature exposed an avoidable clone cost; the branch now removes that cost instead of accepting the regression.

@Ramlaoui Ramlaoui merged commit f59546c into main Jun 5, 2026
5 checks passed
@Ramlaoui Ramlaoui deleted the feat/tensor-property-core branch June 5, 2026 12:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant