Add ANI2x (#48)

* add ani and first successful calculator; still unable to change device

* add torch ani diatomic curves

* update readme about lfs

* reformat legend

.github/README.md

@@ -88,6 +88,8 @@ MLIP Arena is now in pre-alpha. If you're interested in joining the effort, plea
 ### Development
 
 ```
+git lfs fetch --all
+git lfs pull
 streamlit run serve/app.py
 ```
 

mlip_arena/models/externals/ani.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import yaml
+from pathlib import Path
+
+from ase.calculators.calculator import all_changes
+from torchani.ase import Calculator as ANICalculator
+from torchani.models import BuiltinEnsemble
+
+from mlip_arena.models.utils import get_freer_device
+
+
+
+with open(Path(__file__).parents[1] / "registry.yaml", encoding="utf-8") as f:
+    REGISTRY = yaml.safe_load(f)
+
+
+class ANI2x(ANICalculator):
+    def __init__(
+        self,
+        checkpoint=REGISTRY["ANI2x"]["checkpoint"],
+        device: str | None = None,
+        periodic_table_index=False,
+        **kwargs,
+    ):
+        self.device = device or str(get_freer_device())
+        
+        ensemble = BuiltinEnsemble._from_neurochem_resources(
+            checkpoint, periodic_table_index
+        )
+        # TODO: ANICalculator does not offer API to change device
+        # ensemble.species.device = self.device
+        super().__init__(ensemble.species, ensemble, **kwargs)
+    
+    
+    def calculate(
+        self, atoms=None, properties=['energy', 'forces', 'stress'], system_changes=all_changes
+    ):
+        super().calculate(atoms, properties, system_changes)

mlip_arena/models/registry.yaml

@@ -93,25 +93,27 @@ MatterSim:
   npt: true
   license: MIT
 
-ORB:
+ORBv2:
   module: externals
-  class: ORB
+  class: ORBv2
   family: orb
-  package: orb-models==0.3.1
-  checkpoint: orbff-v1-20240827.ckpt
-  username: cyrusyc
-  last-update: 2024-03-25T14:30:00
-  datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
+  package: orb-models==0.4.0
+  checkpoint: orb-v2-20241011.ckpt
+  username: 
+  last-update: 2024-10-29T00:00:00
+  datetime: 2024-10-29T00:00:00 # TODO: Fake datetime
   datasets:
-    - MPTrj # TODO: fake HF dataset repo
+    - MPTrj
     - Alexandria
+  cpu-tasks:
+    - eos_alloy
   gpu-tasks:
     - homonuclear-diatomics
     - combustion
     - stability
   github: https://github.com/orbital-materials/orb-models
   doi: 
-  date: 2024-09-03
+  date: 2024-10-15
   prediction: EFS
   nvt: true
   npt: true
@@ -252,6 +254,27 @@ MACE-OFF(M):
   npt: true
   license: ASL
 
+ANI2x:
+  module: externals
+  class: ANI2x
+  family: ani
+  package: torchani==2.2.4
+  checkpoint: ani-2x_8x.info
+  username: cyrusyc
+  last-update: 2024-12-11T16:00:00
+  datetime: 2024-12-11T16:00:00 # TODO: Fake datetime
+  datasets:
+  cpu-tasks:
+  gpu-tasks:
+    - homonuclear-diatomics
+  github: https://github.com/aiqm/torchani
+  doi: https://www.nature.com/articles/s41598-024-62242-5
+  date: 2024-05-23
+  prediction: EFS
+  nvt: true
+  npt: false
+  license: MIT
+
 ALIGNN:
   module: externals
   class: ALIGNN
@@ -294,27 +317,25 @@ DeepMD:
   npt: true
   license:
 
-ORBv2:
+ORB:
   module: externals
-  class: ORBv2
+  class: ORB
   family: orb
-  package: orb-models==0.4.0
-  checkpoint: orb-v2-20241011.ckpt
-  username: 
-  last-update: 2024-10-29T00:00:00
-  datetime: 2024-10-29T00:00:00 # TODO: Fake datetime
+  package: orb-models==0.3.1
+  checkpoint: orbff-v1-20240827.ckpt
+  username: cyrusyc
+  last-update: 2024-03-25T14:30:00
+  datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
   datasets:
-    - MPTrj
+    - MPTrj # TODO: fake HF dataset repo
     - Alexandria
-  cpu-tasks:
-    - eos_alloy
   gpu-tasks:
     - homonuclear-diatomics
     - combustion
     - stability
   github: https://github.com/orbital-materials/orb-models
   doi: 
-  date: 2024-10-15
+  date: 2024-09-03
   prediction: EFS
   nvt: true
   npt: true

pyproject.toml

@@ -59,6 +59,7 @@ test = [
     "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168",
     "alignn==2024.5.27",
     "mattersim==1.0.0rc9",
+    "torchani==2.2.4",
     "pytest",
     "pytest-xdist",
     "prefect==3.1.11",

serve/tasks/homonuclear-diatomics.py

@@ -30,7 +30,7 @@ valid_models = [
 mlip_methods = container.multiselect(
     "MLIPs",
     valid_models,
-    ["MACE-MP(M)", "CHGNet", "M3GNet", "MatterSim", "SevenNet", "ORBv2", "eqV2(OMat)"],
+    ["MACE-MP(M)", "CHGNet", "M3GNet", "MatterSim", "SevenNet", "ORBv2", "eqV2(OMat)", "ANI2x"],
 )
 dft_methods = container.multiselect("DFT Methods", ["PBE"], ["PBE"])
 
@@ -228,15 +228,15 @@ def get_plots(df, energy_plot: bool, force_plot: bool, method_color_mapping: dic
         fig.update_layout(
             showlegend=True,
             legend=dict(
-                orientation="v",
-                x=0.95,
+                orientation="h",
+                x=1.0,
                 xanchor="right",
                 y=1,
                 yanchor="top",
                 bgcolor="rgba(0, 0, 0, 0)",
                 # traceorder='reversed',
-                # entrywidth=0.3,
-                # entrywidthmode='fraction',
+                entrywidth=0.4,
+                entrywidthmode="fraction",
             ),
             title_text=f"{name}",
             title_x=0.5,