atombatch 0.1.1__py3-none-any.whl

atombatch/__init__.py

@@ -0,0 +1,40 @@
+"""Batch energy and force evaluations from many independent ASE tasks."""
+
+from .backends import (
+    Backend,
+    FunctionBackend,
+    SequentialBackend,
+    SymmetrixBackend,
+    TorchSimBackend,
+)
+from .scheduler import BatchCalculator, Scheduler, TaskContext, TaskResult
+from .tasks import (
+    displacement_forces,
+    equation_of_state,
+    molecular_dynamics,
+    neb,
+    relax,
+    single_point,
+    swap_monte_carlo,
+)
+
+__version__ = "0.1.1"
+
+__all__ = [
+    "Backend",
+    "BatchCalculator",
+    "FunctionBackend",
+    "Scheduler",
+    "SequentialBackend",
+    "SymmetrixBackend",
+    "TaskContext",
+    "TaskResult",
+    "TorchSimBackend",
+    "displacement_forces",
+    "equation_of_state",
+    "molecular_dynamics",
+    "neb",
+    "relax",
+    "single_point",
+    "swap_monte_carlo",
+]

atombatch/backends.py

@@ -0,0 +1,178 @@
+"""Evaluation backends: map a list of structures to a list of results."""
+
+import copy
+import sys
+
+import numpy as np
+from ase.calculators.calculator import all_changes
+from ase.stress import full_3x3_to_voigt_6_stress
+
+from .symmetrix import GraphBatcher, supports_graph_batching
+
+
+class Backend:
+    """Evaluate a list of ASE structures with one model.
+
+    ``evaluate`` returns one result dictionary per structure with at least
+    ``energy`` and ``forces``, plus ``stress`` (ASE Voigt order) when
+    requested. ``max_atoms`` bounds the number of atoms in one call; ``None``
+    means unbounded.
+    """
+
+    max_atoms = None
+
+    def evaluate(self, structures, properties):
+        raise NotImplementedError
+
+
+class SequentialBackend(Backend):
+    """Evaluate structures one at a time with any ASE calculator.
+
+    This provides scheduling for calculators without batched evaluation; it
+    does not reduce the per-call cost.
+    """
+
+    def __init__(self, calculator, max_atoms=None):
+        self.calculator = calculator
+        self.max_atoms = max_atoms
+
+    def evaluate(self, structures, properties):
+        results = []
+        for atoms in structures:
+            self.calculator.calculate(atoms.copy(), list(properties), all_changes)
+            results.append(
+                {
+                    key: np.array(value, copy=True)
+                    if isinstance(value, np.ndarray)
+                    else copy.copy(value)
+                    for key, value in self.calculator.results.items()
+                }
+            )
+        return results
+
+
+class SymmetrixBackend(Backend):
+    """Evaluate structures with a Symmetrix calculator in one native call.
+
+    Releases with ``Symmetrix.calculate_batch`` evaluate the batch themselves.
+    For earlier releases (from Symmetrix-XL 0.1.1), atombatch joins the
+    structures into one disconnected neighbor graph and calls the evaluator
+    directly (:mod:`atombatch.symmetrix`). Models with dispersion corrections
+    or electric-field coupling fall back to one ASE evaluation per structure.
+    """
+
+    def __init__(self, calculator, max_atoms=None, batch_mode="native"):
+        self.calculator = calculator
+        self.max_atoms = max_atoms
+        self.batch_mode = batch_mode
+        self._graph = None
+        if callable(getattr(calculator, "calculate_batch", None)):
+            return
+        if not supports_graph_batching(calculator):
+            raise TypeError(
+                "SymmetrixBackend requires a Symmetrix calculator "
+                "(Symmetrix-XL 0.1.1 or later)."
+            )
+        self._graph = GraphBatcher(calculator)
+
+    def evaluate(self, structures, properties):
+        if self._graph is not None:
+            return self._graph.evaluate(structures, properties)
+        return self.calculator.calculate_batch(
+            structures, properties=tuple(properties), batch_mode=self.batch_mode
+        )
+
+
+class TorchSimBackend(Backend):
+    """Evaluate structures in one call of a TorchSim ``ModelInterface``.
+
+    TorchSim provides batched interfaces for MACE, fairchem, SevenNet, ORB,
+    MatterSim, NequIP, Nequix, metatomic, and classical potentials. The
+    structures are joined into one ``SimState`` on the model's device and
+    dtype, and the outputs are split back per structure. Stress requires a
+    model constructed with ``compute_stress=True``.
+    """
+
+    def __init__(self, model, max_atoms=None):
+        self.model = model
+        self.max_atoms = max_atoms
+
+    def evaluate(self, structures, properties):
+        import torch_sim
+
+        if "stress" in properties and not self.model.compute_stress:
+            raise ValueError(
+                "Stress was requested from a TorchSim model constructed "
+                "without compute_stress=True."
+            )
+        state = torch_sim.io.atoms_to_state(
+            list(structures), device=self.model.device, dtype=self.model.dtype
+        )
+        output = self.model(state)
+        energies = _numpy(output["energy"]).reshape(-1)
+        forces = _numpy(output["forces"])
+        stresses = _numpy(output["stress"]) if "stress" in properties else None
+        results, start = [], 0
+        for index, atoms in enumerate(structures):
+            stop = start + len(atoms)
+            result = {
+                "energy": float(energies[index]),
+                "forces": forces[start:stop].copy(),
+            }
+            if stresses is not None:
+                result["stress"] = full_3x3_to_voigt_6_stress(stresses[index])
+            results.append(result)
+            start = stop
+        return results
+
+
+class FunctionBackend(Backend):
+    """Evaluate structures with a user function of the backend signature.
+
+    ``function(structures, properties)`` returns one result dictionary per
+    structure, as :meth:`Backend.evaluate` does. This connects any model with
+    its own batched evaluation without subclassing :class:`Backend`.
+    """
+
+    def __init__(self, function, max_atoms=None):
+        self.function = function
+        self.max_atoms = max_atoms
+
+    def evaluate(self, structures, properties):
+        return self.function(structures, properties)
+
+
+def _numpy(tensor):
+    return tensor.detach().cpu().numpy()
+
+
+def _is_torch_sim_model(model):
+    # Checked only when TorchSim is already imported, so that ASE-only runs
+    # do not import torch.
+    interface = sys.modules.get("torch_sim.models.interface")
+    return interface is not None and isinstance(model, interface.ModelInterface)
+
+
+def as_backend(model, max_atoms=None):
+    """Return ``model`` as a backend.
+
+    Backends are used as given. Symmetrix calculators and TorchSim models
+    are evaluated in batched calls. Other ASE calculators
+    are evaluated one structure at a time, and plain functions are wrapped in
+    :class:`FunctionBackend`.
+    """
+    if isinstance(model, Backend):
+        if max_atoms is not None:
+            model.max_atoms = max_atoms
+        return model
+    if callable(getattr(model, "calculate_batch", None)) or supports_graph_batching(
+        model
+    ):
+        return SymmetrixBackend(model, max_atoms=max_atoms)
+    if _is_torch_sim_model(model):
+        return TorchSimBackend(model, max_atoms=max_atoms)
+    if callable(getattr(model, "calculate", None)):
+        return SequentialBackend(model, max_atoms=max_atoms)
+    if callable(model):
+        return FunctionBackend(model, max_atoms=max_atoms)
+    raise TypeError(f"Cannot use {type(model).__name__} as a model backend.")

