lorem-jax 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- lorem_jax-0.1.0/.claude/commands/release.md +48 -0
- lorem_jax-0.1.0/.github/dependabot.yml +11 -0
- lorem_jax-0.1.0/.github/workflows/lint.yml +19 -0
- lorem_jax-0.1.0/.github/workflows/release.yml +40 -0
- lorem_jax-0.1.0/.github/workflows/tests.yml +47 -0
- lorem_jax-0.1.0/.gitignore +18 -0
- lorem_jax-0.1.0/LICENSE +28 -0
- lorem_jax-0.1.0/PKG-INFO +226 -0
- lorem_jax-0.1.0/README.md +204 -0
- lorem_jax-0.1.0/examples/calculator/example.py +49 -0
- lorem_jax-0.1.0/examples/eval/README.md +0 -0
- lorem_jax-0.1.0/examples/md-ase/README.md +0 -0
- lorem_jax-0.1.0/examples/md-ipi/.gitignore +5 -0
- lorem_jax-0.1.0/examples/md-ipi/README.md +52 -0
- lorem_jax-0.1.0/examples/md-ipi/input.xml +55 -0
- lorem_jax-0.1.0/examples/md-ipi/run.sh +30 -0
- lorem_jax-0.1.0/examples/md-ipi/start.xyz +5 -0
- lorem_jax-0.1.0/examples/train-bec/.gitignore +5 -0
- lorem_jax-0.1.0/examples/train-bec/README.md +21 -0
- lorem_jax-0.1.0/examples/train-bec/bec_example.xyz +4462 -0
- lorem_jax-0.1.0/examples/train-bec/my_experiment/model.yaml +8 -0
- lorem_jax-0.1.0/examples/train-bec/my_experiment/settings.yaml +20 -0
- lorem_jax-0.1.0/examples/train-bec/prepare.py +64 -0
- lorem_jax-0.1.0/examples/train-mlp/.gitignore +5 -0
- lorem_jax-0.1.0/examples/train-mlp/README.md +34 -0
- lorem_jax-0.1.0/examples/train-mlp/data.xyz +315 -0
- lorem_jax-0.1.0/examples/train-mlp/my_experiment/model.yaml +8 -0
- lorem_jax-0.1.0/examples/train-mlp/my_experiment/settings.yaml +20 -0
- lorem_jax-0.1.0/examples/train-mlp/my_experiment_finetune/model.yaml +8 -0
- lorem_jax-0.1.0/examples/train-mlp/my_experiment_finetune/settings.yaml +21 -0
- lorem_jax-0.1.0/examples/train-mlp/prepare.py +53 -0
- lorem_jax-0.1.0/pyproject.toml +88 -0
- lorem_jax-0.1.0/setup.cfg +4 -0
- lorem_jax-0.1.0/src/lorem/__init__.py +21 -0
- lorem_jax-0.1.0/src/lorem/_install_ipi_driver.py +31 -0
- lorem_jax-0.1.0/src/lorem/_version.py +24 -0
- lorem_jax-0.1.0/src/lorem/batching.py +113 -0
- lorem_jax-0.1.0/src/lorem/calculator.py +245 -0
- lorem_jax-0.1.0/src/lorem/ipi.py +42 -0
- lorem_jax-0.1.0/src/lorem/models/__init__.py +4 -0
- lorem_jax-0.1.0/src/lorem/models/backbone.py +247 -0
- lorem_jax-0.1.0/src/lorem/models/bec.py +378 -0
- lorem_jax-0.1.0/src/lorem/models/mlip.py +299 -0
- lorem_jax-0.1.0/src/lorem/neighborlist.py +111 -0
- lorem_jax-0.1.0/src/lorem/train.py +1034 -0
- lorem_jax-0.1.0/src/lorem/transforms.py +176 -0
- lorem_jax-0.1.0/src/lorem_jax.egg-info/PKG-INFO +226 -0
- lorem_jax-0.1.0/src/lorem_jax.egg-info/SOURCES.txt +59 -0
- lorem_jax-0.1.0/src/lorem_jax.egg-info/dependency_links.txt +1 -0
- lorem_jax-0.1.0/src/lorem_jax.egg-info/entry_points.txt +3 -0
- lorem_jax-0.1.0/src/lorem_jax.egg-info/requires.txt +11 -0
- lorem_jax-0.1.0/src/lorem_jax.egg-info/scm_file_list.json +54 -0
- lorem_jax-0.1.0/src/lorem_jax.egg-info/scm_version.json +8 -0
- lorem_jax-0.1.0/src/lorem_jax.egg-info/top_level.txt +1 -0
- lorem_jax-0.1.0/tests/test_backbone.py +220 -0
- lorem_jax-0.1.0/tests/test_calculator.py +331 -0
- lorem_jax-0.1.0/tests/test_init.py +14 -0
- lorem_jax-0.1.0/tests/test_ipi.py +47 -0
- lorem_jax-0.1.0/tests/test_neighborlist.py +278 -0
- lorem_jax-0.1.0/tests/test_train.py +42 -0
- lorem_jax-0.1.0/tox.ini +54 -0
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
Create a new release of lorem-jax on PyPI. Optional argument: version bump level (`patch`, `minor`, or `major`).
|
|
2
|
+
|
|
3
|
+
## Steps
|
|
4
|
+
|
|
5
|
+
### 1. Pre-flight checks
|
|
6
|
+
|
|
7
|
+
- Work from a clean checkout of `main` at `origin/main`. If the current checkout is on another branch or has changes, do not switch it; create a worktree instead: `git fetch origin main && git worktree add <scratch>/release origin/main`, and run everything below in there. Remove the worktree at the end.
|
|
8
|
+
- Check that CI is passing on the latest commit of `main`: `gh run list --repo lab-cosmo/lorem-jax --branch main --limit 3`
|
|
9
|
+
- If CI is not green, do NOT proceed — investigate and fix first
|
|
10
|
+
- Check that `pyproject.toml` has no direct-URL dependencies (`@ git+...`). PyPI rejects them, and the publish job would fail after the tag is already pushed. In particular `jax-pme` must be a versioned PyPI dependency.
|
|
11
|
+
|
|
12
|
+
### 2. Run full local verification
|
|
13
|
+
|
|
14
|
+
- Run `uvx tox -e lint` — must pass
|
|
15
|
+
- Run `uvx tox -e tests` — must pass
|
|
16
|
+
- Run `uvx tox -e examples` — must pass (calculator, train-mlp, train-finetune, train-bec)
|
|
17
|
+
- Do NOT proceed if any of these fail
|
|
18
|
+
|
|
19
|
+
### 3. Determine version
|
|
20
|
+
|
|
21
|
+
- Find the latest git tag with `git describe --tags --abbrev=0` (or note if there are no tags yet; the first release is then `0.1.0` unless the user says otherwise)
|
|
22
|
+
- If a bump level was given ($ARGUMENTS), compute the new version following semver (e.g., `0.1.0` → `0.1.1` for patch, `0.2.0` for minor, `1.0.0` for major)
|
|
23
|
+
- If NO bump level was given, review the changes (step 4) first, then discuss with the user what the appropriate level should be based on the nature of the changes (breaking → major, new features → minor, fixes/maintenance → patch)
|
|
24
|
+
- Confirm the new version with the user before proceeding
|
|
25
|
+
|
|
26
|
+
### 4. Review changes and write changelog
|
|
27
|
+
|
|
28
|
+
- Run `git log <last-tag>..HEAD --oneline` to see all commits since the last release (or all commits if no prior tag)
|
|
29
|
+
- Write a brief changelog summarising the changes, grouped by category where appropriate (features, fixes, breaking changes, maintenance, etc.)
|
|
30
|
+
- Present the changelog to the user for review and approval
|
|
31
|
+
|
|
32
|
+
### 5. Tag and push
|
|
33
|
+
|
|
34
|
+
- Create an annotated tag: `git tag -a v<version> -m "Release v<version>"`
|
|
35
|
+
- Push the tag: `git push origin v<version>`
|
|
36
|
+
- This triggers `.github/workflows/release.yml`, which builds the package and publishes it to PyPI via trusted publishing from the `release` environment
|
|
37
|
+
- Watch it: `gh run watch` (or `gh run list --repo lab-cosmo/lorem-jax --workflow release.yml --limit 1`), and confirm the publish job succeeded before continuing
|
|
38
|
+
|
|
39
|
+
### 6. Create GitHub release
|
|
40
|
+
|
|
41
|
+
- Use `gh release create v<version> --repo lab-cosmo/lorem-jax --title "v<version>" --notes "<changelog>"` to create a GitHub release with the changelog from step 4
|
|
42
|
+
- Confirm the new version is visible on PyPI: `curl -s https://pypi.org/pypi/lorem-jax/json | python -c "import json,sys; print(json.load(sys.stdin)['info']['version'])"`
|
|
43
|
+
|
|
44
|
+
## Notes
|
|
45
|
+
|
|
46
|
+
- The version is derived from git tags by `setuptools_scm` (written to `src/lorem/_version.py`, which is gitignored). No files need to be modified for a release.
|
|
47
|
+
- The release workflow uses PyPI trusted publishing (no tokens). The one-time setup, already done, is a `release` environment on the GitHub repo and a trusted publisher for `lorem-jax` on pypi.org pointing at `lab-cosmo/lorem-jax`, `release.yml`, environment `release`.
|
|
48
|
+
- Sandbox: `git fetch`, `git push`, and `gh` must run as bare commands in their own Bash call, with the shell already inside the repo (use a separate `cd` call first). Wrapping them in `cd … &&`, `$(…)`, `;` or a pipe keeps them inside the sandbox, where ssh keys and the gh token are unreachable.
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
version: 2
|
|
2
|
+
updates:
|
|
3
|
+
- package-ecosystem: "github-actions"
|
|
4
|
+
directory: "/"
|
|
5
|
+
schedule:
|
|
6
|
+
interval: monthly
|
|
7
|
+
open-pull-requests-limit: 1
|
|
8
|
+
groups:
|
|
9
|
+
action-dependencies:
|
|
10
|
+
patterns:
|
|
11
|
+
- "*" # A wildcard to create one PR for all dependencies in the ecosystem
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
name: Lint
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
pull_request:
|
|
5
|
+
branches: [main]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
lint:
|
|
9
|
+
runs-on: ubuntu-latest
|
|
10
|
+
|
|
11
|
+
steps:
|
|
12
|
+
- uses: actions/checkout@v7
|
|
13
|
+
- name: Set up Python
|
|
14
|
+
uses: actions/setup-python@v7
|
|
15
|
+
with:
|
|
16
|
+
python-version: "3.13"
|
|
17
|
+
- run: python -m pip install tox
|
|
18
|
+
- name: Lint the code
|
|
19
|
+
run: tox -e lint
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
name: Release
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
tags:
|
|
6
|
+
- "v*"
|
|
7
|
+
|
|
8
|
+
jobs:
|
|
9
|
+
build:
|
|
10
|
+
runs-on: ubuntu-22.04
|
|
11
|
+
|
|
12
|
+
steps:
|
|
13
|
+
- uses: actions/checkout@v7
|
|
14
|
+
with:
|
|
15
|
+
fetch-depth: 0
|
|
16
|
+
- name: Set up Python
|
|
17
|
+
uses: actions/setup-python@v7
|
|
18
|
+
with:
|
|
19
|
+
python-version: "3.13"
|
|
20
|
+
- run: python -m pip install build
|
|
21
|
+
- name: Build package
|
|
22
|
+
run: python -m build
|
|
23
|
+
- uses: actions/upload-artifact@v7
|
|
24
|
+
with:
|
|
25
|
+
name: dist
|
|
26
|
+
path: dist/
|
|
27
|
+
|
|
28
|
+
publish:
|
|
29
|
+
needs: build
|
|
30
|
+
runs-on: ubuntu-22.04
|
|
31
|
+
environment: release
|
|
32
|
+
permissions:
|
|
33
|
+
id-token: write
|
|
34
|
+
|
|
35
|
+
steps:
|
|
36
|
+
- uses: actions/download-artifact@v8
|
|
37
|
+
with:
|
|
38
|
+
name: dist
|
|
39
|
+
path: dist/
|
|
40
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
name: Tests
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
pull_request:
|
|
7
|
+
# Check all PR
|
|
8
|
+
|
|
9
|
+
jobs:
|
|
10
|
+
tests:
|
|
11
|
+
runs-on: ${{ matrix.os }}
|
|
12
|
+
strategy:
|
|
13
|
+
matrix:
|
|
14
|
+
include:
|
|
15
|
+
- os: ubuntu-22.04
|
|
16
|
+
python-version: "3.11"
|
|
17
|
+
- os: ubuntu-22.04
|
|
18
|
+
python-version: "3.14"
|
|
19
|
+
- os: macos-14
|
|
20
|
+
python-version: "3.14"
|
|
21
|
+
|
|
22
|
+
steps:
|
|
23
|
+
- uses: actions/checkout@v7
|
|
24
|
+
with:
|
|
25
|
+
fetch-depth: 0
|
|
26
|
+
- name: Set up Python ${{ matrix.python-version }}
|
|
27
|
+
uses: actions/setup-python@v7
|
|
28
|
+
with:
|
|
29
|
+
python-version: ${{ matrix.python-version }}
|
|
30
|
+
- run: python -m pip install tox coverage[toml]
|
|
31
|
+
- name: run Python tests
|
|
32
|
+
run: tox -e tests
|
|
33
|
+
|
|
34
|
+
examples:
|
|
35
|
+
runs-on: ubuntu-22.04
|
|
36
|
+
|
|
37
|
+
steps:
|
|
38
|
+
- uses: actions/checkout@v7
|
|
39
|
+
with:
|
|
40
|
+
fetch-depth: 0
|
|
41
|
+
- name: Set up Python
|
|
42
|
+
uses: actions/setup-python@v7
|
|
43
|
+
with:
|
|
44
|
+
python-version: "3.13"
|
|
45
|
+
- run: python -m pip install tox
|
|
46
|
+
- name: run examples
|
|
47
|
+
run: tox -e examples
|
lorem_jax-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
BSD 3-Clause License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026, lorem-jax developers
|
|
4
|
+
|
|
5
|
+
Redistribution and use in source and binary forms, with or without
|
|
6
|
+
modification, are permitted provided that the following conditions are met:
|
|
7
|
+
|
|
8
|
+
1. Redistributions of source code must retain the above copyright notice, this
|
|
9
|
+
list of conditions and the following disclaimer.
|
|
10
|
+
|
|
11
|
+
2. Redistributions in binary form must reproduce the above copyright notice,
|
|
12
|
+
this list of conditions and the following disclaimer in the documentation
|
|
13
|
+
and/or other materials provided with the distribution.
|
|
14
|
+
|
|
15
|
+
3. Neither the name of the copyright holder nor the names of its
|
|
16
|
+
contributors may be used to endorse or promote products derived from
|
|
17
|
+
this software without specific prior written permission.
|
|
18
|
+
|
|
19
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
20
|
+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
21
|
+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
22
|
+
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
23
|
+
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
24
|
+
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
25
|
+
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
26
|
+
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
27
|
+
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
28
|
+
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
lorem_jax-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,226 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: lorem-jax
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: LOREM: an equivariant long-range machine learning interatomic potential in JAX
|
|
5
|
+
License-Expression: BSD-3-Clause
|
|
6
|
+
Project-URL: Repository, https://github.com/lab-cosmo/lorem-jax
|
|
7
|
+
Requires-Python: >=3.11
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Dist: ase
|
|
11
|
+
Requires-Dist: e3x
|
|
12
|
+
Requires-Dist: flax
|
|
13
|
+
Requires-Dist: ipi
|
|
14
|
+
Requires-Dist: jax
|
|
15
|
+
Requires-Dist: jax-pme>=0.1
|
|
16
|
+
Requires-Dist: jaxtyping
|
|
17
|
+
Requires-Dist: marathon-train[grain]
|
|
18
|
+
Requires-Dist: grain<0.2.17
|
|
19
|
+
Requires-Dist: opsis
|
|
20
|
+
Requires-Dist: optax
|
|
21
|
+
Dynamic: license-file
|
|
22
|
+
|
|
23
|
+
# LOREM-JAX
|
|
24
|
+
|
|
25
|
+
JAX implementation of [LOREM](https://openreview.net/forum?id=pZI9e4SW9P) (Learning Long-Range Representations with Equivariant Messages), a machine learning interatomic potential with equivariant long-range message passing.
|
|
26
|
+
|
|
27
|
+
Built on [JAX](https://github.com/jax-ml/jax), [Flax](https://github.com/google/flax), [e3x](https://github.com/google-research/e3x), and [jax-pme](https://github.com/lab-cosmo/jax-pme).
|
|
28
|
+
|
|
29
|
+
## Installation
|
|
30
|
+
|
|
31
|
+
Requires Python >= 3.11.
|
|
32
|
+
|
|
33
|
+
```bash
|
|
34
|
+
pip install .
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
## Usage
|
|
38
|
+
|
|
39
|
+
### ASE calculator
|
|
40
|
+
|
|
41
|
+
```python
|
|
42
|
+
import jax
|
|
43
|
+
from ase.build import bulk
|
|
44
|
+
from lorem.models.mlip import Lorem
|
|
45
|
+
from lorem.calculator import Calculator
|
|
46
|
+
|
|
47
|
+
model = Lorem(cutoff=5.0)
|
|
48
|
+
params = model.init(jax.random.key(42), *model.dummy_inputs())
|
|
49
|
+
calc = Calculator.from_model(model, params=params)
|
|
50
|
+
|
|
51
|
+
atoms = bulk("Ar") * [2, 2, 2]
|
|
52
|
+
calc.calculate(atoms)
|
|
53
|
+
print(calc.results["energy"], calc.results["forces"].shape)
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
To load a trained model from a checkpoint:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
calc = Calculator.from_checkpoint("path/to/checkpoint")
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
### Training
|
|
63
|
+
|
|
64
|
+
Training a model involves three steps: preparing the data, configuring the model and training settings, and running the training script.
|
|
65
|
+
|
|
66
|
+
#### 1. Prepare data
|
|
67
|
+
|
|
68
|
+
Training data is stored in [marathon](https://github.com/sirmarcel/marathon) format. Convert your extended XYZ dataset using a preparation script (see `examples/train-mlp/prepare.py` for a template):
|
|
69
|
+
|
|
70
|
+
```python
|
|
71
|
+
from marathon.data import datasets, get_splits
|
|
72
|
+
from marathon.grain import prepare
|
|
73
|
+
|
|
74
|
+
# datasets is a Path resolved from the $DATASETS environment variable
|
|
75
|
+
prepare(train_atoms, folder=datasets / "my_project/train", ...)
|
|
76
|
+
prepare(valid_atoms, folder=datasets / "my_project/valid", ...)
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
The `$DATASETS` environment variable sets the root directory where prepared datasets are stored. All dataset paths in `settings.yaml` are resolved relative to this directory.
|
|
80
|
+
|
|
81
|
+
#### 2. Configure the experiment
|
|
82
|
+
|
|
83
|
+
Each experiment lives in its own directory containing two YAML files:
|
|
84
|
+
|
|
85
|
+
**`model.yaml`** defines the model architecture:
|
|
86
|
+
|
|
87
|
+
```yaml
|
|
88
|
+
model:
|
|
89
|
+
lorem.Lorem:
|
|
90
|
+
cutoff: 5.0
|
|
91
|
+
max_degree: 4
|
|
92
|
+
max_degree_lr: 2
|
|
93
|
+
num_features: 128
|
|
94
|
+
num_spherical_features: 4
|
|
95
|
+
num_message_passing: 1
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
Use `lorem.LoremBEC` instead of `lorem.Lorem` to train a model that additionally predicts Born effective charges.
|
|
99
|
+
|
|
100
|
+
**`settings.yaml`** configures training:
|
|
101
|
+
|
|
102
|
+
```yaml
|
|
103
|
+
train: "my_project/train" # path relative to $DATASETS
|
|
104
|
+
valid: "my_project/valid" # path relative to $DATASETS
|
|
105
|
+
seed: 23
|
|
106
|
+
batcher:
|
|
107
|
+
batch_size: 4
|
|
108
|
+
loss_weights: {"energy": 0.5, "forces": 0.5}
|
|
109
|
+
optimizer: adam # adam or muon
|
|
110
|
+
start_learning_rate: 1e-3
|
|
111
|
+
min_learning_rate: 1e-6
|
|
112
|
+
max_epochs: 2000
|
|
113
|
+
valid_every_epoch: 2
|
|
114
|
+
decay_style: linear # linear, exponential, or warmup_cosine
|
|
115
|
+
use_wandb: True
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
<details>
|
|
119
|
+
<summary>All training settings</summary>
|
|
120
|
+
|
|
121
|
+
| Setting | Default | Description |
|
|
122
|
+
|---|---|---|
|
|
123
|
+
| `train` | *required* | Training dataset path (relative to `$DATASETS`) |
|
|
124
|
+
| `valid` | *required* | Validation dataset path (relative to `$DATASETS`) |
|
|
125
|
+
| `test_datasets` | `{}` | Extra test datasets: `{name: [path, save_predictions]}` |
|
|
126
|
+
| `batcher.batch_size` | *required* | Samples per batch |
|
|
127
|
+
| `batcher.size_strategy` | `powers_of_4` | Padding strategy for batch dimensions |
|
|
128
|
+
| `loss_weights` | `{"energy": 0.5, "forces": 0.5}` | Per-target loss weights |
|
|
129
|
+
| `scale_by_variance` | `False` | Scale loss weights by validation set variance |
|
|
130
|
+
| `optimizer` | `adam` | Optimizer (`adam`, `muon`, or any optax optimizer) |
|
|
131
|
+
| `start_learning_rate` | `1e-3` | Initial learning rate |
|
|
132
|
+
| `min_learning_rate` | `1e-6` | Minimum learning rate |
|
|
133
|
+
| `max_epochs` | `2000` | Maximum training epochs |
|
|
134
|
+
| `valid_every_epoch` | `2` | Validate every N epochs |
|
|
135
|
+
| `decay_style` | `linear` | LR schedule: `linear`, `exponential`, or `warmup_cosine` |
|
|
136
|
+
| `start_decay_after` | `10` | Epoch to begin LR decay |
|
|
137
|
+
| `stop_decay_after` | `max_epochs` | Epoch to end LR decay (linear only) |
|
|
138
|
+
| `warmup_epochs` | `0` | Warmup epochs (`warmup_cosine` only) |
|
|
139
|
+
| `gradient_clip` | `0` | Gradient clipping threshold (0 = disabled) |
|
|
140
|
+
| `seed` | `0` | Random seed |
|
|
141
|
+
| `rotational_augmentation` | `False` | Apply random rotations to training data |
|
|
142
|
+
| `filter_mixed_pbc` | `False` | Filter out structures with mixed periodic boundary conditions |
|
|
143
|
+
| `filter_above_num_atoms` | `False` | Filter out structures above this atom count |
|
|
144
|
+
| `checkpointers` | `default` | `default` or `full` (adds RMSE checkpointers) |
|
|
145
|
+
| `use_wandb` | `True` | Log to Weights & Biases |
|
|
146
|
+
| `wandb_project` | auto | W&B project name (default: derived from folder names) |
|
|
147
|
+
| `wandb_name` | auto | W&B run name (default: experiment folder name) |
|
|
148
|
+
| `benchmark_pipeline` | `True` | Benchmark data pipeline before training |
|
|
149
|
+
| `compilation_cache` | `False` | Enable JAX persistent compilation cache |
|
|
150
|
+
| `default_matmul_precision` | `float32` | JAX matmul precision (`default`, `float32`) |
|
|
151
|
+
| `debug_nans` | `False` | Enable JAX NaN debugging (~50% slowdown) |
|
|
152
|
+
| `enable_x64` | `False` | Enable 64-bit floating point |
|
|
153
|
+
| `worker_count` | `4` | Data loading workers (training) |
|
|
154
|
+
| `worker_count_valid` | `worker_count` | Data loading workers (validation) |
|
|
155
|
+
| `worker_buffer_size` | `2` | Prefetch buffer per worker (training) |
|
|
156
|
+
|
|
157
|
+
</details>
|
|
158
|
+
|
|
159
|
+
#### 3. Run training
|
|
160
|
+
|
|
161
|
+
```bash
|
|
162
|
+
cd my_experiment
|
|
163
|
+
DATASETS=/path/to/datasets lorem-train
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
Training writes checkpoints, logs, and plots to a `run/` directory inside the experiment folder. If a `run/` directory already exists, training resumes from the latest checkpoint.
|
|
167
|
+
|
|
168
|
+
See `examples/train-mlp/` and `examples/train-bec/` for complete examples including data preparation and configuration files.
|
|
169
|
+
|
|
170
|
+
### Model variants
|
|
171
|
+
|
|
172
|
+
- **`Lorem`** -- the standard MLIP model (energy + forces + stress)
|
|
173
|
+
- **`LoremBEC`** -- predicts Born effective charges in addition to energy/forces
|
|
174
|
+
|
|
175
|
+
### Key hyperparameters
|
|
176
|
+
|
|
177
|
+
| Parameter | Default | Description |
|
|
178
|
+
|---|---|---|
|
|
179
|
+
| `cutoff` | 5.0 | Short-range cutoff radius (A) |
|
|
180
|
+
| `max_degree` | 6 | Maximum angular momentum for spherical features |
|
|
181
|
+
| `max_degree_lr` | 2 | Maximum angular momentum for long-range charges |
|
|
182
|
+
| `num_features` | 128 | Number of scalar features |
|
|
183
|
+
| `num_spherical_features` | 8 | Number of spherical feature channels |
|
|
184
|
+
| `num_radial` | 32 | Number of radial basis functions |
|
|
185
|
+
| `num_message_passing` | 0 | Number of short-range message passing steps |
|
|
186
|
+
| `lr` | True | Enable long-range (Ewald) interaction |
|
|
187
|
+
|
|
188
|
+
## Installing the i-PI driver
|
|
189
|
+
|
|
190
|
+
After installation of the package, install the i-PI driver via:
|
|
191
|
+
|
|
192
|
+
```bash
|
|
193
|
+
lorem-install-ipi-driver
|
|
194
|
+
```
|
|
195
|
+
|
|
196
|
+
This copies the LOREM driver into the i-PI `pes` directory. You can rerun `lorem-install-ipi-driver` anytime (it is idempotent) if you switch environments or reinstall i-PI.
|
|
197
|
+
|
|
198
|
+
## Development
|
|
199
|
+
|
|
200
|
+
Format and lint:
|
|
201
|
+
|
|
202
|
+
```bash
|
|
203
|
+
ruff format . && ruff check --fix .
|
|
204
|
+
```
|
|
205
|
+
|
|
206
|
+
Run tests:
|
|
207
|
+
|
|
208
|
+
```bash
|
|
209
|
+
python -m pytest tests/ -v --override-ini="addopts="
|
|
210
|
+
```
|
|
211
|
+
|
|
212
|
+
Or use tox:
|
|
213
|
+
|
|
214
|
+
```bash
|
|
215
|
+
tox -e lint # check formatting + linting
|
|
216
|
+
tox -e tests # run unit tests
|
|
217
|
+
tox -e examples # run examples as smoke tests
|
|
218
|
+
tox -e format # auto-format
|
|
219
|
+
```
|
|
220
|
+
|
|
221
|
+
## License
|
|
222
|
+
|
|
223
|
+
BSD-3-Clause
|
|
224
|
+
|
|
225
|
+
|
|
226
|
+
This project is [maintained](https://github.com/lab-cosmo/.github/blob/main/Maintainers.md) by [@PicoCentauri](https://github.com/PicoCentauri) and [@sirmarcel](https://github.com/sirmarcel), who will reply to issues and pull requests opened on this repository as soon as possible. You can mention them directly if you did not receive an answer after a couple of days.
|
|
@@ -0,0 +1,204 @@
|
|
|
1
|
+
# LOREM-JAX
|
|
2
|
+
|
|
3
|
+
JAX implementation of [LOREM](https://openreview.net/forum?id=pZI9e4SW9P) (Learning Long-Range Representations with Equivariant Messages), a machine learning interatomic potential with equivariant long-range message passing.
|
|
4
|
+
|
|
5
|
+
Built on [JAX](https://github.com/jax-ml/jax), [Flax](https://github.com/google/flax), [e3x](https://github.com/google-research/e3x), and [jax-pme](https://github.com/lab-cosmo/jax-pme).
|
|
6
|
+
|
|
7
|
+
## Installation
|
|
8
|
+
|
|
9
|
+
Requires Python >= 3.11.
|
|
10
|
+
|
|
11
|
+
```bash
|
|
12
|
+
pip install .
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
## Usage
|
|
16
|
+
|
|
17
|
+
### ASE calculator
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
import jax
|
|
21
|
+
from ase.build import bulk
|
|
22
|
+
from lorem.models.mlip import Lorem
|
|
23
|
+
from lorem.calculator import Calculator
|
|
24
|
+
|
|
25
|
+
model = Lorem(cutoff=5.0)
|
|
26
|
+
params = model.init(jax.random.key(42), *model.dummy_inputs())
|
|
27
|
+
calc = Calculator.from_model(model, params=params)
|
|
28
|
+
|
|
29
|
+
atoms = bulk("Ar") * [2, 2, 2]
|
|
30
|
+
calc.calculate(atoms)
|
|
31
|
+
print(calc.results["energy"], calc.results["forces"].shape)
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
To load a trained model from a checkpoint:
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
calc = Calculator.from_checkpoint("path/to/checkpoint")
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
### Training
|
|
41
|
+
|
|
42
|
+
Training a model involves three steps: preparing the data, configuring the model and training settings, and running the training script.
|
|
43
|
+
|
|
44
|
+
#### 1. Prepare data
|
|
45
|
+
|
|
46
|
+
Training data is stored in [marathon](https://github.com/sirmarcel/marathon) format. Convert your extended XYZ dataset using a preparation script (see `examples/train-mlp/prepare.py` for a template):
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
from marathon.data import datasets, get_splits
|
|
50
|
+
from marathon.grain import prepare
|
|
51
|
+
|
|
52
|
+
# datasets is a Path resolved from the $DATASETS environment variable
|
|
53
|
+
prepare(train_atoms, folder=datasets / "my_project/train", ...)
|
|
54
|
+
prepare(valid_atoms, folder=datasets / "my_project/valid", ...)
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
The `$DATASETS` environment variable sets the root directory where prepared datasets are stored. All dataset paths in `settings.yaml` are resolved relative to this directory.
|
|
58
|
+
|
|
59
|
+
#### 2. Configure the experiment
|
|
60
|
+
|
|
61
|
+
Each experiment lives in its own directory containing two YAML files:
|
|
62
|
+
|
|
63
|
+
**`model.yaml`** defines the model architecture:
|
|
64
|
+
|
|
65
|
+
```yaml
|
|
66
|
+
model:
|
|
67
|
+
lorem.Lorem:
|
|
68
|
+
cutoff: 5.0
|
|
69
|
+
max_degree: 4
|
|
70
|
+
max_degree_lr: 2
|
|
71
|
+
num_features: 128
|
|
72
|
+
num_spherical_features: 4
|
|
73
|
+
num_message_passing: 1
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
Use `lorem.LoremBEC` instead of `lorem.Lorem` to train a model that additionally predicts Born effective charges.
|
|
77
|
+
|
|
78
|
+
**`settings.yaml`** configures training:
|
|
79
|
+
|
|
80
|
+
```yaml
|
|
81
|
+
train: "my_project/train" # path relative to $DATASETS
|
|
82
|
+
valid: "my_project/valid" # path relative to $DATASETS
|
|
83
|
+
seed: 23
|
|
84
|
+
batcher:
|
|
85
|
+
batch_size: 4
|
|
86
|
+
loss_weights: {"energy": 0.5, "forces": 0.5}
|
|
87
|
+
optimizer: adam # adam or muon
|
|
88
|
+
start_learning_rate: 1e-3
|
|
89
|
+
min_learning_rate: 1e-6
|
|
90
|
+
max_epochs: 2000
|
|
91
|
+
valid_every_epoch: 2
|
|
92
|
+
decay_style: linear # linear, exponential, or warmup_cosine
|
|
93
|
+
use_wandb: True
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
<details>
|
|
97
|
+
<summary>All training settings</summary>
|
|
98
|
+
|
|
99
|
+
| Setting | Default | Description |
|
|
100
|
+
|---|---|---|
|
|
101
|
+
| `train` | *required* | Training dataset path (relative to `$DATASETS`) |
|
|
102
|
+
| `valid` | *required* | Validation dataset path (relative to `$DATASETS`) |
|
|
103
|
+
| `test_datasets` | `{}` | Extra test datasets: `{name: [path, save_predictions]}` |
|
|
104
|
+
| `batcher.batch_size` | *required* | Samples per batch |
|
|
105
|
+
| `batcher.size_strategy` | `powers_of_4` | Padding strategy for batch dimensions |
|
|
106
|
+
| `loss_weights` | `{"energy": 0.5, "forces": 0.5}` | Per-target loss weights |
|
|
107
|
+
| `scale_by_variance` | `False` | Scale loss weights by validation set variance |
|
|
108
|
+
| `optimizer` | `adam` | Optimizer (`adam`, `muon`, or any optax optimizer) |
|
|
109
|
+
| `start_learning_rate` | `1e-3` | Initial learning rate |
|
|
110
|
+
| `min_learning_rate` | `1e-6` | Minimum learning rate |
|
|
111
|
+
| `max_epochs` | `2000` | Maximum training epochs |
|
|
112
|
+
| `valid_every_epoch` | `2` | Validate every N epochs |
|
|
113
|
+
| `decay_style` | `linear` | LR schedule: `linear`, `exponential`, or `warmup_cosine` |
|
|
114
|
+
| `start_decay_after` | `10` | Epoch to begin LR decay |
|
|
115
|
+
| `stop_decay_after` | `max_epochs` | Epoch to end LR decay (linear only) |
|
|
116
|
+
| `warmup_epochs` | `0` | Warmup epochs (`warmup_cosine` only) |
|
|
117
|
+
| `gradient_clip` | `0` | Gradient clipping threshold (0 = disabled) |
|
|
118
|
+
| `seed` | `0` | Random seed |
|
|
119
|
+
| `rotational_augmentation` | `False` | Apply random rotations to training data |
|
|
120
|
+
| `filter_mixed_pbc` | `False` | Filter out structures with mixed periodic boundary conditions |
|
|
121
|
+
| `filter_above_num_atoms` | `False` | Filter out structures above this atom count |
|
|
122
|
+
| `checkpointers` | `default` | `default` or `full` (adds RMSE checkpointers) |
|
|
123
|
+
| `use_wandb` | `True` | Log to Weights & Biases |
|
|
124
|
+
| `wandb_project` | auto | W&B project name (default: derived from folder names) |
|
|
125
|
+
| `wandb_name` | auto | W&B run name (default: experiment folder name) |
|
|
126
|
+
| `benchmark_pipeline` | `True` | Benchmark data pipeline before training |
|
|
127
|
+
| `compilation_cache` | `False` | Enable JAX persistent compilation cache |
|
|
128
|
+
| `default_matmul_precision` | `float32` | JAX matmul precision (`default`, `float32`) |
|
|
129
|
+
| `debug_nans` | `False` | Enable JAX NaN debugging (~50% slowdown) |
|
|
130
|
+
| `enable_x64` | `False` | Enable 64-bit floating point |
|
|
131
|
+
| `worker_count` | `4` | Data loading workers (training) |
|
|
132
|
+
| `worker_count_valid` | `worker_count` | Data loading workers (validation) |
|
|
133
|
+
| `worker_buffer_size` | `2` | Prefetch buffer per worker (training) |
|
|
134
|
+
|
|
135
|
+
</details>
|
|
136
|
+
|
|
137
|
+
#### 3. Run training
|
|
138
|
+
|
|
139
|
+
```bash
|
|
140
|
+
cd my_experiment
|
|
141
|
+
DATASETS=/path/to/datasets lorem-train
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
Training writes checkpoints, logs, and plots to a `run/` directory inside the experiment folder. If a `run/` directory already exists, training resumes from the latest checkpoint.
|
|
145
|
+
|
|
146
|
+
See `examples/train-mlp/` and `examples/train-bec/` for complete examples including data preparation and configuration files.
|
|
147
|
+
|
|
148
|
+
### Model variants
|
|
149
|
+
|
|
150
|
+
- **`Lorem`** -- the standard MLIP model (energy + forces + stress)
|
|
151
|
+
- **`LoremBEC`** -- predicts Born effective charges in addition to energy/forces
|
|
152
|
+
|
|
153
|
+
### Key hyperparameters
|
|
154
|
+
|
|
155
|
+
| Parameter | Default | Description |
|
|
156
|
+
|---|---|---|
|
|
157
|
+
| `cutoff` | 5.0 | Short-range cutoff radius (A) |
|
|
158
|
+
| `max_degree` | 6 | Maximum angular momentum for spherical features |
|
|
159
|
+
| `max_degree_lr` | 2 | Maximum angular momentum for long-range charges |
|
|
160
|
+
| `num_features` | 128 | Number of scalar features |
|
|
161
|
+
| `num_spherical_features` | 8 | Number of spherical feature channels |
|
|
162
|
+
| `num_radial` | 32 | Number of radial basis functions |
|
|
163
|
+
| `num_message_passing` | 0 | Number of short-range message passing steps |
|
|
164
|
+
| `lr` | True | Enable long-range (Ewald) interaction |
|
|
165
|
+
|
|
166
|
+
## Installing the i-PI driver
|
|
167
|
+
|
|
168
|
+
After installation of the package, install the i-PI driver via:
|
|
169
|
+
|
|
170
|
+
```bash
|
|
171
|
+
lorem-install-ipi-driver
|
|
172
|
+
```
|
|
173
|
+
|
|
174
|
+
This copies the LOREM driver into the i-PI `pes` directory. You can rerun `lorem-install-ipi-driver` anytime (it is idempotent) if you switch environments or reinstall i-PI.
|
|
175
|
+
|
|
176
|
+
## Development
|
|
177
|
+
|
|
178
|
+
Format and lint:
|
|
179
|
+
|
|
180
|
+
```bash
|
|
181
|
+
ruff format . && ruff check --fix .
|
|
182
|
+
```
|
|
183
|
+
|
|
184
|
+
Run tests:
|
|
185
|
+
|
|
186
|
+
```bash
|
|
187
|
+
python -m pytest tests/ -v --override-ini="addopts="
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
Or use tox:
|
|
191
|
+
|
|
192
|
+
```bash
|
|
193
|
+
tox -e lint # check formatting + linting
|
|
194
|
+
tox -e tests # run unit tests
|
|
195
|
+
tox -e examples # run examples as smoke tests
|
|
196
|
+
tox -e format # auto-format
|
|
197
|
+
```
|
|
198
|
+
|
|
199
|
+
## License
|
|
200
|
+
|
|
201
|
+
BSD-3-Clause
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
This project is [maintained](https://github.com/lab-cosmo/.github/blob/main/Maintainers.md) by [@PicoCentauri](https://github.com/PicoCentauri) and [@sirmarcel](https://github.com/sirmarcel), who will reply to issues and pull requests opened on this repository as soon as possible. You can mention them directly if you did not receive an answer after a couple of days.
|