Feature/thermodynamic integration

tests/models/test_einstein.py

@@ -0,0 +1,333 @@
+"""Tests for the Einstein model implementation."""
+
+import pytest
+import torch
+from ase.build import bulk
+
+import torch_sim as ts
+from torch_sim.models.einstein import EinsteinModel
+
+
+class TestEinsteinModel:
+    """Test Einstein model implementation."""
+
+    @pytest.fixture
+    def simple_system(self):
+        """Create a simple test system."""
+        device = torch.device("cpu")
+        dtype = torch.float64
+
+        # Create a simple 2-atom system
+        positions = torch.tensor(
+            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=dtype, device=device
+        )
+        masses = torch.tensor([1.0, 1.0], dtype=dtype, device=device)
+        cell = torch.eye(3, dtype=dtype, device=device) * 10.0  # Large cell
+        atomic_numbers = torch.tensor([1, 1], dtype=torch.int64, device=device)
+
+        state = ts.SimState(
+            positions=positions,
+            masses=masses,
+            cell=cell.unsqueeze(0),
+            pbc=True,
+            atomic_numbers=atomic_numbers,
+        )
+
+        return state, device, dtype
+
+    @pytest.fixture
+    def batched_system(self):
+        """Create a batched test system."""
+        device = torch.device("cpu")
+        dtype = torch.float64
+
+        # Create two different 2-atom systems
+        si_atoms = bulk("Si", "diamond", a=5.43, cubic=True)
+        fe_atoms = bulk("Fe", "bcc", a=2.87, cubic=True)
+
+        state = ts.io.atoms_to_state([si_atoms, fe_atoms], device, dtype)
+        return state, device, dtype
+
+    def test_einstein_model_creation(
+        self, simple_system: tuple[ts.SimState, torch.device, torch.dtype]
+    ):
+        """Test basic Einstein model creation."""
+        state, device, dtype = simple_system
+
+        # Create equilibrium positions and frequencies
+        equilibrium_pos = state.positions
+        frequencies = torch.ones(2, dtype=dtype, device=device) * 0.1  # [N]
+
+        model = EinsteinModel(
+            equilibrium_position=equilibrium_pos,
+            frequencies=frequencies,
+            masses=state.masses,
+            device=device,
+            dtype=dtype,
+        )
+
+        assert model.device == device
+        assert model.dtype == dtype
+        assert model.compute_forces is True
+
+    def test_einstein_model_forward_single_system(
+        self, simple_system: tuple[ts.SimState, torch.device, torch.dtype]
+    ):
+        """Test forward pass with single system."""
+        state, device, dtype = simple_system
+
+        equilibrium_pos = state.positions
+        frequencies = torch.ones(2, dtype=dtype, device=device) * 0.1
+
+        model = EinsteinModel(
+            equilibrium_position=equilibrium_pos,
+            frequencies=frequencies,
+            masses=state.masses,
+            device=device,
+            dtype=dtype,
+        )
+
+        # Displace atoms slightly from equilibrium
+        displaced_state = state.clone()
+        displaced_state.positions += 0.1
+
+        results = model(displaced_state)
+
+        assert "energy" in results
+        assert "forces" in results
+        assert results["energy"].shape == (1,)
+        assert results["forces"].shape == (2, 3)  # [N_atoms, 3]
+
+        # Forces should point back toward equilibrium
+        expected_force_direction = -(displaced_state.positions - equilibrium_pos)
+        force_directions = results["forces"]
+        # Check that forces point in the right direction (dot product > 0)
+        for i in range(2):
+            dot_product = torch.dot(force_directions[i], expected_force_direction[i])
+            assert dot_product > 0
+
+    def test_einstein_model_forward_batched_system(
+        self, batched_system: tuple[ts.SimState, torch.device, torch.dtype]
+    ):
+        """Test forward pass with batched system."""
+        state, device, dtype = batched_system
+
+        # Create equilibrium positions for the batched system
+        n_atoms = state.n_atoms
+        equilibrium_pos = state.positions
+        frequencies = torch.ones(n_atoms, dtype=dtype, device=device) * 0.05
+
+        model = EinsteinModel(
+            equilibrium_position=equilibrium_pos,
+            frequencies=frequencies,
+            masses=state.masses,
+            device=device,
+            dtype=dtype,
+        )
+
+        # Displace atoms slightly
+        displaced_state = state.clone()
+        displaced_state.positions += 0.05
+
+        results = model(displaced_state)
+
+        assert "energy" in results
+        assert "forces" in results
+        assert results["energy"].shape == (2,)  # [n_systems]
+        assert results["forces"].shape == (n_atoms, 3)  # [total_atoms, 3]
+
+    def test_einstein_model_from_frequencies(self):
+        """Test creation from frequencies class method."""
+        device = torch.device("cpu")
+        dtype = torch.float64
+
+        # Create ASE atoms
+        atoms = bulk("Si", "diamond", a=5.43, cubic=True)
+        state = ts.io.atoms_to_state([atoms], device, dtype)
+
+        frequencies = torch.ones(len(atoms), dtype=dtype, device=device) * 0.05
+
+        model = EinsteinModel.from_atom_and_frequencies(
+            atom=state,
+            frequencies=frequencies,
+            reference_energy=1.0,
+            device=device,
+            dtype=dtype,
+        )
+
+        assert torch.allclose(model.reference_energy, torch.tensor(1.0, dtype=dtype))
+        assert model.frequencies.shape[0] == len(atoms)
+
+    def test_periodic_boundary_conditions(
+        self, simple_system: tuple[ts.SimState, torch.device, torch.dtype]
+    ):
+        """Test that PBC are handled correctly."""
+        state, device, dtype = simple_system
+
+        # Create model with equilibrium at origin
+        equilibrium_pos = torch.zeros((2, 3), dtype=dtype, device=device)
+        frequencies = torch.ones(2, dtype=dtype, device=device) * 1
+
+        model = EinsteinModel(
+            equilibrium_position=equilibrium_pos,
+            frequencies=frequencies,
+            device=device,
+            dtype=dtype,
+        )
+
+        # Place atoms near opposite faces of the cell
+        test_state = state.clone()
+        test_state.positions = torch.tensor(
+            [
+                [0.1, 0.0, 0.0],  # Near one face
+                [9.9, 0.0, 0.0],  # Near opposite face
+            ],
+            dtype=dtype,
+            device=device,
+        )
+
+        results = model(test_state)
+
+        # Should handle PBC correctly - both atoms far from origin
+        # but should compute minimum image distances
+        assert torch.isfinite(results["energy"])
+        assert torch.isfinite(results["forces"]).all()
+
+        spring = frequencies**2  # since mass=1
+        target_energies = 0.5 * spring * (torch.tensor([0.1, -0.1]) ** 2)
+        target_forces = -spring[:, None] * torch.tensor(
+            [[0.1, 0.0, 0.0], [-0.1, 0.0, 0.0]], dtype=dtype, device=device
+        )
+        assert torch.allclose(results["energy"], target_energies.sum(), atol=1e-6)
+        assert torch.allclose(results["forces"], target_forces, atol=1e-6)
+
+    def test_energy_force_consistency(
+        self, simple_system: tuple[ts.SimState, torch.device, torch.dtype]
+    ):
+        """Test that forces are consistent with energy gradients."""
+        state, device, dtype = simple_system
+
+        equilibrium_pos = state.positions.clone()
+        frequencies = torch.ones(2, dtype=dtype, device=device) * 0.1
+
+        model = EinsteinModel(
+            equilibrium_position=equilibrium_pos,
+            frequencies=frequencies,
+            masses=state.masses,
+            device=device,
+            dtype=dtype,
+        )
+
+        # Create a displaced state with gradients enabled
+        test_positions = state.positions.clone() + 0.1
+        test_positions.requires_grad_(requires_grad=True)
+
+        test_state = state.clone()
+        test_state.positions = test_positions
+
+        results = model(test_state)
+        energy = results["energy"]
+
+        # Compute forces from gradients
+        forces_from_grad = -torch.autograd.grad(
+            energy, test_positions, create_graph=False
+        )[0]
+        forces_direct = results["forces"]
+
+        # Forces should match (within numerical precision)
+        torch.testing.assert_close(forces_direct, forces_from_grad, atol=1e-6, rtol=1e-6)
+
+    def test_get_free_energy(
+        self, simple_system: tuple[ts.SimState, torch.device, torch.dtype]
+    ):
+        """Test free energy calculation."""
+        state, device, dtype = simple_system
+
+        equilibrium_pos = state.positions
+        frequencies = torch.ones(2, dtype=dtype, device=device) * 0.1  # THz
+
+        model = EinsteinModel(
+            equilibrium_position=equilibrium_pos,
+            frequencies=frequencies,
+            masses=state.masses,
+            device=device,
+            dtype=dtype,
+        )
+
+        temperature = 300.0  # K
+        results = model.get_free_energy(temperature)
+
+        # Check that result is a dictionary with free energy
+        assert isinstance(results, dict)
+        assert "free_energy" in results
+        free_energy = results["free_energy"]
+        assert isinstance(free_energy, torch.Tensor)
+        assert free_energy.shape == (1,)  # Single system
+
+        # Free energy should be finite
+        assert torch.isfinite(free_energy).all()
+
+        # At higher temperature, free energy should be lower (more negative)
+        results_high = model.get_free_energy(600.0)
+        free_energy_high = results_high["free_energy"]
+        assert free_energy_high < free_energy
+
+    def test_get_free_energy_batched(
+        self, batched_system: tuple[ts.SimState, torch.device, torch.dtype]
+    ):
+        """Test free energy calculation for batched systems."""
+        state, device, dtype = batched_system
+
+        n_atoms = state.n_atoms
+        equilibrium_pos = state.positions
+        frequencies = torch.ones(n_atoms, dtype=dtype, device=device) * 0.05
+
+        model = EinsteinModel(
+            equilibrium_position=equilibrium_pos,
+            frequencies=frequencies,
+            masses=state.masses,
+            system_idx=state.system_idx,
+            device=device,
+            dtype=dtype,
+        )
+
+        temperature = 300.0
+        results = model.get_free_energy(temperature)
+
+        # Check result format and shape
+        assert isinstance(results, dict)
+        assert "free_energy" in results
+        free_energy = results["free_energy"]
+
+        # Should have one free energy per system
+        assert free_energy.shape == (2,)  # Two systems
+        assert torch.isfinite(free_energy).all()
+
+    def test_sample_method(
+        self, simple_system: tuple[ts.SimState, torch.device, torch.dtype]
+    ):
+        """Test sampling from Einstein model."""
+        state, device, dtype = simple_system
+
+        equilibrium_pos = state.positions
+        frequencies = torch.ones(2, dtype=dtype, device=device) * 0.1
+
+        model = EinsteinModel(
+            equilibrium_position=equilibrium_pos,
+            frequencies=frequencies,
+            masses=state.masses,
+            device=device,
+            dtype=dtype,
+        )
+
+        temperature = 300.0
+        sampled_state = model.sample(state, temperature)
+
+        # Check that sampled state has correct shape and type
+        assert isinstance(sampled_state, ts.SimState)
+        assert sampled_state.positions.shape == state.positions.shape
+        assert sampled_state.positions.dtype == dtype
+        assert sampled_state.positions.device == device
+
+        # Sampled positions should have finite values
+        assert torch.isfinite(sampled_state.positions).all()