Alchemiops KNN+D3, TorchSim, Refactor

[vsimkus]

Porting internal changes to orb-models:

• Refactored codebase, hence why there are so many changes
• Added knn_alchemi for graph construction, and set to default
• Added a D3 dispersion correction module using alchemiops (and added an example in readme)
• Added a TorchSim wrapper module (and added an example in readme, and made the torch-sim dependency optional)
• Removed loss_weights.get(target_name, default) pattern, instead setting the defaults in pretrained.py when loading the model.

Dependency/CI changes:

• uv instead of poetry
• Updated CI to run ruff linter checks and parallelise pytest
• Requires python >=3.12 due to torch_sim
• Requires torch >= 2.8.0 due to nvalchemiops

Checks:

• README.md updated and examples working
• NaClWaterMD.py works
• Tests run
• Finetune runs
• Loading finetuned model works
• Speed test script works
• Check docker image works (Note: had to use torch==2.10 base image since the torch==2.8 images come with python <3.12, and upgrading python would mean that we invalidate the base pytorch installation, hence losing the point of using the base image in the first place.)

Regressions:

• pretrained.orb_v3...() now returns the model and the AtomsAdapter
  • This is a slight regression from Switch from GraphRegressor to DirectForcefieldRegressor #74 (comment), but we need to return atoms adapter now since we don't have SystemConfig anymore

[vsimkus]

• Fix custom ref energies shape in finetune.py #146
• Alchemiops KNN+D3, TorchSim, Refactor #143 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[benrhodes26]

Add a small update about improvements, plugging alchemiops repo?

[vsimkus]

Added this:

**February 2026**: Improved GPU-accelerated graph construction with [ALCHEMI Toolkit-Ops](https://github.com/NVIDIA/nvalchemi-toolkit-ops) and batched simulation with [TorchSim](https://github.com/TorchSim/torch-sim):

* Alchemi-based graph construction (GPU-accelerated, up to 12x faster for large single systems, and sub-linear batch scaling delivering >100x graph construction speed-up for large batches of small systems)
* TorchSim wrapper for batched optimisation and simulation, see [usage with TorchSim](#usage-with-torchsim)
* Alchemi-based D3 dispersion correction module, see [D3 correction](#d3-correction)

What do you think?

[benrhodes26]

🎉🎉🎉