Use ase bulk for ar base state

Author: abhijeetgangan

tests/conftest.py

@@ -65,47 +65,9 @@ def si_double_base_state(si_atoms: Atoms, device: torch.device) -> Any:
 @pytest.fixture
 def ar_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,
-    )
+    # Create FCC Ar using ASE, with 4x4x4 supercell
+    ar_atoms = bulk("Ar", "fcc", a=5.26, cubic=True).repeat([2, 2, 2])
+    return atoms_to_state(ar_atoms, device, torch.float64)
 
 
 @pytest.fixture

tests/models/test_lennard_jones.py

@@ -2,13 +2,15 @@
 
 import pytest
 import torch
+from ase.build import bulk
 
 from torchsim.models.interface import validate_model_outputs
 from torchsim.models.lennard_jones import (
     UnbatchedLennardJonesModel,
     lennard_jones_pair,
     lennard_jones_pair_force,
 )
+from torchsim.runners import atoms_to_state
 from torchsim.state import BaseState
 
 
@@ -130,9 +132,21 @@ def test_lennard_jones_force_energy_consistency() -> None:
     assert torch.allclose(force_direct, force_from_grad, rtol=1e-4, atol=1e-4)
 
 
+# NOTE: This is a large system to test the neighbor list and direct calculation
+#       are consistent. Direct calculation uses minimal image convention, which
+#       is not used in the neighbor list calculation. So to get correct results,
+#       we need a system that is large enough (2*cutoff).
+@pytest.fixture
+def ar_base_state_large(device: torch.device) -> BaseState:
+    """Create a face-centered cubic (FCC) Argon structure."""
+    # Create FCC Ar using ASE, with 4x4x4 supercell
+    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 calculators(
-    ar_base_state: BaseState,
+    ar_base_state_large: BaseState,
 ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
     """Create both neighbor list and direct calculators with Argon parameters."""
     calc_params = {
@@ -152,7 +166,7 @@ def calculators(
         use_neighbor_list=False, cutoff=cutoff, **calc_params
     )
 
-    positions, cell = ar_base_state.positions, ar_base_state.cell.squeeze(0)
+    positions, cell = ar_base_state_large.positions, ar_base_state_large.cell.squeeze(0)
     return calc_nl(positions, cell), calc_direct(positions, cell)
 
 

tests/test_integrators.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import pytest
 import torch
 
 from torchsim.integrators import batched_initialize_momenta, nve, nvt_langevin
@@ -11,23 +10,6 @@
 from torchsim.units import MetalUnits
 
 
-@pytest.fixture
-def ar_double_base_state(ar_base_state: BaseState) -> BaseState:
-    """Create a basic state from ar_fcc_base_state."""
-    batch = torch.repeat_interleave(torch.arange(2), ar_base_state.positions.shape[0])
-
-    return BaseState(
-        positions=torch.cat([ar_base_state.positions, ar_base_state.positions], dim=0),
-        cell=torch.cat([ar_base_state.cell, ar_base_state.cell], dim=0),
-        masses=torch.cat([ar_base_state.masses, ar_base_state.masses], dim=0),
-        atomic_numbers=torch.cat(
-            [ar_base_state.atomic_numbers, ar_base_state.atomic_numbers], dim=0
-        ),
-        batch=batch,
-        pbc=ar_base_state.pbc,
-    )
-
-
 def batched_initialize_momenta_loop(
     positions: torch.Tensor,  # shape: (n_batches, n_atoms_per_batch, 3)
     masses: torch.Tensor,  # shape: (n_batches, n_atoms_per_batch)