atombatch/scheduler.py

@@ -0,0 +1,296 @@
+"""Run many ASE tasks and batch their energy and force requests per model."""
+
+import threading
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from ase.calculators.calculator import Calculator, all_changes
+
+from .backends import as_backend
+
+_REQUIRED = ("energy", "forces")
+
+
+@dataclass
+class TaskResult:
+    """Outcome of one task: its return value, or the exception it raised."""
+
+    name: str
+    value: Any = None
+    error: BaseException | None = None
+
+    @property
+    def ok(self):
+        return self.error is None
+
+
+@dataclass
+class _Request:
+    model: str
+    atoms: Any
+    properties: tuple
+    result: dict | None = None
+    error: BaseException | None = None
+
+
+@dataclass
+class _Call:
+    """Requests one task thread waits on together."""
+
+    requests: list
+    event: threading.Event = field(default_factory=threading.Event)
+
+
+class BatchCalculator(Calculator):
+    """ASE calculator whose evaluations are batched by a :class:`Scheduler`.
+
+    Each instance belongs to one task and one model. ``calculate`` submits a
+    request and blocks the task's thread until the scheduler has evaluated
+    it together with the requests of the other tasks.
+
+    Calculators created together with :meth:`TaskContext.linked_calculators`
+    form a group: a request from one member also evaluates every other member
+    whose structure changed, so a task that queries several structures in
+    turn, such as an NEB band, submits them as one request.
+    """
+
+    implemented_properties = [  # noqa: RUF012 (ASE calculator convention)
+        "energy",
+        "free_energy",
+        "energies",
+        "forces",
+        "stress",
+    ]
+
+    def __init__(self, scheduler, model, properties=_REQUIRED):
+        super().__init__()
+        self._scheduler = scheduler
+        self._model = model
+        self._properties = tuple(dict.fromkeys((*_REQUIRED, *properties)))
+        self._group = None
+
+    def calculate(self, atoms=None, properties=("energy",), system_changes=all_changes):
+        super().calculate(atoms, properties, system_changes)
+        requested = tuple(dict.fromkeys((*self._properties, *properties)))
+        members = [(self, atoms)]
+        for calculator, other in self._group or ():
+            if calculator is self or other.calc is not calculator:
+                continue
+            stale = calculator.check_state(other) or any(
+                prop not in calculator.results for prop in requested
+            )
+            if stale:
+                members.append((calculator, other))
+        results = self._scheduler._evaluate(
+            [(calculator._model, other, requested) for calculator, other in members]
+        )
+        for (calculator, other), result in zip(members, results):
+            if calculator is not self:
+                calculator.atoms = other.copy()
+            calculator.results = result
+
+
+class TaskContext:
+    """Handle given to a task function to request batched evaluations."""
+
+    def __init__(self, scheduler, model):
+        self._scheduler = scheduler
+        self.model = model
+
+    def calculator(self, model=None, properties=_REQUIRED):
+        """Return a new batched calculator for ``model`` (default: the task's).
+
+        Every request from the calculator includes ``properties``, so a task
+        that reads several properties at each state (for example forces and
+        stress under a cell filter) is evaluated once per state.
+        """
+        model = self.model if model is None else model
+        if model not in self._scheduler._backends:
+            raise KeyError(f"Unknown model {model!r}.")
+        return BatchCalculator(self._scheduler, model, properties)
+
+    def evaluate(self, structures, model=None, properties=_REQUIRED):
+        """Evaluate independent structures in one request and return the results.
+
+        Use this when a task needs many evaluations that do not depend on each
+        other, such as displaced supercells or trial configurations: they join
+        the same batched call instead of one call per round. Each result is a
+        dictionary with ``energy``, ``free_energy``, ``forces``, and any further
+        ``properties``.
+        """
+        model = self.model if model is None else model
+        if model not in self._scheduler._backends:
+            raise KeyError(f"Unknown model {model!r}.")
+        if not structures:
+            return []
+        requested = tuple(dict.fromkeys((*_REQUIRED, *properties)))
+        return self._scheduler._evaluate(
+            [(model, atoms, requested) for atoms in structures]
+        )
+
+    def linked_calculators(self, structures, model=None, properties=_REQUIRED):
+        """Attach one grouped calculator to each structure and return them."""
+        calculators = [self.calculator(model, properties) for _ in structures]
+        group = list(zip(calculators, structures))
+        for calculator, atoms in group:
+            calculator._group = group
+            atoms.calc = calculator
+        return calculators
+
+
+@dataclass
+class _Task:
+    name: str
+    function: Callable
+    model: str
+    thread: threading.Thread | None = None
+    result: TaskResult | None = None
+
+
+class Scheduler:
+    """Run ASE tasks in worker threads and batch their model evaluations.
+
+    Each task runs ordinary ASE code with :class:`BatchCalculator` instances.
+    When every running task is waiting for a result, the scheduler evaluates
+    all pending requests in batched calls, one or more per model, and resumes
+    the tasks. At most ``max_active`` tasks run at once; finished tasks are
+    replaced from the queue so that batches stay full.
+    """
+
+    def __init__(self, max_active=None):
+        self.max_active = max_active
+        self._backends = {}
+        self._queue = deque()
+        self._tasks = []
+        self._active = set()
+        self._pending = []
+        self._waiting = 0
+        self._condition = threading.Condition()
+        self.batch_log = []
+
+    def add_model(self, name, model, max_atoms=None):
+        """Register a model under ``name``.
+
+        ``model`` is a :class:`~atombatch.Backend`, a Symmetrix calculator with
+        ``calculate_batch``, a TorchSim model, a function with the
+        :meth:`Backend.evaluate` signature, or any other ASE calculator, which
+        is evaluated one structure at a time. ``max_atoms`` bounds the atoms in
+        one batched call.
+        """
+        self._backends[name] = as_backend(model, max_atoms=max_atoms)
+
+    def submit(self, function, model, name=None):
+        """Queue ``function(context)``; its return value becomes the result."""
+        if model not in self._backends:
+            raise KeyError(f"Unknown model {model!r}; register it with add_model.")
+        task = _Task(name or f"task-{len(self._tasks)}", function, model)
+        self._tasks.append(task)
+        self._queue.append(task)
+        return task.name
+
+    def run(self):
+        """Run all queued tasks and return their results in submission order."""
+        tasks = list(self._queue)
+        self._admit()
+        while True:
+            with self._condition:
+                self._condition.wait_for(lambda: len(self._active) == self._waiting)
+                if not self._active and not self._queue:
+                    break
+                calls, self._pending = self._pending, []
+            if calls:
+                self._serve(calls)
+            self._admit()
+        return [task.result for task in tasks]
+
+    def _admit(self):
+        with self._condition:
+            while self._queue and (
+                self.max_active is None or len(self._active) < self.max_active
+            ):
+                task = self._queue.popleft()
+                self._active.add(task.name)
+                task.thread = threading.Thread(
+                    target=self._run_task, args=(task,), name=task.name, daemon=True
+                )
+                task.thread.start()
+
+    def _run_task(self, task):
+        result = TaskResult(task.name)
+        task.result = result
+        try:
+            result.value = task.function(TaskContext(self, task.model))
+        except Exception as error:  # noqa: BLE001 (reported in the task result)
+            result.error = error
+        finally:
+            with self._condition:
+                self._active.discard(task.name)
+                self._condition.notify_all()
+
+    def _evaluate(self, items):
+        call = _Call(
+            [
+                _Request(model, atoms.copy(), properties)
+                for model, atoms, properties in items
+            ]
+        )
+        with self._condition:
+            self._pending.append(call)
+            self._waiting += 1
+            self._condition.notify_all()
+        call.event.wait()
+        for request in call.requests:
+            if request.error is not None:
+                raise request.error
+        return [request.result for request in call.requests]
+
+    def _serve(self, calls):
+        requests = [request for call in calls for request in call.requests]
+        groups = {}
+        for request in requests:
+            stress = "stress" in request.properties
+            groups.setdefault((request.model, stress), []).append(request)
+        for (model, stress), group in groups.items():
+            backend = self._backends[model]
+            properties = (*_REQUIRED, "stress") if stress else _REQUIRED
+            for chunk in _chunks(group, backend.max_atoms):
+                self.batch_log.append(
+                    (model, len(chunk), sum(len(r.atoms) for r in chunk))
+                )
+                try:
+                    results = backend.evaluate([r.atoms for r in chunk], properties)
+                    if len(results) != len(chunk):
+                        raise RuntimeError(
+                            f"Backend {model!r} returned {len(results)} results "
+                            f"for {len(chunk)} structures."
+                        )
+                except Exception as error:  # noqa: BLE001 (raised in each task)
+                    for request in chunk:
+                        request.error = error
+                    continue
+                for request, result in zip(chunk, results):
+                    result = dict(result)
+                    result.setdefault("free_energy", result["energy"])
+                    request.result = result
+        with self._condition:
+            self._waiting -= len(calls)
+        for call in calls:
+            call.event.set()
+
+
+def _chunks(requests, max_atoms):
+    if max_atoms is None:
+        yield requests
+        return
+    chunk, atoms = [], 0
+    for request in requests:
+        size = len(request.atoms)
+        if chunk and atoms + size > max_atoms:
+            yield chunk
+            chunk, atoms = [], 0
+        chunk.append(request)
+        atoms += size
+    if chunk:
+        yield chunk