rename per_batch to per_graph

Author: curtischong

examples/scripts/1_Introduction/1.2_MACE.py

@@ -64,11 +64,11 @@
 atomic_numbers = torch.tensor(atomic_numbers_numpy, device=device, dtype=torch.int)
 
 # create batch index array of shape (16,) which is 0 for first 8 atoms, 1 for last 8 atoms
-atoms_per_batch = torch.tensor(
+atoms_per_graph = torch.tensor(
     [len(atoms) for atoms in atoms_list], device=device, dtype=torch.int
 )
 batch = torch.repeat_interleave(
-    torch.arange(len(atoms_per_batch), device=device), atoms_per_batch
+    torch.arange(len(atoms_per_graph), device=device), atoms_per_graph
 )
 
 # You can see their shapes are as expected

examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py

@@ -94,11 +94,11 @@
 masses = torch.tensor(masses_numpy, device=device, dtype=dtype)
 
 # Create batch indices tensor for scatter operations
-atoms_per_batch = torch.tensor(
+atoms_per_graph = torch.tensor(
     [len(atoms) for atoms in atoms_list], device=device, dtype=torch.int
 )
 batch_indices = torch.repeat_interleave(
-    torch.arange(len(atoms_per_batch), device=device), atoms_per_batch
+    torch.arange(len(atoms_per_graph), device=device), atoms_per_graph
 )
 """
 

examples/tutorials/state_tutorial.py

@@ -260,7 +260,7 @@
 scope = infer_property_scope(md_state)
 print("Global properties:", scope["global"])
 print("Per-atom properties:", scope["per_atom"])
-print("Per-batch properties:", scope["per_batch"])
+print("Per-batch properties:", scope["per_graph"])
 
 
 # %% [markdown]

tests/test_integrators.py

@@ -372,7 +372,7 @@ def test_compare_single_vs_batched_integrators(
         torch.testing.assert_close(single_state.energy, final_state.energy)
 
 
-def test_compute_cell_force_atoms_per_batch():
+def test_compute_cell_force_atoms_per_graph():
     """Test that compute_cell_force correctly scales by number of atoms per batch.
 
     Covers fix in https://github.com/Radical-AI/torch-sim/pull/153."""

tests/test_state.py

@@ -29,7 +29,7 @@ def test_infer_sim_state_property_scope(si_sim_state: ts.SimState) -> None:
     scope = infer_property_scope(si_sim_state)
     assert set(scope["global"]) == {"pbc"}
     assert set(scope["per_atom"]) == {"positions", "masses", "atomic_numbers", "batch"}
-    assert set(scope["per_batch"]) == {"cell"}
+    assert set(scope["per_graph"]) == {"cell"}
 
 
 def test_infer_md_state_property_scope(si_sim_state: ts.SimState) -> None:
@@ -50,7 +50,7 @@ def test_infer_md_state_property_scope(si_sim_state: ts.SimState) -> None:
         "forces",
         "momenta",
     }
-    assert set(scope["per_batch"]) == {"cell", "energy"}
+    assert set(scope["per_graph"]) == {"cell", "energy"}
 
 
 def test_slice_substate(
@@ -156,9 +156,9 @@ def test_concatenate_si_and_fe_states(
     )
     assert torch.all(concatenated.batch == expected_batch)
 
-    # check n_atoms_per_batch
+    # check n_atoms_per_graph
     assert torch.all(
-        concatenated.n_atoms_per_batch
+        concatenated.n_atoms_per_graph
         == torch.tensor(
             [si_sim_state.n_atoms, fe_supercell_sim_state.n_atoms],
             device=concatenated.device,

tests/test_trajectory.py

@@ -678,9 +678,9 @@ def test_multi_batch_reporter(
 
     # Check that each trajectory has the correct number of atoms
     # (should be half of the total in the double state)
-    atoms_per_batch = si_double_sim_state.positions.shape[0] // 2
-    assert traj0.get_array("positions").shape[1] == atoms_per_batch
-    assert traj1.get_array("positions").shape[1] == atoms_per_batch
+    atoms_per_graph = si_double_sim_state.positions.shape[0] // 2
+    assert traj0.get_array("positions").shape[1] == atoms_per_graph
+    assert traj1.get_array("positions").shape[1] == atoms_per_graph
 
     # Check property data
     assert "ones" in traj0.array_registry

torch_sim/integrators/npt.py

@@ -135,8 +135,8 @@ def _compute_cell_force(
         3, device=state.device, dtype=state.dtype
     ).unsqueeze(0)
 
-    # Correct implementation with scaling by n_atoms_per_batch
-    return virial + e_kin_per_atom * state.n_atoms_per_batch.view(-1, 1, 1)
+    # Correct implementation with scaling by n_atoms_per_graph
+    return virial + e_kin_per_atom * state.n_atoms_per_graph.view(-1, 1, 1)
 
 
 def npt_langevin(  # noqa: C901, PLR0915
@@ -663,13 +663,13 @@ def npt_init(
 
         # Calculate cell masses based on system size and temperature
         # This follows standard NPT barostat mass scaling
-        n_atoms_per_batch = torch.bincount(state.batch)
+        n_atoms_per_graph = torch.bincount(state.batch)
         batch_kT = (
             kT.expand(state.n_batches)
             if isinstance(kT, torch.Tensor) and kT.ndim == 0
             else kT
         )
-        cell_masses = (n_atoms_per_batch + 1) * batch_kT * b_tau * b_tau
+        cell_masses = (n_atoms_per_graph + 1) * batch_kT * b_tau * b_tau
 
         # Create the initial state
         return NPTLangevinState(
@@ -735,8 +735,8 @@ def npt_update(
 
         # Update barostat mass based on current temperature
         # This ensures proper coupling between system and barostat
-        n_atoms_per_batch = torch.bincount(state.batch)
-        state.cell_masses = (n_atoms_per_batch + 1) * batch_kT * b_tau * b_tau
+        n_atoms_per_graph = torch.bincount(state.batch)
+        state.cell_masses = (n_atoms_per_graph + 1) * batch_kT * b_tau * b_tau
 
         # Compute model output for current state
         model_output = model(state)
@@ -1017,8 +1017,8 @@ def update_cell_mass(
         kT_batch = kT.expand(state.n_batches) if kT.ndim == 0 else kT
 
         # Calculate cell masses for each batch
-        n_atoms_per_batch = torch.bincount(state.batch, minlength=state.n_batches)
-        cell_mass = dim * (n_atoms_per_batch + 1) * kT_batch * state.barostat.tau**2
+        n_atoms_per_graph = torch.bincount(state.batch, minlength=state.n_batches)
+        cell_mass = dim * (n_atoms_per_graph + 1) * kT_batch * state.barostat.tau**2
 
         # Update state with new cell masses
         state.cell_mass = cell_mass.to(device=device, dtype=dtype)
@@ -1213,15 +1213,15 @@ def compute_cell_force(
 
         # Compute kinetic energy contribution per batch
         # Split momenta and masses by batch
-        KE_per_batch = torch.zeros(
+        KE_per_graph = torch.zeros(
             n_batches, device=positions.device, dtype=positions.dtype
         )
         for b in range(n_batches):
             batch_mask = batch == b
             if batch_mask.any():
                 batch_momenta = momenta[batch_mask]
                 batch_masses = masses[batch_mask]
-                KE_per_batch[b] = calc_kinetic_energy(batch_momenta, batch_masses)
+                KE_per_graph[b] = calc_kinetic_energy(batch_momenta, batch_masses)
 
         # Get stress tensor and compute trace per batch
         # Handle stress tensor with batch dimension
@@ -1236,7 +1236,7 @@ def compute_cell_force(
         # Compute force on cell coordinate per batch
         # F = alpha * KE - dU/dV - P*V*d
         return (
-            (alpha * KE_per_batch)
+            (alpha * KE_per_graph)
             - (internal_pressure * volume)
             - (external_pressure * volume * dim)
         )
@@ -1285,8 +1285,8 @@ def npt_inner_step(
         model_output = model(state)
 
         # First half step: Update momenta
-        n_atoms_per_batch = torch.bincount(state.batch, minlength=state.n_batches)
-        alpha = 1 + 1 / n_atoms_per_batch  # [n_batches]
+        n_atoms_per_graph = torch.bincount(state.batch, minlength=state.n_batches)
+        alpha = 1 + 1 / n_atoms_per_graph  # [n_batches]
 
         cell_force_val = compute_cell_force(
             alpha=alpha,
@@ -1425,8 +1425,8 @@ def npt_nose_hoover_init(
         kT_batch = kT.expand(n_batches) if kT.ndim == 0 else kT
 
         # Calculate cell masses for each batch
-        n_atoms_per_batch = torch.bincount(state.batch, minlength=n_batches)
-        cell_mass = dim * (n_atoms_per_batch + 1) * kT_batch * b_tau**2
+        n_atoms_per_graph = torch.bincount(state.batch, minlength=n_batches)
+        cell_mass = dim * (n_atoms_per_graph + 1) * kT_batch * b_tau**2
         cell_mass = cell_mass.to(device=device, dtype=dtype)
 
         # Calculate cell kinetic energy (using first batch for initialization)
@@ -1596,19 +1596,19 @@ def npt_nose_hoover_invariant(
     e_pot = state.energy  # Should be scalar or [n_batches]
 
     # Calculate kinetic energy of particles per batch
-    e_kin_per_batch = calc_kinetic_energy(state.momenta, state.masses, batch=state.batch)
+    e_kin_per_graph = calc_kinetic_energy(state.momenta, state.masses, batch=state.batch)
 
     # Calculate degrees of freedom per batch
-    n_atoms_per_batch = torch.bincount(state.batch)
-    DOF_per_batch = (
-        n_atoms_per_batch * state.positions.shape[-1]
+    n_atoms_per_graph = torch.bincount(state.batch)
+    DOF_per_graph = (
+        n_atoms_per_graph * state.positions.shape[-1]
     )  # n_atoms * n_dimensions
 
     # Initialize total energy with PE + KE
     if isinstance(e_pot, torch.Tensor) and e_pot.ndim > 0:
-        e_tot = e_pot + e_kin_per_batch  # [n_batches]
+        e_tot = e_pot + e_kin_per_graph  # [n_batches]
     else:
-        e_tot = e_pot + e_kin_per_batch  # [n_batches]
+        e_tot = e_pot + e_kin_per_graph  # [n_batches]
 
     # Add thermostat chain contributions
     # Note: These are global thermostat variables, so we add them to each batch
@@ -1618,14 +1618,14 @@ def npt_nose_hoover_invariant(
         2 * state.thermostat.masses[0]
     )
 
-    # Ensure kT can broadcast properly with DOF_per_batch
+    # Ensure kT can broadcast properly with DOF_per_graph
     if isinstance(kT, torch.Tensor) and kT.ndim == 0:
-        # Scalar kT - expand to match DOF_per_batch shape
-        kT_expanded = kT.expand_as(DOF_per_batch)
+        # Scalar kT - expand to match DOF_per_graph shape
+        kT_expanded = kT.expand_as(DOF_per_graph)
     else:
         kT_expanded = kT
 
-    thermostat_energy += DOF_per_batch * kT_expanded * state.thermostat.positions[0]
+    thermostat_energy += DOF_per_graph * kT_expanded * state.thermostat.positions[0]
 
     # Add remaining thermostat terms
     for pos, momentum, mass in zip(

torch_sim/integrators/nvt.py

@@ -370,14 +370,14 @@ def nvt_nose_hoover_init(
         KE = calc_kinetic_energy(momenta, state.masses, batch=state.batch)
 
         # Calculate degrees of freedom per batch
-        n_atoms_per_batch = torch.bincount(state.batch)
-        dof_per_batch = (
-            n_atoms_per_batch * state.positions.shape[-1]
+        n_atoms_per_graph = torch.bincount(state.batch)
+        dof_per_graph = (
+            n_atoms_per_graph * state.positions.shape[-1]
         )  # n_atoms * n_dimensions
 
         # For now, sum the per-batch DOF as chain expects a single int
         # This is a limitation that should be addressed in the chain implementation
-        total_dof = int(dof_per_batch.sum().item())
+        total_dof = int(dof_per_graph.sum().item())
 
         # Initialize state
         state = NVTNoseHooverState(
@@ -479,8 +479,8 @@ def nvt_nose_hoover_invariant(
     e_kin = calc_kinetic_energy(state.momenta, state.masses, batch=state.batch)
 
     # Get system degrees of freedom per batch
-    n_atoms_per_batch = torch.bincount(state.batch)
-    dof = n_atoms_per_batch * state.positions.shape[-1]  # n_atoms * n_dimensions
+    n_atoms_per_graph = torch.bincount(state.batch)
+    dof = n_atoms_per_graph * state.positions.shape[-1]  # n_atoms * n_dimensions
 
     # Start with system energy
     e_tot = e_pot + e_kin

torch_sim/io.py

@@ -226,9 +226,9 @@ def atoms_to_state(
     )
 
     # Create batch indices using repeat_interleave
-    atoms_per_batch = torch.tensor([len(a) for a in atoms_list], device=device)
+    atoms_per_graph = torch.tensor([len(a) for a in atoms_list], device=device)
     batch = torch.repeat_interleave(
-        torch.arange(len(atoms_list), device=device), atoms_per_batch
+        torch.arange(len(atoms_list), device=device), atoms_per_graph
     )
 
     # Verify consistent pbc
@@ -298,9 +298,9 @@ def structures_to_state(
     )
 
     # Create batch indices
-    atoms_per_batch = torch.tensor([len(s) for s in struct_list], device=device)
+    atoms_per_graph = torch.tensor([len(s) for s in struct_list], device=device)
     batch = torch.repeat_interleave(
-        torch.arange(len(struct_list), device=device), atoms_per_batch
+        torch.arange(len(struct_list), device=device), atoms_per_graph
     )
 
     return ts.SimState(
@@ -369,9 +369,9 @@ def phonopy_to_state(
     )
 
     # Create batch indices using repeat_interleave
-    atoms_per_batch = torch.tensor([len(a) for a in phonopy_atoms_list], device=device)
+    atoms_per_graph = torch.tensor([len(a) for a in phonopy_atoms_list], device=device)
     batch = torch.repeat_interleave(
-        torch.arange(len(phonopy_atoms_list), device=device), atoms_per_batch
+        torch.arange(len(phonopy_atoms_list), device=device), atoms_per_graph
     )
 
     """

torch_sim/quantities.py

@@ -60,13 +60,13 @@ def calc_kT(  # noqa: N802
 
     # Count degrees of freedom per batch
     batch_sizes = torch.bincount(batch)
-    dof_per_batch = batch_sizes * squared_term.shape[-1]  # multiply by n_dimensions
+    dof_per_graph = batch_sizes * squared_term.shape[-1]  # multiply by n_dimensions
 
     # Calculate temperature per batch
     batch_sums = torch.segment_reduce(
         flattened_squared, reduce="sum", lengths=batch_sizes
     )
-    return batch_sums / dof_per_batch
+    return batch_sums / dof_per_graph
 
 
 def calc_temperature(