Rename batch to system #217

curtischong · 2025-07-09T02:49:43Z

Summary

The usage of the word "batch" is misleading since it really means "sample in a set of systems". Describing multiple systems as a set of systems is so much clearer. By renaming batch to system, we also free up the word "batch" to mean an actual batch (e.g. a set of systems we are processing at once)
See https://github.com/Radical-AI/torch-sim/pull/189/files for more discussion.

This change is quite involved. But I didn't want to split it up into multiple parts because "batch" is so pervasive within the codebase. But it feels like the right thing to do. Notable, I changed these interfaces in the SimState:

SimState.n_atoms_per_batch (renamed to n_atoms_per_system)
SimState.n_batches (renamed to n_systems)
SimState.batch (renamed to system_idx)
SimState.__getitem__(self, batch_indices: int | list[int] | slice | torch.Tensor) -> Self: (renamed to SimState.__getitem__(self, system_indices: int | list[int] | slice | torch.Tensor) -> Self:)
SimState.__init__(self, ..., batch) (renamed to SimState.__init__(self, ..., system_idx))

The changes to getitem and init will cause breaking changes, but I feel like it is worth it while torchsim's adoption is currently minimal.

I have also bumped the minor version from 0.2.2 to 0.3.0

Notes of my changes

I didn't change the autobatcher logic much since it's a bit more involved and would probably require a more experienced eye to properly rename it.
comments that are talking about the "batch dimension" have not been changed. This already makes sense because the system dimension IS the batch dimension.
I probably did not get every rename. But the vast majority have been renamed.
One way to review this PR is to clone it and search for "batch". Here, you can see the things that did NOT get renamed (and the surrounding code that did)
codecov/patch is failing becuase this is just a huge change. However, I do not want to introduce logical changes in this PR so I'm not fixing it (it'll make it harder to review)
http://test_graphpes.py/ failing is probably a flake

Note: this PR is an improved version of this PR: #216 (with Rhys' suggestions)

Summary by CodeRabbit

Refactor
- All references to "batch" and related variables have been renamed to "system" or "system_idx" throughout the codebase, including in models, integrators, optimizers, state management, IO, transforms, neighbors, Monte Carlo, and tests.
- Function signatures, class attributes, and internal variable names updated to use "system" terminology consistently.
- Updated indexing and tensor shape annotations from batch-based to system-based dimensions.
- Error messages, comments, and print statements revised to reflect "system" terminology.
Documentation
- Docstrings, comments, tutorials, and examples updated to replace "batch" with "system" for clarity and consistency.
Tests
- Test code and assertions modified to use "system" or "system_idx" terminology.
- Renamed test functions and updated test fixtures to align with system-based naming.
Chores
- Updated .gitignore to exclude .vscode/ directory.

coderabbitai · 2025-07-09T02:49:54Z

Walkthrough

This change systematically replaces all references to "batch" and "batch indices" with "system" and "system indices" throughout the codebase, including core library modules, models, optimizers, integrators, transforms, IO, and all tests and tutorials. All relevant variable names, attributes, docstrings, comments, and function signatures are updated for consistency.

Changes

