Skip to content

Prepare for release#2

Merged
sirmarcel merged 26 commits intomainfrom
rename-package-and-train-script
Mar 9, 2026
Merged

Prepare for release#2
sirmarcel merged 26 commits intomainfrom
rename-package-and-train-script

Conversation

@sirmarcel
Copy link
Copy Markdown
Contributor

@sirmarcel sirmarcel commented Feb 11, 2026 •
edited
Loading

Summary

  • Rename package loremjax → lorem so import lorem works as expected
  • Move standalone train.py into the package as lorem.train with a lorem-train console script entry point
  • Add train-bec example (BEC prediction training)
  • Factor shared model code into models/backbone.py, deduplicating mlip.py and bec.py
  • Add marathon[grain] extra to dependencies
  • Replace example shell scripts with READMEs; add tox -e examples smoke test environment
  • Add examples CI job
  • Expand README with full training workflow documentation
  • Default matmul precision to float32
  • Remove Windows from CI (uvloop incompatible)
  • Fix isort in ipi.py, add uv.lock to .gitignore

Test plan

  • tox -e lint passes
  • tox -e tests passes (18 tests)
  • tox -e examples passes (calculator, train-mlp, train-bec)
  • CI green on ubuntu (3.11, 3.14) and macOS (3.14)

🤖 Generated with Claude Code

sirmarcel and others added 2 commits February 11, 2026 14:36
@sirmarcel @claude
Rename package loremjax -> lorem, add train script and BEC example
8fb51e3 
- Rename src/loremjax/ to src/lorem/ so `import lorem` works
- Move standalone train.py into package as lorem.train with `lorem-train` console script
- Add train-bec example with prepare.py, model/settings YAML, and run.sh
- Factor shared model code (Initial, MLP, Update, RadialCoefficients, spherical
  helpers) into models/backbone.py to deduplicate mlip.py and bec.py
- Fix __init__.py exports (LOREMCalculator, LoremBEC, comms)
- Add missing dependencies to pyproject.toml (e3x, flax, optax, jaxtyping)
- Adopt ruff rules from marathon (relaxed E741/E731/E402/E501 ignores)
- Update test to use new package name

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Add train-mlp example (energy + forces on cumulene)
a1cc95f 
Small non-periodic C9H4 dataset (21 structures), trains Lorem model
for 2 epochs on energy and forces. Runs in ~45 seconds.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel
Copy link
Copy Markdown
Contributor Author

sirmarcel commented Feb 11, 2026

Added train-mlp example: trains the base Lorem model on energy + forces using a small non-periodic cumulene (C9H4) dataset (21 structures, ~45s to run).

This comment was written by Claude Code on behalf of @sirmarcel.

@sirmarcel @claude
Fix tests for CI and add backbone tests
1264f55 
- Make backbone importable without marathon (lazy import of `masked`)
- Skip test_init when comms is not installed (CI)
- Fix stale loremjax references in _install_ipi_driver.py and test_ipi.py
- Add test_backbone.py: tests for spherical helpers (degree_wise_trace,
  degree_wise_repeat, spherical_norm + custom JVP gradient check),
  MLP, ChemicalEmbedding, RadialEmbedding, and Initial module

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel
Copy link
Copy Markdown
Contributor Author

sirmarcel commented Feb 11, 2026

Note: This PR cannot be fully merged until marathon and comms are released as proper packages (or vendored). Currently:

  • test_init.py is skipped in CI (needs comms)
  • test_ipi.py will fail in CI (needs comms transitively)
  • The backbone tests (spherical helpers, MLP, embeddings) pass in CI since they only need jax/e3x/flax
  • The _masked wrapper in backbone.py is a temporary workaround for lazy-importing marathon.utils.masked — should be reverted to a direct import once marathon is a proper dependency

This comment was written by Claude Code on behalf of @sirmarcel.

@sirmarcel @claude
Fix CI: add jaxpme as git dep, make __init__.py tolerate missing deps
55f1c0d 
jaxpme is public so we can pull it from GitHub directly. The comms
and marathon imports are wrapped in try/except so that backbone tests
can run in CI without them.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel
Copy link
Copy Markdown
Contributor Author

sirmarcel commented Feb 11, 2026

Fix CI failures (55f1c0d, written by Claude)

The backbone tests were failing in CI because from lorem.models.backbone import ... triggered lorem/__init__.py, which unconditionally imported comms (not available in CI).

Changes:

  • Added jaxpme as a git dependency in pyproject.toml — it's public at lab-cosmo/jax-pme, so CI can install it directly
  • Made __init__.py tolerate missing deps: comms import wrapped in try/except (→ None if missing), marathon-dependent imports (models, batching, transforms) wrapped in try/except so backbone-only tests can run

Remaining CI blocker: marathon (and comms) still need to be released before the full test suite runs in CI. The backbone tests (12 tests) and ipi driver test should pass now.

