TST: Add support for python 3.9 (#93)

.github/workflows/build.yml

@@ -2,6 +2,10 @@ name: build with conda
 
 on: [push]
 
+env:
+  TORCH: "1.8.1"
+  CONDA_PREFIX: /usr/share/miniconda
+
 jobs:
   build:
 
@@ -10,11 +14,11 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        version: [3.7, 3.8]
+        version: [3.7, 3.8, 3.9]
 
     steps:
     - name: Cancel Previous Runs
-      uses: styfle/cancel-workflow-action@0.4.0
+      uses: styfle/cancel-workflow-action@0.9.1
       with:
           access_token: ${{ github.token }}
     - uses: actions/checkout@v2
@@ -24,39 +28,34 @@ jobs:
         update-conda: true
         python-version: ${{ matrix.version }}
         conda-channels: anaconda
-    - run: conda --version
-    - run: which python
-    - name: install conda dependencies
-      env:
-        TORCH: "1.8"
-      run: conda install h5py scipy rdkit pytorch==${TORCH} torchvision cpuonly dgl -c pytorch -c conda-forge -c dglteam
-
-    - name: install torch-geometric dependencies
-      env:
-        TORCHG: "1.8.0"
+
+    - name: install dependencies
       run: |
-        pip install --no-index torch-scatter -f https://pytorch-geometric.com/whl/torch-${TORCHG}+cpu.html
-        pip install --no-index torch-sparse -f https://pytorch-geometric.com/whl/torch-${TORCHG}+cpu.html
-        pip install --no-index torch-cluster -f https://pytorch-geometric.com/whl/torch-${TORCHG}+cpu.html
-        pip install --no-index torch-spline-conv -f https://pytorch-geometric.com/whl/torch-${TORCHG}+cpu.html
-        pip install torch-geometric
+        # torch==1.8.1 dependencies
+        conda install scipy rdkit pytorch==${TORCH} torchvision cpuonly dgl -c pytorch -c conda-forge -c dglteam
+        pip install "gpytorch<=1.5.1"  # TODO: Remove when switching to pytorch >=1.9
+
+        # torch-geometric dependencies
+        pip install --no-index torch-scatter -f https://pytorch-geometric.com/whl/torch-${TORCH}+cpu.html
+        pip install --no-index torch-sparse -f https://pytorch-geometric.com/whl/torch-${TORCH}+cpu.html
+        pip install --no-index torch-cluster -f https://pytorch-geometric.com/whl/torch-${TORCH}+cpu.html
+        pip install --no-index torch-spline-conv -f https://pytorch-geometric.com/whl/torch-${TORCH}+cpu.html
+
+        # The package
+        pip install .[test,doc]
 
-    - name: install torch-geometric
-      run:  pip install torch-geometric
+    - name: Conda info
+      run: conda info
 
-    - name: Install the package
-      run: pip install .[test,doc]
-      env:
-        CONDA_PREFIX: /usr/share/miniconda
+    - name: Conda list
+      run: conda list -n base
 
     - name: Test with pytest
-      env:
-        CONDA_PREFIX: /usr/share/miniconda
       run: |
         pytest
 
     - name: coverage
-      uses: codecov/codecov-action@v1
+      uses: codecov/codecov-action@v2
       with:
         file: ./coverage.xml
         name: codecov-umbrella

setup.py

@@ -33,6 +33,7 @@
         'Natural Language :: English',
         'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
     ],
     install_requires=[
         'e3nn@git+https://github.com/e3nn/e3nn@main',

swan/dataset/splitter.py

@@ -1,24 +1,28 @@
-from typing import Generic, NamedTuple, Tuple, TypeVar, Union
+from __future__ import annotations
+
+from typing import NamedTuple, Tuple, TypeVar
 
 import numpy as np
 import torch
 
 from ..state import StateH5
 from ..type_hints import PathLike
 
-T_co = TypeVar('T_co', bound=Union[np.ndarray, torch.Tensor], covariant=True)
-
 
-class SplitDataset(NamedTuple, Generic[T_co]):
+class SplitDataset(NamedTuple):
     indices: np.ndarray       # Shuffled indices to split the data
     ntrain: int               # Number of points used for training
-    features_trainset: T_co   # Features for training
-    features_validset: T_co   # Features for validation
-    labels_trainset: T_co     # Labels for training
-    labels_validset: T_co     # Labels for validation
+    features_trainset: np.ndarray | torch.Tensor   # Features for training
+    features_validset: np.ndarray | torch.Tensor   # Features for validation
+    labels_trainset: np.ndarray | torch.Tensor     # Labels for training
+    labels_validset: np.ndarray | torch.Tensor     # Labels for validation
 
 
-def split_dataset(features: T_co, labels: T_co, frac: Tuple[float, float] = (0.8, 0.2)) -> SplitDataset:
+def split_dataset(
+    features: np.ndarray | torch.Tensor,
+    labels: np.ndarray | torch.Tensor,
+    frac: Tuple[float, float] = (0.8, 0.2),
+) -> SplitDataset:
     """Split the fingerprint dataset into a training and validation set.
 
     Parameters