Rename batch to system (#217)

Author: curtischong

.gitignore

@@ -34,3 +34,6 @@ coverage.xml
 
 # env
 uv.lock
+
+# IDE
+.vscode/

examples/scripts/1_Introduction/1.2_MACE.py

@@ -63,38 +63,38 @@
 cell = torch.tensor(cell_numpy, device=device, dtype=dtype)
 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(
+# create system idx array of shape (16,) which is 0 for first 8 atoms, 1 for last 8 atoms
+atoms_per_system = 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
+system_idx = torch.repeat_interleave(
+    torch.arange(len(atoms_per_system), device=device), atoms_per_system
 )
 
 # You can see their shapes are as expected
 print(f"Positions: {positions.shape}")
 print(f"Cell: {cell.shape}")
 print(f"Atomic numbers: {atomic_numbers.shape}")
-print(f"Batch: {batch.shape}")
+print(f"System indices: {system_idx.shape}")
 
 # Now we can pass them to the model
 results = batched_model(
     dict(
         positions=positions,
         cell=cell,
         atomic_numbers=atomic_numbers,
-        batch=batch,
+        system_idx=system_idx,
         pbc=True,
     )
 )
 
-# The energy has shape (n_batches,) as the structures in a batch
+# The energy has shape (n_systems,) as the structures in a batch
 print(f"Energy: {results['energy'].shape}")
 
 # The forces have shape (n_atoms, 3) same as positions
 print(f"Forces: {results['forces'].shape}")
 
-# The stress has shape (n_batches, 3, 3) same as cell
+# The stress has shape (n_systems, 3, 3) same as cell
 print(f"Stress: {results['stress'].shape}")
 
 # Check if the energy, forces, and stress are the same for the Si system across the batch

examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py

@@ -93,20 +93,20 @@
 masses_numpy = np.concatenate([atoms.get_masses() for atoms in atoms_list])
 masses = torch.tensor(masses_numpy, device=device, dtype=dtype)
 
-# Create batch indices tensor for scatter operations
-atoms_per_batch = torch.tensor(
+# Create system indices tensor for scatter operations
+atoms_per_system = 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
+system_indices = torch.repeat_interleave(
+    torch.arange(len(atoms_per_system), device=device), atoms_per_system
 )
 """
 
 state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype)
 
 print(f"Positions shape: {state.positions.shape}")
 print(f"Cell shape: {state.cell.shape}")
-print(f"Batch indices shape: {state.batch.shape}")
+print(f"System indices shape: {state.system_idx.shape}")
 
 # Run initial inference
 results = batched_model(state)

examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py

@@ -85,7 +85,7 @@
 # Initialize unit cell gradient descent optimizer
 gd_init, gd_update = unit_cell_gradient_descent(
     model=model,
-    cell_factor=None,  # Will default to atoms per batch
+    cell_factor=None,  # Will default to atoms per system
     hydrostatic_strain=False,
     constant_volume=False,
     scalar_pressure=0.0,

examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py

@@ -81,7 +81,7 @@
 # Initialize unit cell gradient descent optimizer
 fire_init, fire_update = unit_cell_fire(
     model=model,
-    cell_factor=None,  # Will default to atoms per batch
+    cell_factor=None,  # Will default to atoms per system
     hydrostatic_strain=False,
     constant_volume=False,
     scalar_pressure=0.0,

examples/scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py

@@ -80,7 +80,7 @@
 # Initialize unit cell gradient descent optimizer
 fire_init, fire_update = ts.optimizers.frechet_cell_fire(
     model=model,
-    cell_factor=None,  # Will default to atoms per batch
+    cell_factor=None,  # Will default to atoms per system
     hydrostatic_strain=False,
     constant_volume=False,
     scalar_pressure=0.0,

examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py

@@ -86,7 +86,7 @@ class HybridSwapMCState(MDState):
 hybrid_state = HybridSwapMCState(
     **vars(md_state),
     last_permutation=torch.zeros(
-        md_state.n_batches, device=md_state.device, dtype=torch.bool
+        md_state.n_systems, device=md_state.device, dtype=torch.bool
     ),
 )
 

examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py

@@ -119,13 +119,15 @@
 for step in range(N_steps):
     if step % 50 == 0:
         temp = (
-            calc_kT(masses=state.masses, momenta=state.momenta, batch=state.batch)
+            calc_kT(
+                masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
+            )
             / Units.temperature
         )
         pressure = get_pressure(
             model(state)["stress"],
             calc_kinetic_energy(
-                masses=state.masses, momenta=state.momenta, batch=state.batch
+                masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
             ),
             torch.linalg.det(state.cell),
         )
@@ -139,15 +141,15 @@
     state = npt_update(state, kT=kT, external_pressure=target_pressure)
 
 temp = (
-    calc_kT(masses=state.masses, momenta=state.momenta, batch=state.batch)
+    calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx)
     / Units.temperature
 )
 print(f"Final temperature: {temp.item():.4f}")
 
 
 stress = model(state)["stress"]
 calc_kinetic_energy = calc_kinetic_energy(
-    masses=state.masses, momenta=state.momenta, batch=state.batch
+    masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
 )
 volume = torch.linalg.det(state.cell)
 pressure = get_pressure(stress, calc_kinetic_energy, volume)

examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py

@@ -68,7 +68,9 @@
 for step in range(N_steps_nvt):
     if step % 10 == 0:
         temp = (
-            calc_kT(masses=state.masses, momenta=state.momenta, batch=state.batch)
+            calc_kT(
+                masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
+            )
             / Units.temperature
         )
         invariant = float(nvt_nose_hoover_invariant(state, kT=kT))
@@ -83,7 +85,9 @@
 for step in range(N_steps_npt):
     if step % 10 == 0:
         temp = (
-            calc_kT(masses=state.masses, momenta=state.momenta, batch=state.batch)
+            calc_kT(
+                masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
+            )
             / Units.temperature
         )
         stress = model(state)["stress"]
@@ -92,7 +96,9 @@
             get_pressure(
                 stress,
                 calc_kinetic_energy(
-                    masses=state.masses, momenta=state.momenta, batch=state.batch
+                    masses=state.masses,
+                    momenta=state.momenta,
+                    system_idx=state.system_idx,
                 ),
                 volume,
             ).item()
@@ -107,7 +113,7 @@
     state = npt_update(state, kT=kT, external_pressure=target_pressure)
 
 final_temp = (
-    calc_kT(masses=state.masses, momenta=state.momenta, batch=state.batch)
+    calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx)
     / Units.temperature
 )
 print(f"Final temperature: {final_temp.item():.4f} K")
@@ -117,7 +123,7 @@
     get_pressure(
         final_stress,
         calc_kinetic_energy(
-            masses=state.masses, momenta=state.momenta, batch=state.batch
+            masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
         ),
         final_volume,
     ).item()

examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py

@@ -77,7 +77,7 @@
 start_time = time.perf_counter()
 for step in range(N_steps):
     total_energy = state.energy + calc_kinetic_energy(
-        masses=state.masses, momenta=state.momenta, batch=state.batch
+        masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
     )
     if step % 10 == 0:
         print(f"Step {step}: Total energy: {total_energy.item():.4f} eV")