sirmarcel and others added 8 commits February 11, 2026 15:29
@sirmarcel @claude
Fix jaxpme dependency name: jax-pme matches package metadata
3c7b1e9 
pip rejects 'jaxpme' because the package metadata declares 'jax-pme'.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Fix remaining marathon imports breaking CI
fd6c894 
- models/__init__.py: wrap bec/mlip imports in try/except
- mlip.py, bec.py: use _masked from backbone instead of direct
  marathon.utils import

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Fix ASE calculator and add usage example
d0d7ce9 
The calculator had broken imports (nonexistent SetUpEwald, get_batch)
and wrong field access patterns. Rewired to use lorem.batching directly,
added from_model() classmethod for easy construction without checkpoints.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Fill in README with installation, usage, and hyperparameter docs
4c81277 
Also fixes the i-PI driver command name (loremjax-install-ipi-driver
-> lorem-install-ipi-driver) to match pyproject.toml.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel
Clarify deps
7a5ff16
@sirmarcel @claude
Allow wandb_project and wandb_name to be set via settings.yaml
4507139 
Previously these were hardcoded to None and always auto-detected from
folder structure. Now settings.yaml can override them while keeping the
folder-based fallback as default.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Apply ruff formatting to calculator and example
869b070 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Add i-PI driven dynamics example with BEC support
57f9265 
Wire up the Calculator → i-PI BEC data path: extract "apt" from
LoremBEC model output, reshape to (3*natoms, 3), and expose as "BEC"
for i-PI's driven_dynamics module. Implement LOREMCalculator as a
proper Calculator subclass and clean up the i-PI driver (remove buggy
compute_structure override that operated on JSON strings).

Add md-ipi example with input.xml (eda-nve, bec mode="driver"),
run.sh, start.xyz, and README.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@PicoCentauri PicoCentauri force-pushed the rename-package-and-train-script branch from 04291dd to 7b5b3f6 Compare February 26, 2026 08:31
@PicoCentauri PicoCentauri force-pushed the rename-package-and-train-script branch from 7b5b3f6 to 1a57216 Compare February 26, 2026 08:47
sirmarcel and others added 11 commits March 3, 2026 14:46
@sirmarcel @claude
Inherit from ASE BaseCalculator for i-PI compatibility
a5cc2d0 
GetPropertiesMixin lacks calculate_properties(), which i-PI's
ASEDriver requires. BaseCalculator provides it (and inherits
from GetPropertiesMixin), so this is a strict superset.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@PicoCentauri
make BEC ase compatible
354785d
@sirmarcel @claude
Implement stress in ASE calculator
7e3c9c1 
The model predicts the virial (stress * volume) as a 3x3 tensor.
Convert to ASE convention: divide by cell volume and return as
Voigt 6-element array.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@PicoCentauri
fix bec wrapper
f306235
@sirmarcel @claude
Remove halfspace kwarg from get_batch calls (jax-pme API update)
b9bac8e 
jax-pme moved halfspace handling into prepare(), so get_batch no
longer accepts this parameter.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Fix remaining test failures after recent refactors
5ac0a71 
- Update BEC tests to use "born_effective_charges" key and (N,3,3) shape
  to match the ASE-compatible calculator output
- Export LOREMCalculator alias from __init__.py
- (ipi entry point was fixed by reinstall after prior flattening commit)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Update imports from marathon.extra.hermes to marathon.grain
90d010f 
hermes is deprecated; marathon.grain is now the canonical module.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Add marathon as git dependency, fix k_grid size counting in ToBatch
888269b 
Add marathon (release/v0.2.0) as a proper git dependency instead of
requiring separate local install. Fix ToBatch.k_size to use kg.shape[0]
directly instead of count_halfspace_kvectors(kg.shape), which expects
grid dimensions (nx, ny, nz) but receives the flat k-vector array shape.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Add tox examples runner, fix marathon dep, modernize example workflow
2c853f4 
- Add marathon[grain] extra to pyproject.toml
- Replace example run.sh scripts with README docs
- Add tox -e examples environment for smoke-testing examples
- Add examples CI job to tests workflow
- Remove stale marathon-dev install note from README
- Add tox commands to README development section
- Fix isort in ipi.py (drop unneeded __future__ import)
- Add uv.lock to .gitignore

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Remove Windows from CI test matrix
a14365a 
uvloop (transitive dep via flax/orbax) does not support Windows.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Expand README training docs, default to float32 matmul precision
bdf55a2 
- Document full training workflow (data prep, config, running)
- Explain $DATASETS environment variable
- Add collapsible reference table for all training settings
- Default matmul precision to float32 instead of default

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel sirmarcel changed the title Rename package loremjax -> lorem, add train script and BEC example Prepare for release Mar 9, 2026
sirmarcel and others added 2 commits March 9, 2026 17:27
@sirmarcel @claude
Remove try/except guards for comms and marathon imports
07b3060 
Both are hard dependencies — the guards were dead code.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel @claude
Add PyPI release workflow (trusted publishing on tag push)
5690468 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sirmarcel sirmarcel merged commit 58b5d07 into main Mar 9, 2026
5 checks passed
@sirmarcel sirmarcel deleted the rename-package-and-train-script branch March 9, 2026 16:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@sirmarcel @PicoCentauri