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.
Files changed (61) hide show
  1. lorem_jax-0.1.0/.claude/commands/release.md +48 -0
  2. lorem_jax-0.1.0/.github/dependabot.yml +11 -0
  3. lorem_jax-0.1.0/.github/workflows/lint.yml +19 -0
  4. lorem_jax-0.1.0/.github/workflows/release.yml +40 -0
  5. lorem_jax-0.1.0/.github/workflows/tests.yml +47 -0
  6. lorem_jax-0.1.0/.gitignore +18 -0
  7. lorem_jax-0.1.0/LICENSE +28 -0
  8. lorem_jax-0.1.0/PKG-INFO +226 -0
  9. lorem_jax-0.1.0/README.md +204 -0
  10. lorem_jax-0.1.0/examples/calculator/example.py +49 -0
  11. lorem_jax-0.1.0/examples/eval/README.md +0 -0
  12. lorem_jax-0.1.0/examples/md-ase/README.md +0 -0
  13. lorem_jax-0.1.0/examples/md-ipi/.gitignore +5 -0
  14. lorem_jax-0.1.0/examples/md-ipi/README.md +52 -0
  15. lorem_jax-0.1.0/examples/md-ipi/input.xml +55 -0
  16. lorem_jax-0.1.0/examples/md-ipi/run.sh +30 -0
  17. lorem_jax-0.1.0/examples/md-ipi/start.xyz +5 -0
  18. lorem_jax-0.1.0/examples/train-bec/.gitignore +5 -0
  19. lorem_jax-0.1.0/examples/train-bec/README.md +21 -0
  20. lorem_jax-0.1.0/examples/train-bec/bec_example.xyz +4462 -0
  21. lorem_jax-0.1.0/examples/train-bec/my_experiment/model.yaml +8 -0
  22. lorem_jax-0.1.0/examples/train-bec/my_experiment/settings.yaml +20 -0
  23. lorem_jax-0.1.0/examples/train-bec/prepare.py +64 -0
  24. lorem_jax-0.1.0/examples/train-mlp/.gitignore +5 -0
  25. lorem_jax-0.1.0/examples/train-mlp/README.md +34 -0
  26. lorem_jax-0.1.0/examples/train-mlp/data.xyz +315 -0
  27. lorem_jax-0.1.0/examples/train-mlp/my_experiment/model.yaml +8 -0
  28. lorem_jax-0.1.0/examples/train-mlp/my_experiment/settings.yaml +20 -0
  29. lorem_jax-0.1.0/examples/train-mlp/my_experiment_finetune/model.yaml +8 -0
  30. lorem_jax-0.1.0/examples/train-mlp/my_experiment_finetune/settings.yaml +21 -0
  31. lorem_jax-0.1.0/examples/train-mlp/prepare.py +53 -0
  32. lorem_jax-0.1.0/pyproject.toml +88 -0
  33. lorem_jax-0.1.0/setup.cfg +4 -0
  34. lorem_jax-0.1.0/src/lorem/__init__.py +21 -0
  35. lorem_jax-0.1.0/src/lorem/_install_ipi_driver.py +31 -0
  36. lorem_jax-0.1.0/src/lorem/_version.py +24 -0
  37. lorem_jax-0.1.0/src/lorem/batching.py +113 -0
  38. lorem_jax-0.1.0/src/lorem/calculator.py +245 -0
  39. lorem_jax-0.1.0/src/lorem/ipi.py +42 -0
  40. lorem_jax-0.1.0/src/lorem/models/__init__.py +4 -0
  41. lorem_jax-0.1.0/src/lorem/models/backbone.py +247 -0
  42. lorem_jax-0.1.0/src/lorem/models/bec.py +378 -0
  43. lorem_jax-0.1.0/src/lorem/models/mlip.py +299 -0
  44. lorem_jax-0.1.0/src/lorem/neighborlist.py +111 -0
  45. lorem_jax-0.1.0/src/lorem/train.py +1034 -0
  46. lorem_jax-0.1.0/src/lorem/transforms.py +176 -0
  47. lorem_jax-0.1.0/src/lorem_jax.egg-info/PKG-INFO +226 -0
  48. lorem_jax-0.1.0/src/lorem_jax.egg-info/SOURCES.txt +59 -0
  49. lorem_jax-0.1.0/src/lorem_jax.egg-info/dependency_links.txt +1 -0
  50. lorem_jax-0.1.0/src/lorem_jax.egg-info/entry_points.txt +3 -0
  51. lorem_jax-0.1.0/src/lorem_jax.egg-info/requires.txt +11 -0
  52. lorem_jax-0.1.0/src/lorem_jax.egg-info/scm_file_list.json +54 -0
  53. lorem_jax-0.1.0/src/lorem_jax.egg-info/scm_version.json +8 -0
  54. lorem_jax-0.1.0/src/lorem_jax.egg-info/top_level.txt +1 -0
  55. lorem_jax-0.1.0/tests/test_backbone.py +220 -0
  56. lorem_jax-0.1.0/tests/test_calculator.py +331 -0
  57. lorem_jax-0.1.0/tests/test_init.py +14 -0
  58. lorem_jax-0.1.0/tests/test_ipi.py +47 -0
  59. lorem_jax-0.1.0/tests/test_neighborlist.py +278 -0
  60. lorem_jax-0.1.0/tests/test_train.py +42 -0
  61. 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
@@ -0,0 +1,18 @@
1
+ *.pyc
2
+ *.ipynb_checkpoints*
3
+ __pycache__
4
+ *.egg-info
5
+ *.swp
6
+ *.swo
7
+ *DS_Store
8
+ *coverage*
9
+ _version.py
10
+ *.gz
11
+
12
+ .tox/
13
+ build/
14
+ dist/
15
+ uv.lock
16
+
17
+ # training run outputs (train.py defaults workdir to ./run)
18
+ run/
@@ -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.
@@ -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.