lorem-jax 0.1.0__py3-none-any.whl

lorem/__init__.py

@@ -0,0 +1,21 @@
+from opsis.comms import Comms
+
+from ._version import __version__
+from .batching import to_batch, to_sample
+from .calculator import Calculator as LOREMCalculator
+from .models.bec import LoremBEC
+from .models.mlip import Lorem
+from .transforms import ToBatch, ToSample
+
+comms = Comms("lorem")
+
+__all__ = [
+    "__version__",
+    "Lorem",
+    "LoremBEC",
+    "LOREMCalculator",
+    "to_sample",
+    "to_batch",
+    "ToSample",
+    "ToBatch",
+]

lorem/_install_ipi_driver.py

@@ -0,0 +1,31 @@
+"""Helper to install the LOREM driver into i-PI."""
+
+import shutil
+from importlib.util import find_spec
+from pathlib import Path
+
+
+def install_ipi_driver():
+    """Copy the installed `lorem.py` into the i-PI `pes` directory."""
+
+    # Locate i-PI installation
+    ipi_spec = find_spec("ipi")
+    if ipi_spec is None or not ipi_spec.submodule_search_locations:
+        raise RuntimeError(
+            "i-PI installation not found. Install i-PI and rerun "
+            "`lorem-install-ipi-driver`.",
+        )
+    pes_dir = Path(ipi_spec.submodule_search_locations[0]) / "pes"
+
+    # Locate source driver module
+    source_spec = find_spec("lorem.ipi")
+    if source_spec is None or source_spec.origin is None:
+        raise FileNotFoundError("Could not locate lorem.ipi module")
+    source_path = Path(source_spec.origin)
+
+    target_path = pes_dir / "lorem.py"
+    shutil.copy(source_path, target_path)
+
+
+if __name__ == "__main__":
+    install_ipi_driver()

lorem/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

lorem/batching.py

@@ -0,0 +1,113 @@
+import numpy as np
+
+from collections import namedtuple
+
+from jaxpme.batched_mixed.batching import get_batch as jaxpme_batcher
+from jaxpme.batched_mixed.batching import prepare as jaxpme_prepare
+from marathon.data.batching import batch_labels
+from marathon.data.properties import DEFAULT_PROPERTIES
+from marathon.data.sample import to_labels
+from marathon.utils import next_size
+
+Batch = namedtuple(
+    "Batch",
+    (
+        "atomic_numbers",
+        "sr",
+        "nopbc",
+        "pbc",
+        "labels",
+    ),
+)
+
+
+Sample = namedtuple(
+    "Sample",
+    (
+        "structure",
+        "labels",
+    ),
+)
+
+
+def to_batch(
+    samples,
+    keys,
+    batch_size=None,
+    strategies={"default": "powers_of_2"},
+    shapes=None,
+    properties=DEFAULT_PROPERTIES,
+):
+    if batch_size is not None:
+        assert batch_size > len(samples)
+    else:
+        batch_size = next_size(len(samples) + 1, strategy="powers_of_2")
+
+    labels, structures = [], []
+
+    for sample in samples:
+        labels.append(sample.labels)
+        structures.append(sample.structure)
+
+    if shapes is None:
+        default = strategies.pop("default", "powers_of_2")
+        _, sr, nopbc, pbc = jaxpme_batcher(
+            structures,
+            strategy=default,
+            num_structures_pbc=strategies.get("fine", default),
+            num_pairs_nonpbc=strategies.get("coarse", default),
+            num_pairs=strategies.get("coarse", default),
+            num_structures=batch_size,
+        )
+    else:
+        kwargs = {
+            "num_structures": batch_size,
+            "num_structures_pbc": shapes["pbc"],
+            "num_atoms": shapes["atoms"],
+            "num_atoms_pbc": shapes["atoms_pbc"],
+            "num_pairs": shapes["pairs"],
+            "num_pairs_nonpbc": shapes["pairs_nonpbc"],
+            "num_k": shapes["k"],
+            "strategy": "multiples",
+        }
+        _, sr, nopbc, pbc = jaxpme_batcher(
+            structures,
+            **kwargs,
+        )
+
+    num_structures = sr.cell.shape[0]
+    num_atoms = sr.positions.shape[0]
+
+    atomic_numbers = np.zeros(num_atoms, dtype=int)
+    Z = np.concatenate([sample.structure["atomic_numbers"] for sample in samples])
+    atomic_numbers[: len(Z)] = Z
+
+    labels = batch_labels(labels, num_structures, num_atoms, keys, properties=properties)
+
+    return Batch(atomic_numbers, sr, nopbc, pbc, labels)
+
+
+def to_sample(
+    atoms,
+    cutoff,
+    keys=None,
+    energy=True,
+    forces=True,
+    stress=False,
+    lr_wavelength=None,
+    smearing=None,
+    properties=DEFAULT_PROPERTIES,
+):
+    structure = jaxpme_prepare(
+        atoms, cutoff, lr_wavelength=lr_wavelength, smearing=smearing, dtype=np.float32
+    )
+    labels = to_labels(
+        atoms,
+        keys=keys,
+        energy=energy,
+        forces=forces,
+        stress=stress,
+        properties=properties,
+    )
+
+    return Sample(structure, labels)

