Rename more batch to system #233
Conversation
WalkthroughThis change performs a comprehensive renaming throughout the codebase, replacing "batch" terminology with "system" in variable names, function signatures, comments, and documentation. The update affects neighbor list computations, model implementations, tutorials, tests, and utility functions, ensuring consistent usage of "system" to refer to collections of atoms or molecules, rather than "batch." Several function and method names are updated to match the new terminology, with no changes to underlying logic or data flow. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Runner
participant Quantities
participant NeighborList
participant Model
User->>Runner: Call integrate/optimize/static (systemwise)
Runner->>NeighborList: Compute neighbor list (system_idx)
Runner->>Model: Forward pass (system_idx, system_mask)
Model->>NeighborList: Use system_mapping for edge construction
Runner->>Quantities: Compute systemwise_max_force
Note over Runner,Quantities: All operations now use "system" terminology
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes Possibly related PRs
Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (7)
examples/scripts/7_Others/7.3_Batched_neighbor_list.py (1)
40-48: Second call mirrors first – same shape caveat
pbc_tensoris reused; ensuring the earlier fix propagates here will resolve any potential mismatch.Consider factoring the neighbour-list call into a small helper to avoid duplication between linked-cell and (N^2) paths.
torch_sim/quantities.py (1)
153-159: Specifyinclude_self=Falseto avoid a silent edge-case
Tensor.scatter_reduceincludes the existing values of the destination tensor in the reduction unlessinclude_self=False.
Right now the destination tensor is initialised with zeros, so the behaviour is correct, but:
- the intent (“pure reduce over
src”) is clearer with the flag set;- future refactors (e.g. different initial fill value) will not change semantics.
- return system_wise_max_force.scatter_reduce( - dim=0, index=state.system_idx, src=max_forces, reduce="amax" - ) + return system_wise_max_force.scatter_reduce( + dim=0, + index=state.system_idx, + src=max_forces, + reduce="amax", + include_self=False, + )torch_sim/models/orb.py (1)
226-236: Minor type/device consistency caveat
system_num_edgesis created ondevicewhich can legitimately beNone(the default).
If the caller passesdevice=Nonewhile the rest of the tensors live on, say, CUDA, this tensor stays on CPU and triggers an implicit device copy insideAtomGraphs.__init__.
Consider defaulting topositions.devicewhendevice is Noneto avoid unnecessary transfers:-device = device # may be None +device = device or positions.devicetests/test_integrators.py (1)
23-30: Out-of-date comment lists “3 systems” but four are generatedThe updated
system_idxcreates four systems (indices 0-3).
Adjust the comment to prevent confusion during future maintenance.torch_sim/neighbors.py (3)
704-706: Inconsistent naming between system_mapping and mapping_system.Inside
strict_nlthe filtered tensor is stored inmapping_system, while the incoming argument issystem_mapping.
Although not wrong, the flip in word order is easy to mis-read and requires mental context switching in the downstream functions.- mapping_system = system_mapping[mask] - ... - return mapping, mapping_system, shifts_idx + filtered_system_mapping = system_mapping[mask] + ... + return mapping, filtered_system_mapping, shifts_idx
753-759: Potential GPU → CPU device hop.
torch.bincount(system_idx)always returns a CPU tensor, regardless of the input device.
Ifpositions(and thereforesystem_idx) lives on the GPU this introduces an unnecessary host–device sync every timetorch_nl_n2is called.Consider keeping the tensor on the same device:
- n_atoms = torch.bincount(system_idx) + n_atoms = torch.bincount(system_idx.cpu()).to(system_idx.device)
810-817: Duplicate filter/rename logic – candidate for helper extraction.
torch_nl_n2andtorch_nl_linked_cellshare the same two-step pattern:
- build_*_neighborhood → (mapping, system_mapping, shifts_idx)
- strict_nl → filtered result
Extracting this into a small private utility would remove duplication and lower the risk of future drift.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
examples/scripts/7_Others/7.3_Batched_neighbor_list.py(1 hunks)examples/tutorials/low_level_tutorial.py(2 hunks)examples/tutorials/state_tutorial.py(5 hunks)tests/test_autobatching.py(1 hunks)tests/test_integrators.py(2 hunks)tests/test_neighbors.py(3 hunks)tests/test_transforms.py(2 hunks)torch_sim/models/graphpes.py(1 hunks)torch_sim/models/mace.py(4 hunks)torch_sim/models/orb.py(3 hunks)torch_sim/models/sevennet.py(1 hunks)torch_sim/neighbors.py(11 hunks)torch_sim/quantities.py(1 hunks)torch_sim/runners.py(10 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (9)
torch_sim/models/graphpes.py (2)
tests/test_state.py (1)
test_deprecated_batch_properties_equal_to_new_system_properties(623-645)torch_sim/state.py (2)
_filter_attrs_by_mask(613-666)n_atoms_per_batch(166-177)
torch_sim/models/sevennet.py (1)
torch_sim/state.py (6)
row_vector_cell(246-248)row_vector_cell(251-257)batch(194-205)n_atoms_per_batch(166-177)SimState(26-373)batch(180-191)
torch_sim/quantities.py (2)
torch_sim/state.py (3)
SimState(26-373)batch(194-205)n_atoms_per_batch(166-177)tests/test_state.py (1)
test_deprecated_batch_properties_equal_to_new_system_properties(623-645)
tests/test_autobatching.py (1)
tests/test_runners.py (3)
test_integrate_with_autobatcher(182-212)test_static_with_autobatcher(601-631)test_static_with_autobatcher_and_reporting(634-739)
tests/test_transforms.py (1)
torch_sim/transforms.py (2)
compute_cell_shifts(537-563)build_linked_cell_neighborhood(965-1039)
torch_sim/runners.py (3)
torch_sim/quantities.py (3)
calc_kinetic_energy(96-130)calc_kT(23-69)systemwise_max_force(144-159)torch_sim/trajectory.py (1)
load_new_trajectories(145-174)tests/test_state.py (1)
test_deprecated_batch_properties_equal_to_new_system_properties(623-645)
torch_sim/models/mace.py (1)
torch_sim/state.py (1)
batch(194-205)
tests/test_neighbors.py (4)
torch_sim/state.py (4)
batch(180-191)batch(194-205)device(142-144)dtype(147-149)torch_sim/transforms.py (1)
compute_cell_shifts(537-563)torch_sim/neighbors.py (1)
strict_nl(640-706)tests/test_state.py (1)
test_deprecated_batch_properties_equal_to_new_system_properties(623-645)
torch_sim/neighbors.py (1)
torch_sim/transforms.py (4)
compute_cell_shifts(537-563)build_naive_neighborhood(614-685)build_linked_cell_neighborhood(965-1039)linked_cell(802-962)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (33)
- GitHub Check: test-examples (examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py)
- GitHub Check: test-examples (examples/scripts/7_Others/7.3_Batched_neighbor_list.py)
- GitHub Check: test-examples (examples/scripts/7_Others/7.1_Soft_sphere_autograd.py)
- GitHub Check: test-examples (examples/scripts/7_Others/7.4_Velocity_AutoCorrelation.py)
- GitHub Check: test-examples (examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descen...
- GitHub Check: test-examples (examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py)
- GitHub Check: test-examples (examples/scripts/5_Workflow/5.3_Elastic.py)
- GitHub Check: test-examples (examples/scripts/5_Workflow/5.1_a2c_silicon_batched.py)
- GitHub Check: test-examples (examples/scripts/4_High_level_api/4.1_high_level_api.py)
- GitHub Check: test-examples (examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py)
- GitHub Check: test-examples (examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py)
- GitHub Check: test-examples (examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py)
- GitHub Check: test-examples (examples/tutorials/low_level_tutorial.py)
- GitHub Check: test-examples (examples/tutorials/metatomic_tutorial.py)
- GitHub Check: test-examples (examples/tutorials/autobatching_tutorial.py)
- GitHub Check: test-examples (examples/tutorials/high_level_tutorial.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, fairchem, tests/models/test_fairchem.py)
- GitHub Check: test-model (macos-14, 3.12, lowest-direct, graphpes, tests/models/test_graphpes.py)
- GitHub Check: test-model (macos-14, 3.11, highest, sevenn, tests/models/test_sevennet.py)
- GitHub Check: test-model (macos-14, 3.11, highest, orb, tests/models/test_orb.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, sevenn, tests/models/test_sevennet.py)
- GitHub Check: test-model (macos-14, 3.11, highest, mattersim, tests/models/test_mattersim.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, orb, tests/models/test_orb.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, orb, tests/models/test_orb.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, mace, tests/test_optimizers_vs_ase.py)
- GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, mace, tests/test_elastic.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, mace, tests/test_optimizers_vs_ase.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, sevenn, tests/models/test_sevennet.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, fairchem, tests/models/test_fairchem.py)
- GitHub Check: test-model (ubuntu-latest, 3.11, highest, mace, tests/test_elastic.py)
- GitHub Check: test-core (macos-14, 3.12, lowest-direct)
- GitHub Check: test-core (ubuntu-latest, 3.12, lowest-direct)
- GitHub Check: build-docs
🔇 Additional comments (22)
torch_sim/models/graphpes.py (1)
70-74: Consistent terminology change – looks good
batch_mask➜system_maskaligns with the project-wide rename and keeps the masking logic intact. No functional impact observed.torch_sim/models/sevennet.py (1)
185-192: Variable rename is correct and completeThe swap to
system_maskon Lines 185-192 is internally consistent (positions, atomic numbers, and later usage). No additional updates required in this block.tests/test_autobatching.py (1)
152-154: Comment wording updated – no code impactChanging “batches” → “systems” keeps the assertion meaning clear and consistent with new terminology.
examples/tutorials/low_level_tutorial.py (1)
110-122: Docstring / print-statement rename confirmed“batchwise” → “systemwise” wording only; runtime behaviour unchanged.
examples/scripts/7_Others/7.3_Batched_neighbor_list.py (1)
28-36: Down-stream uses depend on correctedpbc_tensorOnce the shape issue above is fixed, the renamed
mapping_systemvariables here are correct. No further action required.Please re-run the script locally after the shape fix to ensure
torch_nl_linked_cellandtorch_nl_n2execute withoutRuntimeErroron tensor shape mismatches.tests/test_transforms.py (1)
1186-1189: LGTM – variable rename onlyRenaming
batch_mapping➜system_mappingkeeps the test in sync with the API. No functional issues spotted.examples/tutorials/state_tutorial.py (1)
74-80: Documentation update looks consistentAll occurrences of “batchwise” were correctly replaced with “systemwise”. No further action required.
Also applies to: 145-147, 185-198, 206-207
torch_sim/models/mace.py (4)
187-195: LGTM! Consistent terminology update.The comment updates and method call rename from "batch" to "system_idx" are consistent with the PR objective and maintain the same functionality.
197-199: LGTM! Method renaming is consistent.The method name change from
setup_from_batchtosetup_from_system_idxaligns with the new terminology and maintains the same functionality.
289-298: LGTM! Consistent update in forward method.The comment and method call updates from "batch" to "system_idx" terminology are consistent with the renamed method and maintain the same logic flow.
308-325: LGTM! Consistent variable renaming improves semantic clarity.The renaming from
batch_masktosystem_maskthroughout the neighbor list processing loop is semantically more accurate and consistently applied across all variable usages.torch_sim/runners.py (6)
25-25: LGTM! Import statement correctly updated.The import change from
batchwise_max_forcetosystemwise_max_forcealigns with the function renaming intorch_sim.quantitiesand is necessary for the code to function correctly.
177-184: LGTM! Consistent variable renaming in integrate function.The renaming from
batch_indicestosystem_indicesis semantically more accurate and consistently applied throughout the loop processing. The comment update also reflects the correct terminology.
300-313: LGTM! Convergence function updates are consistent.The docstring updates correctly reflect that the function returns systemwise boolean results, and the function call change to
systemwise_max_forcealigns with the updated import and improved semantics.
440-446: LGTM! Consistent variable renaming in optimize function.The variable renaming from
batch_indicestosystem_indicesmaintains consistency with the integrate function and accurately represents the semantic meaning of the indices.
490-491: LGTM! Docstring updates improve clarity.The docstring changes from "batch" to "system" accurately describe the function's behavior of returning results per system rather than per batch.
550-555: LGTM! Final variable renaming maintains consistency.The variable renaming in the static function matches the pattern used in integrate and optimize functions, ensuring consistent terminology across all runner functions.
tests/test_neighbors.py (3)
345-351: LGTM! Test variable renaming aligns with implementation.The variable renaming from
mapping_batchtomapping_systemcorrectly reflects the updated function signatures and improves semantic clarity in the test code.
499-522: LGTM! Consistent parameter renaming in strict_nl tests.The variable and parameter renaming from
batch_mappingtosystem_mappingcorrectly matches the updatedstrict_nlfunction signature and maintains consistency throughout the test cases.
562-562: LGTM! Performance test variable renaming is consistent.The unpacking variable change from
mapping_batchtomapping_systemmaintains consistency with the updated function signatures in the neighbor list implementations.torch_sim/neighbors.py (2)
692-693: No-op whencell is Noneis still respected – LGTM.
compute_cell_shiftsalready handles theNone-cell case, so the new variable name introduces no functional change.
645-647: Rename Confirmation: no downstreambatch_mappingcalls foundA repository-wide search for the old keyword argument (
batch_mapping=) returned no matches in any.pyfiles. There are no external call-sites using the removed name, so replacingbatch_mappingwithsystem_mappingis safe and won’t break downstream code.
| # Fix: Ensure pbc has the correct shape [n_systems, 3] | ||
| pbc_tensor = torch.tensor([[pbc] * 3] * len(atoms_list), dtype=torch.bool) | ||
|
|
There was a problem hiding this comment.
Broadcast error: pbc_tensor gets wrong shape
[[pbc] * 3] * len(atoms_list) expands both the leading and 2nd dimensions, yielding shape (n_systems, 3, 3) instead of (n_systems, 3).
This will break downstream neighbour-list functions expecting (n_systems, 3) or (3,).
Suggested fix:
-# Fix: Ensure pbc has the correct shape [n_systems, 3]
-pbc_tensor = torch.tensor([[pbc] * 3] * len(atoms_list), dtype=torch.bool)
+# Ensure pbc has shape (n_systems, 3)
+if isinstance(pbc, torch.Tensor) and pbc.dim() == 1:
+ # pbc is a (3,) tensor – broadcast over systems
+ pbc_tensor = pbc.unsqueeze(0).repeat(len(atoms_list), 1)
+else:
+ # pbc is a python bool or already (3,) list
+ pbc_tensor = torch.tensor([pbc] * len(atoms_list), dtype=torch.bool)Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In examples/scripts/7_Others/7.3_Batched_neighbor_list.py around lines 25 to 27,
the construction of pbc_tensor incorrectly expands dimensions resulting in shape
(n_systems, 3, 3) instead of the required (n_systems, 3). To fix this, create a
list of pbc repeated n_systems times without duplicating the inner list three
times, then convert it to a tensor with dtype=torch.bool to ensure the shape is
(n_systems, 3).
Summary
I missed a few batch -> system renames when I did #217
Checklist
Before a pull request can be merged, the following items must be checked:
Run ruff on your code.
We highly recommended installing the pre-commit hooks running in CI locally to speedup the development process. Simply run
pip install pre-commit && pre-commit installto install the hooks which will check your code before each commit.Summary by CodeRabbit
Documentation
Refactor
Bug Fixes
New Features