More robust tests (#12)

* More robust tests for integrators, lj for argon without ASE / Pymatgen

* Add invariant check for nose-hoover

* NPT test

Author: abhijeetgangan

tests/conftest.py

@@ -56,29 +56,70 @@ def fe_fcc_state(device: torch.device) -> Any:
     return atoms_to_state(fe_atoms, device, torch.float64)
 
 
-@pytest.fixture
-def ar_fcc_state(device: torch.device) -> Any:
-    ar_atoms = bulk("Ar", "fcc", a=5.26, cubic=True).repeat([4, 4, 4])
-    return atoms_to_state(ar_atoms, device, torch.float64)
-
-
 @pytest.fixture
 def si_double_base_state(si_atoms: Atoms, device: torch.device) -> Any:
     """Create a basic state from si_structure."""
     return atoms_to_state([si_atoms, si_atoms], device, torch.float64)
 
 
+@pytest.fixture
+def ar_fcc_base_state(device: torch.device) -> BaseState:
+    """Create a face-centered cubic (FCC) Argon structure."""
+    # 5.26 Å is a typical lattice constant for Ar
+    a = 5.26  # Lattice constant
+    N = 4  # Supercell size
+    n_atoms = 4 * N * N * N  # Total number of atoms (4 atoms per unit cell)
+    dtype = torch.float64
+
+    # Create positions tensor directly
+    positions = torch.zeros((n_atoms, 3), device=device, dtype=dtype)
+    idx = 0
+    for i in range(N):
+        for j in range(N):
+            for k in range(N):
+                # Add base FCC positions with offset
+                positions[idx] = torch.tensor([i, j, k], device=device, dtype=dtype) * a
+                positions[idx + 1] = (
+                    torch.tensor([i, j + 0.5, k + 0.5], device=device, dtype=dtype) * a
+                )
+                positions[idx + 2] = (
+                    torch.tensor([i + 0.5, j, k + 0.5], device=device, dtype=dtype) * a
+                )
+                positions[idx + 3] = (
+                    torch.tensor([i + 0.5, j + 0.5, k], device=device, dtype=dtype) * a
+                )
+                idx += 4
+
+    # Create cell tensor with shape (1, 3, 3) to match atoms_to_state format
+    cell = torch.eye(3, device=device, dtype=dtype).unsqueeze(0) * (N * a)
+
+    # Create batch indices
+    batch = torch.zeros(n_atoms, device=device, dtype=torch.long)
+
+    return BaseState(
+        positions=positions,
+        masses=torch.full((n_atoms,), 39.95, device=device, dtype=dtype),  # Ar mass
+        cell=cell,  # Cubic cell
+        pbc=True,
+        atomic_numbers=torch.full(
+            (n_atoms,), 18, device=device, dtype=torch.long
+        ),  # Ar atomic number
+        batch=batch,
+    )
+
+
 @pytest.fixture
 def unbatched_lj_calculator(device: torch.device) -> UnbatchedLennardJonesModel:
-    """Create a Lennard-Jones calculator with reasonable parameters for Si."""
+    """Create a Lennard-Jones calculator with reasonable parameters for Ar."""
     return UnbatchedLennardJonesModel(
-        sigma=2.0,  # Approximate for Si-Si interaction
-        epsilon=0.1,  # Small epsilon for stability during testing
+        use_neighbor_list=True,
+        sigma=3.405,  # Approximate for Ar-Ar interaction
+        epsilon=0.0104,  # Small epsilon for stability during testing
         device=device,
         dtype=torch.float64,
         compute_force=True,
         compute_stress=True,
-        cutoff=5.0,
+        cutoff=2.5 * 3.405,
     )
 
 

tests/models/test_lennard_jones.py

@@ -5,7 +5,7 @@
 
 from torchsim.models.interface import validate_model_outputs
 from torchsim.models.lennard_jones import (
-    LennardJonesModel,
+    UnbatchedLennardJonesModel,
     lennard_jones_pair,
     lennard_jones_pair_force,
 )
@@ -132,7 +132,7 @@ def test_lennard_jones_force_energy_consistency() -> None:
 
 @pytest.fixture
 def calculators(
-    ar_fcc_state: BaseState,
+    ar_fcc_base_state: BaseState,
 ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
     """Create both neighbor list and direct calculators with Argon parameters."""
     calc_params = {
@@ -145,10 +145,14 @@ def calculators(
     }
 
     cutoff = 2.5 * 3.405  # Standard LJ cutoff * sigma
-    calc_nl = LennardJonesModel(use_neighbor_list=True, cutoff=cutoff, **calc_params)
-    calc_direct = LennardJonesModel(use_neighbor_list=False, cutoff=cutoff, **calc_params)
+    calc_nl = UnbatchedLennardJonesModel(
+        use_neighbor_list=True, cutoff=cutoff, **calc_params
+    )
+    calc_direct = UnbatchedLennardJonesModel(
+        use_neighbor_list=False, cutoff=cutoff, **calc_params
+    )
 
-    positions, cell = ar_fcc_state.positions, ar_fcc_state.cell
+    positions, cell = ar_fcc_base_state.positions, ar_fcc_base_state.cell.squeeze(0)
     return calc_nl(positions, cell), calc_direct(positions, cell)
 
 
@@ -197,7 +201,7 @@ def test_stress_tensor_symmetry(
 
 
 def test_validate_model_outputs(
-    lj_calculator: LennardJonesModel,
+    lj_calculator: UnbatchedLennardJonesModel,
     device: torch.device,
 ) -> None:
     """Test that the model outputs are valid."""