Files/Paths	Change Summary
.gitignore	Added `.vscode/` to ignore VS Code settings under a new `# IDE` section.
examples/scripts/, examples/tutorials/	Renamed variables, attributes, comments, and print statements from "batch" to "system" for clarity and accuracy.
tests/**	Replaced all references to "batch" and "n_batches" with "system" and "n_systems" in variables, assertions, and mock classes. Updated function signatures and added tests for deprecated batch properties.
torch_sim/autobatching.py, torch_sim/optimizers.py, torch_sim/monte_carlo.py	Systematic renaming: all batch-related variables, attributes, and comments changed to system-based equivalents.
torch_sim/integrators/*.py	All batch-related terminology, tensor shapes, and function arguments renamed to use system-based nomenclature.
torch_sim/io.py, torch_sim/trajectory.py	Conversion and reporting functions now use system indices and system-based terminology throughout.
torch_sim/math.py, torch_sim/quantities.py	Docstrings and function parameters updated to refer to systems instead of batches.
torch_sim/models/*.py	All model logic, comments, and docstrings updated to use system indices instead of batch indices.
torch_sim/neighbors.py, torch_sim/transforms.py	Function signatures and internal logic updated to use system indices and system mapping.
torch_sim/runners.py	Progress bar and convergence logic updated to track number of systems, not batches.
torch_sim/state.py	All state management, slicing, splitting, and concatenation logic refactored for system indices. Deprecated batch properties retained for compatibility.
torch_sim/typing.py	Type alias `StateKey` now uses `"system_idx"` instead of `"batch"`.
torch_sim/workflows/a2c.py	Logging tensor shapes updated from batch-based to system-based.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant SimState
    participant Model
    participant Optimizer
    participant Integrator

    User->>SimState: Initialize with system_idx tensor
    SimState->>Model: Pass system_idx for per-system computation
    Model->>SimState: Compute outputs (energy, stress) per system
    SimState->>Optimizer: Optimize positions/cells per system
    SimState->>Integrator: Integrate dynamics per system
    Integrator->>SimState: Update state using system_idx

Possibly related PRs

Radical-AI/torch-sim#206: Removes all unbatched code and replaces unbatched models and integrators with their batched counterparts, directly related to the current PR’s systematic renaming and restructuring to support consistent batched/system operations.

Suggested labels

enhancement, pkg

Poem

🐇
Goodbye to batches, hello to systems,
A hop and a skip—no more schisms!
Indices aligned, code now pristine,
Each atom hops in a system scene.
With every run, the logic’s neat—
A rabbit’s work is now complete!
🥕

✨ Finishing Touches

📝 Generate Docstrings

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.

❤️ Share

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

Review comments: Directly reply to a review comment made by CodeRabbit. Example:
- I pushed a fix in commit <commit_id>, please review it.
- Explain this complex logic.
- Open a follow-up GitHub issue for this discussion.
Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
- @coderabbitai explain this code block.
- @coderabbitai modularize this function.
PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
- @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
- @coderabbitai read src/utils.ts and explain its main purpose.
- @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
- @coderabbitai help me debug CodeRabbit configuration file.

Support

Need 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)

@coderabbitai pause to pause the reviews on a PR.
@coderabbitai resume to resume the paused reviews.
@coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
@coderabbitai full review to do a full review from scratch and review all the files again.
@coderabbitai summary to regenerate the summary of the PR.
@coderabbitai generate docstrings to generate docstrings for this PR.
@coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
@coderabbitai resolve resolve all the CodeRabbit review comments.
@coderabbitai configuration to show the current CodeRabbit configuration for the repository.
@coderabbitai help to get help.

Other keywords and placeholders

Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (`.coderabbit.yaml`)

You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
Please see the configuration documentation for more information.
If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

Visit our Documentation for detailed information on how to use CodeRabbit.
Join our Discord Community to get help, request features, and share feedback.
Follow us on X/Twitter for updates and announcements.

coderabbitai

Actionable comments posted: 3

♻️ Duplicate comments (1)

torch_sim/optimizers.py (1)

1144-1150: Same per-system mass aggregation issue as above

The Frechet‐cell initialiser repeats the segment_reduce pattern with the same
assumptions and should be fixed identically. Please apply the scatter-based
solution here as well.

🧹 Nitpick comments (8)

examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py (1)

86-89: Minor wording tweak left behind

Comment now says “atoms per system”, but the surrounding script still labels itself “Batched …” (top-of-file docstring) and uses variables such as positions_lr “for all batches”. For perfect consistency with the new terminology, consider updating those occurrences in a follow-up sweep.

examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py (1)

82-86: Terminology consistency

Same note as in the gradient-descent script: the inline comment was updated, but other text (“Batched …”, etc.) still refers to batch. Recommend aligning those in a later tidy-up commit.

examples/scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py (1)

81-85: Left-over “batch” wording elsewhere

Comment is fixed here, but the file header and logging strings still talk about batches. Harmonising wording across the examples will avoid reader confusion.
torch_sim/math.py (1)
995-1005: Docstring still mixes “batch” and “system” terms

Line 1003 was corrected to n_systems, but the earlier parameter docs still reference batch membership and the argument itself is batch_indices. If renaming the parameter is out of scope, at least update the wording to avoid contradicting the return description:
-        batch_indices: Tensor of shape [N_total_entities] indicating batch membership.
+        batch_indices: Tensor of shape [N_total_entities] indicating system membership
+            (kept as *batch_indices* for backward compatibility).
Tiny polish, but it keeps the documentation coherent.
tests/models/conftest.py (2)
172-172: Consider renaming variable for consistency.

The attribute reference has been correctly updated to sim_state.system_idx. However, consider renaming the variable from og_batch to og_system_idx for better clarity and consistency with the new terminology.
-        og_batch = sim_state.system_idx.clone()
+        og_system_idx = sim_state.system_idx.clone()
180-180: Update assertion to use consistent variable name.

The assertion has been correctly updated to use sim_state.system_idx. If the variable is renamed as suggested above, this assertion should also be updated accordingly.
-        assert torch.allclose(og_batch, sim_state.system_idx)
+        assert torch.allclose(og_system_idx, sim_state.system_idx)
torch_sim/integrators/npt.py (1)
1215-1225: Consider using the existing calc_kinetic_energy function with system_idx.

The manual loop for computing kinetic energy per system could be replaced with the existing calc_kinetic_energy function that supports system_idx parameter:
-        # Compute kinetic energy contribution per system
-        # Split momenta and masses by system
-        KE_per_system = torch.zeros(
-            n_systems, device=positions.device, dtype=positions.dtype
-        )
-        for b in range(n_systems):
-            system_mask = system_idx == b
-            if system_mask.any():
-                system_momenta = momenta[system_mask]
-                system_masses = masses[system_mask]
-                KE_per_system[b] = calc_kinetic_energy(system_momenta, system_masses)
+        # Compute kinetic energy contribution per system
+        KE_per_system = calc_kinetic_energy(momenta, masses, system_idx=system_idx)
This would be more efficient and consistent with the codebase patterns.
torch_sim/optimizers.py (1)
119-135: Guard against device / dtype mismatch when a tensor learning-rate is supplied

If the caller passes lr as an existing tensor on a different device or with a
different dtype than the model, the slice operation
lr[state.system_idx] will raise an error.
Convert external tensors to the optimiser’s (device,dtype) first.
-        # Get per-atom learning rates by mapping batch learning rates to atoms
+        # Get per-atom learning rates; ensure `lr` lives on the correct device/dtype
         if isinstance(lr, float):
             lr = torch.full((state.n_systems,), lr, device=device, dtype=dtype)
+        else:
+            lr = lr.to(device=device, dtype=dtype)
 
         atom_lr = lr[state.system_idx].unsqueeze(-1)  # shape: (total_atoms, 1)

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2bb3bd5 and c96043d.

📒 Files selected for processing (67)

.gitignore (1 hunks)
examples/scripts/1_Introduction/1.2_MACE.py (1 hunks)
examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py (1 hunks)
examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py (1 hunks)
examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py (1 hunks)
examples/scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py (1 hunks)
examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py (1 hunks)
examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py (2 hunks)
examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py (5 hunks)
examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py (1 hunks)
examples/scripts/3_Dynamics/3.2_MACE_NVE.py (1 hunks)
examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py (1 hunks)
examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py (1 hunks)
examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py (1 hunks)
examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py (1 hunks)
examples/scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py (2 hunks)
examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py (4 hunks)
examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py (1 hunks)
examples/scripts/4_High_level_api/4.2_auto_batching_api.py (1 hunks)
examples/scripts/5_Workflow/5.2_In_Flight_WBM.py (1 hunks)
examples/scripts/7_Others/7.3_Batched_neighbor_list.py (2 hunks)
examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py (10 hunks)
examples/tutorials/autobatching_tutorial.py (2 hunks)
examples/tutorials/high_level_tutorial.py (1 hunks)
examples/tutorials/hybrid_swap_tutorial.py (1 hunks)
examples/tutorials/low_level_tutorial.py (1 hunks)
examples/tutorials/state_tutorial.py (6 hunks)
tests/models/conftest.py (1 hunks)
tests/test_autobatching.py (6 hunks)
tests/test_correlations.py (1 hunks)
tests/test_integrators.py (6 hunks)
tests/test_io.py (6 hunks)
tests/test_monte_carlo.py (5 hunks)
tests/test_neighbors.py (3 hunks)
tests/test_optimizers.py (9 hunks)
tests/test_runners.py (9 hunks)
tests/test_state.py (18 hunks)
tests/test_trajectory.py (5 hunks)
tests/test_transforms.py (10 hunks)
torch_sim/autobatching.py (12 hunks)
torch_sim/integrators/md.py (7 hunks)
torch_sim/integrators/npt.py (51 hunks)
torch_sim/integrators/nve.py (5 hunks)
torch_sim/integrators/nvt.py (10 hunks)
torch_sim/io.py (12 hunks)
torch_sim/math.py (1 hunks)
torch_sim/models/fairchem.py (1 hunks)
torch_sim/models/graphpes.py (1 hunks)
torch_sim/models/interface.py (4 hunks)
torch_sim/models/lennard_jones.py (2 hunks)
torch_sim/models/mace.py (9 hunks)
torch_sim/models/metatomic.py (2 hunks)
torch_sim/models/morse.py (2 hunks)
torch_sim/models/orb.py (3 hunks)
torch_sim/models/particle_life.py (2 hunks)
torch_sim/models/sevennet.py (2 hunks)
torch_sim/models/soft_sphere.py (4 hunks)
torch_sim/monte_carlo.py (5 hunks)
torch_sim/neighbors.py (3 hunks)
torch_sim/optimizers.py (38 hunks)
torch_sim/quantities.py (5 hunks)
torch_sim/runners.py (8 hunks)
torch_sim/state.py (23 hunks)
torch_sim/trajectory.py (6 hunks)
torch_sim/transforms.py (15 hunks)
torch_sim/typing.py (1 hunks)
torch_sim/workflows/a2c.py (1 hunks)

🧰 Additional context used

🧬 Code Graph Analysis (31)

examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py (1)

torch_sim/quantities.py (3)

calc_kT (23-69)

get_pressure (133-141)

calc_kinetic_energy (96-130)

examples/scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py (2)

torch_sim/quantities.py (2)

calc_kT (23-69)

calc_kinetic_energy (96-130)

torch_sim/integrators/npt.py (2)

momenta (70-72)

npt_nose_hoover_invariant (1563-1672)

examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py (2)

torch_sim/quantities.py (1)

calc_kT (23-69)

torch_sim/integrators/nvt.py (1)

nvt_nose_hoover_invariant (448-517)

examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py (1)

torch_sim/io.py (1)

atoms_to_state (180-245)

tests/test_monte_carlo.py (1)

torch_sim/monte_carlo.py (1)

validate_permutation (127-141)

torch_sim/workflows/a2c.py (2)

torch_sim/state.py (3)

n_systems (202-204)

device (142-144)

dtype (147-149)

torch_sim/models/interface.py (4)

device (107-109)

device (112-116)

dtype (119-121)

dtype (124-128)

examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py (1)

torch_sim/quantities.py (1)

calc_kT (23-69)

examples/tutorials/autobatching_tutorial.py (1)

torch_sim/state.py (1)

n_systems (202-204)

examples/scripts/5_Workflow/5.2_In_Flight_WBM.py (1)

torch_sim/state.py (1)

n_systems (202-204)

torch_sim/models/soft_sphere.py (1)

torch_sim/state.py (2)

SimState (26-353)

n_systems (202-204)

examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py (1)

torch_sim/state.py (1)

n_systems (202-204)

examples/scripts/4_High_level_api/4.2_auto_batching_api.py (1)

torch_sim/state.py (1)

n_systems (202-204)

examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py (1)

torch_sim/state.py (1)

n_systems (202-204)

torch_sim/models/interface.py (1)

torch_sim/state.py (1)

clone (239-255)

examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py (1)

torch_sim/quantities.py (1)

calc_kT (23-69)

examples/tutorials/hybrid_swap_tutorial.py (1)

torch_sim/state.py (3)

n_systems (202-204)

device (142-144)

dtype (147-149)

examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py (1)

torch_sim/quantities.py (3)

calc_kT (23-69)

get_pressure (133-141)

calc_kinetic_energy (96-130)

torch_sim/models/orb.py (1)

torch_sim/state.py (1)

n_systems (202-204)

torch_sim/runners.py (1)

torch_sim/state.py (1)

n_systems (202-204)

torch_sim/models/graphpes.py (1)

torch_sim/state.py (1)

n_systems (202-204)

tests/test_correlations.py (1)

torch_sim/state.py (2)

n_systems (202-204)

split (281-290)

torch_sim/quantities.py (2)

torch_sim/state.py (4)

momenta (360-366)

n_systems (202-204)

device (142-144)

dtype (147-149)

torch_sim/integrators/md.py (1)

velocities (44-48)

torch_sim/integrators/nve.py (1)

torch_sim/integrators/md.py (1)

calculate_momenta (51-109)

tests/test_autobatching.py (3)

torch_sim/state.py (5)

n_systems (202-204)

device (142-144)

dtype (147-149)

batch (175-181)

batch (184-190)

tests/conftest.py (3)

device (24-25)

dtype (29-30)

lj_model (34-45)

torch_sim/autobatching.py (1)

BinningAutoBatcher (417-699)

torch_sim/autobatching.py (1)

torch_sim/state.py (3)

n_atoms (152-154)

n_systems (202-204)

concatenate_states (809-895)

tests/test_state.py (2)

torch_sim/state.py (12)

_normalize_system_indices (400-437)

SimState (26-353)

_slice_state (762-806)

n_atoms (152-154)

dtype (147-149)

device (142-144)

n_atoms_per_system (157-163)

n_systems (202-204)

batch (175-181)

batch (184-190)

n_batches (193-199)

n_atoms_per_batch (166-172)

tests/conftest.py (7)

si_double_sim_state (323-325)

si_sim_state (142-144)

si_atoms (89-91)

fe_atoms (83-85)

fe_supercell_sim_state (298-302)

dtype (29-30)

device (24-25)

torch_sim/integrators/npt.py (3)

torch_sim/state.py (7)

device (142-144)

dtype (147-149)

n_systems (202-204)

n_atoms_per_system (157-163)

momenta (360-366)

to (321-335)

clone (239-255)

torch_sim/integrators/md.py (2)

velocities (44-48)

calculate_momenta (51-109)

torch_sim/quantities.py (1)

calc_kinetic_energy (96-130)

torch_sim/trajectory.py (2)

torch_sim/state.py (2)

n_systems (202-204)

split (281-290)

tests/test_correlations.py (1)

split (39-42)

torch_sim/monte_carlo.py (1)

torch_sim/state.py (3)

device (142-144)

n_systems (202-204)

n_atoms (152-154)

torch_sim/integrators/nvt.py (4)

torch_sim/integrators/md.py (1)

calculate_momenta (51-109)

torch_sim/quantities.py (1)

calc_kinetic_energy (96-130)

torch_sim/integrators/npt.py (1)

momenta (70-72)

torch_sim/state.py (2)

momenta (360-366)

n_atoms_per_system (157-163)

torch_sim/optimizers.py (2)

torch_sim/state.py (10)

n_systems (202-204)

device (142-144)

dtype (147-149)

to (321-335)

deform_grad (391-397)

clone (239-255)

DeformGradMixin (356-397)

_deform_grad (379-389)

row_vector_cell (226-228)

row_vector_cell (231-237)

torch_sim/math.py (1)

batched_vdot (992-1022)

🪛 GitHub Check: codecov/patch

torch_sim/autobatching.py

[warning] 296-297: torch_sim/autobatching.py#L296-L297
Added lines #L296 - L297 were not covered by tests

[warning] 346-346: torch_sim/autobatching.py#L346
Added line #L346 was not covered by tests

[warning] 936-936: torch_sim/autobatching.py#L936
Added line #L936 was not covered by tests

[warning] 942-942: torch_sim/autobatching.py#L942
Added line #L942 was not covered by tests

[warning] 1022-1023: torch_sim/autobatching.py#L1022-L1023
Added lines #L1022 - L1023 were not covered by tests

[warning] 1028-1029: torch_sim/autobatching.py#L1028-L1029
Added lines #L1028 - L1029 were not covered by tests

[warning] 1060-1060: torch_sim/autobatching.py#L1060
Added line #L1060 was not covered by tests

torch_sim/integrators/md.py

[warning] 83-83: torch_sim/integrators/md.py#L83
Added line #L83 was not covered by tests

[warning] 90-90: torch_sim/integrators/md.py#L90
Added line #L90 was not covered by tests

[warning] 95-96: torch_sim/integrators/md.py#L95-L96
Added lines #L95 - L96 were not covered by tests

torch_sim/integrators/npt.py

[warning] 109-109: torch_sim/integrators/npt.py#L109
Added line #L109 was not covered by tests

[warning] 112-112: torch_sim/integrators/npt.py#L112
Added line #L112 was not covered by tests

[warning] 120-120: torch_sim/integrators/npt.py#L120
Added line #L120 was not covered by tests

[warning] 132-132: torch_sim/integrators/npt.py#L132
Added line #L132 was not covered by tests

[warning] 139-139: torch_sim/integrators/npt.py#L139
Added line #L139 was not covered by tests

[warning] 246-246: torch_sim/integrators/npt.py#L246
Added line #L246 was not covered by tests

[warning] 249-249: torch_sim/integrators/npt.py#L249
Added line #L249 was not covered by tests

[warning] 286-286: torch_sim/integrators/npt.py#L286
Added line #L286 was not covered by tests

[warning] 288-288: torch_sim/integrators/npt.py#L288
Added line #L288 was not covered by tests

[warning] 291-292: torch_sim/integrators/npt.py#L291-L292
Added lines #L291 - L292 were not covered by tests

[warning] 294-294: torch_sim/integrators/npt.py#L294
Added line #L294 was not covered by tests

[warning] 355-355: torch_sim/integrators/npt.py#L355
Added line #L355 was not covered by tests

[warning] 359-359: torch_sim/integrators/npt.py#L359
Added line #L359 was not covered by tests

[warning] 361-361: torch_sim/integrators/npt.py#L361
Added line #L361 was not covered by tests

[warning] 420-420: torch_sim/integrators/npt.py#L420
Added line #L420 was not covered by tests

[warning] 422-422: torch_sim/integrators/npt.py#L422
Added line #L422 was not covered by tests

[warning] 424-424: torch_sim/integrators/npt.py#L424
Added line #L424 was not covered by tests

[warning] 427-428: torch_sim/integrators/npt.py#L427-L428
Added lines #L427 - L428 were not covered by tests

[warning] 498-498: torch_sim/integrators/npt.py#L498
Added line #L498 was not covered by tests

[warning] 505-506: torch_sim/integrators/npt.py#L505-L506
Added lines #L505 - L506 were not covered by tests

⏰ 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). (36)

GitHub Check: test-examples (examples/scripts/5_Workflow/5.3_Elastic.py)
GitHub Check: test-examples (examples/scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py)
GitHub Check: test-examples (examples/scripts/5_Workflow/5.1_a2c_silicon_batched.py)
GitHub Check: test-examples (examples/scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py)
GitHub Check: test-examples (examples/scripts/2_Structural_optimization/2.4_MACE_FIRE.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.3_MACE_Gradient_Descent.py)
GitHub Check: test-examples (examples/scripts/4_High_level_api/4.1_high_level_api.py)
GitHub Check: test-examples (examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.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.13_MACE_NVE_non_pbc.py)
GitHub Check: test-examples (examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py)
GitHub Check: test-examples (examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py)
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.4_Velocity_AutoCorrelation.py)
GitHub Check: test-examples (examples/scripts/7_Others/7.2_Stress_autograd.py)
GitHub Check: test-examples (examples/tutorials/autobatching_tutorial.py)
GitHub Check: test-examples (examples/tutorials/high_level_tutorial.py)
GitHub Check: test-examples (examples/tutorials/low_level_tutorial.py)
GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, sevenn, tests/models/test_sevennet.py)
GitHub Check: test-model (macos-14, 3.12, lowest-direct, metatomic, tests/models/test_metatomic.py)
GitHub Check: test-model (macos-14, 3.12, lowest-direct, mace, tests/test_elastic.py)
GitHub Check: test-model (macos-14, 3.12, lowest-direct, mace, tests/test_optimizers_vs_ase.py)
GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, orb, tests/models/test_orb.py)
GitHub Check: test-model (macos-14, 3.11, highest, mace, tests/models/test_mace.py)
GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, mattersim, tests/models/test_mattersim.py)
GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, mace, tests/test_elastic.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.11, highest, sevenn, tests/models/test_sevennet.py)
GitHub Check: test-model (ubuntu-latest, 3.12, lowest-direct, fairchem, tests/models/test_fairchem.py)
GitHub Check: test-model (ubuntu-latest, 3.11, highest, mace, tests/test_elastic.py)
GitHub Check: test-model (ubuntu-latest, 3.11, highest, orb, tests/models/test_orb.py)
GitHub Check: test-model (ubuntu-latest, 3.11, highest, fairchem, tests/models/test_fairchem.py)
GitHub Check: test-core (ubuntu-latest, 3.12, lowest-direct)
GitHub Check: test-core (macos-14, 3.12, lowest-direct)
GitHub Check: build-docs

🔇 Additional comments (207)

.gitignore (1)

37-39: 👍 Ignoring the .vscode/ workspace folder is sensible

Keeps editor-specific clutter out of the repo.
No further action required.

examples/scripts/5_Workflow/5.2_In_Flight_WBM.py (1)

86-86: LGTM! Consistent with systematic renaming effort.

The change from state.n_batches to state.n_systems correctly aligns with the codebase-wide renaming from "batch" to "system" terminology. The print statement now accurately reflects that it's reporting the number of systems rather than batches.

examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py (1)

89-89: LGTM! Correct tensor dimensionality with updated terminology.

The change from md_state.n_batches to md_state.n_systems correctly updates the tensor initialization to use the new system-based terminology while maintaining the proper dimensionality for tracking permutations across systems.

examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py (1)

80-80: LGTM! Parameter name updated to match new function signature.

The change from batch=state.batch to system_idx=state.system_idx correctly updates the function call to use the new parameter name and attribute, ensuring proper system indexing for kinetic energy calculations.

examples/scripts/4_High_level_api/4.2_auto_batching_api.py (1)

79-79: LGTM! Print statement updated for consistency.

The change from state.n_batches to state.n_systems correctly updates the debug print statement to use the new system-based terminology, maintaining consistency with the codebase-wide renaming effort.

examples/scripts/3_Dynamics/3.2_MACE_NVE.py (1)

91-91: LGTM! Function call updated with correct parameter and attribute names.

The change from batch=state.batch to system_idx=state.system_idx correctly updates both the parameter name and the attribute reference to align with the systematic renaming from "batch" to "system" terminology throughout the codebase.

examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py (1)

68-68: LGTM: Parameter name updated consistently with function signature.

The change from batch=state.batch to system_idx=state.system_idx correctly aligns with the updated calc_kT function signature in torch_sim/quantities.py (lines 23-69), which now accepts system_idx parameter for system-specific temperature calculations.

examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py (1)

74-74: LGTM: Parameter name updated consistently with function signature.

The change from batch=state.batch to system_idx=state.system_idx correctly aligns with the updated calc_kinetic_energy function signature in torch_sim/quantities.py (lines 96-130), which now accepts system_idx parameter for system-specific kinetic energy calculations.

examples/scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py (1)

126-128: LGTM: All parameter names updated consistently with function signatures.

The systematic replacement of batch=state.batch with system_idx=state.system_idx in both calc_kT and calc_kinetic_energy function calls is correct and consistent with the updated function signatures in torch_sim/quantities.py. This maintains the same functionality while using the updated system-based terminology.

Also applies to: 135-135, 150-150, 157-159

examples/tutorials/high_level_tutorial.py (1)

378-378: LGTM: Documentation updated to reflect new terminology.

The change from n_batches to n_systems in the documentation correctly reflects the updated terminology used throughout the codebase. This improves consistency and clarity for users.

examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py (1)

86-88: LGTM: Parameter names updated consistently with function signature.

The changes from batch=state.batch to system_idx=state.system_idx in both calc_kT function calls are correct and consistent with the updated function signature in torch_sim/quantities.py (lines 23-69). This maintains the same temperature calculation functionality while using the updated system-based terminology.

Also applies to: 95-95

examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py (2)

71-73: LGTM! Function call updated correctly.

The calc_kT function call has been correctly updated to use system_idx=state.system_idx instead of batch=state.batch, which aligns with the systematic renaming throughout the codebase.

81-81: LGTM! Function call updated correctly.

The calc_kT function call has been correctly updated to use system_idx=state.system_idx instead of batch=state.batch, which aligns with the systematic renaming throughout the codebase.

torch_sim/typing.py (1)

18-18: LGTM! Type definition updated correctly.

The StateKey type literal has been correctly updated to include "system_idx" instead of "batch", which maintains type safety after the attribute renaming in the state objects.

examples/tutorials/hybrid_swap_tutorial.py (1)

136-136: LGTM! Property reference updated correctly.

The tensor initialization has been correctly updated to use md_state.n_systems instead of md_state.n_batches, which aligns with the property renaming in the state object.

examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py (5)

71-73: LGTM! Function call updated correctly.

The calc_kT function call has been correctly updated to use system_idx=state.system_idx instead of batch=state.batch, which aligns with the systematic renaming throughout the codebase.

88-90: LGTM! Function call updated correctly.

The calc_kT function call has been correctly updated to use system_idx=state.system_idx instead of batch=state.batch, which aligns with the systematic renaming throughout the codebase.

98-102: LGTM! Function call updated correctly.

The calc_kinetic_energy function call has been correctly updated to use system_idx=state.system_idx instead of batch=state.batch, which aligns with the systematic renaming throughout the codebase.

116-116: LGTM! Function call updated correctly.

The calc_kT function call has been correctly updated to use system_idx=state.system_idx instead of batch=state.batch, which aligns with the systematic renaming throughout the codebase.

125-127: LGTM! Function call updated correctly.

The calc_kinetic_energy function call has been correctly updated to use system_idx=state.system_idx instead of batch=state.batch, which aligns with the systematic renaming throughout the codebase.

examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py (4)

122-124: LGTM - Consistent parameter renaming

The calc_kT function call correctly uses the updated system_idx parameter instead of the deprecated batch parameter. This aligns with the function signature changes in torch_sim/quantities.py (lines 23-69).

129-131: LGTM - Consistent parameter renaming

The calc_kinetic_energy function call correctly uses the updated system_idx parameter instead of the deprecated batch parameter. This aligns with the function signature changes in torch_sim/quantities.py (lines 96-130).

143-145: LGTM - Consistent parameter renaming

The final temperature calculation correctly uses the updated system_idx parameter, maintaining consistency with the refactoring pattern.

151-153: LGTM - Consistent parameter renaming

The final kinetic energy calculation correctly uses the updated system_idx parameter, completing the consistent refactoring across all function calls in this script.

torch_sim/models/metatomic.py (2)

77-77: LGTM - Improved docstring clarity

The docstring update from "batch indices" to "system indices" improves semantic clarity and aligns with the broader refactoring to use more descriptive terminology.

203-203: LGTM - Consistent attribute renaming

The logic correctly uses state.system_idx instead of the deprecated state.batch attribute. This maintains the same functionality while using the renamed attribute that better represents the concept of system indexing.

examples/tutorials/autobatching_tutorial.py (3)

252-252: LGTM - Consistent property renaming

The variable assignment correctly uses the updated n_systems property instead of the deprecated n_batches property. This aligns with the property definition in torch_sim/state.py (lines 202-204).

282-282: LGTM - Consistent property renaming

The assertion correctly uses the updated n_systems property to verify all states have been processed. This maintains the same logic while using the renamed property.

286-286: LGTM - Consistent property renaming

The standalone expression correctly uses the updated n_systems property, completing the consistent refactoring across all property references in this tutorial.

examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py (5)

72-74: LGTM - Consistent parameter renaming

The calc_kT function call correctly uses the updated system_idx parameter, aligning with the function signature changes in torch_sim/quantities.py.

91-93: LGTM - Consistent parameter renaming

The calc_kT function call in the NPT loop correctly uses the updated system_idx parameter, maintaining consistency with the refactoring pattern.

101-103: LGTM - Consistent parameter renaming

The calc_kinetic_energy function call correctly uses the updated system_idx parameter, aligning with the function signature changes in torch_sim/quantities.py.

114-114: LGTM - Consistent parameter renaming

The final temperature calculation correctly uses the updated system_idx parameter, completing the consistent refactoring in temperature calculations.

122-124: LGTM - Consistent parameter renaming

The final kinetic energy calculation correctly uses the updated system_idx parameter, completing the consistent refactoring across all function calls in this script.

examples/tutorials/low_level_tutorial.py (1)

230-232: LGTM - Consistent parameter renaming

The calc_kT function call correctly uses the updated system_idx parameter instead of the deprecated batch parameter. This aligns with the function signature changes in torch_sim/quantities.py and maintains consistency with the refactoring pattern across the codebase.

examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py (1)

166-166: LGTM: Consistent with system terminology update.

The parameter change from batch=state.batch to system_idx=state.system_idx correctly aligns with the updated calc_kT function signature shown in the relevant code snippets.

examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py (2)

96-102: LGTM: Improved variable naming for clarity.

The renaming from atoms_per_batch to atoms_per_system and batch_indices to system_indices better reflects the semantic meaning of the data - representing individual atomic systems rather than generic batch processing.

109-109: LGTM: Print statement updated consistently.

The print statement correctly references state.system_idx to match the updated attribute name.

examples/scripts/1_Introduction/1.2_MACE.py (4)

66-72: LGTM: Consistent variable renaming improves clarity.

The variables atoms_per_system and system_idx better convey the semantic meaning of representing individual atomic systems rather than generic batch processing.

78-78: LGTM: Print statement updated consistently.

The print statement correctly uses "System indices" to match the updated variable name.

86-86: LGTM: Function argument updated to match new parameter name.

The argument change from batch=batch to system_idx=system_idx correctly aligns with the updated function signature.

91-91: LGTM: Comments updated to reflect new terminology.

The comments correctly reference "systems" instead of "batches" to match the updated variable semantics.

Also applies to: 97-97

tests/test_monte_carlo.py (3)

53-53: LGTM: Function calls updated to match new API.

The validate_permutation calls correctly use system_idx parameter, which aligns with the updated function signature shown in the relevant code snippets.

Also applies to: 98-100, 110-110

67-69: LGTM: Assertion logic updated consistently.

The test correctly verifies that swaps occur within the same system using the updated system_idx attribute.

152-165: LGTM: Test assertions updated to use new terminology.

The test assertions correctly use system_idx instead of batch to verify system assignments and atomic number distributions, maintaining the same validation logic with improved naming.

torch_sim/models/sevennet.py (2)

184-185: LGTM: Internal implementation updated consistently.

The loop iteration and masking logic correctly use state.system_idx instead of state.batch, maintaining identical functionality while improving naming consistency.

248-248: LGTM: Energy tensor creation updated consistently.

The energy tensor creation correctly uses state.system_idx.max().item() + 1 to maintain the same logic with updated attribute naming.

torch_sim/workflows/a2c.py (1)

733-735: LGTM: Consistent tensor shape update for system-based indexing.

The tensor shapes for energy and stress logging have been correctly updated from state.n_batches to state.n_systems, maintaining the same dimensionality while aligning with the new system-based terminology.

torch_sim/models/orb.py (3)

105-105: LGTM: Correct update to system-based indexing.

The change from state.batch to state.system_idx is consistent with the new system-based indexing scheme and maintains the same functionality.

146-146: LGTM: Proper system count calculation.

Computing the number of systems from state.system_idx.max().item() + 1 is correct and aligns with the new indexing approach.

160-160: LGTM: Correct mask creation for system filtering.

The boolean mask state.system_idx == i properly filters atoms belonging to system i, maintaining the same logic with updated terminology.

torch_sim/models/graphpes.py (1)

70-71: LGTM: Consistent system-based iteration and filtering.

The changes from state.n_batches to state.n_systems and state.batch == i to state.system_idx == i correctly update the iteration and masking logic while maintaining the same functionality.

examples/scripts/7_Others/7.3_Batched_neighbor_list.py (3)

21-21: LGTM: Proper variable renaming in example script.

The variable name change from batch to system_idx correctly reflects the new system-based indexing terminology.

25-26: LGTM: Updated comment reflects new terminology.

The comment correctly describes the tensor shape as corresponding to n_systems instead of batches, maintaining accuracy with the new terminology.

29-29: LGTM: Function calls updated for new API.

The neighbor list function calls have been correctly updated to pass system_idx instead of batch, demonstrating proper usage of the new API.

Also applies to: 41-41

torch_sim/neighbors.py (3)

769-769: LGTM: Function parameter updated for consistency.

The parameter name change from batch to system_idx aligns with the new system-based indexing terminology across the codebase.

787-787: LGTM: Docstring updated to reflect new parameter name.

The docstring correctly describes the system_idx parameter, maintaining documentation accuracy with the API changes.

809-809: LGTM: Internal usage updated consistently.

The internal usage of system_idx in the torch.bincount call is correct and maintains the same functionality with the new parameter name.

torch_sim/runners.py (8)

175-175: LGTM: Progress bar initialization correctly updated.

The change from state.n_batches to state.n_systems properly reflects the new terminology and maintains the correct progress bar behavior for tracking system processing.

197-197: LGTM: Progress bar update correctly updated.

The change from state.n_batches to state.n_systems in the progress bar update call is consistent with the initialization change and maintains proper progress tracking.

310-312: LGTM: Docstring correctly updated to reflect system terminology.

The docstring update from "batch" to "system" terminology accurately describes that the convergence status is returned per system rather than per batch, improving code clarity.

346-348: LGTM: Docstring consistently updated to system terminology.

The docstring update maintains consistency with the broader refactoring, correctly describing convergence status per system.

375-375: LGTM: Function parameter docstring updated consistently.

The docstring update for the convergence_fn parameter correctly describes it as returning a boolean tensor of length n_systems, maintaining consistency with the terminology change.

437-437: LGTM: Progress bar initialization updated consistently.

The change to state.n_systems in the optimize function's progress bar initialization is consistent with the other progress bar updates in the file.

548-548: LGTM: Progress bar initialization updated consistently.

The change to state.n_systems in the static function's progress bar initialization maintains consistency with the terminology changes throughout the file.

571-571: LGTM: Progress bar update updated consistently.

The change to sub_state.n_systems in the progress bar update maintains consistency with the terminology changes and ensures proper progress tracking.

tests/test_trajectory.py (6)

35-37: LGTM: Test fixture attribute updated consistently.

The change from batch_idx to system_idx in the test fixture aligns with the broader terminology refactoring and maintains test functionality.

681-683: LGTM: Variable name updated for clarity.

The change from atoms_per_batch to atoms_per_system improves code readability and aligns with the system terminology used throughout the refactoring.

701-714: LGTM: Loop variable renamed consistently.

The change from batch to system_idx in the loop variable and related usage maintains test functionality while improving semantic clarity.

730-738: LGTM: Variable names updated consistently.

The changes from batch to system_idx and system_props improve code clarity and align with the broader refactoring terminology.

770-776: LGTM: Variable names updated consistently.

The systematic renaming from batch to system_props and system_idx maintains test functionality while improving semantic consistency.

783-798: LGTM: Variable names updated consistently.

The changes from batch to system_idx and system_props in the test function maintain test logic while improving code clarity through consistent terminology.

tests/test_correlations.py (2)

34-37: LGTM: MockState attributes updated consistently.

The changes from n_batches to n_systems and batch to system_idx align with the broader terminology refactoring while maintaining the mock class functionality for testing.

40-42: LGTM: Docstring and comment updated consistently.

The updates to the docstring and comment reflect the new system terminology, improving code clarity while maintaining the mock class behavior.

tests/test_io.py (6)

51-55: LGTM: Test assertion updated consistently.

The change from state.batch to state.system_idx in the assertion maintains test functionality while aligning with the broader terminology refactoring.

69-70: LGTM: Test assertions updated consistently.

The changes from state.batch to state.system_idx in the assertions maintain test functionality while improving semantic clarity.

84-88: LGTM: Test assertions updated consistently.

The changes from state.batch to state.system_idx in the assertions align with the broader refactoring and maintain proper test validation.

176-180: LGTM: Test assertions updated consistently.

The changes from state.batch to state.system_idx in the assertions maintain test functionality while improving semantic consistency.

241-245: LGTM: Test logic updated consistently.

The changes from sim_state.batch to sim_state.system_idx in the round-trip test logic maintain test functionality while aligning with the new terminology.

254-254: LGTM: Test assertion updated consistently.

The change from sim_state.batch to sim_state.system_idx in the assertion maintains test validation while improving semantic clarity.

tests/test_optimizers.py (9)

131-131: LGTM: SimState construction updated consistently.

The change from batch to system_idx in the SimState construction maintains test functionality while aligning with the broader terminology refactoring.

232-232: LGTM: Variable name updated consistently.

The change from initial_dt_batch to reflect system terminology maintains test functionality while improving semantic clarity.

265-271: LGTM: Variable indexing updated consistently.

The changes from batch to system_idx in the variable indexing maintain test functionality while aligning with the new attribute naming convention.

320-322: LGTM: Variable indexing updated consistently.

The changes from batch to system_idx in the velocity indexing maintain test functionality while improving semantic consistency.

348-348: LGTM: SimState construction updated consistently.

The change from batch to system_idx in the SimState construction aligns with the broader refactoring and maintains test functionality.

432-432: LGTM: Tensor shape parameter updated consistently.

The change from n_batches to n_systems in the tensor shape parameter aligns with the broader terminology refactoring and maintains correct tensor dimensions.

455-461: LGTM: Comment and method call updated consistently.

The changes from n_batches to n_systems in the comment and from batch to system_idx in the method call maintain test functionality while improving semantic clarity.

528-528: LGTM: SimState construction updated consistently.

The change from batch to system_idx in the SimState construction maintains test functionality while aligning with the broader terminology refactoring.

877-879: LGTM: Comment updated consistently.

The change from "batch" to "system" terminology in the comment improves code clarity while maintaining the test's intent.

torch_sim/models/fairchem.py (1)

352-362: LGTM: Consistent renaming from batch to system_idx.

The changes correctly replace state.batch with state.system_idx throughout the forward method. The logic is preserved and the usage of torch.bincount(state.system_idx) and tensor sizing with state.system_idx.size(0) is appropriate for the new system-based indexing.

torch_sim/integrators/nve.py (5)

40-42: LGTM: Docstring updates reflect new system terminology.

The docstring changes from [n_batches] to [n_systems] correctly reflect the updated terminology for tensor shapes.

75-75: LGTM: Docstring update for kT parameter.

The docstring correctly specifies that kT can be scalar or have shape [n_systems].

91-91: LGTM: Function call updated to use system_idx.

The call to calculate_momenta correctly uses state.system_idx instead of state.batch, matching the function signature shown in the relevant code snippets from torch_sim/integrators/md.py.

102-102: LGTM: MDState initialization updated.

The MDState initialization correctly uses system_idx=state.system_idx instead of the deprecated batch attribute.

119-119: LGTM: Docstring update for dt parameter.

The docstring correctly specifies that dt can be scalar or have shape [n_systems].

torch_sim/models/lennard_jones.py (5)

360-360: LGTM: Docstring updated to use system terminology.

The docstring correctly updates from "batch" to "system" terminology in the method description.

369-373: LGTM: Return value documentation updated.

The docstring correctly reflects that energy and stress tensors have shapes [n_systems] and [n_systems, 3, 3] respectively.

380-380: LGTM: Error description updated.

The docstring correctly updates the error description to use "system" instead of "batch".

388-390: LGTM: Example documentation updated.

The example correctly shows that energy and stress have shapes [n_systems] and [n_systems, 3, 3] respectively.

397-400: LGTM: Runtime validation and iteration updated.

The runtime check correctly uses state.system_idx and the error message properly references "System" instead of "batch". The iteration over systems correctly uses state.n_systems.

tests/test_autobatching.py (7)

118-122: LGTM: Test assertions updated for new system terminology.

The test correctly uses state[1].n_systems == 1 and state[1].system_idx == 0 to validate split state properties. The comment appropriately explains that system indices are reset to 0 in split states.

475-482: LGTM: Convergence function updated to use system-based indexing.

The convergence function correctly:

Uses system_wise_max_force instead of batch_wise_max_force

Creates tensor with shape state.n_systems instead of n_batches

Uses state.system_idx for scatter_reduce indexing

Returns system-wise convergence criteria

517-517: LGTM: Variable name updated for consistency.

The variable name optimal_n_systems correctly reflects the new terminology while maintaining the same logic.

525-527: LGTM: System counting updated.

The system counting logic correctly uses n_systems instead of n_batches while preserving the test logic.

538-538: LGTM: Assertion updated for new variable names.

The assertion correctly compares n_systems == optimal_n_systems using the updated variable names.

564-564: LGTM: Convergence tensor creation updated.

The convergence tensor creation correctly uses state.n_systems instead of n_batches for sizing.

577-577: LGTM: Convergence tensor update consistent.

The convergence tensor update correctly uses state.n_systems for consistency with the new system-based indexing.

torch_sim/models/particle_life.py (4)

226-226: LGTM: Docstring updated for energy tensor shape.

The docstring correctly specifies that energy has shape [n_systems].

229-229: LGTM: Docstring updated for stress tensor shape.

The docstring correctly specifies that stress has shape [n_systems, 3, 3].

242-245: LGTM: Runtime validation updated with clearer error message.

The runtime check correctly uses state.system_idx instead of state.batch, and the error message clearly indicates that system_idx can only be inferred for single-system states.

247-247: LGTM: System iteration updated.

The iteration correctly uses state.n_systems instead of state.n_batches to process each system.

torch_sim/models/soft_sphere.py (5)

384-385: LGTM: Improved documentation clarity

The documentation update correctly clarifies that the method handles "each system" rather than generic "batched states," which better reflects the atomistic simulation context.

394-394: LGTM: Consistent shape documentation updates

The documentation correctly updates tensor shapes to use n_systems instead of n_batches, maintaining consistency with the new terminology while accurately describing the expected output dimensions.

Also applies to: 397-397, 410-410

402-402: LGTM: Consistent error handling and attribute access

The error message and state attribute access are correctly updated to use system_idx instead of batch, maintaining consistency with the new SimState API.

Also applies to: 418-421

423-423: LGTM: Correct loop iteration update

The loop correctly iterates over state.n_systems instead of state.n_batches, maintaining the same functional behavior with the updated terminology.

821-821: LGTM: Consistent terminology updates in SoftSphereMultiModel

All changes in the SoftSphereMultiModel class correctly apply the same terminology updates, ensuring consistency across both model implementations.

Also applies to: 824-824, 860-863, 865-865

examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py (3)

144-144: LGTM: Consistent state property updates

All references to n_batches are correctly updated to n_systems, maintaining the same functionality while using the improved terminology.

Also applies to: 171-171, 222-222, 233-233

332-332: LGTM: Consistent variable renaming

Variable names are consistently updated from batch-related to system-related terminology (batch_idx → system_idx, all_batches_for_gd → all_systems_for_gd), improving code readability and semantic accuracy.

Also applies to: 337-337, 353-355, 364-364, 378-378

396-397: LGTM: Consistent parameter and comment updates

The comment and parameter name updates (concatenated_batch_indices → concatenated_system_indices, batch → system_idx) maintain consistency with the new terminology while preserving functionality.

Also applies to: 416-416

tests/test_integrators.py (2)

112-112: LGTM: Consistent function call parameter updates

All calls to calc_kT are correctly updated to use system_idx=state.system_idx instead of batch=state.batch, maintaining test functionality with the updated API.

Also applies to: 175-175, 216-216, 276-276

375-376: LGTM: Consistent test function and parameter updates

The test function name and parameter updates (test_compute_cell_force_atoms_per_batch → test_compute_cell_force_atoms_per_system, batch → system_idx) ensure the test remains accurate with the refactored API.

Also applies to: 392-392

examples/tutorials/state_tutorial.py (2)

34-34: LGTM: Consistent documentation and comment updates

Comments and documentation are correctly updated to use "system" terminology instead of "batch," improving clarity and alignment with the codebase refactoring.

Also applies to: 78-78, 263-263

61-61: LGTM: Consistent property and attribute access updates

All property and attribute access are correctly updated to use the new API (n_systems instead of n_batches, system_idx instead of batch), ensuring the tutorial remains functional and educational.

Also applies to: 69-69, 112-112, 119-119, 151-151, 157-157, 162-162, 256-256

tests/test_neighbors.py (6)

40-40: LGTM: Docstring updated to reflect new terminology.

The docstring correctly updates the return value description from "batch index" to "system index" which better describes the semantic meaning of this tensor.

52-52: LGTM: Variable renamed consistently.

The variable system_idx is more descriptive than batch as it represents the system index for each atom.

56-56: LGTM: Variable usage updated consistently.

The assignment correctly uses the renamed system_idx variable.

62-62: LGTM: Return value updated to use renamed variable.

The function return statement correctly uses the renamed system_idx variable.

559-559: LGTM: Variable renamed consistently in test function.

The variable naming is consistent with the broader refactoring from "batch" to "system_idx" terminology.

563-563: LGTM: Function call updated with renamed parameter.

The function call correctly passes the renamed system_idx parameter instead of the old batch parameter.

tests/test_transforms.py (10)

253-254: LGTM: Variable names updated to use new terminology.

The mask variables are correctly renamed to use system_idx instead of batch, making the code more semantically clear.

271-273: LGTM: Function call updated with renamed parameter.

The function call correctly uses the new system_idx parameter name instead of the old batch parameter.

322-322: LGTM: Function parameter updated consistently.

The function call correctly passes the system_idx parameter with the appropriate tensor.

349-349: LGTM: Variable renamed for consistency.

The variable system_idx is more descriptive than the previous naming convention.

352-352: LGTM: Function call updated with renamed parameter.

The function call correctly uses the updated parameter name.

372-372: LGTM: Variable renamed consistently.

The variable naming follows the new terminology convention.

377-377: LGTM: Function parameters updated consistently.

All function calls correctly use the new system_idx parameter name, maintaining consistency throughout the test suite.

Also applies to: 385-385, 393-393

405-405: LGTM: Variable usage updated throughout the test.

The variable references are consistently updated to use system_idx instead of batch, maintaining the test logic while using more descriptive naming.

Also applies to: 410-410, 421-421, 425-425, 428-428, 432-432

439-439: LGTM: Comments and variable names updated consistently.

The comments and variable names are correctly updated to reflect the new "system" terminology while preserving the test logic.

Also applies to: 450-450, 455-455, 456-456, 463-463

849-849: LGTM: Test documentation and logic updated consistently.

The test names, comments, and assertions are correctly updated to use "system" terminology instead of "batch", making the tests more semantically accurate.

Also applies to: 852-852, 865-865, 874-874, 875-875

torch_sim/models/morse.py (6)

359-359: LGTM: Documentation updated to reflect new terminology.

The docstring correctly updates the energy tensor shape description from [n_batches] to [n_systems], which is more semantically accurate.

362-362: LGTM: Documentation updated consistently.

The stress tensor shape description is correctly updated from [n_batches, 3, 3] to [n_systems, 3, 3].

375-375: LGTM: Example documentation updated.

The example comment correctly reflects the new energy tensor shape using [n_systems] instead of [n_batches].

382-382: LGTM: Code updated to use new attribute name.

The conditional check correctly uses state.system_idx instead of state.batch, maintaining the same logic with the updated attribute name.

383-385: LGTM: Error message updated to reflect new terminology.

The error message correctly mentions system_idx instead of batch, providing clear guidance to users about the updated API.

387-387: LGTM: Iteration updated to use new property name.

The loop correctly uses state.n_systems instead of state.n_batches, maintaining the same functionality with the updated property name.

torch_sim/autobatching.py (8)

237-237: LGTM: Consistent terminology update

The change from state.n_batches to state.n_systems aligns with the broader codebase refactoring to use more semantically accurate terminology.

296-297: LGTM: Variable renaming maintains logic

The renaming from n_batches to n_systems and the corresponding usage in concatenate_states is consistent with the terminology update while preserving the original functionality.

346-346: LGTM: Consistent attribute access

The change from state.n_batches to state.n_systems maintains consistency with the SimState class updates referenced in the relevant code snippets.

408-409: LGTM: Updated error messaging

The print statement correctly uses the new state.n_systems attribute, maintaining consistency with the terminology changes while preserving the debugging information.

431-431: LGTM: Consistent docstring updates

The parameter descriptions have been updated from "batch" to "system" terminology, improving semantic clarity while maintaining accurate documentation.

Also applies to: 480-480, 719-719, 779-779

936-936: LGTM: Consistent variable usage

The variable renaming from n_batches to n_systems maintains the same logic flow while using more semantically appropriate terminology.

Also applies to: 942-942

974-974: LGTM: Updated shape annotation

The docstring update from [n_batches] to [n_systems] correctly reflects the tensor shape expectations after the terminology change.

1022-1023: LGTM: Consistent error handling updates

The error validation and conditional checks have been updated to use the new n_systems attribute, maintaining the same validation logic with improved terminology.

Also applies to: 1028-1029, 1060-1060

tests/test_runners.py (5)

152-152: LGTM: Consistent test updates

The trajectory file generation logic has been updated to use n_systems instead of n_batches, maintaining the same test behavior while aligning with the terminology changes.

Also applies to: 234-234, 343-343, 417-417

801-801: LGTM: Updated assertion

The assertion correctly uses the new n_systems attribute, ensuring the test validates the expected tensor shape after the terminology change.

809-809: LGTM: Mock state fixture updates

The mock state creation has been consistently updated to use system_idx instead of batch terminology, maintaining the same test fixture functionality while using more appropriate naming.

Also applies to: 815-816, 820-820, 823-823

861-861: LGTM: Test assertion updates

The test assertions have been updated to use the new n_systems attribute, ensuring proper validation of the state objects after the terminology change.

Also applies to: 880-880, 891-891

933-933: LGTM: Controlled mock state updates

The controlled mock state creation has been updated to use system_idx terminology and related variable names, maintaining the same test logic while using more semantically appropriate naming.

Also applies to: 935-936, 940-948

torch_sim/models/interface.py (3)

68-68: LGTM: Updated documentation examples

The shape annotations in the docstring examples have been updated from [n_batches] to [n_systems], maintaining consistency with the terminology changes while providing accurate documentation.

Also applies to: 70-70

177-177: LGTM: Interface specification updates

The interface documentation has been updated to use "system_idx" instead of "batch" and updated tensor shape annotations to use [n_systems]. This improves semantic clarity while maintaining the same interface contract.

Also applies to: 178-178, 184-184, 186-186

259-259: LGTM: Validation function updates

The validation function variable has been renamed from og_batch to og_system_idx with corresponding updates to the validation logic, maintaining the same validation behavior while using consistent terminology.

Also applies to: 269-270

torch_sim/io.py (7)

36-36: LGTM: Updated docstring descriptions

The docstring descriptions have been updated from "one per batch" to "one per system", providing more semantically accurate documentation for the conversion functions.

Also applies to: 82-82, 138-138

53-53: LGTM: Consistent variable renaming in state_to_atoms

The variable renaming from batch to system_idx and related updates maintain the same conversion logic while using more appropriate terminology. The iteration and masking logic remains functionally identical.

Also applies to: 55-55, 58-59, 60-62

101-101: LGTM: Consistent updates in state_to_structures

The conversion function has been updated to use system_idx terminology with corresponding variable renaming throughout the function, maintaining the same Pymatgen Structure conversion logic.

Also applies to: 103-103, 105-106, 109-114

155-155: LGTM: Consistent updates in state_to_phonopy

The PhonopyAtoms conversion function has been updated to use system_idx terminology, maintaining the same conversion logic while using more semantically appropriate variable names.

Also applies to: 157-157, 160-164

228-232: LGTM: Updated atoms_to_state function

The function has been updated to use system_idx terminology in variable names and comments, with the SimState constructor call updated accordingly. The conversion logic remains unchanged.

Also applies to: 244-244

300-304: LGTM: Updated structures_to_state function

The Pymatgen Structure to SimState conversion has been updated to use system_idx terminology, maintaining the same conversion logic while using more appropriate variable names.

Also applies to: 312-312

371-375: LGTM: Updated phonopy_to_state function

The PhonopyAtoms to SimState conversion has been updated to use system_idx terminology in variable names and comments, with the SimState constructor call updated accordingly.

Also applies to: 390-390

tests/test_state.py (4)

12-12: LGTM - Import updated correctly

The import statement correctly updates from _normalize_batch_indices to _normalize_system_indices to match the function rename.

31-37: LGTM - Property scope assertions updated correctly

The test assertions correctly reflect the rename from batch-related properties to system-related properties (system_idx and per_system scope).

409-463: LGTM - Function rename and test logic are correct

The test function has been properly renamed from test_normalize_batch_indices to test_normalize_system_indices and all function calls and assertions have been updated consistently. The test logic remains intact while using the new terminology.

623-646: LGTM - Excellent backward compatibility test

This new test ensures that deprecated batch properties continue to work correctly and are properly aliased to the new system properties. The test covers both getter and setter behavior, which is crucial for maintaining API compatibility.

torch_sim/trajectory.py (2)

211-242: LGTM - Clear documentation and validation logic updates

The documentation and validation logic have been correctly updated to use system terminology. The validation logic properly checks state.n_systems against the number of trajectory files, maintaining the same functional behavior.

676-725: LGTM - Parameter rename is consistent and well-documented

The parameter rename from batch_index to system_index is consistent throughout the method signature, documentation, and implementation. The logic for handling system indices remains functionally identical.

torch_sim/models/mace.py (4)

94-94: LGTM - Attribute documentation updated correctly

The attribute documentation properly reflects the rename from batch to system_idx while maintaining the same semantic meaning.

115-135: LGTM - Constructor parameter rename is consistent

The constructor parameter has been correctly renamed from batch to system_idx with corresponding documentation updates. The parameter meaning and usage remain identical.

197-199: LGTM - Method name and signature updated appropriately

The method name change from setup_from_batch to setup_from_batch and parameter rename to system_idx accurately reflects the new terminology while maintaining the same functionality.

337-337: LGTM - Model interface maintained correctly

The call to the underlying MACE model continues to use the batch parameter name as expected by the MACE interface, while the internal TorchSim code uses the clearer system_idx terminology.

torch_sim/monte_carlo.py (4)

44-56: LGTM - Documentation updated to reflect system-based swaps

The documentation correctly explains that swaps occur within the same system, and the return type documentation properly reflects that swaps are generated per system.

57-104: LGTM - System indexing logic is correct

The variable renames from batch to system are consistent throughout the function. The logic for generating swaps within systems remains functionally identical, with proper handling of system boundaries and indices.

127-141: LGTM - Validation function updated correctly

The validation function parameter and logic have been properly updated to use system_idx terminology. The validation ensures swaps only occur within the same system, which is the correct behavior.

236-236: LGTM - State initialization and validation updated

The state initialization and validation calls have been properly updated to use system_idx instead of batch, maintaining the same functionality with clearer terminology.

Also applies to: 268-268

torch_sim/quantities.py (5)

27-40: LGTM - Parameter rename and documentation updated correctly

The parameter rename from batch to system_idx is consistent and the documentation clearly explains that it indicates system membership of each particle. The terminology better reflects the semantic meaning.

54-69: LGTM - System-based temperature calculation logic is correct

The calculation logic has been properly updated to use system terminology while maintaining the same mathematical operations. The bincount and segment_reduce operations correctly group by system indices.

76-92: LGTM - Function signature and documentation updated consistently

The calc_temperature function properly forwards the renamed parameter to calc_kT, maintaining the same interface with clearer terminology.

100-130: LGTM - Kinetic energy calculation updated correctly

The kinetic energy calculation function has been properly updated to use system terminology in parameters and documentation. The logic for per-system calculations remains mathematically correct.

144-159: LGTM - Batch-wise function correctly renamed to system-wise

The function rename from computing maximum force per batch to per system is semantically correct. The implementation properly uses state.system_idx and state.n_systems to group forces by system.

torch_sim/integrators/md.py (4)

24-34: LGTM! Terminology improvements enhance code clarity.

The renaming from "batch" to "system" terminology provides better semantic clarity about what these tensors represent. The shape annotations and attribute names are consistently updated throughout the docstring.

54-68: LGTM! Parameter renaming improves semantic clarity.

The parameter name change from batch to system_idx and corresponding docstring updates make the function signature more descriptive and consistent with the broader terminology changes.

82-107: LGTM! Implementation consistently uses new terminology.

The variable renaming from batchwise_momenta to systemwise_momenta and consistent use of system_idx throughout the function maintains the same logic while improving code readability.

Note: Static analysis indicates some lines (83, 90, 95-96) lack test coverage. Consider adding tests for these conditional branches to ensure the temperature indexing logic is properly tested.

121-152: LGTM! Docstring updates and function calls are consistent.

The shape annotations in docstrings are properly updated from [n_batches] to [n_systems], and the function call to pbc_wrap_batched correctly uses the new system_idx attribute.

torch_sim/integrators/nvt.py (6)

48-52: LGTM! Docstring shape annotations are consistently updated.

The shape annotations are properly updated from [n_batches] to [n_systems] throughout the function documentation, maintaining consistency with the broader terminology changes.

96-122: LGTM! Function consistently uses new system indexing.

The docstring and implementation are properly updated to use state.system_idx instead of state.batch, maintaining the same indexing logic while improving code clarity.

150-181: LGTM! Function calls and object construction use consistent terminology.

The calculate_momenta function call and MDState constructor are properly updated to use system_idx instead of batch, maintaining functionality while improving code clarity.

205-209: LGTM! Docstring consistently updated.

The shape annotations in the docstring are properly updated from [n_batches] to [n_systems], maintaining consistency with the terminology changes throughout the codebase.

366-380: LGTM! Function consistently uses system terminology.

The function call to calculate_momenta and the subsequent calculations are properly updated to use system_idx and system-related terminology, maintaining the same logic while improving code clarity.

434-483: LGTM! Functions consistently use system terminology.

The calc_kinetic_energy function calls and related calculations are properly updated to use system_idx and system-related terminology, maintaining the same computational logic while improving code readability.

torch_sim/transforms.py (4)

26-59: LGTM! Improved terminology enhances code clarity.

The shape annotations and error messages are updated to use "system" terminology instead of "batch", providing better semantic clarity about what these dimensions represent. The error message is also more descriptive about the requirements for multiple systems.

158-172: LGTM! Function signature and docstring are consistently updated.

The parameter name change from batch to system_idx and the corresponding docstring updates make the function signature more descriptive and consistent with the broader terminology changes.

185-211: LGTM! Implementation consistently uses new terminology.

The variable names and comments are properly updated to use system terminology (unique_systems, n_systems, system_idx) while maintaining the same computational logic. The code is more readable with the improved naming.

538-1037: LGTM! All remaining functions consistently use new terminology.

The functions compute_cell_shifts, build_naive_neighborhood, and build_linked_cell_neighborhood are properly updated to use system_mapping instead of batch_mapping in parameter names, variable names, and comments. The computational logic remains unchanged while improving code clarity.

torch_sim/integrators/npt.py (3)

22-74: LGTM! Comprehensive docstring updates for system terminology.

The NPTLangevinState class documentation has been thoroughly updated to reflect the batch→system renaming, with correct tensor shape annotations.

109-139: Consider adding test coverage for system-specific cell force calculations.

The implementation correctly uses system-based indexing, but static analysis indicates these critical lines lack test coverage:

Volume calculation per system (line 109)

Pressure tensor expansion (line 120)

Temperature expansion (line 132)

Scaling by atoms per system (line 139)

These calculations are essential for NPT dynamics correctness.

806-898: Well-structured NPTNoseHooverState with proper system dimensions.

The class correctly implements system-based indexing with appropriate tensor shapes for cell variables. The logarithmic cell coordinate parameterization is clearly documented.

torch_sim/state.py (3)

25-164: Excellent refactoring of SimState to use system terminology.

The core class has been thoroughly updated with:

Clear documentation using system-based terminology

Proper validation of system indices as consecutive integers

Correct tensor shape expectations

620-639: Excellent implementation of system index remapping.

The code correctly handles the remapping of system indices to ensure they remain consecutive after filtering. This is crucial for maintaining the invariant that system indices must be consecutive integers starting from 0.

809-896: Well-implemented state concatenation with proper system index management.

The concatenate_states function correctly:

Maintains system index continuity across concatenated states

Handles per-system properties appropriately

Updates indices with proper offsets

torch_sim/integrators/npt.py

torch_sim/state.py

torch_sim/optimizers.py

add .vscode to gitignore rename per_batch to per_graph wip convert many batches to graphs convert more batches to graphs rename all of batch indices rename all of batch indices wip fix tests fix more tests more test renaming more test renaming again fix test_deform_grad_batched rename more batch to use graphs update mock state to use graphs precommit fix bug rename in calculate_momenta fix more precommit rename more batch to graphs more renames more renaming fix deprecated text in docstring added test for deprecated batch properties in simstate rename batch to graph more renames cleanup minor typo code rabbit reviewgs rename batch to system bump minor version

coderabbitai

Actionable comments posted: 0

♻️ Duplicate comments (2)

torch_sim/integrators/npt.py (1)

498-502: Redundant L_n shape check remains from previous review.

This conditional recalculation of L_n is unnecessary as identified in the previous review. The check is redundant because L_n is always freshly computed with the correct shape at lines 751-753 immediately before calling langevin_position_step.

torch_sim/state.py (1)

165-219: Deprecated properties properly implemented with warnings.

The deprecated properties n_atoms_per_batch, batch, and n_batches have been correctly implemented with DeprecationWarning warnings and proper redirection to the new properties. This addresses the past review comment requesting these warnings.

🧹 Nitpick comments (1)

torch_sim/state.py (1)

122-126: Consider improving system index validation logic.

The current validation logic uses torch.unique_consecutive and torch.bincount to ensure system indices are unique consecutive integers. However, the TODO comment indicates this approach may not be reliable.

Consider implementing a more robust validation approach:

-# TODO(curtis): I feel like this logic is not reliable.
-# I'll come up with something better later.
-_, counts = torch.unique_consecutive(self.system_idx, return_counts=True)
-if not torch.all(counts == torch.bincount(self.system_idx)):
-    raise ValueError("System indices must be unique consecutive integers")
+# Validate system indices are consecutive integers starting from 0
+sorted_indices = torch.sort(self.system_idx)[0]
+expected_indices = torch.arange(self.n_systems, device=self.device, dtype=torch.int64)
+expected_indices = expected_indices.repeat_interleave(torch.bincount(self.system_idx))
+if not torch.equal(sorted_indices, expected_indices):
+    raise ValueError("System indices must be unique consecutive integers starting from 0")

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c96043d and 657ea73.

📒 Files selected for processing (68)

.gitignore (1 hunks)
examples/scripts/1_Introduction/1.2_MACE.py (1 hunks)
examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py (1 hunks)
examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py (1 hunks)
examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py (1 hunks)
examples/scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py (1 hunks)
examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py (1 hunks)
examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py (2 hunks)
examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py (5 hunks)
examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py (1 hunks)
examples/scripts/3_Dynamics/3.2_MACE_NVE.py (1 hunks)
examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py (1 hunks)
examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py (1 hunks)
examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py (1 hunks)
examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py (1 hunks)
examples/scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py (2 hunks)
examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py (4 hunks)
examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py (1 hunks)
examples/scripts/4_High_level_api/4.2_auto_batching_api.py (1 hunks)
examples/scripts/5_Workflow/5.2_In_Flight_WBM.py (1 hunks)
examples/scripts/7_Others/7.3_Batched_neighbor_list.py (2 hunks)
examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py (10 hunks)
examples/tutorials/autobatching_tutorial.py (2 hunks)
examples/tutorials/high_level_tutorial.py (1 hunks)
examples/tutorials/hybrid_swap_tutorial.py (1 hunks)
examples/tutorials/low_level_tutorial.py (1 hunks)
examples/tutorials/state_tutorial.py (6 hunks)
pyproject.toml (1 hunks)
tests/models/conftest.py (1 hunks)
tests/test_autobatching.py (6 hunks)
tests/test_correlations.py (1 hunks)
tests/test_integrators.py (6 hunks)
tests/test_io.py (6 hunks)
tests/test_monte_carlo.py (5 hunks)
tests/test_neighbors.py (3 hunks)
tests/test_optimizers.py (9 hunks)
tests/test_runners.py (9 hunks)
tests/test_state.py (18 hunks)
tests/test_trajectory.py (5 hunks)
tests/test_transforms.py (10 hunks)
torch_sim/autobatching.py (12 hunks)
torch_sim/integrators/md.py (7 hunks)
torch_sim/integrators/npt.py (51 hunks)
torch_sim/integrators/nve.py (5 hunks)
torch_sim/integrators/nvt.py (10 hunks)
torch_sim/io.py (12 hunks)
torch_sim/math.py (1 hunks)
torch_sim/models/fairchem.py (1 hunks)
torch_sim/models/graphpes.py (1 hunks)
torch_sim/models/interface.py (4 hunks)
torch_sim/models/lennard_jones.py (2 hunks)
torch_sim/models/mace.py (9 hunks)
torch_sim/models/metatomic.py (2 hunks)
torch_sim/models/morse.py (2 hunks)
torch_sim/models/orb.py (3 hunks)
torch_sim/models/particle_life.py (2 hunks)
torch_sim/models/sevennet.py (2 hunks)
torch_sim/models/soft_sphere.py (4 hunks)
torch_sim/monte_carlo.py (5 hunks)
torch_sim/neighbors.py (3 hunks)
torch_sim/optimizers.py (38 hunks)
torch_sim/quantities.py (5 hunks)
torch_sim/runners.py (8 hunks)
torch_sim/state.py (23 hunks)
torch_sim/trajectory.py (6 hunks)
torch_sim/transforms.py (15 hunks)
torch_sim/typing.py (1 hunks)
torch_sim/workflows/a2c.py (1 hunks)

✅ Files skipped from review due to trivial changes (6)

examples/scripts/5_Workflow/5.2_In_Flight_WBM.py
pyproject.toml
examples/tutorials/high_level_tutorial.py
examples/scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py
examples/tutorials/state_tutorial.py
torch_sim/optimizers.py

🚧 Files skipped from review as they are similar to previous changes (58)

examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py
.gitignore
examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py
examples/scripts/4_High_level_api/4.2_auto_batching_api.py
examples/scripts/3_Dynamics/3.2_MACE_NVE.py
examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py
torch_sim/math.py
examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py
tests/models/conftest.py
torch_sim/typing.py
examples/tutorials/low_level_tutorial.py
examples/tutorials/autobatching_tutorial.py
torch_sim/models/orb.py
torch_sim/models/sevennet.py
examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py
torch_sim/workflows/a2c.py
tests/test_autobatching.py
torch_sim/models/graphpes.py
examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py
tests/test_integrators.py
torch_sim/integrators/nve.py
torch_sim/neighbors.py
tests/test_optimizers.py
tests/test_correlations.py
examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py
tests/test_trajectory.py
examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py
examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py
torch_sim/models/lennard_jones.py
torch_sim/models/fairchem.py
examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py
tests/test_neighbors.py
torch_sim/io.py
examples/tutorials/hybrid_swap_tutorial.py
examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py
examples/scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py
torch_sim/models/metatomic.py
torch_sim/quantities.py
examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py
examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py
examples/scripts/1_Introduction/1.2_MACE.py
tests/test_monte_carlo.py
examples/scripts/7_Others/7.3_Batched_neighbor_list.py
tests/test_runners.py
examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py
torch_sim/runners.py
tests/test_io.py
torch_sim/trajectory.py
tests/test_transforms.py
torch_sim/models/morse.py
torch_sim/models/particle_life.py
torch_sim/models/mace.py
torch_sim/models/soft_sphere.py
torch_sim/monte_carlo.py
torch_sim/models/interface.py
torch_sim/transforms.py
tests/test_state.py
torch_sim/integrators/nvt.py

🧰 Additional context used

🧬 Code Graph Analysis (2)

torch_sim/autobatching.py (1)

torch_sim/state.py (2)

n_systems (222-224)

concatenate_states (829-915)

torch_sim/integrators/npt.py (7)

torch_sim/state.py (8)

device (142-144)

dtype (147-149)

n_systems (222-224)

n_atoms_per_system (157-163)

volume (227-229)

momenta (380-386)

to (341-355)

clone (259-275)

torch_sim/models/interface.py (4)

device (107-109)

device (112-116)

dtype (119-121)

dtype (124-128)

tests/conftest.py (2)

device (24-25)

dtype (29-30)

tests/test_neighbors.py (1)

dtype (15-16)

torch_sim/integrators/md.py (2)

velocities (44-48)

calculate_momenta (51-109)

torch_sim/integrators/nvt.py (1)

velocities (269-273)

torch_sim/quantities.py (1)

calc_kinetic_energy (96-130)

🪛 GitHub Check: codecov/patch

torch_sim/autobatching.py

[warning] 1023-1023: torch_sim/autobatching.py#L1023
Added line #L1023 was not covered by tests

[warning] 1029-1029: torch_sim/autobatching.py#L1029
Added line #L1029 was not covered by tests

torch_sim/integrators/md.py

[warning] 83-83: torch_sim/integrators/md.py#L83
Added line #L83 was not covered by tests

[warning] 90-90: torch_sim/integrators/md.py#L90
Added line #L90 was not covered by tests

[warning] 95-96: torch_sim/integrators/md.py#L95-L96
Added lines #L95 - L96 were not covered by tests

torch_sim/integrators/npt.py

[warning] 246-246: torch_sim/integrators/npt.py#L246
Added line #L246 was not covered by tests

[warning] 249-249: torch_sim/integrators/npt.py#L249
Added line #L249 was not covered by tests

[warning] 286-286: torch_sim/integrators/npt.py#L286
Added line #L286 was not covered by tests

[warning] 294-294: torch_sim/integrators/npt.py#L294
Added line #L294 was not covered by tests

[warning] 511-511: torch_sim/integrators/npt.py#L511
Added line #L511 was not covered by tests

[warning] 514-514: torch_sim/integrators/npt.py#L514
Added line #L514 was not covered by tests

[warning] 585-585: torch_sim/integrators/npt.py#L585
Added line #L585 was not covered by tests

[warning] 588-588: torch_sim/integrators/npt.py#L588
Added line #L588 was not covered by tests

[warning] 891-893: torch_sim/integrators/npt.py#L891-L893
Added lines #L891 - L893 were not covered by tests

[warning] 966-968: torch_sim/integrators/npt.py#L966-L968
Added lines #L966 - L968 were not covered by tests

[warning] 979-979: torch_sim/integrators/npt.py#L979
Added line #L979 was not covered by tests

[warning] 1017-1017: torch_sim/integrators/npt.py#L1017
Added line #L1017 was not covered by tests

[warning] 1020-1021: torch_sim/integrators/npt.py#L1020-L1021
Added lines #L1020 - L1021 were not covered by tests

[warning] 1089-1089: torch_sim/integrators/npt.py#L1089
Added line #L1089 was not covered by tests

[warning] 1154-1154: torch_sim/integrators/npt.py#L1154
Added line #L1154 was not covered by tests

[warning] 1212-1212: torch_sim/integrators/npt.py#L1212
Added line #L1212 was not covered by tests

[warning] 1216-1216: torch_sim/integrators/npt.py#L1216
Added line #L1216 was not covered by tests

[warning] 1219-1224: torch_sim/integrators/npt.py#L1219-L1224
Added lines #L1219 - L1224 were not covered by tests

[warning] 1234-1234: torch_sim/integrators/npt.py#L1234
Added line #L1234 was not covered by tests

[warning] 1273-1275: torch_sim/integrators/npt.py#L1273-L1275
Added lines #L1273 - L1275 were not covered by tests

[warning] 1288-1289: torch_sim/integrators/npt.py#L1288-L1289
Added lines #L1288 - L1289 were not covered by tests

[warning] 1413-1413: torch_sim/integrators/npt.py#L1413
Added line #L1413 was not covered by tests

[warning] 1417-1418: torch_sim/integrators/npt.py#L1417-L1418
Added lines #L1417 - L1418 were not covered by tests

[warning] 1425-1425: torch_sim/integrators/npt.py#L1425
Added line #L1425 was not covered by tests

[warning] 1428-1429: torch_sim/integrators/npt.py#L1428-L1429
Added lines #L1428 - L1429 were not covered by tests

[warning] 1438-1438: torch_sim/integrators/npt.py#L1438
Added line #L1438 was not covered by tests

[warning] 1448-1448: torch_sim/integrators/npt.py#L1448
Added line #L1448 was not covered by tests

[warning] 1545-1545: torch_sim/integrators/npt.py#L1545
Added line #L1545 was not covered by tests

🔇 Additional comments (26)

torch_sim/autobatching.py (8)

237-237: LGTM: Consistent terminology update in print statement.

The print statement now correctly uses "systems" instead of "batches" to align with the new terminology.

296-297: LGTM: Variable renaming improves clarity.

The variable name change from n_batches to n_systems better reflects the actual concept being represented.

346-347: LGTM: Property access updated correctly.

The change from n_batches to n_systems is consistent with the SimState class property renaming.

408-410: LGTM: Print statement terminology consistently updated.

The print statement now correctly references "systems" instead of "batches" in both contexts, maintaining consistency with the new terminology.

431-431: LGTM: Docstring terminology updated consistently.

The docstring updates correctly change "batch" to "system" throughout the parameter descriptions, improving clarity about what the code actually handles.

Also applies to: 480-480, 719-719, 779-779

936-942: LGTM: Variable names updated consistently.

The variable name change from n_batches to n_systems is consistent with the overall terminology update and maintains the same logic.

1022-1023: LGTM: Error messages updated with consistent terminology.

The error messages now correctly reference n_systems instead of n_batches, maintaining consistency with the updated property names.

Also applies to: 1028-1029

1060-1060: LGTM: Conditional check updated correctly.

The property access is correctly updated to use n_systems instead of n_batches.

torch_sim/integrators/md.py (7)

24-24: LGTM: MDState class documentation updated consistently.

The tensor shape annotations and attribute descriptions now correctly use "system" terminology instead of "batch", improving clarity about what the class represents.

Also applies to: 27-27, 29-29, 34-34

54-54: LGTM: Function signature and docstring updated correctly.

The parameter rename from batch to system_idx and the corresponding docstring updates improve clarity about the indexing scheme.

Also applies to: 67-68

82-83: LGTM: Code logic updated with consistent terminology.

The comment and indexing operation correctly use system_idx instead of batch, maintaining the same functionality with clearer terminology.

90-91: LGTM: Variable names updated consistently throughout function.

The variable names are systematically updated to use "system" terminology while maintaining the same computational logic for momenta calculation.

Also applies to: 94-96

107-107: LGTM: Indexing operation updated correctly.

The indexing operation now uses system_idx instead of batch, maintaining consistency with the parameter renaming.

121-121: LGTM: Docstring tensor shape annotations updated.

The docstring comments now correctly reference [n_systems] instead of [n_batches] for the timestep tensor shape.

Also applies to: 141-141

150-152: LGTM: PBC wrapping function call updated correctly.

The function call now uses state.system_idx instead of state.batch, maintaining consistency with the attribute renaming in the MDState class.

torch_sim/integrators/npt.py (1)

1-1673: Excellent systematic refactoring of batch-to-system terminology.

The renaming from "batch" to "system" terminology throughout this NPT integrator is comprehensive and well-executed. The changes consistently update:

Class attributes and properties (n_batches → n_systems, batch → system_idx)

Function parameters and variable names

Docstrings and comments to reflect the new semantics

Tensor shape annotations ([n_batches] → [n_systems])

The mathematical logic and algorithmic implementation remain unchanged, which is exactly what's needed for this semantic refactoring. This aligns well with the PR objective to clarify that "batch" actually refers to "systems" in the codebase.

torch_sim/state.py (10)

32-84: Systematic documentation and attribute renaming looks good.

The documentation updates throughout the class docstring and attribute definitions are consistent and accurately reflect the conceptual shift from "batch" to "system" terminology. The attribute renaming from batch to system_idx maintains semantic clarity.

312-373: Method parameter renaming is consistent.

The parameter renaming in methods pop and __getitem__ from batch_indices to system_indices is consistent with the overall refactoring and maintains the same functionality.

420-457: Function renaming and implementation is correct.

The function _normalize_batch_indices has been properly renamed to _normalize_system_indices with all parameter names, documentation, and internal logic updated accordingly. The function handles various input types correctly.

496-574: Property scope inference updates are accurate.

The infer_property_scope function has been correctly updated to use "per_system" instead of "per_batch" terminology throughout, including parameter names, documentation, and return type annotations.

613-666: Attribute filtering logic properly updated.

The _filter_attrs_by_mask function has been correctly updated to handle "per_system" attributes instead of "per_batch", including the system index remapping logic for maintaining consecutive indices.

669-724: State splitting logic correctly refactored.

The _split_state function has been properly updated to use "system" terminology throughout, including variable names, comments, and logic for handling per-system attributes.

727-779: State popping implementation is consistent.

The _pop_states function has been correctly updated with "system" terminology, including parameter names, documentation, and the logic for creating system masks.

782-826: State slicing functionality properly updated.

The _slice_state function has been correctly refactored to use "system" terminology throughout, including parameter names, documentation, and the implementation logic.

829-915: State concatenation logic is correctly updated.

The concatenate_states function has been properly updated to use "system" terminology, including variable names, comments, and the logic for handling per-system properties and system index offsets.

946-950: State initialization validation updated correctly.

The validation logic in initialize_state has been properly updated to check for n_systems == 1 instead of n_batches == 1, maintaining the same functionality with the new terminology.

orionarcher · 2025-07-09T20:02:08Z

I think this is broadly a good change from a usability and intuitiveness perspective, thanks @curtischong. I couldn't find any missed batchs.

The only tension I see arises from how batch is used in other contexts. Nequip, for example, uses batch to refer to the batch indices and this is also common in other ML packages.

The question is then, "to what extent do we adopt prior norms vs make choices that improve usability for the materials community?" In this case, I think batch -> system is a good change and I'm supportive of merging the PR, just wanted to highlight the reason we originally chose batch and open up the discussion, maybe @abhijeetgangan wants to chime in?

curtischong · 2025-07-10T00:51:41Z

Mace also uses something similar. However, .batch is used in a very minimal way (mainly to determine the scatter_sum index)

I also think it's fine for MLIPs to not follow this convention right now because they are more like individual repos with few dependents. I think for a framework, it's more important to have variable names that accurately reflect what they are.

I think we should do it because the surface area of using "systems" in our codebase is might wider. Additionally, we have a higher chance of ambiguity (due to the autobatchers).

abhijeetgangan · 2025-07-10T01:07:39Z

In the context of ML potentials, batch is used in many independent packages because it is part of torch_geometric as seen from your screenshot or from here. So the question I would be asking is that if it is worth changing the nomenclature for people more focused on materials side and possibly breaking code or keeping it as it is. This is because the batch during training and inference would mean the same thing and would be associated with similar operations (scatter). We use that for similar operations anyway...

curtischong · 2025-07-10T15:41:47Z

I agree that we should not be disruptive. But I also think that renaming it as system really helps the readability in our codebase.

How about I make the .batch getter/setter aliases to be non-deprecated so people can get/set .batch as they please? But we keep the system terminology inside torchsim.

In an upcoming PR I will also make it so that we'll have multiple constructors for SimState so when they init a sim object, they can pass in either batch or system_idx (whichever they find more intuitive).

abhijeetgangan · 2025-07-13T05:59:06Z

I do not have a strong opinion on the suggestions you made. As I don't have much experience in maintaining packages, I will leave this decision to others.

orionarcher · 2025-07-15T20:57:22Z

How about I make the .batch getter/setter aliases to be non-deprecated so people can get/set .batch as they please? But we keep the system terminology inside torchsim.

I think this is a good middle ground, would be nice to not immediately break some workflows. Let's just throw a deprecation notice and pull the batch getters/setters in a month.

SimState so when they init a sim object, they can pass in either batch or system_idx (whichever they find more intuitive).

If we are going to deprecate I don't see too much of a need to maintain a separate constructor long term. I suspect most people are instantiating from pymatgen or ase anyway.

cla-bot bot added the cla-signed Contributor license agreement signed label Jul 9, 2025

curtischong marked this pull request as draft July 9, 2025 02:49

coderabbitai bot reviewed Jul 9, 2025

View reviewed changes

torch_sim/integrators/npt.py Show resolved Hide resolved

torch_sim/state.py Show resolved Hide resolved

torch_sim/optimizers.py Show resolved Hide resolved

curtischong mentioned this pull request Jul 9, 2025

Rename batch to graph #216

Closed

3 tasks

curtischong force-pushed the rename-batch-to-system branch from 14a9b23 to c335511 Compare July 9, 2025 03:10

curtischong force-pushed the rename-batch-to-system branch from c335511 to 9a1b87e Compare July 9, 2025 03:11

add warnings to deprecated attributes

657ea73

curtischong assigned CompRhys, orionarcher and janosh Jul 9, 2025

curtischong marked this pull request as ready for review July 9, 2025 15:04

coderabbitai bot reviewed Jul 9, 2025

View reviewed changes

curtischong assigned abhijeetgangan Jul 10, 2025

orionarcher approved these changes Jul 15, 2025

View reviewed changes

curtischong merged commit 317985c into main Jul 16, 2025
91 of 92 checks passed

curtischong deleted the rename-batch-to-system branch July 16, 2025 16:44

curtischong mentioned this pull request Aug 3, 2025

Rename more batch to system #233

Merged

2 tasks

This was referenced Aug 8, 2025

Improve Typing of ModelInterface #215

Merged

Define attribute scopes in SimStates #228

Merged

curtischong added the breaking Breaking changes label Aug 12, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Rename batch to system #217

Rename batch to system #217

Uh oh!

curtischong commented Jul 9, 2025 •

edited

Loading

Uh oh!

coderabbitai bot commented Jul 9, 2025 •

edited

Loading

Chat

Support

CodeRabbit Commands (Invoked using PR comments)

Other keywords and placeholders

CodeRabbit Configuration File (`.coderabbit.yaml`)

Documentation and Community

Uh oh!

coderabbitai bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot left a comment

Uh oh!

orionarcher commented Jul 9, 2025 •

edited

Loading

Uh oh!

curtischong commented Jul 10, 2025 •

edited

Loading

Uh oh!

abhijeetgangan commented Jul 10, 2025

Uh oh!

curtischong commented Jul 10, 2025

Uh oh!

abhijeetgangan commented Jul 13, 2025

Uh oh!

orionarcher commented Jul 15, 2025

Uh oh!

Uh oh!

Uh oh!

Rename batch to system #217

Rename batch to system #217

Uh oh!

Conversation

curtischong commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Notes of my changes

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Possibly related PRs

Suggested labels

Poem

Chat

Support

CodeRabbit Commands (Invoked using PR comments)

Other keywords and placeholders

CodeRabbit Configuration File (.coderabbit.yaml)

Documentation and Community

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

orionarcher commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

curtischong commented Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

abhijeetgangan commented Jul 10, 2025

Uh oh!

curtischong commented Jul 10, 2025

Uh oh!

abhijeetgangan commented Jul 13, 2025

Uh oh!

orionarcher commented Jul 15, 2025

Uh oh!

Uh oh!

Uh oh!

curtischong commented Jul 9, 2025 •

edited

Loading

coderabbitai bot commented Jul 9, 2025 •

edited

Loading

CodeRabbit Configuration File (`.coderabbit.yaml`)

orionarcher commented Jul 9, 2025 •

edited

Loading

curtischong commented Jul 10, 2025 •

edited

Loading