cleanup

Author: curtischong

tests/test_state.py

@@ -109,7 +109,7 @@ def test_concatenate_two_si_states(
     assert concatenated.graph_idx.shape == si_double_sim_state.graph_idx.shape
 
     # Check graph indices
-    expected_graphs = torch.cat(
+    expected_graph_indices = torch.cat(
         [
             torch.zeros(
                 si_sim_state.n_atoms, dtype=torch.int64, device=si_sim_state.device
@@ -119,7 +119,7 @@ def test_concatenate_two_si_states(
             ),
         ]
     )
-    assert torch.all(concatenated.graph_idx == expected_graphs)
+    assert torch.all(concatenated.graph_idx == expected_graph_indices)
 
     # Check that positions match (accounting for graph indices)
     for graph_idx in range(2):
@@ -153,13 +153,13 @@ def test_concatenate_si_and_fe_states(
     # Check graph indices
     si_atoms = si_sim_state.n_atoms
     fe_atoms = fe_supercell_sim_state.n_atoms
-    expected_graphs = torch.cat(
+    expected_graph_indices = torch.cat(
         [
             torch.zeros(si_atoms, dtype=torch.int64, device=si_sim_state.device),
             torch.ones(fe_atoms, dtype=torch.int64, device=fe_supercell_sim_state.device),
         ]
     )
-    assert torch.all(concatenated.graph_idx == expected_graphs)
+    assert torch.all(concatenated.graph_idx == expected_graph_indices)
 
     # check n_atoms_per_graph
     assert torch.all(
@@ -203,15 +203,15 @@ def test_concatenate_double_si_and_fe_states(
     fe_atoms = fe_supercell_sim_state.n_atoms
 
     # The double Si state already has graphs 0 and 1, so Ar should be graph 2
-    expected_graphs = torch.cat(
+    expected_graph_indices = torch.cat(
         [
             si_double_sim_state.graph_idx,
             torch.full(
                 (fe_atoms,), 2, dtype=torch.int64, device=fe_supercell_sim_state.device
             ),
         ]
     )
-    assert torch.all(concatenated.graph_idx == expected_graphs)
+    assert torch.all(concatenated.graph_idx == expected_graph_indices)
     assert torch.unique(concatenated.graph_idx).shape[0] == 3
 
     # Check that we can slice back to the original states
@@ -638,6 +638,7 @@ def test_deprecated_batch_properties_equal_to_new_graph_properties(
     assert state.n_batches == state.n_graphs
     assert torch.allclose(state.n_atoms_per_batch, state.n_atoms_per_graph)
 
+    # now test that assigning the old .batch property behaves the same
     new_graph_idx = torch.arange(4, device=device)
     state.batch = new_graph_idx
     assert torch.allclose(state.graph_idx, new_graph_idx)