remove debug print statements from tests and replace them with assertions, adjust ruff config to flag prints in package code

Author: janosh

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.12
+    rev: v0.11.13
     hooks:
       - id: ruff
         args: [--fix]

pyproject.toml

@@ -114,10 +114,8 @@ ignore = [
     "S301",    # pickle and modules that wrap it can be unsafe, possible security issue
     "S311",    # Standard pseudo-random generators are not suitable for cryptographic purposes
     "SIM105",  # Use contextlib.suppress instead of try-except-pass
-    "T201",    # print found
     "TD",      # flake8-todos
     "TRY003",  # Avoid specifying long messages outside the exception class
-    "TRY301",  # Abstract raise to an inner function
 ]
 pydocstyle.convention = "google"
 isort.split-on-trailing-comma = false
@@ -126,7 +124,7 @@ pep8-naming.ignore-names = ["get_kT", "kT"]
 
 [tool.ruff.lint.per-file-ignores]
 "**/tests/*" = ["ANN201", "D", "S101"]
-"examples/**/*" = ["B018"]
+"examples/**/*" = ["B018", "T201"]
 "examples/tutorials/**/*" = ["ALL"]
 
 [tool.ruff.format]

tests/models/test_morse.py

@@ -28,7 +28,6 @@ def test_morse_pair_asymptotic() -> None:
     dr = torch.tensor([[1.0]])  # Large distance
     epsilon = 5.0
     energy = morse_pair(dr, epsilon=epsilon)
-    print(energy, -epsilon * torch.ones_like(energy))
     torch.testing.assert_close(
         energy, -epsilon * torch.ones_like(energy), rtol=1e-2, atol=1e-5
     )
@@ -55,8 +54,6 @@ def test_morse_force_energy_consistency() -> None:
     force_from_grad = -torch.autograd.grad(energy.sum(), dr, create_graph=True)[0]
 
     # Compare forces
-    print(force_direct)
-    print(force_from_grad)
     assert torch.allclose(force_direct, force_from_grad, rtol=1e-4, atol=1e-4)
 
 

tests/test_autobatching.py

@@ -457,10 +457,7 @@ def convergence_fn(state: ts.SimState) -> bool:
 
     all_completed_states, convergence_tensor = [], None
     while True:
-        print(f"Starting new batch of {state.n_batches} states.")
-
         state, completed_states = batcher.next_batch(state, convergence_tensor)
-        print("Number of completed states", len(completed_states))
 
         all_completed_states.extend(completed_states)
         if state is None:

tests/test_optimizers.py

