Alchemiops KNN+D3, TorchSim, Refactor (#143)

* Switch to uv

* Update code

* Update tests

* Readmes, examples, fixes

* TorchSim and D3 examples, bug fixes

* Use uv in CI

* Clean up

* Dockerfile

* Made torch-sim optional

* Removed default loss_weights from conservative/direct_regressor, added them in pretrained.py instead

* Added changelog to README.md

Author: vsimkus

.flake8

@@ -2,4 +2,4 @@
 # the repo. Unless a later match takes precedence,
 # @global-owner1 and @global-owner2 will be requested for
 # review when someone opens a pull request.
-*       @DeNeutoy @benrhodes26 @jg8610 @reactiv @zhiyil1230
+*       @benrhodes26 @vsimkus @jg8610 @reactiv

.github/CODEOWNERS

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.12", "3.13", "3.14"]
     if: startsWith(github.ref, 'refs/tags/') 
     steps:
     - uses: actions/checkout@v4
@@ -25,13 +25,20 @@ jobs:
         python-version: ${{ matrix.python-version }}
         cache: pip
         cache-dependency-path: pyproject.toml
-    - name: Install dependencies
+    - name: Set up uv
       run: |
-        pip install '.[test]'
+        pip install uv
+    - name: Install dependencies
+      run: uv sync --group dev
 
-    - name: Run tests
+    - name: Run linters
       run: |
-        pytest
+        uv run ruff check ./orb_models ./tests ./scripts
+        uv run ruff format --check ./orb_models ./tests ./scripts
+        uv run mypy ./orb_models ./tests ./scripts
+    - name: Run tests
+      run: uv run pytest ./tests -n auto
+
   deploy:
     runs-on: ubuntu-latest
     needs: [test]

.github/workflows/publish.yml

@@ -5,12 +5,13 @@ on: [push, pull_request]
 permissions:
   contents: read
 
+
 jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.12", "3.13", "3.14"]
     steps:
     - uses: actions/checkout@v4
     - name: Update version
@@ -22,10 +23,16 @@ jobs:
         python-version: ${{ matrix.python-version }}
         cache: pip
         cache-dependency-path: pyproject.toml
-    - name: Install dependencies
+    - name: Set up uv
       run: |
-        pip install '.[test]'
+        pip install uv
+    - name: Install dependencies
+      run: uv sync --group dev
 
-    - name: Run tests
+    - name: Run linters
       run: |
-        pytest
+        uv run ruff check ./orb_models ./tests ./scripts
+        uv run ruff format --check ./orb_models ./tests ./scripts
+        uv run mypy ./orb_models ./tests ./scripts
+    - name: Run tests
+      run: uv run pytest ./tests -n auto

.github/workflows/test.yml

@@ -1,3 +1,6 @@
+# MacOS
+.DS_Store
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -101,6 +104,10 @@ ipython_config.py
 #   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
 #poetry.lock
 
+# uv
+#   https://docs.astral.sh/uv/concepts/projects/sync/#upgrading-locked-package-versions
+uv.lock
+
 # pdm
 #   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
 #pdm.lock
@@ -166,5 +173,7 @@ datasets/
 
 # VS Code
 .devcontainer/
+.vscode/
 
-wandb/
+wandb/
+ckpts/

.gitignore

@@ -0,0 +1,4 @@
+# Python version, ensures consistent development environment 
+# If the version changes in this file, running `uv sync` or `uv sync --group dev` 
+# will download the relevant Python version
+3.12

.python-version

@@ -1,26 +1,30 @@
 ### Setting up a development environment
 
-The `orb_models` package uses [Poetry](https://python-poetry.org/) for dependency management. To install the package and its dependencies, run the following command:
+The `orb_models` repository uses [uv](https://docs.astral.sh/uv/) for dependency management. To install the package and its dependencies, run the following command:
 
 ```bash
-pip install poetry  # Install Poetry if you don't have it
-poetry install
+pipx install uv  # If you don't have uv, we recommend installing it into an isolated environment with pipx: https://docs.astral.sh/uv/getting-started/installation/#pypi
+uv sync --group dev  # Install orb-models and development packages
 ```
 
-Optionally, also install [cuML](https://docs.rapids.ai/install/) (requires CUDA):
+### Running linters
+
+The `orb_models` repository uses `ruff` for formatting and linting, and `mypy` for type checking. To run the linters, use the following commands:
+
 ```bash
-pip install "cuml-cu11==25.2.*"  # For cuda versions >=11.4, <11.8
-pip install "cuml-cu12==25.2.*"  # For cuda versions >=12.0, <13.0
+ruff format .   # Format code
+ruff check .    # Check for linting errors
+mypy .          # Run type checking
 ```
 
 ### Running tests
 
-The `orb_models` package uses `pytest` for testing. To run the tests, navigate to the root directory of the package and run the following command:
+The `orb_models` repository uses `pytest` for testing. To run the tests, navigate to the root directory of the package and run the following command:
 
 ```bash
-pytest
+pytest -n auto ./tests/
 ```
 
 ### Publishing
 
-The `orb_models` package is published using [trusted publishers](https://docs.pypi.org/trusted-publishers/). Whenever a new release is created on GitHub, the package is automatically published to PyPI using GitHub Actions.
+The `orb_models` package is published using [trusted publishers](https://docs.pypi.org/trusted-publishers/). Whenever a new release is created on GitHub, the package is automatically published to PyPI using GitHub Actions.

CONTRIBUTING.md

@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime
+FROM pytorch/pytorch:2.10.0-cuda12.6-cudnn9-runtime
 
 ENV DEBIAN_FRONTEND=noninteractive
 
@@ -10,14 +10,12 @@ RUN apt-get update && \
         git \
         sudo \
         gcc \
-        g++
-
-# Help Numba find lubcudart.so
-ENV LD_LIBRARY_PATH=/opt/conda/lib/python3.11/site-packages/nvidia/cuda_runtime/lib:$LD_LIBRARY_PATH
-RUN ln -s \
-    /opt/conda/lib/python3.11/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12 \
-    /opt/conda/lib/python3.11/site-packages/nvidia/cuda_runtime/lib/libcudart.so
+        g++ \ 
+    # Cleanup
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
 ## Install Python requirements
-RUN pip install orb-models && \
-    pip install "cuml-cu12==25.2.*"
+# For details on --break-system-packages, see: https://veronneau.org/python-311-pip-and-breaking-system-packages.html 
+RUN pip install --break-system-packages orb-models && \
+    pip install --break-system-packages "cuml-cu12==25.2.*"

Dockerfile

@@ -243,9 +243,11 @@ Lines starting with `#` are treated as comments and ignored.
 
 ### Loss Weights
 
-- `--energy_loss_weight`: Weight for energy loss (default: uses model default, usually 1.0)
-- `--forces_loss_weight`: Weight for forces loss (automatically uses correct key for model type)
+- `--energy_loss_weight`: Weight for energy loss (default: 1.0)
+- `--forces_loss_weight`: Weight for forces loss (default: 1.0)
 - `--stress_loss_weight`: Weight for stress loss (set to 0 to disable)
+- `--equigrad_loss_weight`: Weight for the Equigrad loss (turned off by default). Only available for the conservative models.
+  - NOTE: We've found that Equigrad loss should be ≳1000x smaller than the other losses
 
 ### Reference Energies
 
@@ -274,7 +276,7 @@ If you prefer to write your own finetuning script, you can use the clean API dir
 from orb_models.forcefield import pretrained
 
 # Load model with custom configuration
-model = pretrained.orb_v3_conservative_omol(
+model, atoms_adapter = pretrained.orb_v3_conservative_omol(
     device='cuda',
     precision='float32-high',
     train=True,
@@ -287,7 +289,7 @@ model = pretrained.orb_v3_conservative_omol(
 )
 
 # For direct models, use 'forces' and 'stress' keys:
-model = pretrained.orb_v3_direct_omol(
+model, atoms_adapter = pretrained.orb_v3_direct_omol(
     device='cuda',
     train=True,
     loss_weights={
@@ -300,8 +302,6 @@ model = pretrained.orb_v3_direct_omol(
 # The model is now ready for training with your custom configuration!
 ```
 
-This approach is more "pythonic" and clearly documents what configuration options are available. It also encapsulates the implementation details, making your code less fragile to internal changes.
-
 ## How It Works
 
 ### Reference Energies
@@ -319,7 +319,7 @@ import torch
 from orb_models.forcefield import pretrained
 
 # Load model architecture (set train=False for inference)
-model = pretrained.orb_v3_conservative_omol(train=False)
+model, atoms_adapter = pretrained.orb_v3_conservative_omol(train=False)
 
 # Load your finetuned checkpoint
 model.load_state_dict(torch.load('path/to/finetuned_checkpoint.pt'))
@@ -331,7 +331,7 @@ You can also specify loss weights when loading for further finetuning:
 
 ```python
 # Load for continued finetuning with different loss weights
-model = pretrained.orb_v3_conservative_omol(
+model, atoms_adapter = pretrained.orb_v3_conservative_omol(
     train=True,
     loss_weights={'energy': 0.5, 'grad_forces': 20.0}
 )
@@ -370,7 +370,7 @@ python finetune.py \
 from orb_models.forcefield import pretrained
 import torch
 
-model = pretrained.orb_v3_conservative_omol(train=False)
+model, atoms_adapter = pretrained.orb_v3_conservative_omol(train=False)
 model.load_state_dict(torch.load('checkpoints/my_finetuned_model.pt'))
 # Reference energies from my_refs.json are now loaded!
 ```
@@ -384,7 +384,7 @@ from orb_models.forcefield import pretrained
 from orb_models.dataset.ase_sqlite_dataset import AseSqliteDataset
 
 # Load model with configuration
-model = pretrained.orb_v3_conservative_omol(
+model, atoms_adapter = pretrained.orb_v3_conservative_omol(
     device='cuda',
     train=True,
     train_reference_energies=False,  # Fixed reference energies

FINETUNING_GUIDE.md

@@ -4,7 +4,7 @@ We provide several pretrained models that can be used to calculate energies, for
 
 ### OrbMol Models
 
-These models are a continuation of the [`orb-v3`](#v3-models) series trained on the [Open Molecules 2025 (OMol25)](https://arxiv.org/pdf/2505.08762) dataset—over 100M high-accuracy DFT calculations (ωB97M-V/def2-TZVPD) on diverse molecular systems including metal complexes, biomolecules, and electrolytes. Note: The training data does not contain periodic systmems and these models have not been carefully tested on periodic systems. 
+These models are a continuation of the [`orb-v3`](#v3-models) series trained on the [Open Molecules 2025 (OMol25)](https://arxiv.org/pdf/2505.08762) dataset—over 100M high-accuracy DFT calculations (ωB97M-V/def2-TZVPD) on diverse molecular systems including metal complexes, biomolecules, and electrolytes. Note: The training data does not contain periodic systems and these models have not been carefully tested on periodic systems. 
 
 There are two options:
 * `orb-v3-conservative-omol`
@@ -15,11 +15,11 @@ See below for more explanation of this naming convention. Both models have `inf`
 ### [V3 Models](https://arxiv.org/abs/2504.06231)
 
 V3 models use the following naming convention: ```orb-v3-X-Y-Z``` where:
-- `X`: Model type - `direct` or `conservative`. Conservative models compute forces and stress via backpropagation, which is a physically motivated choice that appears necessary for certain types of simulation such as NVE Molecular dynamics. Conservative models are signficantly slower and use more memory than their direct counterparts.
+- `X`: Model type - `direct` or `conservative`. Conservative models compute forces and stress via backpropagation, which is a physically motivated choice that appears necessary for certain types of simulation such as NVE Molecular dynamics. Conservative models are significantly slower and use more memory than their direct counterparts.
 
-- `Y`: Maximum neighbors per atom: `20` or `inf`. A finite cutoff of `20` induces discontinuties in the PES, which can lead to significant inaccuracies for certain types of highly sensitive calculations (e.g. calculations involving Hessians). However, finite cutoffs reduce the amount of edge processing in the network, reducing latency and memory use.
+- `Y`: Maximum neighbors per atom: `20` or `inf`. A finite cutoff of `20` induces discontinuities in the PES, which can lead to significant inaccuracies for certain types of highly sensitive calculations (e.g. calculations involving Hessians). However, finite cutoffs reduce the amount of edge processing in the network, reducing latency and memory use.
 
-- `Z`: Training dataset - `omat` or `mpa`. Both of these dataset consist of small bulk crystal structures. We find that models trained on such data can generalise reasonably well to non-periodic systems (organic molecules) or partially periodic systems (slabs), but caution is advised in these scenarios.
+- `Z`: Training dataset - `omat` or `mpa`. Both of these datasets consist of small bulk crystal structures. We find that models trained on such data can generalise reasonably well to non-periodic systems (organic molecules) or partially periodic systems (slabs), but caution is advised in these scenarios.
 
 #### Features:
 
@@ -35,7 +35,7 @@ V3 models use the following naming convention: ```orb-v3-X-Y-Z``` where:
 #### Advice / Caveats
 
 - Consider using `orb-v3-conservative-120-omat` for initial testing, specifying `precision='float32-highest'` when loading the model. This is the most computational expensive but accurate configuration. If this level of accuracy meets your needs, then other models and precisions can be investigated to improve speed and system-size scalability.
-- We do not advise using the `-mpa` models unless they are required for compatability with benchmarks (for example, Matbench Discovery). They are generally less performant.
+- We do not advise using the `-mpa` models unless they are required for compatibility with benchmarks (for example, Matbench Discovery). They are generally less performant.
 - Orb-v3 models are **compiled** by default and use Pytorch's dynamic batching, which means that they do not need to recompile as graph sizes change. However, the first call to the model will be slower, as the graph is compiled by torch.
 
 ### [V2 Models](https://arxiv.org/abs/2410.22570)
@@ -53,4 +53,4 @@ V3 models use the following naming convention: ```orb-v3-X-Y-Z``` where:
 
 ### [V1 Models](https://arxiv.org/abs/2410.22570)
 
-Our initial release. These models were state of the art performance on the matbench discovery dataset at time of release, but have since been superceeded and removed.
+Our initial release. These models were state of the art performance on the matbench discovery dataset at time of release, but have since been superseded and removed.