fix tests

.github/workflows/testing.yml

@@ -10,6 +10,9 @@ jobs:
       matrix:
         os: [ubuntu-latest, macos-latest]
         python-version: [3.7, 3.8]
+        exclude:
+            - os: macos-latest
+              python-version: 3.7
 
     runs-on: ${{ matrix.os }}
 

megnet/data/molecule.py

@@ -25,11 +25,11 @@
     pybel = None
 
 try:
-    from rdkit import Chem
+    from rdkit import Chem  # type: ignore
 except ImportError:
     Chem = None
 
-from typing import List, Dict, Union
+from typing import Sequence, Dict, Union, List
 
 __date__ = '12/01/2018'
 
@@ -216,7 +216,7 @@ def convert(self,
                 'index1': index1,
                 'index2': index2}
 
-    def _create_pair_feature_vector(self, bond: Dict) -> List[float]:
+    def _create_pair_feature_vector(self, bond: Dict) -> List[int]:
         """Generate the feature vector from the bond feature dictionary
 
         Handles the binarization of categorical variables, and performing the distance conversion
@@ -226,7 +226,7 @@ def _create_pair_feature_vector(self, bond: Dict) -> List[float]:
         Returns:
             ([float]) Values converted to a vector
             """
-        bond_temp = []
+        bond_temp: List[int] = []
         for i in self.bond_features:
             # Some features require conversion (e.g., binarization)
             if i in bond:
@@ -550,7 +550,7 @@ def __del__(self):
         if self.pool is not None:
             self.pool.close()  # Kill thread pool if generator is deleted
 
-    def _generate_inputs(self, batch_index: int) -> np.ndarray:
+    def _generate_inputs(self, batch_index: list) -> np.ndarray:
         # Get the molecules for this batch
         mols = self.mols[batch_index]
 
@@ -587,7 +587,13 @@ def create_cached_generator(self) -> GraphBatchGenerator:
         graphs = self._generate_graphs(self.mols)
 
         # Turn them into a fat array
-        inputs = self.converter.get_flat_data(graphs, self.targets)
-
-        return GraphBatchGenerator(*inputs, is_shuffle=self.is_shuffle,
+        atom_features, bond_features, state_features, index1_list, index2_list, targets = \
+            self.converter.get_flat_data(graphs, self.targets)  # type: ignore
+        return GraphBatchGenerator(atom_features=atom_features,
+                                   bond_features=bond_features,
+                                   state_features=state_features,
+                                   index1_list=index1_list,
+                                   index2_list=index2_list,
+                                   targets=targets,
+                                   is_shuffle=self.is_shuffle,
                                    batch_size=self.batch_size)

megnet/data/tests/test_crystal.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 from megnet.data.graph import GaussianDistance
 from megnet.utils.general import to_list

megnet/data/tests/test_graph.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 from megnet.data.graph import GaussianDistance, GraphBatchGenerator, GraphBatchDistanceConvert,\
     EmbeddingMap

megnet/data/tests/test_molecule.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 import os
 import json

megnet/layers/featurizer/tests/test_gaussian.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 
 import numpy as np

megnet/layers/graph/tests/test_cgcnn.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 
 import numpy as np

megnet/layers/graph/tests/test_megnet.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 
 import numpy as np

megnet/layers/graph/tests/test_schnet.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 
 import numpy as np

megnet/layers/readout/tests/test_linear.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 from megnet.layers import LinearWithIndex
 import numpy as np
 from tensorflow.keras.layers import Input

megnet/layers/readout/tests/test_set2set.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 from megnet.layers import Set2Set
 import numpy as np
 from tensorflow.keras.layers import Input

megnet/models/base.py

@@ -186,7 +186,7 @@ def train_from_graphs(self,
                 val_generator = None  # type: ignore
                 steps_per_val = None  # type: ignore
         else:
-            val_generator = None
+            val_generator = None  # type: ignore
             steps_per_val = None  # type: ignore
 
         train_inputs = self.graph_converter.get_flat_data(train_graphs, train_targets)

megnet/models/megnet.py

@@ -129,7 +129,7 @@ def __init__(self,
         super().__init__(model=model, target_scaler=target_scaler, graph_converter=graph_converter)
 
     @classmethod
-    def from_url(cls, url: str) -> 'MEGNetModel':
+    def from_url(cls, url: str) -> GraphModel:
         """
         Download and load a model from a URL. E.g.
         https://github.com/materialsvirtuallab/megnet/blob/master/mvl_models/mp-2019.4.1/formation_energy.hdf5
@@ -147,7 +147,7 @@ def from_url(cls, url: str) -> 'MEGNetModel':
         return cls.from_file(fname)
 
     @classmethod
-    def from_mvl_models(cls, name: str) -> 'MEGNetModel':
+    def from_mvl_models(cls, name: str) -> GraphModel:
         from megnet.utils.models import load_model
         return load_model(name)
 
@@ -331,7 +331,7 @@ def one_block(a, b, c, has_ff=True, block_index=0):
     if is_classification:
         final_act = 'sigmoid'
     else:
-        final_act = None
+        final_act = None  # type: ignore
     out = Dense(ntarget, activation=final_act, name='readout_2')(final_vec)
     model = Model(inputs=[x1, x2, x3, x4, x5, x6, x7], outputs=out)
     return model

megnet/tests/test_activations.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 from megnet.activations import softplus2
 import numpy as np

megnet/tests/test_callbacks.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import os
 import unittest
 

megnet/tests/test_config.py

@@ -1,6 +1,7 @@
 """
 Test the data types
 """
+import tensorflow as tf
 import unittest
 
 import numpy as np

megnet/tests/test_losses.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 import numpy as np
 

megnet/tests/test_models.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 import numpy as np
 

megnet/utils/descriptor.py

@@ -4,7 +4,7 @@
 """
 
 import os
-from typing import Union
+from typing import Union, Dict
 
 import numpy as np
 from tensorflow.keras.models import Model
@@ -69,7 +69,7 @@ def __init__(self,
         model.model = full_model
         self.model = model
         self.valid_names = valid_names
-        self._cache = {}
+        self._cache: Dict[str, float] = {}
         self.use_cache = use_cache
 
     def _predict_structure(self, structure: StructureOrMolecule) -> np.ndarray:
@@ -111,7 +111,7 @@ def _get_updated_prefix_level(self, prefix: str, level: int):
                    "set2_set": ["set2set_atom" if level == 1 else "set2set_bond", None],
                    "concatenate": ["concatenate", None]}
         if self.version == "v2":
-            return mapping[prefix][0], mapping[prefix][1]
+            return mapping[prefix][0], mapping[prefix][1]  # type: ignore
         return prefix, level
 
     def get_atom_features(self, structure: StructureOrMolecule,

megnet/utils/tests/test_data.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 from pymatgen import Structure, Lattice
 from megnet.utils.data import get_graphs_within_cutoff

megnet/utils/tests/test_descriptor.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 from pymatgen import Structure, Lattice
 

megnet/utils/tests/test_layer.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 from megnet.utils.layer import repeat_with_index, _repeat
 import tensorflow as tf

megnet/utils/tests/test_models_util.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import unittest
 
 from megnet.utils.models import load_model

megnet/utils/tests/test_molecule_util.py

@@ -1,3 +1,4 @@
+import tensorflow as tf
 import os
 import unittest
 from pymatgen import Molecule