Skip to content

Commit 56da4dc

Browse files
LuthafLuthaf
committed
Split ASE integration into its own package
1 parent a06e5b2 commit 56da4dc

27 files changed

Lines changed: 3078 additions & 2616 deletions

‎.github/workflows/ase-tests.yml‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
name: ASE integration tests
2+
3+
on:
4+
push:
5+
branches: [main]
6+
pull_request:
7+
# Check all PR
8+
9+
concurrency:
10+
group: ase-tests-${{ github.ref }}
11+
cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
12+
13+
jobs:
14+
tests:
15+
runs-on: ubuntu-24.04
16+
name: ASE
17+
steps:
18+
- uses: actions/checkout@v6
19+
with:
20+
fetch-depth: 0
21+
22+
- name: setup Python
23+
uses: actions/setup-python@v6
24+
with:
25+
python-version: "3.13"
26+
27+
- name: Setup sccache
28+
uses: mozilla-actions/sccache-action@v0.0.9
29+
with:
30+
version: "v0.10.0"
31+
32+
- name: Setup sccache environnement variables
33+
run: |
34+
echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV
35+
echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV
36+
echo "CMAKE_C_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV
37+
echo "CMAKE_CXX_COMPILER_LAUNCHER=sccache" >> $GITHUB_ENV
38+
39+
- name: install tests dependencies
40+
run: |
41+
python -m pip install --upgrade pip
42+
python -m pip install tox coverage
43+
44+
- name: run tests
45+
run: tox -e ase-tests
46+
env:
47+
PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
48+
49+
- name: combine Python coverage files
50+
shell: bash
51+
run: |
52+
coverage combine .tox/*/.coverage
53+
coverage xml
54+
55+
- name: upload to codecov.io
56+
uses: codecov/codecov-action@v5
57+
with:
58+
fail_ci_if_error: true
59+
files: coverage.xml
60+
token: ${{ secrets.CODECOV_TOKEN }}

‎.github/workflows/torch-tests.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ jobs:
6464
python -m pip install tox coverage
6565
6666
- name: run tests
67-
run: tox
67+
run: tox -e lint,torch-tests,torch-tests-cxx,torch-install-tests-cxx,docs-tests
6868
env:
6969
PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
7070
METATOMIC_TESTS_TORCH_VERSION: ${{ matrix.torch-version }}

‎docs/src/engines/ase.rst‎

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ ASE
1414
Supported model outputs
1515
^^^^^^^^^^^^^^^^^^^^^^^
1616

17-
.. py:currentmodule:: metatomic.torch.ase_calculator
17+
.. py:currentmodule:: metatomic_ase
1818
1919
- the :ref:`energy <energy-output>`, non-conservative :ref:`forces
2020
<non-conservative-forces-output>` and :ref:`stress <non-conservative-stress-output>`
@@ -26,16 +26,37 @@ Supported model outputs
2626
- for non-equivariant architectures like
2727
`PET <https://docs.metatensor.org/metatrain/latest/architectures/pet.html>`_,
2828
rotationally-averaged energies, forces, and stresses can be computed using
29-
:py:class:`metatomic.torch.ase_calculator.SymmetrizedCalculator`.
29+
:py:class:`metatomic_ase.SymmetrizedCalculator`.
3030

3131
How to install the code
3232
^^^^^^^^^^^^^^^^^^^^^^^
3333

34-
The code is available in the ``metatomic-torch`` package, in the
35-
:py:class:`metatomic.torch.ase_calculator.MetatomicCalculator` class.
34+
The code is available in the ``metatomic-ase`` package, which can be installed
35+
using ``pip install metatomic-ase``.
3636

3737
How to use the code
3838
^^^^^^^^^^^^^^^^^^^
3939

40-
See the :ref:`corresponding tutorial <atomistic-tutorial-md>`, and API
41-
documentation of the :py:class:`MetatomicCalculator` class.
40+
We offer two ASE calculators: :py:class:`metatomic_ase.MetatomicCalculator` is
41+
the default one, and support all the features described above, while
42+
:py:class:`metatomic_ase.SymmetrizedCalculator` is a wrapper around the former
43+
that allows to compute rotationally-averaged energies, forces, and stresses for
44+
non-equivariant architectures. Both calculators are designed to be used as
45+
drop-in replacements for any ASE calculator, and can be used in any ASE
46+
workflow. You can also check the :ref:`corresponding tutorial
47+
<atomistic-tutorial-md>`.
48+
49+
.. _ase-integration-api:
50+
51+
API documentation
52+
-----------------
53+
54+
.. _calculator: https://ase-lib.org/ase/calculators/calculators.html
55+
56+
.. autoclass:: metatomic_ase.MetatomicCalculator
57+
:show-inheritance:
58+
:members:
59+
60+
.. autoclass:: metatomic_ase.SymmetrizedCalculator
61+
:show-inheritance:
62+
:members:

‎docs/src/torch/reference/ase.rst‎

Lines changed: 9 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,12 @@
11
Atomic Simulation Environment (ASE) integration
22
===============================================
33

4-
.. py:currentmodule:: metatomic.torch
5-
6-
The code in ``metatomic.torch.ase_calculator`` defines a class that
7-
allows using a :py:class:`AtomisticModel` which predicts the energy and forces of a
8-
system as an ASE `calculator`_; enabling the use of machine learning interatomic
9-
potentials to drive calculations compatible with ASE calculators.
10-
11-
Additionally, it allows using arbitrary models with prediction targets which are
12-
not just the energy, through the
13-
:py:meth:`ase_calculator.MetatomicCalculator.run_model` function.
14-
15-
.. _calculator: https://ase-lib.org/ase/calculators/calculators.html
16-
17-
.. autoclass:: metatomic.torch.ase_calculator.MetatomicCalculator
18-
:show-inheritance:
19-
:members:
20-
21-
.. autoclass:: metatomic.torch.ase_calculator.SymmetrizedCalculator
22-
:show-inheritance:
23-
:members:
4+
The integration of metatomic with the Atomic Simulation Environment (ASE) was
5+
moved into it's own package, ``metatomic-ase``, which is available on PyPI. The
6+
documentation for this package can be found in the :ref:`corresponding section
7+
of the documentation <ase-integration-api>`.
8+
9+
Both calculators classes are re-exported from the
10+
``metatomic.torch.ase_calculator`` module for baclwards compatibility, but users
11+
are encouraged to import them from the ``metatomic_ase`` package instead. The
12+
old import paths will be removed in a future release.

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ ignore = ["B018", "B904"]
7777

7878
[tool.ruff.lint.isort]
7979
lines-after-imports = 2
80-
known-first-party = ["metatomic"]
80+
known-first-party = ["metatomic", "metatomic_ase"]
8181
known-third-party = ["torch"]
8282

8383
[tool.ruff.format]

‎python/examples/2-running-ase-md.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@
4242
)
4343

4444
# Integration with ASE for metatomic models
45-
from metatomic.torch.ase_calculator import MetatomicCalculator
45+
from metatomic_ase import MetatomicCalculator
4646

4747

4848
# %%

‎python/examples/3-atomistic-model-with-nl.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@
5959
NeighborListOptions,
6060
System,
6161
)
62-
from metatomic.torch.ase_calculator import MetatomicCalculator
62+
from metatomic_ase import MetatomicCalculator
6363

6464

6565
# %%

‎python/examples/4-profiling.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
ModelOutput,
3030
System,
3131
)
32-
from metatomic.torch.ase_calculator import MetatomicCalculator
32+
from metatomic_ase import MetatomicCalculator
3333

3434

3535
# %%

‎python/metatomic_ase/AUTHORS‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
../../AUTHORS

‎python/metatomic_ase/README.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
``metatomic_ase``
2+
=================
3+
4+
TODO

0 commit comments

Comments
 (0)