Add scikit-learn as dependence and update README

README.md

@@ -15,6 +15,13 @@ git clone https://github.com/wengroup/matten.git
 pip install -e ./matten
 ```
 
+If you get package version conflicts, try the below command to install the dependencies
+with strict version requirements.
+
+```
+pip install -e "./matten[strict]"
+```
+
 ## Use the pretrained model
 
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wengroup/matten/blob/main/notebooks/predict_colab.ipynb)

pyproject.toml

@@ -6,21 +6,34 @@ build-backend = "setuptools.build_meta"
 name = "matten"
 version = "0.0.1"
 readme = "README.md"
-requires-python = ">=3.8, <3.11"
+requires-python = ">=3.8, <3.12"
 dependencies = [
-    "torch>=2.0.0",
     "pytorch-lightning>=2.0.0",
     "torchmetrics>=0.10.0, <1.0.0",
     "torch_geometric>=2.3.0",
     "torch_scatter",
     "e3nn",
+    "scikit-learn",
     "ase",
     "pymatgen",
-    "loguru",
     "torchtyping",
+    "loguru",
 ]
 
 [project.optional-dependencies]
+strict = [
+    "pytorch-lightning==2.4.0",
+    "torchmetrics==0.11.4",
+    "torch_geometric==2.6.1",
+    "torch_scatter==2.1.2",
+    "e3nn==0.5.1",
+    "scikit-learn==1.5.2",
+    "ase==3.23.0",
+    "pymatgen==2024.10.3",
+    "torchtyping==0.1.5",
+    "loguru==0.7.2",
+]
+
 test = ["pytest"]
 
 [tool.black]