General improvements and first steps to enable VAE in ModelMerger.

docs/source/models.rst

@@ -119,14 +119,14 @@ predicting energy and atomic forces.
 
     features = Gaussian(cutoff=6.5, normalized=True, save_preprocessor="features.scaler")
 
-In the code snippet above we are building Gaussian type using the
+In the code snippet above we are building Gaussian type features using the
 :class:`ml4chem.fingerprints.gaussian.Gaussian` class. We use a ``cutoff``
-radius of :math:`6.5` angstrom, normalized, and the preprocessing is saved to
-the file ``features.scaler`` (by default the scaler is ``MinMaxScaler`` in a
-range :math:`(-1, 1)` as implemented in ``scikit-learn``). The ``angular``
-symmetry function used by default is :math:`G_i^3`, if you are interested in
-using :math:`G_i^4`, then you need to pass ``angular_type`` keyword
-argument::
+radius of :math:`6.5` angstrom, we normalized by the squared cutoff raidous,
+and the preprocessing is saved to the file ``features.scaler`` (by default
+the preprocessing used is ``MinMaxScaler`` in a range :math:`(-1, 1)` as
+implemented in ``scikit-learn``). The ``angular`` symmetry functions used by
+default are :math:`G_i^3`, if you are interested in using :math:`G_i^4`, then
+you need to pass ``angular_type`` keyword argument::
 
     features = Gaussian(cutoff=6.5, normalized=True,
                         save_preprocessor="features.scaler", angular_type="G4")
@@ -160,7 +160,7 @@ In ML4Chem, a neural network can be instantiated as shown below:
     nn = NeuralNetwork(hiddenlayers=(n, n), activation=activation)
     nn.prepare_model()
 
-Here, weare building a NN with the
+Here, we are building a NN with the
 :class:`ml4chem.models.neuralnetwork.NeuralNetwork` class with two
 hidden-layers composed 10 neurons each, and a ReLu activation function.
 
@@ -187,13 +187,13 @@ reconstruct the input data.
     autoencoder.prepare_model(input_dimension, output_dimension, data=data_handler)
 
 
-ML4Chem also provides access to variational autoencoders (VAE). These
-architectures differ from an AE in that the encoder codes a distribution with
-mean and variance (two vectors with the desired latent space dimension)
-instead of a single latent vector. Subsequently, this distribution is sampled
-and used by the decoder to reconstruct the input. This creates a generative
-model because now we will generate a latent distribution that allows a
-continuous change from one class to another.
+ML4Chem also provides access to variational autoencoders (VAE)[Kingma2013]_.
+These architectures differ from an AE in that the encoder codes a
+distribution with mean and variance (two vectors with the desired latent
+space dimension) instead of a single latent vector. Subsequently, this
+distribution is sampled and used by the decoder to reconstruct the input.
+This creates a generative model because now we will generate a latent
+distribution that allows a continuous change from one class to another.
 
 .. image:: _static/vae.png
    :alt: VAE
@@ -240,4 +240,5 @@ uncertainty of each prediction.
 
 .. [Behler2007] Behler, J. & Parrinello, M. Generalized Neural-Network Representation of High-Dimensional Potential-Energy Surfaces. Phys. Rev. Lett. 98, 146401 (2007).
 .. [Behler2015] Behler, J. Constructing high-dimensional neural network potentials: A tutorial review. Int. J. Quantum Chem. 115, 1032–1050 (2015).
+.. [Kingma2013] Kingma, D. P. & Welling, M. Auto-Encoding Variational Bayes. arXiv Prepr. arXiv1312.6114 (2013).
 .. [Rupp2015] Rupp, M. Machine learning for quantum mechanics in a nutshell. Int. J. Quantum Chem. 115, 1058–1073 (2015).

examples/merged_models/training.py

@@ -104,6 +104,6 @@ def hybrid():
 if __name__ == "__main__":
     logger()
     cluster = LocalCluster(n_workers=5, threads_per_worker=2, dashboard_address=8798)
-    client = Client(cluster, asyncronous=True)
+    client = Client(cluster)
     # Let's do this
     hybrid()

ml4chem/models/autoencoders.py

@@ -880,7 +880,7 @@ def closure(
         """Closure
 
         This method clears previous gradients, iterates over chunks, accumulate
-        the gradiends, update model params, and return loss.
+        the gradients, update model params, and return loss.
         """
 
         outputs_ = []
@@ -1014,7 +1014,6 @@ def train_batches(
             args.update(latent)
 
         if loss_name == "EncoderMapLoss":
-            # This is for the EncoderMap
             latent = {
                 "latent": model.get_latent_space(
                     inputs, svm=False, purpose="preprocessing"

ml4chem/models/merger.py

@@ -32,6 +32,8 @@ class ModelMerger(torch.nn.Module):
 
     NAME = "Merged"
 
+    autoencoders = ["AutoEncoder", "VAE"]
+
     @classmethod
     def name(cls):
         """Returns name of class"""
@@ -65,7 +67,7 @@ def forward(self, X, models):
             if inspect.ismethod(x):
                 x = X[i - 1]
 
-            if name == "AutoEncoder":
+            if name in ModelMerger.autoencoders:
                 x = OrderedDict(x)
             elif name == "PytorchPotentials":
                 x = X[i](OrderedDict(x))
@@ -226,7 +228,7 @@ def train(
                 ]
                 targets[index] = _targets
                 del _targets
-            elif model.name() == "AutoEncoder":
+            elif model.name() in ModelMerger.autoencoders:
                 targets[index] = lod_to_list(targets[index])
 
         # Data scattering
@@ -392,7 +394,7 @@ def closure(self, index, model, independent_loss, name=None):
             )
             return loss, outputs_
 
-        elif name == "AutoEncoder" and independent_loss:
+        elif name in ModelMerger.autoencoders and independent_loss:
             train = dynamic_import("train", "ml4chem.models", alt_name="autoencoders")
             targets = self.targets[index]