lorem/calculator.py

@@ -0,0 +1,245 @@
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+jax.config.update("jax_default_matmul_precision", "float32")
+
+from ase.calculators.calculator import (
+    BaseCalculator,
+    PropertyNotImplementedError,
+)
+from marathon.emit.checkpoint import read_msgpack
+from marathon.io import from_dict, read_yaml
+
+from lorem.neighborlist import NeighborListCache
+
+
+class Calculator(BaseCalculator):
+    name = "lorem"
+    parameters = {}
+
+    def todict(self):
+        return self.parameters
+
+    # default; rebuilt per-instance in __init__ from what the model predicts
+    implemented_properties = ["energy", "forces"]
+
+    def __init__(
+        self,
+        pred_fn,
+        species_weights,
+        params,
+        cutoff,
+        atoms=None,
+        stress=False,
+        bec=False,
+        add_offset=True,
+        double_precision=False,
+        skin=0.25,
+    ):
+        self.params = params
+        self.cutoff = cutoff
+        self.skin = skin
+        self.add_offset = add_offset
+        self.double_precision = double_precision
+
+        self._nl_cache = NeighborListCache(skin=skin)
+
+        # only advertise what the model actually produces
+        self.implemented_properties = ["energy", "forces"]
+        if stress:
+            self.implemented_properties.append("stress")
+        if bec:
+            self.implemented_properties.append("born_effective_charges")
+
+        predict_fn = lambda params, batch: pred_fn(params, batch, stress=stress)
+
+        self.predict_fn = jax.jit(predict_fn)
+        self.species_weights = species_weights
+
+        self.atoms = None
+        self.batch = None
+        self.results = {}
+        if atoms is not None:
+            self.setup(atoms)
+
+    @classmethod
+    def from_model(cls, model, params=None, species_weights=None, **kwargs):
+        """Create a Calculator from a Lorem model instance."""
+        if params is None:
+            params = model.init(jax.random.key(0), *model.dummy_inputs())
+        if species_weights is None:
+            species_weights = {}
+            kwargs.setdefault("add_offset", False)
+        kwargs.setdefault("bec", _model_predicts_bec(model))
+        return cls(model.predict, species_weights, params, model.cutoff, **kwargs)
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        folder,
+        **kwargs,
+    ):
+        from pathlib import Path
+
+        folder = Path(folder)
+
+        model = from_dict(read_yaml(folder / "model/model.yaml"))
+
+        _ = model.init(jax.random.key(1), *model.dummy_inputs())
+
+        baseline = read_yaml(folder / "model/baseline.yaml")
+        species_to_weight = baseline["elemental"]
+
+        params = read_msgpack(folder / "model/model.msgpack")
+
+        kwargs.setdefault("bec", _model_predicts_bec(model))
+        return cls(model.predict, species_to_weight, params, model.cutoff, **kwargs)
+
+    def update(self, atoms):
+        if self._nl_cache.needs_update(atoms):
+            # Structural change or combined displacement beyond skin
+            self.results = {}
+            self.atoms = atoms.copy()
+            self.setup(atoms)
+        elif self.atoms is None or not self._geometry_unchanged(atoms):
+            # Positions and/or cell changed but within skin budget
+            self.results = {}
+            self.atoms = atoms.copy()
+            self._update_geometry(atoms)
+
+    def _geometry_unchanged(self, atoms):
+        return np.array_equal(
+            atoms.get_positions(), self.atoms.get_positions()
+        ) and np.array_equal(atoms.get_cell()[:], self.atoms.get_cell()[:])
+
+    def setup(self, atoms):
+        from lorem.batching import to_batch, to_sample
+
+        nl_cutoff = self.cutoff + self.skin
+
+        # Derive Ewald parameters from physical cutoff so the long-range
+        # decomposition is unchanged when using the extended cutoff.
+        lr_wavelength = self.cutoff / 8.0
+        smearing = lr_wavelength * 2.0
+
+        sample = to_sample(
+            atoms,
+            nl_cutoff,
+            lr_wavelength=lr_wavelength,
+            smearing=smearing,
+            energy=False,
+            forces=False,
+            stress=False,
+        )
+        batch = to_batch([sample], [])
+        self.batch = jax.tree.map(lambda x: jnp.array(x), batch)
+
+        max_cell_shift = int(np.abs(np.array(self.batch.sr.cell_shifts)).max())
+        self._nl_cache.save_reference(atoms, max_cell_shift=max_cell_shift)
+
+    def _update_geometry(self, atoms):
+        """Update positions and cell in cached batch without rebuilding
+        the neighbor list. The model recomputes R_ij from the updated
+        sr.positions and sr.cell, and the Ewald calculator recomputes
+        k-vectors from sr.cell (pbc.k_grid stores only integer frequency
+        indices). So forces, energy, and stress remain correct."""
+        sr = self.batch.sr
+        n_atoms = len(atoms)
+
+        positions = np.zeros(np.array(sr.positions).shape, dtype=np.float32)
+        positions[:n_atoms] = atoms.get_positions()
+
+        cell = np.array(sr.cell)
+        new_cell = atoms.get_cell()[:].astype(np.float32)
+        if atoms.get_pbc().sum() == 2:
+            from jaxpme.batched_mixed.batching import shrink_2d_cell
+
+            new_cell = shrink_2d_cell(new_cell, atoms.get_pbc(), positions[:n_atoms])
+        cell[0] = new_cell
+
+        new_sr = sr._replace(
+            positions=jnp.array(positions),
+            cell=jnp.array(cell),
+        )
+        self.batch = self.batch._replace(sr=new_sr)
+
+    def calculate(
+        self,
+        atoms=None,
+        properties=None,
+        system_changes=None,
+        **kwargs,
+    ):
+        self.update(atoms)
+
+        results = self.predict_fn(self.params, self.batch)
+
+        actual_results = {}
+        for key in self.implemented_properties:
+            if key == "energy":
+                actual_results[key] = float(
+                    results[key][self.batch.sr.structure_mask].squeeze()
+                )
+            elif key == "forces":
+                actual_results[key] = np.array(
+                    results[key][self.batch.sr.atom_mask].reshape(-1, 3),
+                    dtype=np.float32,
+                )
+            elif key == "stress":
+                virial = np.array(
+                    results[key][self.batch.sr.structure_mask].reshape(3, 3),
+                    dtype=np.float32,
+                )
+                volume = atoms.get_volume()
+                from ase.stress import full_3x3_to_voigt_6_stress
+
+                actual_results[key] = full_3x3_to_voigt_6_stress(virial / volume)
+
+        # BEC passthrough: when model outputs "apt" (e.g. LoremBEC), expose as
+        # "born_effective_charges" in (natoms, 3, 3) layout for ase compatibility
+        if "apt" in results:
+            actual_results["born_effective_charges"] = np.array(
+                results["apt"][self.batch.sr.atom_mask].reshape(-1, 3, 3),
+                dtype=np.float32,
+            )
+
+        if self.add_offset:
+            energy_offset = np.sum(
+                [self.species_weights[Z] for Z in atoms.get_atomic_numbers()]
+            )
+            actual_results["energy"] += energy_offset
+
+        self.results = actual_results
+        return actual_results
+
+    def get_property(self, name, atoms=None, allow_calculation=True):
+        if name not in self.implemented_properties:
+            raise PropertyNotImplementedError(f"{name} property not implemented")
+
+        self.update(atoms)
+
+        if name not in self.results:
+            if not allow_calculation:
+                return None
+            self.calculate(atoms=atoms)
+
+        if name not in self.results:
+            # For some reason the calculator was not able to do what we want,
+            # and that is OK.
+            raise PropertyNotImplementedError(f"{name} property not present in results!")
+
+        result = self.results[name]
+        if isinstance(result, np.ndarray):
+            result = result.copy()
+        return result
+
+    def get_potential_energy(self, atoms=None, force_consistent=True):
+        # force_consistent is ignored; we are always consistent
+        return self.get_property(name="energy", atoms=atoms)
+
+
+def _model_predicts_bec(model):
+    from lorem.models.bec import LoremBEC
+
+    return isinstance(model, LoremBEC)

