From 472a9a1134a8c137e2c330324be662774d0bf461 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 6 Jan 2026 15:21:08 +0000 Subject: [PATCH] release: create release-0.1.8 branch --- .github/workflows/tests_and_linters.yaml | 1 + CHANGELOG | 7 +++++ pyproject.toml | 2 +- src/mlip/data/graph_dataset_builder.py | 1 + src/mlip/data/helpers/neighborhood.py | 2 +- .../simulation/ase/ase_simulation_engine.py | 5 +++- .../test_batched_inference.py | 19 +++++++++++++ tests/simulation/test_ase_simulation.py | 27 +++++++++++++++++++ 8 files changed, 61 insertions(+), 3 deletions(-) diff --git a/.github/workflows/tests_and_linters.yaml b/.github/workflows/tests_and_linters.yaml index 373ef4f..dd2d26c 100644 --- a/.github/workflows/tests_and_linters.yaml +++ b/.github/workflows/tests_and_linters.yaml @@ -15,6 +15,7 @@ jobs: run: | apt-get update && apt-get install -y \ coreutils \ + gcc \ git git config --system --add safe.directory $GITHUB_WORKSPACE diff --git a/CHANGELOG b/CHANGELOG index 03f0bef..11054b0 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -1,5 +1,12 @@ # Changelog +## Release 0.1.8 + +- Fixing bug in ASE simulation engine to allow for passing Periodic Boundary + Conditions via the box config value. +- Adapting setup line for zero shifts array to use `np.zeros` to guarantee correct + array shape, see issue #36 for reference. + ## Release 0.1.7 - Fixing issues with Periodic Boundary Conditions (PBCs) during inference. diff --git a/pyproject.toml b/pyproject.toml index 475a7ed..d308fef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mlip" -version = "0.1.7" +version = "0.1.8" description = "Machine Learning Interatomic Potentials in JAX" license-files = [ "LICENSE" diff --git a/src/mlip/data/graph_dataset_builder.py b/src/mlip/data/graph_dataset_builder.py index 84785f8..bdff854 100644 --- a/src/mlip/data/graph_dataset_builder.py +++ b/src/mlip/data/graph_dataset_builder.py @@ -103,6 +103,7 @@ def __init__( self._reader = reader self._config = dataset_config self._dataset_info: Optional[DatasetInfo] = dataset_info + self._graphs: Optional[dict[str, list[jraph.GraphsTuple]]] = None self._datasets: Optional[dict[str, Optional[GraphDataset]]] = None # Sanity check when DatasetInfo is passed from the outside cutoff = self._config.graph_cutoff_angstrom diff --git a/src/mlip/data/helpers/neighborhood.py b/src/mlip/data/helpers/neighborhood.py index c00c6ef..6bb733c 100644 --- a/src/mlip/data/helpers/neighborhood.py +++ b/src/mlip/data/helpers/neighborhood.py @@ -104,7 +104,7 @@ def get_neighborhood( ) # If we are not having PBCs, then use shifts of zero - shifts = senders_unit_shifts if any(pbc) else np.array([[0] * 3] * len(senders)) + shifts = senders_unit_shifts if any(pbc) else np.zeros((len(senders), 3)) # See docstring of functions get_edge_relative_vectors() and # get_edge_vectors() on how these return values are used diff --git a/src/mlip/simulation/ase/ase_simulation_engine.py b/src/mlip/simulation/ase/ase_simulation_engine.py index 1030754..60277b4 100644 --- a/src/mlip/simulation/ase/ase_simulation_engine.py +++ b/src/mlip/simulation/ase/ase_simulation_engine.py @@ -92,7 +92,10 @@ def _initialize( def _init_box(self) -> None: """Update the PBC parameters of the underlying `ase.Atoms`""" # Pass if atoms already have PBC and cell, best source of truth - if self.atoms.cell is not None and self.atoms.pbc is not None: + if np.any(self.atoms.cell) or np.any(self.atoms.pbc): + logger.warning( + "Ignoring `box` parameter as `atoms` already has PBC configured." + ) return # Support cubic periodic box from config for Jax-MD consistency. # To be discouraged once both engines support arbitrary lattices. diff --git a/tests/models_inference/test_batched_inference.py b/tests/models_inference/test_batched_inference.py index f9fa50b..ddf79e1 100644 --- a/tests/models_inference/test_batched_inference.py +++ b/tests/models_inference/test_batched_inference.py @@ -74,3 +74,22 @@ def test_batched_inference_works_correctly( assert result[0].pressure is None assert result[0].energy == pytest.approx(-0.11254195, abs=1e-3) assert result[0].forces[0][0] == pytest.approx(0.04921325, abs=1e-3) + + +def test_batched_inference_with_graph_without_edges(setup_system_and_mace_model): + atoms, _, _, mace_ff = setup_system_and_mace_model + + positions = np.array([[0, 0, 0], [0, 0, 10]]) + atomic_numbers = np.array([6, 6]) + atoms_without_edges = ase.Atoms(positions=positions, numbers=atomic_numbers) + + structures = [atoms_without_edges, atoms] + result = run_batched_inference(structures, mace_ff, batch_size=2) + + # For first structure, no energy if no edges + assert result[0].energy == 0.0 + assert not result[0].forces.any() + + # For second structure, normal result + assert result[1].energy == pytest.approx(-0.11254195, abs=1e-3) + assert result[1].forces[0][0] == pytest.approx(0.04921325, abs=1e-3) diff --git a/tests/simulation/test_ase_simulation.py b/tests/simulation/test_ase_simulation.py index 75c0e67..7165add 100644 --- a/tests/simulation/test_ase_simulation.py +++ b/tests/simulation/test_ase_simulation.py @@ -105,3 +105,30 @@ def test_md_can_be_restarted_from_velocities_with_ase_backend( for i in range(1, 5): assert not np.allclose(engine.state.velocities[i], velocities_to_restore) + + +@pytest.mark.parametrize("atoms_cell", [None, np.eye(3) * 10.0]) +def test_ase_engine_sets_cell_from_config( + setup_system_and_mace_model, atoms_cell +) -> None: + atoms, _, _, mace_ff = setup_system_and_mace_model + config_box_length = 25.0 + + _atoms = deepcopy(atoms) + if atoms_cell is None: # If atoms have no cell, use config.box + _atoms.set_cell(None) + _atoms.set_pbc(None) + target_cell = np.eye(3) * config_box_length + else: # If atoms have a cell, ignore config.box + _atoms.set_cell(atoms_cell) + _atoms.set_pbc(True) + target_cell = atoms_cell + + md_config = ASESimulationEngine.Config( + simulation_type=SimulationType.MD, + num_steps=1, + box=config_box_length, + ) + engine = ASESimulationEngine(_atoms, mace_ff, md_config) + assert (engine.atoms.get_cell() == target_cell).all() + assert engine.atoms.get_pbc().all()