""" Copyright (c) Meta Platforms, Inc. and affiliates. This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree. """ from __future__ import annotations import os import sys import tempfile import hydra import numpy as np import pytest from ase import Atoms pytest.importorskip("lammps") from fairchem.lammps import lammps_fc # noqa: E402 from fairchem.lammps.lammps_fc import restricted_cell_from_lammps_box # noqa: E402 def create_lammps_data_file(filepath, positions, cell, atom_types, masses): """ Create a LAMMPS data file for a triclinic box. Args: filepath: Path to write the data file positions: Nx3 array of atom positions in Cartesian coordinates cell: 3x3 cell matrix (rows are lattice vectors) atom_types: List of atom type IDs (1-indexed for LAMMPS) masses: Dict mapping atom type ID to mass """ n_atoms = len(positions) n_types = len(masses) # Extract cell parameters for LAMMPS triclinic box # Cell rows are: a = cell[0], b = cell[1], c = cell[2] a_vec = cell[0] b_vec = cell[1] c_vec = cell[2] # LAMMPS restricted triclinic parameters # See: https://docs.lammps.org/Howto_triclinic.html xlo, ylo, zlo = 0.0, 0.0, 0.0 xhi = a_vec[0] # lx xy = b_vec[0] yhi = b_vec[1] # ly xz = c_vec[0] yz = c_vec[1] zhi = c_vec[2] # lz with open(filepath, "w") as f: f.write("LAMMPS data file\n\n") f.write(f"{n_atoms} atoms\n") f.write(f"{n_types} atom types\n\n") f.write(f"{xlo} {xhi} xlo xhi\n") f.write(f"{ylo} {yhi} ylo yhi\n") f.write(f"{zlo} {zhi} zlo zhi\n") f.write(f"{xy} {xz} {yz} xy xz yz\n\n") f.write("Masses\n\n") for type_id, mass in masses.items(): f.write(f"{type_id} {mass}\n") f.write("\n") f.write("Atoms\n\n") for i, (pos, atype) in enumerate(zip(positions, atom_types), start=1): f.write(f"{i} {atype} {pos[0]} {pos[1]} {pos[2]}\n") @pytest.mark.parametrize( "cell_name,cell,fractional_positions", [ ( "cubic", np.array( [[5.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 5.0]], dtype=np.float64 ), np.array( [[0.1, 0.2, 0.3], [0.7, 0.8, 0.1], [0.5, 0.5, 0.5]], dtype=np.float64 ), ), ( "orthorhombic", np.array( [[3.0, 0.0, 0.0], [0.0, 4.0, 0.0], [0.0, 0.0, 5.0]], dtype=np.float64 ), np.array( [[0.1, 0.2, 0.3], [0.7, 0.8, 0.1], [0.5, 0.5, 0.5]], dtype=np.float64 ), ), ( "monoclinic", np.array( [[4.0, 0.0, 0.0], [-1.5, 4.5, 0.0], [0.0, 0.0, 6.0]], dtype=np.float64 ), np.array( [[0.1, 0.2, 0.3], [0.7, 0.8, 0.1], [0.5, 0.5, 0.5]], dtype=np.float64 ), ), ( "triclinic", np.array( [[5.0, 0.0, 0.0], [1.2, 4.8, 0.0], [0.8, 0.5, 5.5]], dtype=np.float64 ), np.array( [[0.1, 0.2, 0.3], [0.7, 0.8, 0.1], [0.5, 0.5, 0.5]], dtype=np.float64 ), ), ( "triclinic_with_boundary_atoms", np.array( [[5.0, 0.0, 0.0], [1.0, 4.0, 0.0], [0.5, 0.3, 6.0]], dtype=np.float64 ), np.array( [ [0.0, 0.0, 0.0], [0.5, 0.0, 0.0], [0.0, 0.5, 0.0], [0.0, 0.0, 0.5], [0.5, 0.5, 0.5], [0.25, 0.75, 0.25], ], dtype=np.float64, ), ), ], ) def test_scaled_positions_lammps_vs_ase(cell_name, cell, fractional_positions): """ Test that scaled atomic positions computed by ASE match those from LAMMPS. This test: 1. Creates a LAMMPS simulation with atoms in a triclinic box 2. Extracts box parameters and Cartesian positions from LAMMPS 3. Uses restricted_cell_from_lammps_box to get the ASE cell 4. Creates an ASE Atoms object with the positions and cell 5. Verifies that scaled (fractional) positions match """ lammps = pytest.importorskip("lammps") # Convert fractional to Cartesian: pos = frac @ cell cartesian_positions = fractional_positions @ cell # Atom types (all type 1 = Carbon for simplicity) atom_types = [1] * len(cartesian_positions) masses = {1: 12.011} # Carbon mass with tempfile.TemporaryDirectory() as tmpdir: data_file = os.path.join(tmpdir, "test.data") create_lammps_data_file( data_file, cartesian_positions, cell, atom_types, masses ) # Create LAMMPS instance and read the data file lmp = lammps.lammps(cmdargs=["-screen", "none", "-log", "none"]) lmp.command("units metal") lmp.command("atom_style atomic") lmp.command("boundary p p p") lmp.command(f"read_data {data_file}") # Extract box parameters from LAMMPS boxlo, boxhi, xy_lmp, yz_lmp, xz_lmp, periodicity, box_change = ( lmp.extract_box() ) # Extract atom positions from LAMMPS nlocal = lmp.get_natoms() x_lammps = lmp.numpy.extract_atom("x")[:nlocal].copy() # Get ASE cell from LAMMPS box parameters cell_from_lammps = restricted_cell_from_lammps_box( boxlo, boxhi, xy_lmp, yz_lmp, xz_lmp ) cell_np = cell_from_lammps.squeeze().numpy() # Create ASE Atoms object ase_atoms = Atoms( symbols=["C"] * nlocal, positions=x_lammps, cell=cell_np, pbc=True ) # Get scaled positions from ASE (wrap=False to avoid [0,1) wrapping issues) ase_scaled_positions = ase_atoms.get_scaled_positions(wrap=False) # The key validation: scaled positions from ASE should match our original # fractional positions. We compare modulo 1 to handle periodic boundary effects. def normalize_fractional(frac): """Normalize fractional coordinates to [0, 1) handling numerical precision.""" normalized = frac % 1.0 # Handle values very close to 1.0 that should wrap to 0.0 normalized = np.where(np.abs(normalized - 1.0) < 1e-6, 0.0, normalized) return normalized wrapped_original = normalize_fractional(fractional_positions) wrapped_ase = normalize_fractional(ase_scaled_positions) assert np.allclose( wrapped_ase, wrapped_original, atol=1e-5 ), f"Cell {cell_name}: Scaled positions don't match.\nASE: {wrapped_ase}\nOriginal: {wrapped_original}" lmp.close() @pytest.mark.parametrize( "box_name,boxlo,boxhi,xy,yz,xz", [ ("cubic", [0.0, 0.0, 0.0], [5.0, 5.0, 5.0], 0.0, 0.0, 0.0), ("orthorhombic", [0.0, 0.0, 0.0], [3.0, 4.0, 5.0], 0.0, 0.0, 0.0), ("xy_tilt", [0.0, 0.0, 0.0], [4.0, 5.0, 6.0], 1.0, 0.0, 0.0), ("yz_tilt", [0.0, 0.0, 0.0], [4.0, 5.0, 6.0], 0.0, 0.5, 0.0), ("xz_tilt", [0.0, 0.0, 0.0], [4.0, 5.0, 6.0], 0.0, 0.0, 0.3), ("full_triclinic", [0.0, 0.0, 0.0], [4.0, 5.0, 6.0], 1.0, 0.5, 0.3), ("nonzero_boxlo", [1.0, 2.0, 3.0], [6.0, 7.0, 8.0], 0.5, 0.3, 0.2), ], ) def test_cell_conversion_preserves_volume(box_name, boxlo, boxhi, xy, yz, xz): """ Test that restricted_cell_from_lammps_box preserves cell volume. """ cell = restricted_cell_from_lammps_box(boxlo, boxhi, xy, yz, xz) cell_np = cell.squeeze().numpy() # Expected volume: lx * ly * lz (for restricted triclinic) lx = boxhi[0] - boxlo[0] ly = boxhi[1] - boxlo[1] lz = boxhi[2] - boxlo[2] expected_volume = lx * ly * lz actual_volume = np.abs(np.linalg.det(cell_np)) assert np.isclose(expected_volume, actual_volume, atol=1e-6), ( f"Volume mismatch for {box_name}: boxlo={boxlo}, boxhi={boxhi}, xy={xy}, yz={yz}, xz={xz}.\n" f"Expected: {expected_volume}, Actual: {actual_volume}" ) def test_lammps_hydra_entrypoint_startup(monkeypatch, tmp_path): """ Test that the Hydra-decorated LAMMPS entry point composes config and starts. """ hydra.core.global_hydra.GlobalHydra.instance().clear() lammps_state = {} run_args = {} def fake_instantiate(_cfg): return object() def fake_run_lammps_with_fairchem( predictor, lammps_input_path, task_name, charge=0, spin=0, ): run_args.update( { "predictor": predictor, "lammps_input_path": lammps_input_path, "task_name": task_name, "charge": charge, "spin": spin, } ) class DummyLammps: pass lmp = DummyLammps() lmp._predictor = predictor lammps_state["lmp"] = lmp return lmp monkeypatch.setattr(lammps_fc.hydra.utils, "instantiate", fake_instantiate) monkeypatch.setattr( lammps_fc, "run_lammps_with_fairchem", fake_run_lammps_with_fairchem ) old_argv = sys.argv[:] try: sys.argv = [ "lammps_fc.py", f"hydra.run.dir={tmp_path}", "hydra.output_subdir=null", ] lammps_fc.main() finally: sys.argv = old_argv hydra.core.global_hydra.GlobalHydra.instance().clear() assert run_args["lammps_input_path"] == "lammps_in_example.file" assert run_args["task_name"] == "omol" assert run_args["charge"] == 0 assert run_args["spin"] == 0 assert not hasattr(lammps_state["lmp"], "_predictor")