lorem/ipi.py

@@ -0,0 +1,42 @@
+"""LOREM driver for i-PI."""
+
+import numpy as np
+
+import json
+
+from ipi.pes._ase import ASEDriver
+
+from lorem.calculator import Calculator
+
+__DRIVER_NAME__ = "lorem"
+__DRIVER_CLASS__ = "LOREM_driver"
+
+
+class LOREM_driver(ASEDriver):
+    def __init__(self, template, model_path, *args, skin=0.25, **kwargs):
+        self.model_path = model_path
+        self.skin = skin
+        super().__init__(template, *args, **kwargs)
+        # only advertise BEC if the model predicts it
+        if "born_effective_charges" in self.ase_calculator.implemented_properties:
+            self.capabilities.append("born_effective_charges")
+
+    def check_parameters(self):
+        super().check_parameters()
+        has_stress = "stress" in self.capabilities
+        self.ase_calculator = Calculator.from_checkpoint(
+            self.model_path, stress=has_stress, skin=self.skin
+        )
+
+    def compute_structure(self, cell, pos):
+        pot_ipi, force_ipi, vir_ipi, extras = super().compute_structure(cell, pos)
+
+        # BEC passthrough: when model outputs "born_effective_charges" (e.g. LoremBEC),
+        # expose as "BEC" in (3*natoms, 3) layout for i-PI compatibility
+        if extras and "born_effective_charges" in extras:
+            extras = json.loads(extras)
+            BEC = np.reshape(extras.pop("born_effective_charges"), (-1, 3))
+            extras["BEC"] = BEC.tolist()
+            extras = json.dumps(extras)
+
+        return pot_ipi, force_ipi, vir_ipi, extras

lorem/models/__init__.py

@@ -0,0 +1,4 @@
+from .bec import LoremBEC
+from .mlip import Lorem
+
+__all__ = ["Lorem", "LoremBEC"]