From d1af82f1f5923d602595c1ba255a426aef010da4 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Sat, 30 Nov 2024 13:12:02 -0500
Subject: [PATCH 01/15] feat: add sevennet support to mlp recipes

---
 pyproject.toml                                    |  2 +-
 src/quacc/recipes/mlp/_base.py                    | 15 +++++++++++++--
 src/quacc/recipes/mlp/core.py                     |  9 +++++----
 src/quacc/recipes/mlp/phonons.py                  |  2 +-
 .../core/recipes/mlp_recipes/test_core_recipes.py |  6 ++++++
 tests/requirements-mlp.txt                        |  1 +
 6 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 69b3dc69dd..e963f1a087 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,7 @@ covalent = ["covalent>=0.234.1-rc.0; platform_system!='Windows'", "covalent-clou
 dask = ["dask[distributed]>=2023.12.1", "dask-jobqueue>=0.8.2"]
 defects = ["pymatgen-analysis-defects>=2023.8.22", "shakenbreak>=3.2.0"]
 jobflow = ["jobflow[fireworks]>=0.1.14", "jobflow-remote>=0.1.0"]
-mlp = ["matgl>=1.1.2", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0"]
+mlp = ["matgl>=1.1.2", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0", "sevenn>=0.10.1"]
 mp = ["atomate2>=0.0.14"]
 newtonnet = ["newtonnet>=1.1"]
 parsl = ["parsl[monitoring]>=2024.5.27; platform_system!='Windows'"]
diff --git a/src/quacc/recipes/mlp/_base.py b/src/quacc/recipes/mlp/_base.py
index 1d3e4dcffa..db231d70f9 100644
--- a/src/quacc/recipes/mlp/_base.py
+++ b/src/quacc/recipes/mlp/_base.py
@@ -16,7 +16,7 @@
 
 @lru_cache
 def pick_calculator(
-    method: Literal["mace-mp-0", "m3gnet", "chgnet"], **kwargs
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet"], **kwargs
 ) -> Calculator:
     """
     Adapted from `matcalc.util.get_universal_calculator`.
@@ -29,7 +29,8 @@ def pick_calculator(
         Custom kwargs for the underlying calculator. Set a value to
         `quacc.Remove` to remove a pre-existing key entirely. For a list of available
         keys, refer to the `mace.calculators.mace_mp`, `chgnet.model.dynamics.CHGNetCalculator`,
-        or `matgl.ext.ase.M3GNetCalculator` calculators.
+        `matgl.ext.ase.M3GNetCalculator`, or `sevenn.sevennet_calculator.SevenNetCalculator`
+        calculators.
 
     Returns
     -------
@@ -63,6 +64,16 @@ def pick_calculator(
         if "default_dtype" not in kwargs:
             kwargs["default_dtype"] = "float64"
         calc = mace_mp(**kwargs)
+    
+    elif method.lower() == "sevennet":
+        from sevenn import __version__
+        from sevenn.sevennet_calculator import SevenNetCalculator
+        
+        if "model" not in kwargs:
+            kwargs["model"] = "7net-0"
+        if "device" not in kwargs:
+            kwargs["device"] = "cpu"
+        calc = SevenNetCalculator(**kwargs)
 
     else:
         raise ValueError(f"Unrecognized {method=}.")
diff --git a/src/quacc/recipes/mlp/core.py b/src/quacc/recipes/mlp/core.py
index 635fd75925..50e5338322 100644
--- a/src/quacc/recipes/mlp/core.py
+++ b/src/quacc/recipes/mlp/core.py
@@ -21,7 +21,7 @@
 @job
 def static_job(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet"],
     properties: list[str] | None = None,
     additional_fields: dict[str, Any] | None = None,
     **calc_kwargs,
@@ -43,7 +43,7 @@ def static_job(
         Custom kwargs for the underlying calculator. Set a value to
         `quacc.Remove` to remove a pre-existing key entirely. For a list of available
         keys, refer to the `mace.calculators.mace_mp`, `chgnet.model.dynamics.CHGNetCalculator`,
-        or `matgl.ext.ase.M3GNetCalculator` calculators.
+        `matgl.ext.ase.M3GNetCalculator`, or `sevenn.sevennet_calculator.SevenNetCalculator` calculators.
 
     Returns
     -------
@@ -63,7 +63,7 @@ def static_job(
 @job
 def relax_job(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet"],
     relax_cell: bool = False,
     opt_params: OptParams | None = None,
     additional_fields: dict[str, Any] | None = None,
@@ -89,7 +89,8 @@ def relax_job(
         Custom kwargs for the underlying calculator. Set a value to
         `quacc.Remove` to remove a pre-existing key entirely. For a list of available
         keys, refer to the `mace.calculators.mace_mp`, `chgnet.model.dynamics.CHGNetCalculator`,
-        or `matgl.ext.ase.M3GNetCalculator` calculators.
+        `matgl.ext.ase.M3GNetCalculator`, or `sevenn.sevennet_calculator.SevenNetCalculator`
+        calculators.
 
     Returns
     -------
diff --git a/src/quacc/recipes/mlp/phonons.py b/src/quacc/recipes/mlp/phonons.py
index 87f67225b6..09c9feced0 100644
--- a/src/quacc/recipes/mlp/phonons.py
+++ b/src/quacc/recipes/mlp/phonons.py
@@ -33,7 +33,7 @@
 )
 def phonon_flow(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet"],
     symprec: float = 1e-4,
     min_lengths: float | tuple[float, float, float] | None = 20.0,
     supercell_matrix: (
diff --git a/tests/core/recipes/mlp_recipes/test_core_recipes.py b/tests/core/recipes/mlp_recipes/test_core_recipes.py
index 0c3ae96ecd..de42214e4d 100644
--- a/tests/core/recipes/mlp_recipes/test_core_recipes.py
+++ b/tests/core/recipes/mlp_recipes/test_core_recipes.py
@@ -21,6 +21,9 @@
 if has_chgnet := find_spec("chgnet"):
     methods.append("chgnet")
 
+if has_sevennet := find_spec("sevenn"):
+    methods.append("sevennet")
+
 
 @pytest.mark.skipif(has_chgnet is None, reason="chgnet not installed")
 def test_bad_method():
@@ -48,6 +51,7 @@ def test_static_job(tmp_path, monkeypatch, method):
         "chgnet": -4.083308219909668,
         "m3gnet": -4.0938973,
         "mace-mp-0": -4.083906650543213,
+        "sevennet": -4.096191883087158,
     }
     atoms = bulk("Cu")
     output = static_job(atoms, method=method)
@@ -68,6 +72,7 @@ def test_relax_job(tmp_path, monkeypatch, method):
         "chgnet": -32.665428161621094,
         "m3gnet": -32.75003433227539,
         "mace-mp-0": -32.6711566550002,
+        "sevennet": -32.76924133300781,
     }
 
     atoms = bulk("Cu") * (2, 2, 2)
@@ -108,6 +113,7 @@ def test_relax_cell_job(tmp_path, monkeypatch, method):
         "chgnet": -32.66698455810547,
         "m3gnet": -32.750858306884766,
         "mace-mp-0": -32.67840391814377,
+        "sevennet": -32.76963806152344,
     }
 
     atoms = bulk("Cu") * (2, 2, 2)
diff --git a/tests/requirements-mlp.txt b/tests/requirements-mlp.txt
index 6f157ccd87..bfb21bc2f8 100644
--- a/tests/requirements-mlp.txt
+++ b/tests/requirements-mlp.txt
@@ -1,3 +1,4 @@
 chgnet==0.4.0
 mace-torch==0.3.8
 torch-dftd==0.5.1
+sevenn==0.10.1

From ca05be0341e701a6f5a192a17ef60797d7869bd8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 30 Nov 2024 18:21:42 +0000
Subject: [PATCH 02/15] pre-commit auto-fixes

---
 src/quacc/recipes/mlp/_base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/quacc/recipes/mlp/_base.py b/src/quacc/recipes/mlp/_base.py
index db231d70f9..7434db21cb 100644
--- a/src/quacc/recipes/mlp/_base.py
+++ b/src/quacc/recipes/mlp/_base.py
@@ -64,11 +64,11 @@ def pick_calculator(
         if "default_dtype" not in kwargs:
             kwargs["default_dtype"] = "float64"
         calc = mace_mp(**kwargs)
-    
+
     elif method.lower() == "sevennet":
         from sevenn import __version__
         from sevenn.sevennet_calculator import SevenNetCalculator
-        
+
         if "model" not in kwargs:
             kwargs["model"] = "7net-0"
         if "device" not in kwargs:

From d63362137feee5c95b3021c3a7a9bbbd254b05f4 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Sat, 30 Nov 2024 21:16:27 -0500
Subject: [PATCH 03/15] PR comments addressed: removed explicit declaration of
 sevennet default variables

---
 src/quacc/recipes/mlp/_base.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/quacc/recipes/mlp/_base.py b/src/quacc/recipes/mlp/_base.py
index 7434db21cb..4c6922d0b5 100644
--- a/src/quacc/recipes/mlp/_base.py
+++ b/src/quacc/recipes/mlp/_base.py
@@ -69,10 +69,6 @@ def pick_calculator(
         from sevenn import __version__
         from sevenn.sevennet_calculator import SevenNetCalculator
 
-        if "model" not in kwargs:
-            kwargs["model"] = "7net-0"
-        if "device" not in kwargs:
-            kwargs["device"] = "cpu"
         calc = SevenNetCalculator(**kwargs)
 
     else:

From 6fbc0a53c4b8c1a3299867205e462045faf65488 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Thu, 19 Dec 2024 13:25:28 -0500
Subject: [PATCH 04/15] orb-models support added

---
 pyproject.toml                                    |  2 +-
 src/quacc/recipes/mlp/_base.py                    | 15 ++++++++++++---
 src/quacc/recipes/mlp/core.py                     | 11 ++++++-----
 src/quacc/recipes/mlp/phonons.py                  |  2 +-
 .../core/recipes/mlp_recipes/test_core_recipes.py | 12 +++++++++---
 tests/requirements-mlp.txt                        |  1 +
 6 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e963f1a087..a8b93156e0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,7 @@ covalent = ["covalent>=0.234.1-rc.0; platform_system!='Windows'", "covalent-clou
 dask = ["dask[distributed]>=2023.12.1", "dask-jobqueue>=0.8.2"]
 defects = ["pymatgen-analysis-defects>=2023.8.22", "shakenbreak>=3.2.0"]
 jobflow = ["jobflow[fireworks]>=0.1.14", "jobflow-remote>=0.1.0"]
-mlp = ["matgl>=1.1.2", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0", "sevenn>=0.10.1"]
+mlp = ["matgl>=1.1.2", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0", "sevenn>=0.10.1", "orb-models>=4.1.0"]
 mp = ["atomate2>=0.0.14"]
 newtonnet = ["newtonnet>=1.1"]
 parsl = ["parsl[monitoring]>=2024.5.27; platform_system!='Windows'"]
diff --git a/src/quacc/recipes/mlp/_base.py b/src/quacc/recipes/mlp/_base.py
index 4c6922d0b5..d999b884f9 100644
--- a/src/quacc/recipes/mlp/_base.py
+++ b/src/quacc/recipes/mlp/_base.py
@@ -16,7 +16,7 @@
 
 @lru_cache
 def pick_calculator(
-    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet"], **kwargs
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb-models"], **kwargs
 ) -> Calculator:
     """
     Adapted from `matcalc.util.get_universal_calculator`.
@@ -29,8 +29,8 @@ def pick_calculator(
         Custom kwargs for the underlying calculator. Set a value to
         `quacc.Remove` to remove a pre-existing key entirely. For a list of available
         keys, refer to the `mace.calculators.mace_mp`, `chgnet.model.dynamics.CHGNetCalculator`,
-        `matgl.ext.ase.M3GNetCalculator`, or `sevenn.sevennet_calculator.SevenNetCalculator`
-        calculators.
+        `matgl.ext.ase.M3GNetCalculator`, `sevenn.sevennet_calculator.SevenNetCalculator`, or
+        `orb_models.forcefield.calculator.ORBCalculator` calculators.
 
     Returns
     -------
@@ -71,6 +71,15 @@ def pick_calculator(
 
         calc = SevenNetCalculator(**kwargs)
 
+    elif method.lower() == "orb-models":
+        from orb_models import __version__
+        from orb_models.forcefield import pretrained
+        from orb_models.forcefield.calculator import ORBCalculator
+
+        orb_model = kwargs.pop("model", "orb_v2")
+        orbff = getattr(pretrained, orb_model)()
+        calc = ORBCalculator(model=orbff, **kwargs)
+
     else:
         raise ValueError(f"Unrecognized {method=}.")
 
diff --git a/src/quacc/recipes/mlp/core.py b/src/quacc/recipes/mlp/core.py
index 50e5338322..192e486c98 100644
--- a/src/quacc/recipes/mlp/core.py
+++ b/src/quacc/recipes/mlp/core.py
@@ -21,7 +21,7 @@
 @job
 def static_job(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb-models"],
     properties: list[str] | None = None,
     additional_fields: dict[str, Any] | None = None,
     **calc_kwargs,
@@ -43,7 +43,8 @@ def static_job(
         Custom kwargs for the underlying calculator. Set a value to
         `quacc.Remove` to remove a pre-existing key entirely. For a list of available
         keys, refer to the `mace.calculators.mace_mp`, `chgnet.model.dynamics.CHGNetCalculator`,
-        `matgl.ext.ase.M3GNetCalculator`, or `sevenn.sevennet_calculator.SevenNetCalculator` calculators.
+        `matgl.ext.ase.M3GNetCalculator`, `sevenn.sevennet_calculator.SevenNetCalculator`, or
+        `orb_models.forcefield.calculator.ORBCalculator` calculators.
 
     Returns
     -------
@@ -63,7 +64,7 @@ def static_job(
 @job
 def relax_job(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb-models"],
     relax_cell: bool = False,
     opt_params: OptParams | None = None,
     additional_fields: dict[str, Any] | None = None,
@@ -89,8 +90,8 @@ def relax_job(
         Custom kwargs for the underlying calculator. Set a value to
         `quacc.Remove` to remove a pre-existing key entirely. For a list of available
         keys, refer to the `mace.calculators.mace_mp`, `chgnet.model.dynamics.CHGNetCalculator`,
-        `matgl.ext.ase.M3GNetCalculator`, or `sevenn.sevennet_calculator.SevenNetCalculator`
-        calculators.
+        `matgl.ext.ase.M3GNetCalculator`, `sevenn.sevennet_calculator.SevenNetCalculator`, or
+        `orb_models.forcefield.calculator.ORBCalculator` calculators.
 
     Returns
     -------
diff --git a/src/quacc/recipes/mlp/phonons.py b/src/quacc/recipes/mlp/phonons.py
index 09c9feced0..7348f1029b 100644
--- a/src/quacc/recipes/mlp/phonons.py
+++ b/src/quacc/recipes/mlp/phonons.py
@@ -33,7 +33,7 @@
 )
 def phonon_flow(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb-models"],
     symprec: float = 1e-4,
     min_lengths: float | tuple[float, float, float] | None = 20.0,
     supercell_matrix: (
diff --git a/tests/core/recipes/mlp_recipes/test_core_recipes.py b/tests/core/recipes/mlp_recipes/test_core_recipes.py
index de42214e4d..5102aa48bf 100644
--- a/tests/core/recipes/mlp_recipes/test_core_recipes.py
+++ b/tests/core/recipes/mlp_recipes/test_core_recipes.py
@@ -24,6 +24,9 @@
 if has_sevennet := find_spec("sevenn"):
     methods.append("sevennet")
 
+if has_orb_models := find_spec("orb_models"):
+    methods.append("orb-models")
+
 
 @pytest.mark.skipif(has_chgnet is None, reason="chgnet not installed")
 def test_bad_method():
@@ -52,10 +55,11 @@ def test_static_job(tmp_path, monkeypatch, method):
         "m3gnet": -4.0938973,
         "mace-mp-0": -4.083906650543213,
         "sevennet": -4.096191883087158,
+        "orb-models": -4.093477725982666,
     }
     atoms = bulk("Cu")
     output = static_job(atoms, method=method)
-    assert output["results"]["energy"] == pytest.approx(ref_energy[method])
+    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-5)
     assert np.shape(output["results"]["forces"]) == (1, 3)
     assert output["atoms"] == atoms
 
@@ -73,12 +77,13 @@ def test_relax_job(tmp_path, monkeypatch, method):
         "m3gnet": -32.75003433227539,
         "mace-mp-0": -32.6711566550002,
         "sevennet": -32.76924133300781,
+        "orb-models": -32.7361946105957,
     }
 
     atoms = bulk("Cu") * (2, 2, 2)
     atoms[0].position += 0.1
     output = relax_job(atoms, method=method)
-    assert output["results"]["energy"] == pytest.approx(ref_energy[method])
+    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-5)
     assert np.shape(output["results"]["forces"]) == (8, 3)
     assert output["atoms"] != atoms
     assert output["atoms"].get_volume() == pytest.approx(atoms.get_volume())
@@ -114,12 +119,13 @@ def test_relax_cell_job(tmp_path, monkeypatch, method):
         "m3gnet": -32.750858306884766,
         "mace-mp-0": -32.67840391814377,
         "sevennet": -32.76963806152344,
+        "orb-models": -32.73428726196289,
     }
 
     atoms = bulk("Cu") * (2, 2, 2)
     atoms[0].position += 0.1
     output = relax_job(atoms, method=method, relax_cell=True)
-    assert output["results"]["energy"] == pytest.approx(ref_energy[method])
+    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-5)
     assert np.shape(output["results"]["forces"]) == (8, 3)
     assert output["atoms"] != atoms
     assert output["atoms"].get_volume() != pytest.approx(atoms.get_volume())
diff --git a/tests/requirements-mlp.txt b/tests/requirements-mlp.txt
index bfb21bc2f8..873efbdf08 100644
--- a/tests/requirements-mlp.txt
+++ b/tests/requirements-mlp.txt
@@ -2,3 +2,4 @@ chgnet==0.4.0
 mace-torch==0.3.8
 torch-dftd==0.5.1
 sevenn==0.10.1
+orb-models=4.1.0

From df930ec50b398ec57c84689bd8f4cfd02783f331 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 19 Dec 2024 18:37:26 +0000
Subject: [PATCH 05/15] pre-commit auto-fixes

---
 tests/core/recipes/mlp_recipes/test_core_recipes.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/core/recipes/mlp_recipes/test_core_recipes.py b/tests/core/recipes/mlp_recipes/test_core_recipes.py
index c34c63d31f..5102aa48bf 100644
--- a/tests/core/recipes/mlp_recipes/test_core_recipes.py
+++ b/tests/core/recipes/mlp_recipes/test_core_recipes.py
@@ -27,6 +27,7 @@
 if has_orb_models := find_spec("orb_models"):
     methods.append("orb-models")
 
+
 @pytest.mark.skipif(has_chgnet is None, reason="chgnet not installed")
 def test_bad_method():
     atoms = bulk("Cu")

From 132a78c86cafb9471b79ad1c727afd62093cb413 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Thu, 19 Dec 2024 13:40:50 -0500
Subject: [PATCH 06/15] requirement-mlp.txt is fixed

---
 tests/requirements-mlp.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/requirements-mlp.txt b/tests/requirements-mlp.txt
index bc23f3ccb4..d35f692663 100644
--- a/tests/requirements-mlp.txt
+++ b/tests/requirements-mlp.txt
@@ -2,4 +2,4 @@ chgnet==0.4.0
 mace-torch==0.3.9
 torch-dftd==0.5.1
 sevenn==0.10.3
-orb-models=4.1.0
+orb-models==4.1.0

From cc4afa2f625ac444c95a1aeeb0501611ca794a09 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Thu, 19 Dec 2024 13:55:21 -0500
Subject: [PATCH 07/15] pynanoflann requirements added

---
 pyproject.toml             | 2 +-
 tests/requirements-mlp.txt | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8311d67fc8..c2979d9e1f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,7 @@ covalent = ["covalent>=0.234.1-rc.0; platform_system!='Windows'", "covalent-clou
 dask = ["dask[distributed]>=2023.12.1", "dask-jobqueue>=0.8.2"]
 defects = ["pymatgen-analysis-defects>=2024.10.22", "shakenbreak>=3.2.0"]
 jobflow = ["jobflow[fireworks]>=0.1.14", "jobflow-remote>=0.1.0"]
-mlp = ["matgl>=1.1.2", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0", "sevenn>=0.10.1", "orb-models>=4.1.0"]
+mlp = ["matgl>=1.1.2", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0", "sevenn>=0.10.1", "orb-models>=4.1.0", "pynanoflann @ git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"]
 mp = ["atomate2>=0.0.14"]
 newtonnet = ["newtonnet>=1.1"]
 parsl = ["parsl[monitoring]>=2024.5.27; platform_system!='Windows'"]
diff --git a/tests/requirements-mlp.txt b/tests/requirements-mlp.txt
index d35f692663..f275001482 100644
--- a/tests/requirements-mlp.txt
+++ b/tests/requirements-mlp.txt
@@ -3,3 +3,4 @@ mace-torch==0.3.9
 torch-dftd==0.5.1
 sevenn==0.10.3
 orb-models==4.1.0
+pynanoflann @ git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168

From 20fff6db67e2b014edec6700b1ac8ca7e3bd5548 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Thu, 19 Dec 2024 14:00:55 -0500
Subject: [PATCH 08/15] turned out rel=1e-5 was not enough for the pytest
 comparison

---
 tests/core/recipes/mlp_recipes/test_core_recipes.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/core/recipes/mlp_recipes/test_core_recipes.py b/tests/core/recipes/mlp_recipes/test_core_recipes.py
index 5102aa48bf..4ae24d4965 100644
--- a/tests/core/recipes/mlp_recipes/test_core_recipes.py
+++ b/tests/core/recipes/mlp_recipes/test_core_recipes.py
@@ -59,7 +59,7 @@ def test_static_job(tmp_path, monkeypatch, method):
     }
     atoms = bulk("Cu")
     output = static_job(atoms, method=method)
-    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-5)
+    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-4)
     assert np.shape(output["results"]["forces"]) == (1, 3)
     assert output["atoms"] == atoms
 
@@ -83,7 +83,7 @@ def test_relax_job(tmp_path, monkeypatch, method):
     atoms = bulk("Cu") * (2, 2, 2)
     atoms[0].position += 0.1
     output = relax_job(atoms, method=method)
-    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-5)
+    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-4)
     assert np.shape(output["results"]["forces"]) == (8, 3)
     assert output["atoms"] != atoms
     assert output["atoms"].get_volume() == pytest.approx(atoms.get_volume())
@@ -125,7 +125,7 @@ def test_relax_cell_job(tmp_path, monkeypatch, method):
     atoms = bulk("Cu") * (2, 2, 2)
     atoms[0].position += 0.1
     output = relax_job(atoms, method=method, relax_cell=True)
-    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-5)
+    assert output["results"]["energy"] == pytest.approx(ref_energy[method], rel=1e-4)
     assert np.shape(output["results"]["forces"]) == (8, 3)
     assert output["atoms"] != atoms
     assert output["atoms"].get_volume() != pytest.approx(atoms.get_volume())

From 4b3dd59cbcaedc6c8065aa7be9e828c9a0192108 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Thu, 19 Dec 2024 23:18:15 -0500
Subject: [PATCH 09/15] PR comments addressed

---
 pyproject.toml                                |  2 +-
 src/quacc/recipes/mlp/_base.py                | 20 +++++++++++++++----
 .../recipes/mlp_recipes/test_core_recipes.py  | 10 +++++-----
 tests/requirements-mlp.txt                    |  2 +-
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index c2979d9e1f..8311d67fc8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,7 @@ covalent = ["covalent>=0.234.1-rc.0; platform_system!='Windows'", "covalent-clou
 dask = ["dask[distributed]>=2023.12.1", "dask-jobqueue>=0.8.2"]
 defects = ["pymatgen-analysis-defects>=2024.10.22", "shakenbreak>=3.2.0"]
 jobflow = ["jobflow[fireworks]>=0.1.14", "jobflow-remote>=0.1.0"]
-mlp = ["matgl>=1.1.2", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0", "sevenn>=0.10.1", "orb-models>=4.1.0", "pynanoflann @ git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"]
+mlp = ["matgl>=1.1.2", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0", "sevenn>=0.10.1", "orb-models>=4.1.0"]
 mp = ["atomate2>=0.0.14"]
 newtonnet = ["newtonnet>=1.1"]
 parsl = ["parsl[monitoring]>=2024.5.27; platform_system!='Windows'"]
diff --git a/src/quacc/recipes/mlp/_base.py b/src/quacc/recipes/mlp/_base.py
index d999b884f9..85146cc3bd 100644
--- a/src/quacc/recipes/mlp/_base.py
+++ b/src/quacc/recipes/mlp/_base.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from functools import lru_cache
+from importlib.util import find_spec
 from logging import getLogger
 from typing import TYPE_CHECKING
 
@@ -16,15 +17,20 @@
 
 @lru_cache
 def pick_calculator(
-    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb-models"], **kwargs
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb"], **kwargs
 ) -> Calculator:
     """
     Adapted from `matcalc.util.get_universal_calculator`.
 
+    !!! Note
+
+        To use `orb` method, `pynanoflann` must be installed. To install `pynanoflann`,
+        run `pip install git+https://github.com/u1234x1234/pynanoflann.git`.
+
     Parameters
     ----------
     method
-        Name of the calculator to use
+        Name of the calculator to use.
     **kwargs
         Custom kwargs for the underlying calculator. Set a value to
         `quacc.Remove` to remove a pre-existing key entirely. For a list of available
@@ -71,12 +77,18 @@ def pick_calculator(
 
         calc = SevenNetCalculator(**kwargs)
 
-    elif method.lower() == "orb-models":
+    elif method.lower() == "orb":
+        if not find_spec("pynanoflann"):
+            raise ImportError(
+                """orb-models requires pynanoflann.
+                Install pynanoflann with `pip install git+https://github.com/u1234x1234/pynanoflann.git`.
+                """
+            )
         from orb_models import __version__
         from orb_models.forcefield import pretrained
         from orb_models.forcefield.calculator import ORBCalculator
 
-        orb_model = kwargs.pop("model", "orb_v2")
+        orb_model = kwargs.get("model", "orb_v2")
         orbff = getattr(pretrained, orb_model)()
         calc = ORBCalculator(model=orbff, **kwargs)
 
diff --git a/tests/core/recipes/mlp_recipes/test_core_recipes.py b/tests/core/recipes/mlp_recipes/test_core_recipes.py
index 4ae24d4965..5dcc7947f5 100644
--- a/tests/core/recipes/mlp_recipes/test_core_recipes.py
+++ b/tests/core/recipes/mlp_recipes/test_core_recipes.py
@@ -24,8 +24,8 @@
 if has_sevennet := find_spec("sevenn"):
     methods.append("sevennet")
 
-if has_orb_models := find_spec("orb_models"):
-    methods.append("orb-models")
+if has_orb := find_spec("orb_models"):
+    methods.append("orb")
 
 
 @pytest.mark.skipif(has_chgnet is None, reason="chgnet not installed")
@@ -55,7 +55,7 @@ def test_static_job(tmp_path, monkeypatch, method):
         "m3gnet": -4.0938973,
         "mace-mp-0": -4.083906650543213,
         "sevennet": -4.096191883087158,
-        "orb-models": -4.093477725982666,
+        "orb": -4.093477725982666,
     }
     atoms = bulk("Cu")
     output = static_job(atoms, method=method)
@@ -77,7 +77,7 @@ def test_relax_job(tmp_path, monkeypatch, method):
         "m3gnet": -32.75003433227539,
         "mace-mp-0": -32.6711566550002,
         "sevennet": -32.76924133300781,
-        "orb-models": -32.7361946105957,
+        "orb": -32.7361946105957,
     }
 
     atoms = bulk("Cu") * (2, 2, 2)
@@ -119,7 +119,7 @@ def test_relax_cell_job(tmp_path, monkeypatch, method):
         "m3gnet": -32.750858306884766,
         "mace-mp-0": -32.67840391814377,
         "sevennet": -32.76963806152344,
-        "orb-models": -32.73428726196289,
+        "orb": -32.73428726196289,
     }
 
     atoms = bulk("Cu") * (2, 2, 2)
diff --git a/tests/requirements-mlp.txt b/tests/requirements-mlp.txt
index f275001482..a7486eb78c 100644
--- a/tests/requirements-mlp.txt
+++ b/tests/requirements-mlp.txt
@@ -3,4 +3,4 @@ mace-torch==0.3.9
 torch-dftd==0.5.1
 sevenn==0.10.3
 orb-models==4.1.0
-pynanoflann @ git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168
+pynanoflann @ git+https://github.com/u1234x1234/pynanoflann.git

From 276ab8b3e7a72f8fd254968593381e10fa8c9233 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Thu, 19 Dec 2024 23:21:50 -0500
Subject: [PATCH 10/15] orb-models changed to orb in core

---
 src/quacc/recipes/mlp/core.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/quacc/recipes/mlp/core.py b/src/quacc/recipes/mlp/core.py
index 192e486c98..d6c47d47d1 100644
--- a/src/quacc/recipes/mlp/core.py
+++ b/src/quacc/recipes/mlp/core.py
@@ -21,7 +21,7 @@
 @job
 def static_job(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb-models"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb"],
     properties: list[str] | None = None,
     additional_fields: dict[str, Any] | None = None,
     **calc_kwargs,
@@ -64,7 +64,7 @@ def static_job(
 @job
 def relax_job(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb-models"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb"],
     relax_cell: bool = False,
     opt_params: OptParams | None = None,
     additional_fields: dict[str, Any] | None = None,

From 333e9de888ce8f9f1013e56979bbc68d5f6fd5d3 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Thu, 19 Dec 2024 23:52:38 -0500
Subject: [PATCH 11/15] changed orb-models in phonons.py to orb

---
 src/quacc/recipes/mlp/phonons.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/quacc/recipes/mlp/phonons.py b/src/quacc/recipes/mlp/phonons.py
index 7348f1029b..8223432bd9 100644
--- a/src/quacc/recipes/mlp/phonons.py
+++ b/src/quacc/recipes/mlp/phonons.py
@@ -33,7 +33,7 @@
 )
 def phonon_flow(
     atoms: Atoms,
-    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb-models"],
+    method: Literal["mace-mp-0", "m3gnet", "chgnet", "sevennet", "orb"],
     symprec: float = 1e-4,
     min_lengths: float | tuple[float, float, float] | None = 20.0,
     supercell_matrix: (

From 79567d1a45c4ab1a62cfecb45dcb18be0fe343a2 Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Sun, 22 Dec 2024 16:39:44 -0500
Subject: [PATCH 12/15] pin pynanoflann dependency to version 0.0.8 to resolve
 uv installation issue

---
 tests/requirements-mlp.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/requirements-mlp.txt b/tests/requirements-mlp.txt
index a7486eb78c..4300e05054 100644
--- a/tests/requirements-mlp.txt
+++ b/tests/requirements-mlp.txt
@@ -3,4 +3,4 @@ mace-torch==0.3.9
 torch-dftd==0.5.1
 sevenn==0.10.3
 orb-models==4.1.0
-pynanoflann @ git+https://github.com/u1234x1234/pynanoflann.git
+pynanoflann @ git+https://github.com/u1234x1234/pynanoflann.git@0.0.8

From eded2af6710cd42807ca9a77877e38392cf7c9ab Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Sun, 22 Dec 2024 17:10:37 -0500
Subject: [PATCH 13/15] switch pynanoflann dependency from u1234x1234 to
 dwastberg's fork which works on my local computer with uv pip install

---
 tests/requirements-mlp.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/requirements-mlp.txt b/tests/requirements-mlp.txt
index 4300e05054..583a8262e2 100644
--- a/tests/requirements-mlp.txt
+++ b/tests/requirements-mlp.txt
@@ -3,4 +3,4 @@ mace-torch==0.3.9
 torch-dftd==0.5.1
 sevenn==0.10.3
 orb-models==4.1.0
-pynanoflann @ git+https://github.com/u1234x1234/pynanoflann.git@0.0.8
+pynanoflann@git+https://github.com/dwastberg/pynanoflann

From a3240a65a50ee70dbcbf9072736b798a443b8d4d Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Sun, 22 Dec 2024 17:15:53 -0500
Subject: [PATCH 14/15] Changed the docstring to replace pynanoflann from
 u1234x1234 version to dwastberg's fork for consistency

---
 src/quacc/recipes/mlp/_base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/quacc/recipes/mlp/_base.py b/src/quacc/recipes/mlp/_base.py
index 85146cc3bd..c540d09909 100644
--- a/src/quacc/recipes/mlp/_base.py
+++ b/src/quacc/recipes/mlp/_base.py
@@ -25,7 +25,7 @@ def pick_calculator(
     !!! Note
 
         To use `orb` method, `pynanoflann` must be installed. To install `pynanoflann`,
-        run `pip install git+https://github.com/u1234x1234/pynanoflann.git`.
+        run `pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann"`.
 
     Parameters
     ----------

From af4bf514fb39c5580daa1a0724738b7cd59720bf Mon Sep 17 00:00:00 2001
From: Osman Mamun <mamun.che06@gmail.com>
Date: Sun, 22 Dec 2024 19:01:31 -0500
Subject: [PATCH 15/15] Test added for pick calculator when missing pynanoflann

---
 src/quacc/recipes/mlp/_base.py                    |  2 +-
 .../core/recipes/mlp_recipes/test_core_recipes.py | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/quacc/recipes/mlp/_base.py b/src/quacc/recipes/mlp/_base.py
index c540d09909..2207560f55 100644
--- a/src/quacc/recipes/mlp/_base.py
+++ b/src/quacc/recipes/mlp/_base.py
@@ -81,7 +81,7 @@ def pick_calculator(
         if not find_spec("pynanoflann"):
             raise ImportError(
                 """orb-models requires pynanoflann.
-                Install pynanoflann with `pip install git+https://github.com/u1234x1234/pynanoflann.git`.
+                Install pynanoflann with `pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann"`.
                 """
             )
         from orb_models import __version__
diff --git a/tests/core/recipes/mlp_recipes/test_core_recipes.py b/tests/core/recipes/mlp_recipes/test_core_recipes.py
index 5dcc7947f5..b25abe99e1 100644
--- a/tests/core/recipes/mlp_recipes/test_core_recipes.py
+++ b/tests/core/recipes/mlp_recipes/test_core_recipes.py
@@ -64,6 +64,21 @@ def test_static_job(tmp_path, monkeypatch, method):
     assert output["atoms"] == atoms
 
 
+def test_relax_job_missing_pynanoflann(monkeypatch):
+    def mock_find_spec(name):
+        if name == "pynanoflann":
+            return None
+        return find_spec(name)
+
+    import quacc.recipes.mlp._base
+
+    quacc.recipes.mlp._base.pick_calculator.cache_clear()
+    monkeypatch.setattr("importlib.util.find_spec", mock_find_spec)
+    monkeypatch.setattr("quacc.recipes.mlp._base.find_spec", mock_find_spec)
+    with pytest.raises(ImportError, match=r"orb-models requires pynanoflann"):
+        relax_job(bulk("Cu"), method="orb")
+
+
 @pytest.mark.parametrize("method", methods)
 def test_relax_job(tmp_path, monkeypatch, method):
     monkeypatch.chdir(tmp_path)