Rename and reorder some variables

.github/PULL_REQUEST_TEMPLATE.md

@@ -11,4 +11,4 @@ Before a pull request can be merged, the following items must be checked:
 * [ ] Run [ruff](https://beta.ruff.rs/docs/rules/#pydocstyle-d) on your code.
 * [ ] Tests have been added for any new functionality or bug fixes.
 
-We highly recommended installing the `prek` hooks running in CI locally to speedup the development process. Simply run `pip install prek && prek install` to install the hooks which will check your code before each commit.
+<!-- We highly recommended installing the `prek` hooks running in CI locally to speedup the development process. Simply run `pip install prek && prek install` to install the hooks which will check your code before each commit. -->

CONTRIBUTING.md

@@ -2,6 +2,10 @@
 
 TorchSim welcomes contributions and feedback from the community.
 
+## Developer Installation
+
+To install TorchSim, follow the instructions in the [Developer Installation](https://torchsim.github.io/torch-sim/dev/dev_install.html) guide.
+
 ## Code Reviews
 
 All submissions require review by project maintainers before merging:

examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py

@@ -65,7 +65,9 @@
 
 # Create the cell tensor
 cell = torch.tensor(
-    [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], device=device, dtype=dtype
+    [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]],
+    device=device,
+    dtype=dtype,
 )
 
 # Create the atomic numbers tensor (Argon = 18)
@@ -106,7 +108,16 @@
     torch.tensor(10_000, device=device, dtype=dtype) * Units.pressure
 )  # Target pressure (10 kbar)
 
-state = ts.npt_langevin_init(state=state, model=model, dt=dt, kT=kT, seed=1)
+state = ts.npt_langevin_init(
+    state=state,
+    model=model,
+    dt=dt,
+    kT=kT,
+    seed=1,
+    alpha=1.0 / (100 * dt),
+    cell_alpha=1.0 / (100 * dt),
+    b_tau=1 / (1000 * dt),
+)
 
 # Run the simulation
 for step in range(N_steps):
@@ -137,9 +148,6 @@
         dt=dt,
         kT=kT,
         external_pressure=target_pressure,
