Enforcing derived deep learning classes and general improvements.

These changes are needed for ML4Chem publication.

- All deep learning models are now inheriting from DeepLearningModel
  base class.
- Visualization module moved from `data.visualization` to
 `.visualization`.

bin/ml4chem

@@ -5,7 +5,7 @@ import sys
 
 path = os.path.dirname(os.path.abspath(__file__)).strip("bin")
 sys.path.append(path)
-from ml4chem.data.visualization import read_log, plot_atomic_features
+from ml4chem.visualization import read_log, plot_atomic_features
 
 
 @click.command()

docs/source/data.rst

@@ -40,15 +40,15 @@ For more information please refer to :mod:`ml4chem.data.handler`.
 Visualization
 ===================
 
-We also offer a :mod:`ml4chem.data.visualization` module to plot interesting
+We also offer a :mod:`ml4chem.visualization` module to plot interesting
 graphics about your model, features, or even monitor the progress of the loss
 function and error minimization.
 
 Two backends are supported to plot in ML4Chem: Seaborn and Plotly.
 
 An example is shown below::
 
-    from ml4chem.data.visualization import plot_atomic_features
+    from ml4chem.visualization import plot_atomic_features
     fig = plot_atomic_features("latent_space.db",
                                method="pca",
                                dimensions=3,

docs/source/ml4chem.data.rst

@@ -44,10 +44,10 @@ ml4chem.data.utils module
     :undoc-members:
     :show-inheritance:
 
-ml4chem.data.visualization module
+ml4chem.visualization module
 ---------------------------------
 
-.. automodule:: ml4chem.data.visualization
+.. automodule:: ml4chem.visualization
     :members:
     :undoc-members:
     :show-inheritance:

ml4chem/models/autoencoders.py

@@ -7,6 +7,7 @@
 import numpy as np
 from collections import OrderedDict
 from ml4chem.metrics import compute_rmse
+from ml4chem.models.base import DeepLearningModel
 from ml4chem.models.loss import MSELoss
 from ml4chem.optim.handler import get_optimizer, get_lr_scheduler
 from ml4chem.utils import convert_elapsed_time, get_chunks, lod_to_list
@@ -16,7 +17,7 @@
 logger = logging.getLogger()
 
 
-class AutoEncoder(torch.nn.Module):
+class AutoEncoder(DeepLearningModel):
     """Fully connected atomic autoencoder
 
 
@@ -67,7 +68,7 @@ def name(cls):
     def __init__(
         self, hiddenlayers=None, activation="relu", one_for_all=False, **kwargs
     ):
-        super(AutoEncoder, self).__init__()
+        super(DeepLearningModel, self).__init__()
 
         self.hiddenlayers = hiddenlayers
         self.activation = activation
@@ -956,6 +957,8 @@ def trainer(self):
             if self.anneal:
                 annealing = annealer.update(epoch)
                 print(annealing)
+            else: 
+                annealing = None
 
             self.optimizer.zero_grad()  # clear previous gradients
 

ml4chem/models/base.py

@@ -1,7 +1,8 @@
 from abc import ABC, abstractmethod
+import torch
 
 
-class DeepLearningModel(ABC):
+class DeepLearningModel(ABC, torch.nn.Module):
     @abstractmethod
     def name(cls):
         """Return name of the class"""

ml4chem/models/neuralnetwork.py

@@ -7,6 +7,7 @@
 import numpy as np
 from collections import OrderedDict
 from ml4chem.metrics import compute_rmse
+from ml4chem.models.base import DeepLearningModel
 from ml4chem.models.loss import AtomicMSELoss
 from ml4chem.optim.handler import get_optimizer, get_lr_scheduler
 from ml4chem.utils import convert_elapsed_time, get_chunks, get_number_of_parameters
@@ -17,7 +18,7 @@
 logger = logging.getLogger()
 
 
-class NeuralNetwork(torch.nn.Module):
+class NeuralNetwork(DeepLearningModel):
     """Atom-centered Neural Network Regression with Pytorch
 
     This model is based on Ref. 1 by Behler and Parrinello.
@@ -48,7 +49,7 @@ def name(cls):
         return cls.NAME
 
     def __init__(self, hiddenlayers=(3, 3), activation="relu", **kwargs):
-        super(NeuralNetwork, self).__init__()
+        super(DeepLearningModel, self).__init__()
         self.hiddenlayers = hiddenlayers
         self.activation = activation
 

ml4chem/data/visualization.py → ml4chem/visualization.py

@@ -199,7 +199,7 @@ def plot_atomic_features(
     """
     method = method.lower()
     backend = backend.lower()
-    dot_size = kwargs["dot_size"]
+    dot_size = kwargs.get("dot_size", 2)
 
     supported_methods = ["pca", "tsne"]