atombatch 0.1.1__tar.gz

atombatch-0.1.1/.github/workflows/publish.yml

@@ -0,0 +1,40 @@
+name: publish
+
+# Publishes to PyPI when a GitHub release is published, using PyPI trusted
+# publishing (no stored token). The PyPI project must list this repository,
+# this workflow file, and the "pypi" environment as a trusted publisher.
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          python-version: "3.12"
+      - name: Check that the tag matches the package version
+        run: |
+          version=$(uv run --no-project python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])")
+          test "v${version}" = "${GITHUB_REF_NAME}"
+      - run: uv build
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

atombatch-0.1.1/.github/workflows/tests.yml

@@ -0,0 +1,34 @@
+name: tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+          activate-environment: true
+      - run: uv pip install -e ".[test,phonons]"
+      - run: uvx ruff check src tests benchmarks
+      - run: pytest -q
+
+  docs:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          python-version: "3.12"
+          activate-environment: true
+      - run: uv pip install -e ".[docs]"
+      - run: sphinx-build -W docs docs/_build/html

atombatch-0.1.1/.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.venv/
+build/
+dist/
+.pytest_cache/
+.ruff_cache/
+docs/_build/

atombatch-0.1.1/AGENTS.md

@@ -0,0 +1,43 @@
+# AGENTS.md
+
+Guidance for coding agents working in this repository. Read `README.md` for the user-facing description and `docs/design.md` for the design decisions behind the code.
+
+## Project
+
+atombatch batches the energy and force requests of many independent ASE tasks (single points, relaxations, NEB, molecular dynamics) into one model call per model. The ASE optimizers and integrators run unmodified, and every task receives exactly the result it would get from an independent run. Preserving that equivalence is the central invariant of the project.
+
+## Layout
+
+- `src/atombatch/scheduler.py`: `Scheduler`, `BatchCalculator`, `TaskContext`, `TaskResult`. Each task runs in its own worker thread. `BatchCalculator.calculate` submits a request and blocks. The main thread waits until every active task is blocked or finished, then `_serve` groups requests by `(model, needs stress)`, splits them into chunks under the backend's `max_atoms`, evaluates each chunk, and wakes the tasks. Finished tasks are replaced from the queue by `_admit`, bounded by `max_active`.
+- `src/atombatch/backends.py`: `Backend` interface (`evaluate(structures, properties) -> list[dict]`), `SequentialBackend` for any ASE calculator, `SymmetrixBackend`, which uses `Symmetrix.calculate_batch` when the installed Symmetrix has it and otherwise the graph batching in `src/atombatch/symmetrix.py` (adapted from Symmetrix-XL, MIT, and limited to evaluator calls present in Symmetrix-XL 0.1.1), `TorchSimBackend` wrapping a TorchSim `ModelInterface` (MACE, fairchem, SevenNet, ORB, MatterSim, NequIP, Nequix, metatomic), and `FunctionBackend` for a user function with the `evaluate` signature. `as_backend` picks the backend by type and duck typing. It recognizes TorchSim models only when `torch_sim` is already imported, so ASE-only runs never import torch.
+- `src/atombatch/tasks.py`: task factories `single_point`, `relax`, `neb`, `molecular_dynamics`, and the batched workflows `displacement_forces` (phonopy and phono3py), `equation_of_state`, and `swap_monte_carlo`. Each returns a closure `task(context)` that copies its input structures and runs plain ASE code with calculators from the context. Workflows with independent evaluations submit them together through `TaskContext.evaluate`.
+- `tests/test_units.py`: fast unit tests of chunking, backend selection (with a fake `torch_sim` module), result splitting, calculator caching, linked calculators, error paths, and Monte Carlo proposals. They use small fake backends and run in under a second.
+- `tests/test_scheduler.py`: tests against EMT, which needs no MLIP. `tests/test_workflows.py` covers `context.evaluate` and the batched workflows (phonopy and phono3py tests are skipped when those packages are absent). `tests/test_torch_sim_backend.py` uses TorchSim's Lennard-Jones model and is skipped when TorchSim is not installed. `tests/test_mace_backend.py` compares batched MACE through TorchSim with MACE's own ASE calculator and runs only when `ATOMBATCH_MACE_MODEL` names a MACE-MP model (for example `small`) and `mace-torch` is installed. `tests/test_symmetrix_backend.py` is skipped unless `ATOMBATCH_SYMMETRIX_MODEL` points to a Symmetrix JSON model and the installed Symmetrix provides `calculate_batch`.
+
+## Development
+
+```bash
+uv venv && source .venv/bin/activate
+uv pip install -e ".[test]"          # add ,torchsim on Python 3.12+ for the TorchSim tests
+pytest
+```
+
+Python 3.10 or later, with `numpy` and `ase>=3.23` as the only runtime dependencies. Symmetrix (`.[symmetrix]`) and TorchSim (`.[torchsim]`, Python 3.12 or later) are optional extras and must stay optional: import them inside the backend, never at module level. Ruff is configured with line length 88 in `pyproject.toml`, and existing `noqa` comments carry the reason in parentheses. Follow that convention.
+
+## Documentation
+
+Sphinx with the Furo theme and MyST Markdown, in `docs/` (`conf.py`, guide pages in Markdown, `api.rst` for autodoc). `docs/design.md` is part of the site. Build with `sphinx-build -W docs docs/_build/html`, which must pass without warnings. Code examples in `index.md`, `quickstart.md`, `workflows.md`, and `custom-tasks.md` use EMT so that they run without an MLIP. Run them after changing the API, and keep any printed output quoted in the pages in sync. Docstrings use Sphinx roles (`:class:`, `:meth:`) and double-backtick literals. The logo is `docs/_static/logo.svg`, a flat 3x3 grid of atoms in indigo tones with one amber atom.
+
+## Conventions
+
+- Validate every new task type or scheduling change by comparing with an independent ASE run of the same algorithm on EMT: identical step counts and convergence, and positions, energies, and forces equal within `atol=1e-9` (see `assert_same` in `tests/test_scheduler.py`). The tolerance is nonzero only because EMT reuses its neighbor list in an independent run.
+- A request must be evaluated at exactly the state the task asked for. `_evaluate` copies `atoms` at submission time, so later mutation by the task cannot leak into the batch.
+- A workflow whose evaluations do not depend on each other must submit them through `context.evaluate` in one request rather than one at a time, otherwise it contributes one structure per round.
+- Tasks that query several structures in turn, such as an NEB band, use `TaskContext.linked_calculators` so that one request covers every stale member of the group. Tasks that read stress (cell filters) pass `properties` including `"stress"` so that each state is evaluated once.
+- Exceptions raised in a task are captured in its `TaskResult.error` and do not stop other tasks. A backend failure is raised in every task of the failing chunk.
+- New MLIP backends are added one at a time, each validated against that model's own ASE calculator. A backend returns one dict per structure with `energy`, `forces`, and `stress` in ASE Voigt order when requested.
+- Public API changes must be reflected in `src/atombatch/__init__.py` (`__all__`), the README example, and `docs/design.md` where the design changes.
+
+## Status
+
+Early development (version 0.1.1). The version is set in both `pyproject.toml` and `src/atombatch/__init__.py`, and `CHANGELOG.md` records each release. Correctness is tested. `benchmarks/bench_batching.py` measures model throughput against batch size and end-to-end relaxations against a sequential ASE loop, and its results are quoted in `docs/performance.md` with the JSON in `benchmarks/results/`. When quoting new numbers, state the hardware, model, precision, structure size, and `max_active`, and repeat CPU runs, which vary by about 20 %.