-        alpha=1.0 / (100 * dt),
-        cell_alpha=1.0 / (100 * dt),
-        b_tau=1 / (1000 * dt),
     )
 
 temp = (

examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py

@@ -71,7 +71,16 @@
         print(f"{step=}: Temperature: {temp.item():.4f}: {invariant=:.4f}, ")
     state = ts.nvt_nose_hoover_step(model=model, state=state, dt=dt, kT=kT)
 
-state = ts.npt_langevin_init(state=state, model=model, kT=kT, dt=dt, seed=1)
+state = ts.npt_langevin_init(
+    state=state,
+    model=model,
+    kT=kT,
+    dt=dt,
+    seed=1,
+    alpha=1.0 / (100 * dt),
+    cell_alpha=1.0 / (100 * dt),
+    b_tau=1 / (1000 * dt),
+)
 
 for step in range(N_steps_npt):
     if step % 10 == 0:
@@ -107,9 +116,6 @@
         dt=dt,
         kT=kT,
         external_pressure=target_pressure,
-        alpha=1.0 / (100 * dt),
-        cell_alpha=1.0 / (100 * dt),
-        b_tau=1 / (1000 * dt),
     )
 
 final_temp = (

tests/test_autobatching.py

@@ -550,13 +550,13 @@ def test_in_flight_max_iterations(
     # Create states that won't naturally converge
     states = [si_sim_state.clone(), fe_supercell_sim_state.clone()]
 
-    # Set max_attempts to a small value to ensure quick termination
-    max_attempts = 3
+    # Set max_iterations to a small value to ensure quick termination
+    max_iterations = 3
     batcher = InFlightAutoBatcher(
         model=lj_model,
         memory_scales_with="n_atoms",
         max_memory_scaler=800.0,
-        max_iterations=max_attempts,
+        max_iterations=max_iterations,
     )
     batcher.load_states(states)
 
@@ -580,15 +580,15 @@ def test_in_flight_max_iterations(
         if state is not None:
             convergence_tensor = torch.zeros(state.n_systems, dtype=torch.bool)
 
-        if iteration_count > max_attempts + 4:
+        if iteration_count > max_iterations + 4:
             raise ValueError("Should have terminated by now")
 
     # Verify all states were processed
     assert len(all_completed_states) == len(states)
 
-    # Verify we didn't exceed max_attempts + 1 iterations (first call doesn't count)
+    # Verify we didn't exceed max_iterations + 1 iterations (first call doesn't count)
     assert iteration_count == 3
 
-    # Verify swap_attempts tracking
+    # Verify iteration_count tracking
     for idx in range(len(states)):
-        assert batcher.swap_attempts[idx] == max_attempts
+        assert batcher.iteration_count[idx] == max_iterations

tests/test_integrators.py

@@ -90,7 +90,14 @@ def test_npt_langevin(
 
     # Initialize integrator using new direct API
     state = ts.npt_langevin_init(
-        state=ar_double_sim_state, model=lj_model, dt=dt, kT=kT, alpha=alpha, seed=42
+        state=ar_double_sim_state,
+        model=lj_model,
+        dt=dt,
+        kT=kT,
+        alpha=alpha,
+        cell_alpha=cell_alpha,
+        b_tau=b_tau,
+        seed=42,
     )
 
     # Run dynamics for several steps
@@ -103,9 +110,6 @@ def test_npt_langevin(
             dt=dt,
             kT=kT,
             external_pressure=external_pressure,
-            alpha=alpha,
-            cell_alpha=cell_alpha,
-            b_tau=b_tau,
         )
 
         # Calculate instantaneous temperature from kinetic energy
@@ -161,7 +165,14 @@ def test_npt_langevin_multi_kt(
 
     # Initialize integrator using new direct API
     state = ts.npt_langevin_init(
-        state=ar_double_sim_state, model=lj_model, dt=dt, kT=kT, alpha=alpha, seed=42
+        state=ar_double_sim_state,
+        model=lj_model,
+        dt=dt,
+        kT=kT,
+        alpha=alpha,
+        cell_alpha=cell_alpha,
+        b_tau=b_tau,
+        seed=42,
     )
 
     # Run dynamics for several steps
@@ -174,9 +185,6 @@ def test_npt_langevin_multi_kt(
             dt=dt,
             kT=kT,
             external_pressure=external_pressure,
-            alpha=alpha,
-            cell_alpha=cell_alpha,
-            b_tau=b_tau,
         )
 
         # Calculate instantaneous temperature from kinetic energy
@@ -613,7 +621,11 @@ def test_npt_nose_hoover_multi_equivalent_to_single(
     )
     for _step in range(n_steps):
         state = ts.npt_nose_hoover_step(
-            state=state, model=lj_model, dt=dt, kT=kT, external_pressure=external_pressure
+            state=state,
+            model=lj_model,
+            dt=dt,
+            kT=kT,
+            external_pressure=external_pressure,
         )
 
         # Calculate instantaneous temperature from kinetic energy
@@ -787,6 +799,9 @@ def test_compute_cell_force_atoms_per_system():
         cell_positions=torch.ones((2, 3, 3)),
         cell_velocities=torch.zeros((2, 3, 3)),
         cell_masses=torch.ones(2),
+        alpha=torch.ones(2),
+        cell_alpha=torch.ones(2),
+        b_tau=torch.ones(2),
     )
 
     # Get forces and compare ratio

torch_sim/autobatching.py

@@ -32,7 +32,9 @@
 from torch_sim.typing import MemoryScaling
 
 
-def to_constant_volume_bins[T: dict[int, float] | list[float] | list[tuple[T, ...]]](  # noqa: C901, PLR0915
+def to_constant_volume_bins[  # noqa: C901, PLR0915
+    T: dict[int, float] | list[float] | list[tuple[T, ...]]
+](
     items: T,
     max_volume: float,
     *,
@@ -704,7 +706,7 @@ class InFlightAutoBatcher[T: SimState]:
         completed_idx (list[int]): Indices of states that have been processed.
         completed_idx_og_order (list[int]): Original indices of completed states.
         current_scalers (list[float]): Memory metrics for states in current batch.
-        swap_attempts (dict[int, int]): Count of iterations for each state.
+        iteration_count (dict[int, int]): Number of iterations for each state.
 
     Example::
 
@@ -772,7 +774,7 @@ def __init__(
         self.max_memory_scaler = max_memory_scaler or None
         self.max_atoms_to_try = max_atoms_to_try
         self.memory_scaling_factor = memory_scaling_factor
-        self.max_attempts = max_iterations  # TODO: change to max_iterations
+        self.max_iterations = max_iterations
         self.max_memory_padding = max_memory_padding
 
     def load_states(self, states: Sequence[T] | Iterator[T] | T) -> None:
@@ -818,7 +820,7 @@ def load_states(self, states: Sequence[T] | Iterator[T] | T) -> None:
         self.current_scalers = []
         self.current_idx = []
         self.iterator_idx = 0
-        self.swap_attempts = []  # Track attempts for each state
+        self.iteration_count = []  # Track attempts for each state
 
         self.completed_idx_og_order = []
 
@@ -857,7 +859,7 @@ def _get_next_states(self) -> list[T]:
             new_idx.append(self.iterator_idx)
             new_states.append(state)
             # Initialize attempt counter for new state
-            self.swap_attempts.append(0)
+            self.iteration_count.append(0)
             self.iterator_idx += 1
 
         self.current_scalers.extend(new_metrics)
@@ -898,7 +900,7 @@ def _get_first_batch(self) -> T:
         first_metric = calculate_memory_scaler(first_state, self.memory_scales_with)
         self.current_scalers += [first_metric]
         self.current_idx += [0]
-        self.swap_attempts.append(0)  # Initialize attempt counter for first state
+        self.iteration_count.append(0)  # Initialize attempt counter for first state
         self.iterator_idx += 1
 
         # if max_metric is not set, estimate it
@@ -998,9 +1000,9 @@ def next_batch(  # noqa: C901
 
         # Increment attempt counters and check for max attempts in a single loop
         for cur_idx, abs_idx in enumerate(self.current_idx):
-            self.swap_attempts[abs_idx] += 1
-            if self.max_attempts is not None and (
-                self.swap_attempts[abs_idx] >= self.max_attempts
+            self.iteration_count[abs_idx] += 1
+            if self.max_iterations is not None and (
+                self.iteration_count[abs_idx] >= self.max_iterations
             ):
                 # Force convergence for states that have reached max attempts
                 convergence_tensor[cur_idx] = torch.tensor(True)  # noqa: FBT003