Merge pull request #3 from wengroup/devel

Simplify installation and downgrade support to py3.8

.github/workflows/testing.yml

@@ -4,60 +4,52 @@ on: [push, pull_request]
 
 jobs:
   lint:
-    runs-on: ubuntu-latest
-
+    strategy:
+      max-parallel: 3
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.10"]
+    runs-on: ${{ matrix.os }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v3.5.3
 
-      - name: Set up Python
+      - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v4
         with:
-          python-version: "3.10"
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+          cache-dependency-path: pyproject.toml
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           pip install pre-commit
+
       - name: Lint
-        run: pre-commit run --show-diff-on-failure
+        run: pre-commit run --all-files --show-diff-on-failure
 
   test:
-    runs-on: ${{ matrix.os }}
     strategy:
       max-parallel: 3
       matrix:
         os: [ubuntu-latest]
-        python-version: ["3.10"]
-
+        python-version: ["3.8", "3.9", "3.10"]
+    runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v3
 
-      - uses: conda-incubator/setup-miniconda@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
         with:
-          auto-update-conda: true
           python-version: ${{ matrix.python-version }}
-          # activate-environment should be the same name as in environment.yml
-          activate-environment: matten
-          environment-file: environment.yml
+          cache: pip
+          cache-dependency-path: pyproject.toml
 
       - name: Install package
-        shell: bash -l {0}
-        run: |
-          # Multiple channels having pytorch. Specifying it in environment.yml
-          mamba install pytorch>=2.0.0 -c pytorch
-          mamba install pyg>=2.0.0 pytorch-scatter -c pyg -c conda-forge
-          mamba install lightning torchmetrics -c conda-forge
-          pip install -e .
-
-      # - name: Lint with flake8
-      #   shell: bash -l {0}
-      #   run: |
-      #     # stop the build if there are Python syntax errors or undefined names
-      #     flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-      #     # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-      #     flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
-
-      - name: Test with pytest
-        shell: bash -l {0}
         run: |
-          cd tests
-          pytest
+          python -m pip install --upgrade pip
+          pip install torch
+          pip install -e .[test]
+
+      - name: Test
+        run: pytest

README.md

@@ -7,14 +7,12 @@ MatTen is an equivariant graph neural network built using [e3nn](https://github.
 
 ## Install
 
-Follow the official documentation to install the dependencies: [pytorch>=2.0.0](https://pytorch.org/get-started/locally/),
-[pyg>=2.3.0](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html),
-and [pytorch-lightning>=2.0.0](https://lightning.ai/docs/pytorch/latest/).
+Follow the official documentation to install [pytorch>=2.0.0](https://pytorch.org/get-started/locally/).
 Then
 
 ```
 git clone https://github.com/wengroup/matten.git
-pip install -e matten
+pip install -e ./matten
 ```
 
 ## Use the pretrained model

devtools/check_prediction.py

@@ -5,6 +5,7 @@
 the MAE.
 """
 from pathlib import Path
+from typing import List
 
 import pandas as pd
 import torch
@@ -16,7 +17,7 @@
 
 def get_data():
     filename = (
-        Path(__file__).parent.parent
+        Path(__file__).resolve().parent.parent
         / "datasets"
         / "example_crystal_elasticity_tensor_n100.json"
     )
@@ -30,7 +31,7 @@ def get_data():
     return structures, targets
 
 
-def get_spherical_tensor(t: list[torch.Tensor]):
+def get_spherical_tensor(t: List[torch.Tensor]):
     """Convert a Cartesian tensor to a spherical tensor."""
     t = torch.stack(t)
     converter = CartesianTensorWrapper("ijkl=jikl=klij")

pyproject.toml

@@ -6,12 +6,11 @@ build-backend = "setuptools.build_meta"
 name = "matten"
 version = "0.0.1"
 readme = "README.md"
-requires-python = ">=3.10, <3.11"
+requires-python = ">=3.8, <3.11"
 dependencies = [
-    "pytorch-lightning>1.7.0, <1.9.0",
-    "lightning-bolts>=0.6.0",
-    "torchmetrics>=0.10.0,<1.0.0",
-    "torch_geometric",
+    "pytorch-lightning>=2.0.0",
+    "torchmetrics>=0.10.0, <1.0.0",
+    "torch_geometric>=2.3.0",
     "torch_scatter",
     "e3nn",
     "ase",
@@ -20,9 +19,12 @@ dependencies = [
     "torchtyping",
 ]
 
+[project.optional-dependencies]
+test = ["pytest"]
+
 [tool.black]
 line-length = 88
 
-# Explicitly add src_paths so that import of current package will be placed at the bottom.
+# Add src_paths so that import of current package will be placed at the bottom.
 [tool.isort]
 src_paths = ["src", "tests"]

scripts/train_materials_tensor.py

@@ -1,6 +1,7 @@
 """Script to train the materials tensor model."""
 
 from pathlib import Path
+from typing import Dict, List, Union
 
 import yaml
 from loguru import logger
@@ -13,7 +14,7 @@
 from matten.model_factory.tfn_scalar_tensor import ScalarTensorModel
 
 
-def instantiate_class(d: dict | list):
+def instantiate_class(d: Union[Dict, List]):
     args = tuple()  # no positional args
     if isinstance(d, dict):
         return lit_instantiate_class(args, d)
@@ -30,7 +31,7 @@ def get_args(path: Path):
     return config
 
 
-def main(config: dict):
+def main(config: Dict):
     dm = TensorDataModule(**config["data"])
     dm.prepare_data()
     dm.setup()

src/matten/dataset/structure_scalar_tensor.py

@@ -78,7 +78,7 @@ def __init__(
         scalar_target_names: List[str] = None,
         log_scalar_targets: List[bool] = None,
         normalize_scalar_targets: List[bool] = None,
-        tensor_target_weight: dict[str, dict[str, float]] = None,
+        tensor_target_weight: Dict[str, Dict[str, float]] = None,
         global_featurizer: None = None,
         normalize_global_features: bool = False,
         atom_featurizer: None = None,
@@ -379,7 +379,7 @@ def __init__(
         self,
         filename: str,  # this can be None
         r_cut: float,
-        structures: list[Structure],
+        structures: List[Structure],
         **kwargs,
     ):
         self.structures = structures
@@ -410,13 +410,13 @@ def __init__(
         tensor_target_formula: str = "ijkl=jikl=klij",
         tensor_target_scale: float = 1.0,
         normalize_tensor_target: bool = False,
-        tensor_target_weight: dict[str, dict[str, float]] = None,
+        tensor_target_weight: Dict[str, Dict[str, float]] = None,
         scalar_target_names: List[str] = None,
         log_scalar_targets: List[bool] = None,
         normalize_scalar_targets: List[bool] = None,
-        global_featurizer: Union[str, list[str]] = None,
+        global_featurizer: Union[str, List[str]] = None,
         normalize_global_features: bool = False,
-        atom_featurizer: Union[str, list[str]] = None,
+        atom_featurizer: Union[str, List[str]] = None,
         normalize_atom_features: bool = False,
         root: Union[str, Path] = ".",
         reuse: bool = True,

src/matten/model_factory/tfn_scalar_tensor.py

@@ -11,7 +11,7 @@
 """
 
 from collections import OrderedDict
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 from e3nn.io import CartesianTensor
@@ -34,7 +34,7 @@ def init_backbone(
         self,
         backbone_hparams: Dict[str, Any],
         dataset_hparams: Optional[Dict[str, Any]] = None,
-    ) -> tuple[torch.nn.Module, dict]:
+    ) -> Tuple[torch.nn.Module, Dict]:
         backbone = create_model(backbone_hparams, dataset_hparams)
 
         formula = (backbone_hparams["output_formula"]).lower()

src/matten/predict.py

@@ -1,6 +1,7 @@
 import tempfile
 import warnings
 from pathlib import Path
+from typing import List, Union
 
 import numpy as np
 import torch
@@ -47,7 +48,7 @@ def get_pretrained_config(identifier: str, config_filename: str = "config_final.
 
 
 def get_data_loader(
-    structures: list[Structure], identifier: str, batch_size: int = 200
+    structures: List[Structure], identifier: str, batch_size: int = 200
 ) -> DataLoader:
     # config contains info for dataset and data loader, we only use the dataset part,
     # and adjust some parameters
@@ -85,7 +86,7 @@ def get_data_loader(
     return DataLoader(dataset=dataset, shuffle=False, batch_size=batch_size)
 
 
-def check_species(model, structures: list[Structure]):
+def check_species(model, structures: List[Structure]):
     """
     Check if the species in the structures are support by the model.
     """
@@ -110,7 +111,7 @@ def evaluate(
     model,
     loader,
     target_name: str = "elastic_tensor_full",
-) -> list[torch.Tensor]:
+) -> List[torch.Tensor]:
     """
     Evaluate the model to generate predictions.
 
@@ -139,11 +140,11 @@ def evaluate(
 
 
 def predict(
-    structure: Structure | list[Structure],
+    structure: Union[Structure, List[Structure]],
     model_identifier="20230627",
     batch_size: int = 200,
     logger_level: str = "ERROR",
-) -> ElasticTensor | list[ElasticTensor]:
+) -> Union[ElasticTensor, List[ElasticTensor]]:
     f"""
     Predict the property of a structure or a list of structures.