Skip to content

feat(state): make pbc a per-system attribute#594

Open
lil-lon wants to merge 5 commits into
TorchSim:mainfrom
lil-lon:feat/per-system-pbc
Open

feat(state): make pbc a per-system attribute#594
lil-lon wants to merge 5 commits into
TorchSim:mainfrom
lil-lon:feat/per-system-pbc

Conversation

@lil-lon

@lil-lon lil-lon commented Jul 25, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Make SimState.pbc a per-system attribute of shape (n_systems, 3); a bool, list[bool], or (3,) tensor is broadcast to all systems. Moving it from _global_attributes to _system_attributes lets slicing, splitting, popping, and concatenation handle it like cell.
  • atoms_to_state / structures_to_state accept inputs with different pbc instead of raising, and state_to_atoms / state_to_structures / trajectory files store each system's own pbc row.
  • Ewald and PME models evaluate the periodic systems of a mixed batch and zero-fill systems with no periodic axis.
  • Add require_full_pbc; elastic tensor calculation now rejects systems that are not fully periodic.
  • The orb / NequIP / metatomic wrappers convert pbc back to the legacy (3,) row on a shallow copy (_to_legacy_pbc_state) before delegating, because those integrations broadcast state.pbc themselves assuming the old shape; mixed-pbc batches raise a clear error there.
  • Declare pbc after system_idx so __setattr__ can broadcast a (3,) row already during __init__, leaving a single broadcast site (__post_init__ only validates the shape). Simplify pbc handling in autobatching and io now that pbc is always a tensor.

Related issue: #584

Compatibility

Breaking:

  • External code reading state.pbc as a global (3,) tensor must index it per system. The bundled orb / NequIP / metatomic wrappers already adjusted this.
  • get_attrs_for_scope: pbc moved from "global" to "per-system".

Unchanged:

  • bool / list[bool] / (3,) inputs are still accepted and broadcast to all systems.
  • Trajectory files still store a (3,) pbc array; existing files stay readable.

Discussion points

  • _to_legacy_pbc_state hands third-party forward() a SimState whose pbc temporarily violates the (n_systems, 3) invariant. The preferred fix is upstream support for per-system pbc in the integrations; until then the shim keeps the current orb / NequIP / metatomic releases working unchanged, since they only read plain attributes off the state.
  • state.pbc = <1-D value> after construction broadcasts any length without validation (__post_init__ runs only at construction). Missing validation is pre-existing, but should __setattr__ enforce a length-3 row?

Checklist

Before a pull request can be merged, the following items must be checked:

  • Doc strings have been added in the Google docstring format.
  • Run ruff on your code.
  • Tests have been added for any new functionality or bug fixes.

lil-lon added 2 commits July 25, 2026 12:24
During __init__ the dataclass assigned pbc before system_idx, so the (3,) to (n_systems, 3) broadcast in __setattr__ could not run at construction time and had to be duplicated in __post_init__. Declaring pbc after system_idx makes n_systems available when pbc is assigned, leaving a single broadcast site in __setattr__ and only shape validation in __post_init__. kw_only=True keeps the public API unchanged.
@lil-lon
lil-lon marked this pull request as draft July 25, 2026 06:11
@lil-lon
lil-lon marked this pull request as ready for review July 25, 2026 09:57
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