@@ -152,8 +152,9 @@ def test_fire_optimization(
         energies.append(state.energy.item())
         steps_taken += 1
 
-    if steps_taken == max_steps:
-        print(f"FIRE optimization for {md_flavor=} did not converge in {max_steps} steps")
+    assert steps_taken < max_steps, (
+        f"FIRE optimization for {md_flavor=} did not converge in {max_steps=}"
+    )
 
     energies = energies[1:]
 
@@ -327,7 +328,6 @@ def test_unit_cell_fire_optimization(
     ar_supercell_sim_state: ts.SimState, lj_model: torch.nn.Module, md_flavor: MdFlavor
 ) -> None:
     """Test that the Unit Cell FIRE optimizer actually minimizes energy."""
-    print(f"\n--- Starting test_unit_cell_fire_optimization for {md_flavor=} ---")
 
     # Add random displacement to positions and cell
     current_positions = (
@@ -347,48 +347,33 @@ def test_unit_cell_fire_optimization(
         atomic_numbers=ar_supercell_sim_state.atomic_numbers.clone(),
         batch=ar_supercell_sim_state.batch.clone(),
     )
-    print(f"[{md_flavor}] Initial SimState created.")
 
     initial_state_positions = current_sim_state.positions.clone()
     initial_state_cell = current_sim_state.cell.clone()
 
     # Initialize FIRE optimizer
-    print(f"Initializing {md_flavor} optimizer...")
     init_fn, update_fn = unit_cell_fire(
         model=lj_model,
         dt_max=0.3,
         dt_start=0.1,
         md_flavor=md_flavor,
     )
-    print(f"[{md_flavor}] Optimizer functions obtained.")
 
     state = init_fn(current_sim_state)
-    energy = float(getattr(state, "energy", "nan"))
-    print(f"[{md_flavor}] Initial state created by init_fn. {energy=:.4f}")
 
     # Run optimization for a few steps
     energies = [1000.0, state.energy.item()]
     max_steps = 1000
     steps_taken = 0
-    print(f"[{md_flavor}] Entering optimization loop (max_steps: {max_steps})...")
 
     while abs(energies[-2] - energies[-1]) > 1e-6 and steps_taken < max_steps:
         state = update_fn(state)
         energies.append(state.energy.item())
         steps_taken += 1
 
-    print(f"[{md_flavor}] Loop finished after {steps_taken} steps.")
-
-    if steps_taken == max_steps and abs(energies[-2] - energies[-1]) > 1e-6:
-        print(
-            f"WARNING: Unit Cell FIRE {md_flavor=} optimization did not converge "
-            f"in {max_steps} steps. Final energy: {energies[-1]:.4f}"
-        )
-    else:
-        print(
-            f"Unit Cell FIRE {md_flavor=} optimization converged in {steps_taken} "
-            f"steps. Final energy: {energies[-1]:.4f}"
-        )
+    assert steps_taken < max_steps, (
+        f"Unit Cell FIRE {md_flavor=} optimization did not converge in {max_steps=}"
+    )
 
     energies = energies[1:]
 
@@ -522,7 +507,6 @@ def test_frechet_cell_fire_optimization(
 ) -> None:
     """Test that the Frechet Cell FIRE optimizer actually minimizes energy for different
     md_flavors."""
-    print(f"\n--- Starting test_frechet_cell_fire_optimization for {md_flavor=} ---")
 
     # Add random displacement to positions and cell
     # Create a fresh copy for each test run to avoid interference
@@ -543,48 +527,33 @@ def test_frechet_cell_fire_optimization(
         atomic_numbers=ar_supercell_sim_state.atomic_numbers.clone(),
         batch=ar_supercell_sim_state.batch.clone(),
     )
-    print(f"[{md_flavor}] Initial SimState created for Frechet test.")
 
     initial_state_positions = current_sim_state.positions.clone()
     initial_state_cell = current_sim_state.cell.clone()
 
     # Initialize FIRE optimizer
-    print(f"Initializing Frechet {md_flavor} optimizer...")
     init_fn, update_fn = frechet_cell_fire(
         model=lj_model,
         dt_max=0.3,
         dt_start=0.1,
         md_flavor=md_flavor,
     )
-    print(f"[{md_flavor}] Frechet optimizer functions obtained.")
 
     state = init_fn(current_sim_state)
-    energy = float(getattr(state, "energy", "nan"))
-    print(f"[{md_flavor}] Initial state created by Frechet init_fn. {energy=:.4f}")
 
     # Run optimization for a few steps
     energies = [1000.0, state.energy.item()]  # Ensure float for comparison
     max_steps = 1000
     steps_taken = 0
-    print(f"[{md_flavor}] Entering Frechet optimization loop (max_steps: {max_steps})...")
 
     while abs(energies[-2] - energies[-1]) > 1e-6 and steps_taken < max_steps:
         state = update_fn(state)
         energies.append(state.energy.item())
         steps_taken += 1
 
-    print(f"[{md_flavor}] Frechet loop finished after {steps_taken} steps.")
-
-    if steps_taken == max_steps and abs(energies[-2] - energies[-1]) > 1e-6:
-        print(
-            f"WARNING: Frechet Cell FIRE {md_flavor=} optimization did not converge "
-            f"in {max_steps} steps. Final energy: {energies[-1]:.4f}"
-        )
-    else:
-        print(
-            f"Frechet Cell FIRE {md_flavor=} optimization converged in {steps_taken} "
-            f"steps. Final energy: {energies[-1]:.4f}"
-        )
+    assert steps_taken < max_steps, (
+        f"Frechet FIRE {md_flavor=} optimization did not converge in {max_steps=}"
+    )
 
     energies = energies[1:]
 
@@ -600,8 +569,7 @@ def test_frechet_cell_fire_optimization(
     pressure = torch.trace(state.stress.squeeze(0)) / 3.0
 
     # Adjust tolerances if needed, Frechet might behave slightly differently
-    pressure_tol = 0.01
-    force_tol = 0.2
+    pressure_tol, force_tol = 0.01, 0.2
 
     assert torch.abs(pressure) < pressure_tol, (
         f"{md_flavor=} pressure should be below {pressure_tol=} after Frechet "

tests/test_runners.py

@@ -768,12 +768,12 @@ def test_readme_example(lj_model: LennardJonesModel, tmp_path: Path) -> None:
 
     cu_atoms = bulk("Cu", "fcc", a=3.58, cubic=True).repeat((2, 2, 2))
     many_cu_atoms = [cu_atoms] * 5
-    trajectory_files = [tmp_path / f"Cu_traj_{i}" for i in range(len(many_cu_atoms))]
+    trajectory_files = [tmp_path / f"Cu_traj_{i}.h5md" for i in range(len(many_cu_atoms))]
 
     # run them all simultaneously with batching
     final_state = ts.integrate(
         system=many_cu_atoms,
-        model=lj_model,
+        model=lj_model,  # using LJ instead of MACE for testing
         n_steps=50,
         timestep=0.002,
         temperature=1000,
@@ -788,17 +788,17 @@ def test_readme_example(lj_model: LennardJonesModel, tmp_path: Path) -> None:
         with ts.TorchSimTrajectory(filename) as traj:
             final_energies.append(traj.get_array("potential_energy")[-1])
 
-    print(final_energies)
+    assert len(final_energies) == len(trajectory_files)
 
     # relax all of the high temperature states
     relaxed_state = ts.optimize(
         system=final_state,
         model=lj_model,
         optimizer=ts.frechet_cell_fire,
-        # autobatcher=True,
+        # autobatcher=True,  # disabled for CPU-based LJ model in test
     )
 
-    print(relaxed_state.energy)
+    assert relaxed_state.energy.shape == (final_state.n_batches,)
 
 
 @pytest.fixture

torch_sim/math.py

@@ -984,7 +984,7 @@ def matrix_log_33(
             "Falling back to scipy"
         )
         if fallback_warning:
-            print(msg)
+            print(msg)  # noqa: T201
         # Fall back to scipy implementation
         return matrix_log_scipy(matrix).to(sim_dtype)
 

torch_sim/models/fairchem.py

@@ -14,6 +14,8 @@
     pretrained model checkpoints.
 """
 
+# ruff: noqa: T201
+
 from __future__ import annotations
 
 import copy

torch_sim/models/mace.py

@@ -169,7 +169,7 @@ def __init__(
             self.model = self.model.to(dtype=self.dtype)
 
         if enable_cueq:
-            print("Converting models to CuEq for acceleration")
+            print("Converting models to CuEq for acceleration")  # noqa: T201
             self.model = run_e3nn_to_cueq(self.model)
 
         # Set model properties

torch_sim/optimizers.py

@@ -1599,7 +1599,7 @@ def _ase_fire_step(  # noqa: C901, PLR0915
         volumes = torch.linalg.det(state.cell).view(n_batches, 1, 1)
         if torch.any(volumes <= 0):
             bad_indices = torch.where(volumes <= 0)[0].tolist()
-            print(
+            print(  # noqa: T201
                 f"WARNING: Non-positive volume(s) detected during _ase_fire_step: "
                 f"{volumes[bad_indices].tolist()} at {bad_indices=} ({is_frechet=})"
             )