atombatch-0.1.1/CHANGELOG.md

@@ -0,0 +1,19 @@
+# Changelog
+
+## 0.1.1 (2026-10-15)
+
+- The `symmetrix` extra installs matscipy. Without it, Symmetrix uses ASE's Python neighbor list, and batched relaxations with the `metal-preview-1` wheel ran three times slower than a sequential loop. atombatch now warns when Symmetrix has fallen back to the ASE neighbor list.
+- Tested with the `metal-preview-1` Symmetrix-XL wheels (Metal, no `calculate_batch`): batched results match the Symmetrix calculator, and 64 FIRE relaxations ran 2.0 times faster than a sequential loop with identical step counts.
+- Releases are published to PyPI from GitHub Actions with trusted publishing.
+
+## 0.1.0 (2026-10-15)
+
+First public release.
+
+- Scheduler that runs unmodified ASE tasks in worker threads and evaluates their pending energy and force requests in one batched call per model, with refill from a queue (`max_active`) and a per-model atom budget (`max_atoms`).
+- Tasks: single points, relaxations with any ASE optimizer and cell filter, NEB (regular and climbing image), and molecular dynamics with any ASE integrator.
+- Batched workflows: `TaskContext.evaluate` for independent structures, force sets for phonopy and phono3py (`displacement_forces`), equations of state (`equation_of_state`), and swap Monte Carlo with replica exchange (`swap_monte_carlo`).
+- Backends: Symmetrix-XL from 0.1.1 (native `calculate_batch` when available, otherwise one joined neighbor graph per call), TorchSim models (MACE, fairchem, SevenNet, ORB, MatterSim, NequIP, Nequix, metatomic), user functions, and any ASE calculator evaluated sequentially.
+- Results agree with independent ASE runs and with the models' own ASE calculators to rounding.
+- Benchmarks: on an Apple M1 Max GPU with Symmetrix, 6.2 times higher model throughput at 64 structures per call and 3.0 times faster FIRE relaxations than a sequential ASE loop.
+- Documentation (Sphinx), unit and integration tests, and a benchmark script.