replace assert statements with ValueError in multiple files

Author: janosh

examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py

@@ -484,7 +484,8 @@ def run_optimization_ase(  # noqa: C901, PLR0915
     final_state_opt: SimState | GDState | None = None
 
     if optimizer_type_val == "torch_sim":
-        assert ts_md_flavor_val is not None, "ts_md_flavor must be provided for torch_sim"
+        if ts_md_flavor_val is None:
+            raise ValueError(f"{ts_md_flavor_val=} must be provided for torch_sim")
         steps, final_state_opt = run_optimization_ts(
             initial_state=state.clone(),
             ts_md_flavor=ts_md_flavor_val,

pyproject.toml

@@ -102,14 +102,11 @@ ignore = [
     "ISC001",  # avoid conflicts with the formatter
     "N803",    # Variable name should be lowercase
     "N806",    # Uppercase letters in variable names
-    "PD010",   # .pivot_table is preferred to .pivot or .unstack; provides same functionality
-    "PD015",   # pandas-use-of-pd-merge
     "PLR0912", # too many branches
     "PLR0913", # too many function arguments
     "PLR2004", # Magic value used in comparison, consider replacing {value} with a constant variable
     "PLW2901", # Outer for loop variable overwritten by inner assignment target
     "PTH",     # flake8-use-pathlib
-    "S101",    # Use of assertion statements
     "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

torch_sim/autobatching.py

@@ -955,10 +955,8 @@ def _get_first_batch(self) -> SimState:
             self.max_memory_scaler = self.max_memory_scaler * self.max_memory_padding
         return concatenate_states([first_state, *states])
 
-    def next_batch(
-        self,
-        updated_state: SimState | None,
-        convergence_tensor: torch.Tensor | None,
+    def next_batch(  # noqa: C901
+        self, updated_state: SimState | None, convergence_tensor: torch.Tensor | None
     ) -> (
         tuple[SimState | None, list[SimState]]
         | tuple[SimState | None, list[SimState], list[int]]
@@ -1022,10 +1020,14 @@ def next_batch(
 
         # assert statements helpful for debugging, should be moved to validate fn
         # the first two are most important
-        assert len(convergence_tensor) == updated_state.n_batches
-        assert len(self.current_idx) == len(self.current_scalers)
-        assert len(convergence_tensor.shape) == 1
-        assert updated_state.n_batches > 0
+        if len(convergence_tensor) != updated_state.n_batches:
+            raise ValueError(f"{len(convergence_tensor)=} != {updated_state.n_batches=}")
+        if len(self.current_idx) != len(self.current_scalers):
+            raise ValueError(f"{len(self.current_idx)=} != {len(self.current_scalers)=}")
+        if len(convergence_tensor.shape) != 1:
+            raise ValueError(f"{len(convergence_tensor.shape)=} != 1")
+        if updated_state.n_batches <= 0:
+            raise ValueError(f"{updated_state.n_batches=} <= 0")
 
         # Increment attempt counters and check for max attempts in a single loop
         for cur_idx, abs_idx in enumerate(self.current_idx):

torch_sim/models/fairchem.py

@@ -175,7 +175,8 @@ def __init__(  # noqa: C901, PLR0915
             )
 
         # Either the config path or the checkpoint path needs to be provided
-        assert config_yml or model is not None
+        if not config_yml and model is None:
+            raise ValueError("Either config_yml or model must be provided")
 
         checkpoint = None
         if config_yml is not None:

torch_sim/models/interface.py

@@ -199,10 +199,8 @@ def forward(self, state: SimState | StateDict, **kwargs) -> dict[str, torch.Tens
         """
 
 
-def validate_model_outputs(
-    model: ModelInterface,
-    device: torch.device,
-    dtype: torch.dtype,
+def validate_model_outputs(  # noqa: C901, PLR0915
+    model: ModelInterface, device: torch.device, dtype: torch.dtype
 ) -> None:
     """Validate the outputs of a model implementation against the interface requirements.
 
@@ -233,10 +231,9 @@ def validate_model_outputs(
     """
     from ase.build import bulk
 
-    assert model.dtype is not None
-    assert model.device is not None
-    assert model.compute_stress is not None
-    assert model.compute_forces is not None
+    for attr in ("dtype", "device", "compute_stress", "compute_forces"):
+        if not hasattr(model, attr):
+            raise ValueError(f"model.{attr} is not set")
 
     try:
         if not model.compute_stress:
@@ -265,52 +262,56 @@ def validate_model_outputs(
     model_output = model.forward(sim_state)
 
     # assert model did not mutate the input
-    assert torch.allclose(og_positions, sim_state.positions)
-    assert torch.allclose(og_cell, sim_state.cell)
-    assert torch.allclose(og_batch, sim_state.batch)
-    assert torch.allclose(og_atomic_numbers, sim_state.atomic_numbers)
+    if not torch.allclose(og_positions, sim_state.positions):
+        raise ValueError(f"{og_positions=} != {sim_state.positions=}")
+    if not torch.allclose(og_cell, sim_state.cell):
+        raise ValueError(f"{og_cell=} != {sim_state.cell=}")
+    if not torch.allclose(og_batch, sim_state.batch):
+        raise ValueError(f"{og_batch=} != {sim_state.batch=}")
+    if not torch.allclose(og_atomic_numbers, sim_state.atomic_numbers):
+        raise ValueError(f"{og_atomic_numbers=} != {sim_state.atomic_numbers=}")
 
     # assert model output has the correct keys
-    assert "energy" in model_output
-    assert "forces" in model_output if force_computed else True
-    assert "stress" in model_output if stress_computed else True
+    if "energy" not in model_output:
+        raise ValueError("energy not in model output")
+    if force_computed and "forces" not in model_output:
+        raise ValueError("forces not in model output")
+    if stress_computed and "stress" not in model_output:
+        raise ValueError("stress not in model output")
 
     # assert model output shapes are correct
-    assert model_output["energy"].shape == (2,)
-    assert model_output["forces"].shape == (20, 3) if force_computed else True
-    assert model_output["stress"].shape == (2, 3, 3) if stress_computed else True
+    if model_output["energy"].shape != (2,):
+        raise ValueError(f"{model_output['energy'].shape=} != (2,)")
+    if force_computed and model_output["forces"].shape != (20, 3):
+        raise ValueError(f"{model_output['forces'].shape=} != (20, 3)")
+    if stress_computed and model_output["stress"].shape != (2, 3, 3):
+        raise ValueError(f"{model_output['stress'].shape=} != (2, 3, 3)")
 
     si_state = ts.io.atoms_to_state([si_atoms], device, dtype)
     fe_state = ts.io.atoms_to_state([fe_atoms], device, dtype)
 
     si_model_output = model.forward(si_state)
-    assert torch.allclose(
+    if not torch.allclose(
         si_model_output["energy"], model_output["energy"][0], atol=10e-3
-    )
-    assert torch.allclose(
-        si_model_output["forces"],
-        model_output["forces"][: si_state.n_atoms],
+    ):
+        raise ValueError(f"{si_model_output['energy']=} != {model_output['energy'][0]=}")
+    if not torch.allclose(
+        forces := si_model_output["forces"],
+        expected_forces := model_output["forces"][: si_state.n_atoms],
         atol=10e-3,
-    )
-    # assert torch.allclose(
-    #     si_model_output["stress"],
-    #     model_output["stress"][0],
-    #     atol=10e-3,
-    # )
+    ):
+        raise ValueError(f"{forces=} != {expected_forces=}")
 
     fe_model_output = model.forward(fe_state)
     si_model_output = model.forward(si_state)
 
-    assert torch.allclose(
+    if not torch.allclose(
         fe_model_output["energy"], model_output["energy"][1], atol=10e-2
-    )
-    assert torch.allclose(
-        fe_model_output["forces"],
-        model_output["forces"][si_state.n_atoms :],
+    ):
+        raise ValueError(f"{fe_model_output['energy']=} != {model_output['energy'][1]=}")
+    if not torch.allclose(
+        forces := fe_model_output["forces"],
+        expected_forces := model_output["forces"][si_state.n_atoms :],
         atol=10e-2,
-    )
-    # assert torch.allclose(
-    #     arr_model_output["stress"],
-    #     model_output["stress"][1],
-    #     atol=10e-3,
-    # )
+    ):
+        raise ValueError(f"{forces=} != {expected_forces=}")

torch_sim/models/soft_sphere.py

@@ -642,9 +642,10 @@ def __init__(
         )
 
         # Ensure parameter matrices are symmetric (required for energy conservation)
-        assert torch.allclose(self.sigma_matrix, self.sigma_matrix.T)
-        assert torch.allclose(self.epsilon_matrix, self.epsilon_matrix.T)
-        assert torch.allclose(self.alpha_matrix, self.alpha_matrix.T)
+        for matrix_name in ("sigma_matrix", "epsilon_matrix", "alpha_matrix"):
+            matrix = getattr(self, matrix_name)
+            if not torch.allclose(matrix, matrix.T):
+                raise ValueError(f"{matrix_name} is not symmetric")
 
         # Set interaction cutoff distance
         self.cutoff = torch.tensor(

torch_sim/neighbors.py

@@ -109,7 +109,8 @@ def primitive_neighbor_list(  # noqa: C901, PLR0915
             1 / l3 if l3 > 0 else pytorch_scalar_1,
         ]
     )
-    assert face_dist_c.shape == (3,)
+    if face_dist_c.shape != (3,):
+        raise ValueError(f"{face_dist_c.shape=} != (3,)")
 
     # we don't handle other fancier cutoffs
     max_cutoff: torch.Tensor = cutoff
@@ -214,8 +215,10 @@ def primitive_neighbor_list(  # noqa: C901, PLR0915
         bin_index_i = bin_index_i[mask]
 
     # Make sure that all atoms have been sorted into bins.
-    assert len(atom_i) == 0
-    assert len(bin_index_i) == 0
+    if len(atom_i) != 0:
+        raise ValueError(f"{len(atom_i)=} != 0")
+    if len(bin_index_i) != 0:
+        raise ValueError(f"{len(bin_index_i)=} != 0")
 
     # Now we construct neighbor pairs by pairing up all atoms within a bin or
     # between bin and neighboring bin. atom_pairs_pn is a helper buffer that

torch_sim/optimizers.py

@@ -1273,7 +1273,8 @@ def _vv_fire_step(  # noqa: C901, PLR0915
     if is_cell_optimization:
         cell_factor_reshaped = state.cell_factor.view(n_batches, 1, 1)
         if is_frechet:
-            assert isinstance(state, FrechetCellFIREState)
+            if not isinstance(state, expected_cls := FrechetCellFIREState):
+                raise ValueError(f"{type(state)=} must be a {expected_cls.__name__}")
             cur_deform_grad = state.deform_grad()
             deform_grad_log = torch.zeros_like(cur_deform_grad)
             for b in range(n_batches):
@@ -1291,7 +1292,8 @@ def _vv_fire_step(  # noqa: C901, PLR0915
             state.row_vector_cell = new_row_vector_cell
             state.cell_positions = cell_positions_log_scaled_new
         else:
-            assert isinstance(state, UnitCellFireState)
+            if not isinstance(state, expected_cls := UnitCellFireState):
+                raise ValueError(f"{type(state)=} must be a {expected_cls.__name__}")
             cur_deform_grad = state.deform_grad()
             cell_factor_expanded = state.cell_factor.expand(n_batches, 3, 1)
             current_cell_positions_scaled = (
@@ -1329,7 +1331,8 @@ def _vv_fire_step(  # noqa: C901, PLR0915
             ).unsqueeze(0).expand(n_batches, -1, -1)
 
         if is_frechet:
-            assert isinstance(state, FrechetCellFIREState)
+            if not isinstance(state, expected_cls := FrechetCellFIREState):
+                raise ValueError(f"{type(state)=} must be a {expected_cls.__name__}")
             ucf_cell_grad = torch.bmm(
                 virial, torch.linalg.inv(torch.transpose(deform_grad_new, 1, 2))
             )
@@ -1353,7 +1356,8 @@ def _vv_fire_step(  # noqa: C901, PLR0915
                 new_cell_forces[b] = forces_flat.reshape(3, 3)
             state.cell_forces = new_cell_forces / cell_factor_reshaped
         else:
-            assert isinstance(state, UnitCellFireState)
+            if not isinstance(state, expected_cls := UnitCellFireState):
+                raise ValueError(f"{type(state)=} must be a {expected_cls.__name__}")
             state.cell_forces = virial / cell_factor_reshaped
 
     state.velocities += 0.5 * atom_wise_dt * state.forces / state.masses.unsqueeze(-1)
@@ -1564,15 +1568,17 @@ def _ase_fire_step(  # noqa: C901, PLR0915
         )
 
         if is_frechet:
-            assert isinstance(state, FrechetCellFIREState)
+            if not isinstance(state, expected_cls := FrechetCellFIREState):
+                raise ValueError(f"{type(state)=} must be a {expected_cls.__name__}")
             new_logm_F_scaled = state.cell_positions + dr_cell
             state.cell_positions = new_logm_F_scaled
             logm_F_new = new_logm_F_scaled / (state.cell_factor + eps)
             F_new = torch.matrix_exp(logm_F_new)
             new_row_vector_cell = torch.bmm(state.reference_row_vector_cell, F_new.mT)
             state.row_vector_cell = new_row_vector_cell
         else:
-            assert isinstance(state, UnitCellFireState)
+            if not isinstance(state, expected_cls := UnitCellFireState):
+                raise ValueError(f"{type(state)=} must be a {expected_cls.__name__}")
             F_current = state.deform_grad()
             cell_factor_exp_mult = state.cell_factor.expand(n_batches, 3, 1)
             current_F_scaled = F_current * cell_factor_exp_mult
@@ -1619,13 +1625,16 @@ def _ase_fire_step(  # noqa: C901, PLR0915
             ).unsqueeze(0).expand(n_batches, -1, -1)
 
         if is_frechet:
-            assert isinstance(state, FrechetCellFIREState)
-            assert F_new is not None, (
-                "F_new should be defined for Frechet cell force calculation"
-            )
-            assert logm_F_new is not None, (
-                "logm_F_new should be defined for Frechet cell force calculation"
-            )
+            if not isinstance(state, expected_cls := FrechetCellFIREState):
+                raise ValueError(f"{type(state)=} must be a {expected_cls.__name__}")
+            if F_new is None:
+                raise ValueError(
+                    "F_new should be defined for Frechet cell force calculation"
+                )
+            if logm_F_new is None:
+                raise ValueError(
+                    "logm_F_new should be defined for Frechet cell force calculation"
+                )
             ucf_cell_grad = torch.bmm(
                 virial, torch.linalg.inv(torch.transpose(F_new, 1, 2))
             )
@@ -1649,7 +1658,8 @@ def _ase_fire_step(  # noqa: C901, PLR0915
                 new_cell_forces_log_space[b_idx] = forces_flat.reshape(3, 3)
             state.cell_forces = new_cell_forces_log_space / (state.cell_factor + eps)
         else:
-            assert isinstance(state, UnitCellFireState)
+            if not isinstance(state, expected_cls := UnitCellFireState):
+                raise ValueError(f"{type(state)=} must be a {expected_cls.__name__}")
             state.cell_forces = virial / state.cell_factor
 
     return state

torch_sim/transforms.py

@@ -521,8 +521,10 @@ def compute_distances_with_cell_shifts(
         torch.Tensor: A tensor of shape (n_pairs,) containing the
             computed distances for each pair.
     """
-    assert mapping.dim() == 2
-    assert mapping.shape[0] == 2
+    if mapping.dim() != 2:
+        raise ValueError(f"Mapping must be a 2D tensor, got {mapping.shape}")
+    if mapping.shape[0] != 2:
+        raise ValueError(f"Mapping must have 2 rows, got {mapping.shape[0]}")
 
     if cell_shifts is None:
         dr = pos[mapping[1]] - pos[mapping[0]]