Update public repo from private (#61)

* Bump torch, ase, scipy versions
* Remove pynanoflann
* Added model compilation
* PropertyDefinition improvements
* Adaptive supercell creation
* Various model implementation improvements
* New tests

Author: vsimkus

.github/workflows/publish.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
     if: startsWith(github.ref, 'refs/tags/') 
     steps:
     - uses: actions/checkout@v4
@@ -27,7 +27,7 @@ jobs:
         cache-dependency-path: pyproject.toml
     - name: Install dependencies
       run: |
-        pip install '.[test]' "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"
+        pip install '.[test]'
 
     - name: Run tests
       run: |

.github/workflows/test.yml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
     steps:
     - uses: actions/checkout@v4
     - name: Update version
@@ -24,7 +24,7 @@ jobs:
         cache-dependency-path: pyproject.toml
     - name: Install dependencies
       run: |
-        pip install '.[test]' "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"
+        pip install '.[test]'
 
     - name: Run tests
       run: |

.gitignore

@@ -160,3 +160,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Private
+datasets/

CONTRIBUTING.md

@@ -1,4 +1,13 @@
-### Tests
+### 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:
+
+```bash
+pip install poetry  # Install Poetry if you don't have it
+poetry install
+```
+
+### 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:
 

README.md

@@ -18,9 +18,7 @@
 
 ```bash
 pip install orb-models
-pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168",
 ```
-Pynanoflann is not available on PyPI, so you must install it from the git repository.
 
 Orb models are expected to work on MacOS and Linux. Windows support is not guaranteed.
 
@@ -74,9 +72,9 @@ result = orbff.predict(graph)
 # Convert to ASE atoms (unbatches the results and transfers to cpu if necessary)
 atoms = atomic_system.atom_graphs_to_ase_atoms(
     graph,
-    energy=result["graph_pred"],
-    forces=result["node_pred"],
-    stress=result["stress_pred"]
+    energy=result["energy"],
+    forces=result["forces"],
+    stress=result["stress"]
 )
 ```
 
@@ -94,7 +92,7 @@ orbff = pretrained.orb_v2(device=device) # or choose another model using ORB_PRE
 calc = ORBCalculator(orbff, device=device)
 atoms = bulk('Cu', 'fcc', a=3.58, cubic=True)
 
-atoms.set_calculator(calc)
+atoms.calc = calc
 atoms.get_potential_energy()
 ```
 

examples/NaClWaterMD.py

@@ -1,13 +1,11 @@
 import torch
-import numpy as np
 from ase.io import read, write
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.md.langevin import Langevin
 from ase import units
-from orb_models.forcefield import atomic_system, pretrained
-from ase.calculators.calculator import Calculator, all_properties
-from ase.build import molecule, make_supercell
 from ase.md import MDLogger
+
+from orb_models.forcefield import pretrained
 from orb_models.forcefield.calculator import ORBCalculator
 
 
@@ -34,7 +32,7 @@ def run_md_simulation(
     log_interval: int = 1,
 ):
     """Run molecular dynamics simulation with specified parameters.
-    
+
     Args:
         input_file: Path to input XYZ file
         cell_size: Size of cubic simulation cell
@@ -47,17 +45,14 @@ def run_md_simulation(
     """
     # Set up device
     device = setup_device()
-    
+
     # Read in the system from file and set the cell size and pbc
     atoms = read(input_file)
     atoms.set_cell([cell_size] * 3)
     atoms.set_pbc([True] * 3)
 
     # Set the calculator
-    atoms.calc = ORBCalculator(
-        model=pretrained.orb_d3_v2(),
-        device=device
-    )
+    atoms.calc = ORBCalculator(model=pretrained.orb_d3_v2(), device=device)
 
     # Set the initial velocities
     MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_K)
@@ -67,8 +62,7 @@ def run_md_simulation(
 
     # Define output functions and attach to dynamics
     dyn.attach(
-        lambda: write('NaClWaterMD.xyz', atoms, append=True), 
-        interval=traj_interval
+        lambda: write("NaClWaterMD.xyz", atoms, append=True), interval=traj_interval
     )
     dyn.attach(MDLogger(dyn, atoms, "md_nvt.log"), interval=log_interval)
 
@@ -82,4 +76,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()