Update featomic (#8)

* updating to featomic
- fixes test
- fixes imports
- fixes calculator checks

updates tinyurl to models which comply with featomic

renames all models from ShiftML to ShiftMLrev

removes 1.0 model

* formatting

.gitignore

@@ -4,3 +4,4 @@ dist/
 __pycache__/
 *.egg-info/
 .tox/
+*.pt

README.md

@@ -19,7 +19,7 @@ from ase.build import bulk
 from shiftml.ase import ShiftML
 
 frame = bulk("C", "diamond", a=3.566)
-calculator = ShiftML("ShiftML1.0")
+calculator = ShiftML("ShiftML1.1rev")
 
 cs_iso = calculator.get_cs_iso(frame)
 
@@ -29,21 +29,20 @@ print(cs_iso)
 
 ## IMPORTANT: Install pre-instructions before PiPy release
 
-Rascaline-torch, one of the main dependence of ShiftML, requires CXX and Rust compilers to be built from source.
+Featomic-torch, one of the main dependence of ShiftML, requires CXX and Rust compilers to be built from source.
 Most systems come already with configured C/C++ compilers (make sure that some environment variables CC and CXX are set
 and gcc can be found), but Rust typically needs to be installed manually.
 For ease of use we strongly recommend to use some sort of package manager to install Rust, such as conda and a fresh environment.
 
 
 ```bash
 
-conda create -n shiftml python=3.10
+conda create -n shiftml python=3.12
 conda activate shiftml
 conda install -c conda-forge rust
 
 ```
 
-
 ## Installation
 
 To install ShiftML, you can use clone this repository and install it using pip, a pipy release will follow soon:
@@ -56,8 +55,8 @@ pip install --extra-index-url https://download.pytorch.org/whl/cpu .
 
 This project would not have been possible without the following packages:
 
-- Metadata and model handling: [metatensor](https://github.com/lab-cosmo/metatensor)
-- Atomic descriptor engine: [rascaline](https://github.com/Luthaf/rascaline)
+- Metadata and model handling: [metatensor](https://github.com/metatensor/metatensor)
+- Atomic descriptor engine: [featomic](https://github.com/metatensor/featomic)
 
 ## Documentation
 

pyproject.toml

@@ -8,10 +8,10 @@ authors = [
 dependencies = [
     "numpy<2.0.0",
     "ase==3.22.1",
-    "metatensor-operations<0.2.0,>=0.1.0",
+    "metatensor-operations<0.4.0,>=0.3.0",
     "metatensor[torch]",
     "platformdirs",
-    "rascaline-torch@git+https://github.com/luthaf/rascaline@e215461#subdirectory=python/rascaline-torch",
+    "featomic-torch @git+https://github.com/metatensor/featomic@937679e#subdirectory=python/featomic-torch",
 ]
 
 readme = "README.md"

src/shiftml/ase/calculator.py

@@ -12,25 +12,24 @@
 
 
 url_resolve = {
-    "ShiftML1.0": "https://tinyurl.com/3xwec68f",
-    "ShiftML1.1": "https://tinyurl.com/53ymkhvd",
-    "ShiftML2.0": "https://tinyurl.com/bdcp647w",
+    "ShiftML1.1rev": "https://tinyurl.com/msnss4ds",
+    "ShiftML2.0rev": "https://tinyurl.com/3axupmsd",
 }
 
 resolve_outputs = {
-    "ShiftML1.0": {"mtt::cs_iso": ModelOutput(quantity="", unit="ppm", per_atom=True)},
-    "ShiftML1.1": {"mtt::cs_iso": ModelOutput(quantity="", unit="ppm", per_atom=True)},
-    "ShiftML2.0": {
+    "ShiftML1.1rev": {
+        "mtt::cs_iso": ModelOutput(quantity="", unit="ppm", per_atom=True)
+    },
+    "ShiftML2.0rev": {
         "mtt::cs_iso": ModelOutput(quantity="", unit="ppm", per_atom=True),
         "mtt::cs_iso_std": ModelOutput(quantity="", unit="ppm", per_atom=True),
         "mtt::cs_iso_ensemble": ModelOutput(quantity="", unit="ppm", per_atom=True),
     },
 }
 
 resolve_fitted_species = {
-    "ShiftML1.0": set([1, 6, 7, 8, 16]),
-    "ShiftML1.1": set([1, 6, 7, 8, 16]),
-    "ShiftML2.0": set([1, 6, 7, 8, 9, 11, 12, 15, 16, 17, 19, 20]),
+    "ShiftML1.1rev": set([1, 6, 7, 8, 16]),
+    "ShiftML2.0rev": set([1, 6, 7, 8, 9, 11, 12, 15, 16, 17, 19, 20]),
 }
 
 
@@ -57,7 +56,7 @@ def __init__(self, model_version, force_download=False):
         ----------
         model_version : str
             The version of the ShiftML model to use. Supported versions are
-            "ShiftML1.0".
+            "ShiftML1.1rev" and "ShiftML2.0rev".
         force_download : bool, optional
             If True, the model will be downloaded even if it is already in the cache.
             The chache-dir will be determined via the platformdirs library and should
@@ -68,21 +67,21 @@ def __init__(self, model_version, force_download=False):
         try:
             # The rascline import is necessary because
             # it is required for the scripted model
-            import rascaline.torch
+            import featomic.torch
 
-            logging.info("rascaline version: {}".format(rascaline.torch.__version__))
-            logging.info("rascaline-torch is installed, importing rascaline-torch")
+            logging.info("featomic version: {}".format(featomic.torch.__version__))
+            logging.info("featomic-torch is installed, importing featomic-torch")
 
             assert (
-                rascaline.torch.__version__ == "0.1.0.dev558"
-            ), "wrong rascaline-torch installed"
+                featomic.torch.__version__ == "0.1.0.dev597"
+            ), "wrong featomic-torch installed"
 
         except ImportError:
             raise ImportError(
-                "rascaline-torch is required for ShiftML calculators,\
+                "featomic-torch is required for ShiftML calculators,\
                  please install it using\
-                 pip install git+https://github.com/luthaf/rascaline#subdirectory\
-                 =python/rascaline-torch"
+                 pip install git+https://github.com/metatensor/featomic#subdirectory\
+                 =python/featomic-torch"
             )
 
         try:

tests/test_ase.py

@@ -7,6 +7,7 @@
 from shiftml.ase import ShiftML
 
 expected_output = np.array([137.5415, 137.5415])
+
 expected_ensemble_v2 = np.array(
     [
         [
@@ -148,10 +149,10 @@
 
 
 def test_shiftml1_regression():
-    """Regression test for the ShiftML1.0 model."""
+    """Regression test for the ShiftML1.1rev model."""
 
     frame = bulk("C", "diamond", a=3.566)
-    model = ShiftML("ShiftML1.0", force_download=True)
+    model = ShiftML("ShiftML1.1rev", force_download=True)
     out = model.get_cs_iso(frame)
 
     assert np.allclose(
@@ -160,10 +161,10 @@ def test_shiftml1_regression():
 
 
 def test_shiftml1_rotational_invariance():
-    """Rotational invariance test for the ShiftML1.0 model."""
+    """Rotational invariance test for the ShiftML1.1rev model."""
 
     frame = bulk("C", "diamond", a=3.566)
-    model = ShiftML("ShiftML1.0")
+    model = ShiftML("ShiftML1.1rev")
     out = model.get_cs_iso(frame)
 
     assert np.allclose(
@@ -181,10 +182,10 @@ def test_shiftml1_rotational_invariance():
 
 
 def test_shiftml1_size_extensivity_test():
-    """Test ShiftML1.0 for translational invariance."""
+    """Test ShiftML1.1rev for translational invariance."""
 
     frame = bulk("C", "diamond", a=3.566)
-    model = ShiftML("ShiftML1.0")
+    model = ShiftML("ShiftML1.1rev")
     out = model.get_cs_iso(frame)
 
     assert np.allclose(
@@ -200,10 +201,10 @@ def test_shiftml1_size_extensivity_test():
 
 
 def test_shftml1_fail_invalid_species():
-    """Test ShiftML1.0 for non-fitted species"""
+    """Test ShiftML1.1rev for non-fitted species"""
 
     frame = bulk("Si", "diamond", a=3.566)
-    model = ShiftML("ShiftML1.0")
+    model = ShiftML("ShiftML1.1rev")
     with pytest.raises(ValueError) as exc_info:
         model.get_cs_iso(frame)
 
@@ -214,10 +215,10 @@ def test_shftml1_fail_invalid_species():
 
 
 def test_shiftml2_regression_mean():
-    """Regression test for the ShiftML2.0 model."""
+    """Regression test for the ShiftML2.0rev model."""
 
     frame = bulk("C", "diamond", a=3.566)
-    model = ShiftML("ShiftML2.0", force_download=True)
+    model = ShiftML("ShiftML2.0rev", force_download=True)
     out_mean = model.get_cs_iso(frame)
     out_std = model.get_cs_iso_std(frame)
     out_ensemble = model.get_cs_iso_ensemble(frame)

tests/test_regression_pretrained.py

@@ -4,29 +4,14 @@
 from shiftml.ase import ShiftML
 
 expected_outputs = {
-    "ShiftML1.0": np.array([137.5415, 137.5415]),
-    "ShiftML1.1": np.array([163.07251, 163.07251]),
+    "ShiftML1.1rev": np.array([137.5415, 137.5415]),
 }
 
 
-def test_shiftml1_regression():
-    """Regression test for the ShiftML1.0 model."""
-
-    MODEL = "ShiftML1.0"
-
-    frame = bulk("C", "diamond", a=3.566)
-    model = ShiftML(MODEL, force_download=True)
-    out = model.get_cs_iso(frame)
-
-    assert np.allclose(
-        out.flatten(), expected_outputs[MODEL]
-    ), "ShiftML1.0 failed regression test"
-
-
 def test_shiftml1_1_regression():
     """Regression test for the ShiftML1.1 model."""
 
-    MODEL = "ShiftML1.1"
+    MODEL = "ShiftML1.1rev"
 
     frame = bulk("C", "diamond", a=3.566)
     model = ShiftML(MODEL, force_download=True)