Merge class consistency was checked.

- This commit finally closes #10. The module can do both independent and
  dependent loss functions.
- Addition of a `merged_models` directory with an example.
- Improved docstrings of models.merger module.

examples/merged_models/training.py

@@ -0,0 +1,109 @@
+from ase.io import Trajectory
+from dask.distributed import Client, LocalCluster
+import sys
+
+sys.path.append("../ml4chem")
+from ml4chem.data.handler import DataSet
+from ml4chem.fingerprints import Cartesian
+from ml4chem.models.autoencoders import AutoEncoder
+from ml4chem.models.neuralnetwork import NeuralNetwork
+from ml4chem.models.merger import ModelMerger
+from ml4chem.models.loss import MSELoss
+from ml4chem import Potentials
+from ml4chem.utils import logger
+from ml4chem.models.loss import AtomicMSELoss
+from ml4chem.data.serialization import dump
+
+
+def hybrid():
+    # Load the images with ASE, and prepare data handler
+    images = Trajectory("cu_training.traj")
+    purpose = "training"
+
+    latent_dimension = 32
+    data_handler = DataSet(images, purpose=purpose)
+    data_handler.get_unique_element_symbols(images, purpose=purpose)
+    training_set, energy_targets = data_handler.get_data(purpose=purpose)
+
+    # Preprocessor setup
+    preprocessor = ("MinMaxScaler", {"feature_range": (-1, 1)})
+
+    """
+    Preparing the input
+    """
+    features = Cartesian(
+        preprocessor=preprocessor, save_preprocessor="cartesian.scaler"
+    )
+    _inputs = features.calculate_features(training_set, data=data_handler)
+
+    """
+    Building AutoEncoder Model1
+    """
+    # Arguments for building the model
+    hiddenlayers = {
+        "encoder": (144, 72, latent_dimension),
+        "decoder": (latent_dimension, 72, 144),
+    }
+    # hiddenlayers = {"encoder": (2, 2, 2), "decoder": (2, 2, 2)}
+    activation = "tanh"
+    autoencoder = AutoEncoder(hiddenlayers=hiddenlayers, activation=activation)
+    autoencoder.prepare_model(3, 3, data=data_handler)
+
+    """
+    Building the ml potential model
+    """
+
+    # Arguments for building the model
+    n = 40
+    activation = "tanh"
+
+    nn = NeuralNetwork(hiddenlayers=(n, n), activation=activation)
+    nn.prepare_model(latent_dimension, data=data_handler)
+
+    models = [autoencoder, nn]
+    losses = [MSELoss, AtomicMSELoss]
+    # losses = [EncoderMapLoss, AtomicMSELoss]
+
+    merged = ModelMerger(models)
+    # Arguments for training the potential
+    convergence = {"rmse": [1.5e-1, 1.0e-1]}
+    lr = 1e-4
+    weight_decay = 1e-5
+    regularization = None
+
+    # Optimizer
+    optimizer = ("adam", {"lr": lr, "weight_decay": weight_decay, "amsgrad": True})
+    lr_scheduler = None
+
+    inputs = [_inputs, autoencoder.get_latent_space]
+    targets = [_inputs, energy_targets]
+    batch_size = 2
+
+    merged.train(
+        inputs=inputs,
+        targets=targets,
+        data=data_handler,
+        regularization=regularization,
+        convergence=convergence,
+        optimizer=optimizer,
+        device="cpu",
+        batch_size=batch_size,
+        lr_scheduler=lr_scheduler,
+        lossfxn=losses,
+        independent_loss=True,
+    )
+
+    for index, model in enumerate(merged.models):
+        label = "{}_{}".format(index, model.name())
+        Potentials.save(model, label=label)
+
+    dump_ls = merged.models[0].get_latent_space(inputs[0])
+    dump(dump_ls, filename="checkme.latent")
+
+
+if __name__ == "__main__":
+    logger()
+    cluster = LocalCluster(n_workers=5, threads_per_worker=2, dashboard_address=8798)
+    client = Client(cluster, asyncronous=True)
+    # Let's do this
+    hybrid()

ml4chem/models/merger.py

@@ -210,6 +210,9 @@ def train(
                     "train", "ml4chem.models", alt_name="autoencoders"
                 )
                 self.inputs_chunk_vals = train.get_inputs_chunks(chunks[index])
+            else:
+                self.inputs_chunk_vals = None
+
 
         parameters = []
         for index, model in enumerate(self.models):
@@ -295,12 +298,13 @@ def train(
 
             if independent_loss:
                 losses = []
+                outputs = []
                 for model_index, model in enumerate(self.models):
-                    name = model.name()
-                    loss, outputs = self.closure(
-                        model_index, model, independent_loss, name=name
+                    loss, output = self.closure(
+                        model_index, model, independent_loss, name=model.name()
                     )
                     losses.append(loss)
+                    outputs.append(output)
 
             else:
                 loss, outputs = self.closure(index, self.models, independent_loss)
@@ -373,8 +377,10 @@ def closure(self, index, model, independent_loss, name=None):
             train = dynamic_import("train", "ml4chem.models", alt_name="neuralnetwork")
 
             inputs = []
-            for chunk in self.chunks[index - 1]:
-                inputs_ = self.chunks[index](OrderedDict(chunk))
+            # FIXME this is not scaling to n number of models.
+            for chunk_index, chunk in enumerate(self.chunks[index - 1]):
+
+                inputs_ = self.chunks[index][chunk_index](OrderedDict(chunk.result()))
                 inputs.append(client.scatter(inputs_))
 
             loss, outputs_ = train.closure(
@@ -389,8 +395,6 @@ def closure(self, index, model, independent_loss, name=None):
 
         elif name == "AutoEncoder" and independent_loss:
             train = dynamic_import("train", "ml4chem.models", alt_name="autoencoders")
-            # The indexing [0] is needed because of the way the array is built
-            # above.
             targets = self.targets[index]
 
             loss, outputs_ = train.closure(
@@ -429,6 +433,7 @@ def closure(self, index, model, independent_loss, name=None):
             grads = {}
             outputs_ = {}
             losses = {}
+
             for model_index, (outputs, loss, grad) in enumerate(accumulation):
                 for model_index in range(len(self.models)):
                     if model_index not in grads.keys():
@@ -461,18 +466,17 @@ def train_batches(
         losses = []
         for model_index, model in enumerate(models):
             # _losses = []
-            name = model.name()
             # for j, output in enumerate(outputs[i]):
 
             output = outputs[model_index]
-            if name == "PytorchPotentials":
+            if model.name() == "PytorchPotentials":
                 loss = lossfxn[model_index](
                     output,
                     targets[model_index][chunk_index],
                     atoms_per_image[chunk_index],
                 )
 
-            elif name == "AutoEncoder":
+            elif model.name() == "AutoEncoder":
                 loss = lossfxn[model_index](output, targets[model_index][chunk_index])
 
             batch_loss += loss * self.loss_weights[model_index]