diff --git a/examples/scripts/3_dynamics.py b/examples/scripts/3_dynamics.py index 39357a2d..dd790b85 100644 --- a/examples/scripts/3_dynamics.py +++ b/examples/scripts/3_dynamics.py @@ -374,7 +374,7 @@ ) ) stress = mace_model_stress(state)["stress"] - volume = torch.det(state.current_cell) + volume = state.volume e_kin = ts.calc_kinetic_energy( masses=state.masses, momenta=state.momenta, system_idx=state.system_idx ) @@ -400,7 +400,7 @@ log.info(f"Final temperature: {final_temp.item():.4f} K") final_stress = mace_model_stress(state)["stress"] -final_volume = torch.det(state.current_cell) +final_volume = state.volume final_pressure = ts.get_pressure( final_stress, ts.calc_kinetic_energy( diff --git a/tests/conftest.py b/tests/conftest.py index af15b7e1..68f60b91 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -165,3 +165,17 @@ def mixed_double_sim_state( [ar_supercell_sim_state, si_sim_state], device=ar_supercell_sim_state.device, ) + + +def make_left_handed(state: ts.SimState) -> ts.SimState: + """Describe the same crystal with a left-handed cell. + + Negating the third lattice vector leaves Cartesian positions, the lattice and the + order of the a and b vectors untouched (some algorithms depend on that order) but + flips the sign of det(cell). `state.cell` stores lattice vectors as columns. + """ + left = state.clone() + cell = state.cell.clone() + cell[:, :, 2] *= -1 + left.cell = cell + return left diff --git a/tests/models/test_dispersion.py b/tests/models/test_dispersion.py index 60519318..2cc63588 100644 --- a/tests/models/test_dispersion.py +++ b/tests/models/test_dispersion.py @@ -8,7 +8,7 @@ import torch import torch_sim as ts -from tests.conftest import DEVICE, DTYPE +from tests.conftest import DEVICE, DTYPE, make_left_handed from tests.models.conftest import make_validate_model_outputs_test @@ -290,3 +290,15 @@ def spy_nl(*args: object, **nl_kwargs: object) -> object: ) model(state) assert seen == [None] + + +def test_d3_stress_is_independent_of_cell_handedness( + d3_model_pbe: D3DispersionModel, si_sim_state: ts.SimState +) -> None: + """The Cartesian D3 stress must not depend on the handedness of the cell.""" + right = d3_model_pbe(si_sim_state) + left = d3_model_pbe(make_left_handed(si_sim_state)) + + torch.testing.assert_close(left["energy"], right["energy"]) + torch.testing.assert_close(left["forces"], right["forces"]) + torch.testing.assert_close(left["stress"], right["stress"]) diff --git a/tests/test_autobatching.py b/tests/test_autobatching.py index 65b7143a..99c9bb96 100644 --- a/tests/test_autobatching.py +++ b/tests/test_autobatching.py @@ -100,7 +100,7 @@ def test_calculate_scaling_metric(si_sim_state: ts.SimState) -> None: # Test n_atoms_x_density metric density_metric = calculate_memory_scalers(si_sim_state, "n_atoms_x_density") - volume = torch.abs(torch.linalg.det(si_sim_state.cell[0])) / 1000 + volume = ts.quantities.cell_volume(si_sim_state.cell[0]) / 1000 expected = si_sim_state.n_atoms * (si_sim_state.n_atoms / volume.item()) assert pytest.approx(density_metric[0], rel=1e-5) == expected diff --git a/tests/test_integrators.py b/tests/test_integrators.py index 8fec9423..54ed2989 100644 --- a/tests/test_integrators.py +++ b/tests/test_integrators.py @@ -1,9 +1,11 @@ +from collections.abc import Callable + import pytest import torch from ase.build import bulk import torch_sim as ts -from tests.conftest import DEVICE, DTYPE +from tests.conftest import DEVICE, DTYPE, make_left_handed from torch_sim.integrators import initialize_momenta from torch_sim.integrators.npt import _npt_langevin_anisotropic_compute_cell_force from torch_sim.models.lennard_jones import LennardJonesModel @@ -719,7 +721,7 @@ def test_npt_crescale_default_coupling_is_stable(integrator_name: str) -> None: external_pressure = torch.tensor(0.0, dtype=DTYPE) state = ts.npt_crescale_init(state=state, model=model, dt=dt, kT=kT) - initial_volume = torch.det(state.cell).abs().clone() + initial_volume = state.volume.clone() for step in range(100): state = step_fn( @@ -730,7 +732,7 @@ def test_npt_crescale_default_coupling_is_stable(integrator_name: str) -> None: "barostat coupling" ) - ratio = (torch.det(state.cell).abs() / initial_volume).tolist() + ratio = (state.volume / initial_volume).tolist() assert all(0.5 < r < 2.0 for r in ratio), ( f"{integrator_name}: default barostat changed the volume by {ratio} in " "100 steps at the system's own pressure" @@ -1374,3 +1376,57 @@ def _run(seed: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: assert not torch.allclose(pos_a, pos_c) assert not torch.allclose(mom_a, mom_c) assert not torch.allclose(cell_a, cell_c) + + +@pytest.mark.parametrize( + ("init_fn", "step_fn"), + [ + (ts.npt_crescale_init, ts.npt_crescale_isotropic_step), + (ts.npt_crescale_init, ts.npt_crescale_triclinic_step), + (ts.npt_nose_hoover_isotropic_init, ts.npt_nose_hoover_isotropic_step), + (ts.npt_langevin_isotropic_init, ts.npt_langevin_isotropic_step), + ], + ids=["crescale_isotropic", "crescale_triclinic", "nose_hoover", "langevin"], +) +def test_npt_is_independent_of_cell_handedness( + init_fn: Callable, + step_fn: Callable, + ar_double_sim_state: ts.SimState, + lj_model: LennardJonesModel, +) -> None: + """NPT dynamics of a left-handed cell match the right-handed description. + + The barostats use the cell volume (e.g. sqrt(V)), which must be |det(cell)|: + a signed determinant gives NaN or the wrong pressure for a left-handed cell. + """ + dt = torch.tensor(0.001, dtype=DTYPE) + kT = torch.tensor(300.0, dtype=DTYPE) * MetalUnits.temperature + external_pressure = torch.tensor(10.0, dtype=DTYPE) * MetalUnits.pressure + + final = {} + for name, start in ( + ("right", ar_double_sim_state.clone()), + ("left", make_left_handed(ar_double_sim_state)), + ): + start.rng = 42 + state = init_fn(state=start, model=lj_model, dt=dt, kT=kT) + for _ in range(10): + state = step_fn( + state=state, + model=lj_model, + dt=dt, + kT=kT, + external_pressure=external_pressure, + ) + final[name] = state + + assert torch.isfinite(final["left"].energy).all() + # The triclinic barostat re-orthonormalises its random deformation, so the two + # trajectories agree to ~1e-6 + for attr in ("energy", "volume", "positions"): + torch.testing.assert_close( + getattr(final["left"], attr), + getattr(final["right"], attr), + rtol=1e-5, + atol=1e-5, + ) diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py index 7e77146f..aee83391 100644 --- a/tests/test_optimizers.py +++ b/tests/test_optimizers.py @@ -8,6 +8,7 @@ import torch import torch_sim as ts +from tests.conftest import make_left_handed from torch_sim.models.interface import ModelInterface from torch_sim.optimizers import BFGSState, FireFlavor, FireState, LBFGSState, OptimState from torch_sim.optimizers.cell_filters import CellLBFGSState, CellOptimState, deform_grad @@ -976,7 +977,7 @@ def test_unit_cell_fire_ase_non_positive_volume_warning( torch.randn_like(perturbed_state.cell) * 0.5 ) # Large perturbation # Also ensure no PBC issues by slightly expanding cell if it got too small - if torch.linalg.det(perturbed_state.cell[0]) < 1.0: + if ts.quantities.cell_volume(perturbed_state.cell[0]) < 1.0: perturbed_state.cell[0] *= 2.0 state = ts.fire_init( @@ -1562,3 +1563,56 @@ def test_lbfgs_prev_cell_positions_stored_before_step(lj_model: ModelInterface) "prev_positions should be fractional coords in the adjusted cell frame. " f"max diff = {(opt_state.prev_positions - expected_prev).abs().max():.2e}" ) + + +@pytest.mark.parametrize("cell_filter", [ts.CellFilter.unit, ts.CellFilter.frechet]) +@pytest.mark.parametrize("include_cell_forces", [False, True]) +def test_cell_optimization_is_independent_of_cell_handedness( + ar_supercell_sim_state: SimState, + lj_model: ModelInterface, + cell_filter: ts.CellFilter, + *, + include_cell_forces: bool, +) -> None: + """A left-handed description of a crystal relaxes exactly like the right-handed. + + The cell filters turn the stress into cell forces through the cell volume, which + must be |det(cell)|: a signed volume flips the cell forces, so the cell relaxes + the wrong way. A force criterion only sees their norm, so a left-handed cell would + still be reported as converged (`include_cell_forces=True`). + """ + strained = ar_supercell_sim_state.clone() + strained.cell = strained.cell * 1.05 + strained.positions = strained.positions * 1.05 + initial_volume = strained.volume.clone() + convergence_fn = ( + ts.generate_force_convergence_fn(force_tol=1e-3, include_cell_forces=True) + if include_cell_forces + else None # ts.optimize falls back to its default energy criterion + ) + + initial_cell_forces, final = {}, {} + for name, start in (("right", strained), ("left", make_left_handed(strained))): + state = ts.fire_init(state=start, model=lj_model, cell_filter=cell_filter) + initial_cell_forces[name] = state.cell_forces + final[name] = ts.optimize( + system=start, + model=lj_model, + optimizer=ts.Optimizer.fire, + convergence_fn=convergence_fn, + init_kwargs={"cell_filter": cell_filter}, + max_steps=1000, + ) + if convergence_fn is not None: + assert convergence_fn(final[name]).all() + + # The cell forces come from the (Cartesian) stress and the physical volume, so they + # must not depend on the handedness, and for a stretched crystal they point inward. + torch.testing.assert_close(initial_cell_forces["left"], initial_cell_forces["right"]) + assert (torch.diagonal(initial_cell_forces["right"], dim1=-2, dim2=-1) < 0).all() + + # The strained Ar crystal is under tension, so the volume must shrink. + assert (final["right"].volume < initial_volume).all() + torch.testing.assert_close(final["left"].volume, final["right"].volume) + torch.testing.assert_close(final["left"].energy, final["right"].energy) + torch.testing.assert_close(final["left"].positions, final["right"].positions) diff --git a/tests/test_quantities.py b/tests/test_quantities.py index 3eec4ad9..61952459 100644 --- a/tests/test_quantities.py +++ b/tests/test_quantities.py @@ -11,6 +11,7 @@ calc_kinetic_energy, calc_kT, calc_temperature, + cell_volume, ) from torch_sim.units import MetalUnits @@ -276,3 +277,12 @@ def test_calc_temperature(single_system_data: dict[str, Tensor]) -> None: velocities=single_system_data["velocities"], ) assert torch.allclose(temp, kt / MetalUnits.temperature) + + +def test_cell_volume_handedness_invariant() -> None: + """Volume is positive and identical for left- and right-handed cells.""" + cell = torch.tensor([[[2.0, 0, 0], [0, 3.0, 0], [0, 0, 4.0]]] * 2) + cell[1, 2] *= -1 # flip handedness of the second cell + volumes = cell_volume(cell) + torch.testing.assert_close(volumes, torch.full((2,), 24.0)) + assert cell_volume(cell[0]).shape == () diff --git a/tests/test_state.py b/tests/test_state.py index 7c73a649..ffe484c8 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -5,7 +5,7 @@ import torch import torch_sim as ts -from tests.conftest import DEVICE +from tests.conftest import DEVICE, make_left_handed from torch_sim.integrators import MDState from torch_sim.integrators.md import NoseHooverChain, NoseHooverChainFns from torch_sim.integrators.npt import ( @@ -1239,3 +1239,11 @@ def test_lbfgsstate_instantiation() -> None: max_atoms=torch.tensor([n_atoms]), ) _check_coercion(state) + + +def test_volume_is_positive_for_left_handed_cell(si_sim_state: SimState) -> None: + left = make_left_handed(si_sim_state) + + assert torch.det(left.cell).item() < 0 + assert (left.volume > 0).all() + torch.testing.assert_close(left.volume, si_sim_state.volume) diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index 6d997981..4b2c7c86 100644 --- a/torch_sim/autobatching.py +++ b/torch_sim/autobatching.py @@ -406,7 +406,7 @@ def calculate_memory_scalers( ) ) if system_pbc_all: - volume = torch.abs(torch.linalg.det(system_state.cell[0])) / 1000 + volume = system_state.volume[0] / 1000 else: bbox = ( system_state.positions.max(dim=0).values diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index bc0ef634..b028f74f 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -125,11 +125,6 @@ def current_cell(self) -> torch.Tensor: scale = torch.exp(self.cell_positions) # (n_systems, 3) return scale.unsqueeze(-1) * self.reference_cell - @property - def volume(self) -> torch.Tensor: - """Current volume from cell determinant.""" - return torch.linalg.det(self.cell) - def _npt_langevin_particle_beta( state: "NPTLangevinAnisotropicState | NPTLangevinIsotropicState", @@ -697,7 +692,7 @@ def current_cell(self) -> torch.Tensor: def volume(self) -> torch.Tensor: """Current volume V = V₀ · exp(d·ε).""" dim = self.positions.shape[1] - V_0 = torch.linalg.det(self.reference_cell) + V_0 = ts.quantities.cell_volume(self.reference_cell) return V_0 * torch.exp(dim * self.cell_positions) @@ -1205,7 +1200,7 @@ def current_cell(self) -> torch.Tensor: [n_systems, n_dimensions, n_dimensions] """ dim = self.positions.shape[1] - V_0 = torch.det(self.reference_cell) # [n_systems] + V_0 = ts.quantities.cell_volume(self.reference_cell) # [n_systems] V = V_0 * torch.exp(dim * self.cell_position) # [n_systems] scale = (V / V_0) ** (1.0 / dim) # [n_systems] # Expand scale to [n_systems, 1, 1] for broadcasting @@ -1244,7 +1239,7 @@ def _npt_nose_hoover_isotropic_cell_info( """ dim = state.positions.shape[1] ref = state.reference_cell # [n_systems, dim, dim] - V_0 = torch.det(ref) # [n_systems] - Reference volume + V_0 = ts.quantities.cell_volume(ref) # [n_systems] - Reference volume V = V_0 * torch.exp(dim * state.cell_position) # [n_systems] - Current volume def volume_to_cell(V: torch.Tensor) -> torch.Tensor: @@ -2026,7 +2021,7 @@ def npt_nose_hoover_isotropic_invariant( multiple systems. """ # Calculate volume and potential energy - volume = torch.det(state.current_cell) # [n_systems] + volume = ts.quantities.cell_volume(state.current_cell) # [n_systems] e_pot = state.energy # Should be scalar or [n_systems] # Calculate kinetic energy of particles per system @@ -2189,7 +2184,7 @@ def _crescale_triclinic_barostat_step( dt: torch.Tensor, external_pressure: torch.Tensor, ) -> NPTCRescaleState: - volume = torch.det(state.cell) # shape: (n_systems,) + volume = state.volume # shape: (n_systems,) P_int = ts.quantities.compute_instantaneous_pressure_tensor( momenta=state.momenta, masses=state.masses, @@ -2285,7 +2280,7 @@ def _crescale_anisotropic_barostat_step( dt: torch.Tensor, external_pressure: torch.Tensor, ) -> NPTCRescaleState: - volume = torch.det(state.cell) # shape: (n_systems,) + volume = state.volume # shape: (n_systems,) P_int = ts.quantities.compute_instantaneous_pressure_tensor( momenta=state.momenta, masses=state.masses, @@ -2379,7 +2374,7 @@ def _crescale_triclinic_average_barostat_step( dt: torch.Tensor, external_pressure: torch.Tensor, ) -> NPTCRescaleState: - volume = torch.det(state.cell) # shape: (n_systems,) + volume = state.volume # shape: (n_systems,) P_int = compute_average_pressure_tensor( degrees_of_freedom=state.get_number_of_degrees_of_freedom() / 3, kT=kT, @@ -2458,7 +2453,7 @@ def _crescale_isotropic_barostat_step( dt: torch.Tensor, external_pressure: torch.Tensor, ) -> NPTCRescaleState: - volume = torch.det(state.cell) # shape: (n_systems,) + volume = state.volume # shape: (n_systems,) P_int = ts.quantities.compute_instantaneous_pressure_tensor( momenta=state.momenta, masses=state.masses, @@ -2971,7 +2966,7 @@ def npt_crescale_init( # Store initial cell for deviatoric correction initial_cell = state.cell.clone() initial_cell_inv = torch.inverse(initial_cell) - initial_volume = torch.det(initial_cell) + initial_volume = state.volume # Create the initial state npt_state = NPTCRescaleState.from_state( diff --git a/torch_sim/models/pair_potential.py b/torch_sim/models/pair_potential.py index fe3ea9e4..e2ed7101 100644 --- a/torch_sim/models/pair_potential.py +++ b/torch_sim/models/pair_potential.py @@ -74,6 +74,7 @@ def bmhtf_pair(dr, zi, zj, A, B, C, D, sigma): from torch_sim.models.interface import ModelInterface from torch_sim.neighbors import torchsim_nl +from torch_sim.quantities import cell_volume from torch_sim.transforms import compute_cell_shifts @@ -224,7 +225,7 @@ def _virial_stress( ``[n_systems, 3, 3]``, stress_per_pair ``[n_pairs, 3, 3]``, and volumes ``[n_systems]``. """ - volumes = torch.abs(torch.linalg.det(row_cell)) + volumes = cell_volume(row_cell) stress_per_pair = torch.einsum("...i,...j->...ij", dr_vec, force_vectors) stress = torch.zeros((n_systems, 3, 3), dtype=dtype, device=device) stress = stress.index_add(0, system_mapping, -stress_per_pair) diff --git a/torch_sim/optimizers/cell_filters.py b/torch_sim/optimizers/cell_filters.py index 25b1b60f..cc6a8eea 100644 --- a/torch_sim/optimizers/cell_filters.py +++ b/torch_sim/optimizers/cell_filters.py @@ -202,7 +202,7 @@ def unit_cell_filter_init[T: AnyCellState]( # Calculate initial cell forces stress = _get_constrained_stress(model_output, state) - volumes = torch.linalg.det(state.cell).view(n_systems, 1, 1) + volumes = state.volume.view(n_systems, 1, 1) virial = -volumes * (stress + pressure) virial = _apply_constraints( virial, hydrostatic_strain=hydrostatic_strain, constant_volume=constant_volume @@ -255,7 +255,7 @@ def frechet_cell_filter_init[T: AnyCellState]( # Calculate initial cell forces using Frechet approach stress = _get_constrained_stress(model_output, state) - volumes = torch.linalg.det(state.cell).view(n_systems, 1, 1) + volumes = state.volume.view(n_systems, 1, 1) virial = -volumes * (stress + pressure) virial = _apply_constraints( virial, hydrostatic_strain=hydrostatic_strain, constant_volume=constant_volume @@ -355,7 +355,7 @@ def compute_cell_forces[T: AnyCellState]( ) -> None: """Compute cell forces for both unit and frechet methods.""" stress = _get_constrained_stress(model_output, state) - volumes = torch.linalg.det(state.cell).view(state.n_systems, 1, 1) + volumes = state.volume.view(state.n_systems, 1, 1) virial = -volumes * (stress + state.pressure) virial = _apply_constraints( virial, diff --git a/torch_sim/quantities.py b/torch_sim/quantities.py index aa2b4a2a..ab6d04bf 100644 --- a/torch_sim/quantities.py +++ b/torch_sim/quantities.py @@ -12,6 +12,22 @@ from torch_sim.optimizers import OptimState +def cell_volume(cell: torch.Tensor) -> torch.Tensor: + """Volume of one or more cells, positive for left- and right-handed cells. + + ``det(cell)`` is negative for left-handed cells, which would flip the sign of + anything derived from it (pressure, barostat forces, scaled cells). Use this + helper rather than calling ``torch.det`` on a cell directly. + + Args: + cell: Cell matrices with shape ``[..., 3, 3]``. + + Returns: + Volumes with shape ``[...]``. + """ + return torch.linalg.det(cell).abs() + + # @torch.jit.script def calc_kT( # noqa: N802 *, diff --git a/torch_sim/state.py b/torch_sim/state.py index 402cb990..4c8387d8 100644 --- a/torch_sim/state.py +++ b/torch_sim/state.py @@ -442,7 +442,7 @@ def n_systems(self) -> int: @property def volume(self) -> torch.Tensor: """Volume of the system.""" - return torch.det(self.cell) + return ts.quantities.cell_volume(self.cell) @property def attributes(self) -> dict[str, Any]: