Improved forward pass efficiency and general enhancements

These changes improve efficiency and represent a step forward for GPU
training as well as the addition of force training.

- Updated reference to cite ML4Chem.
- Bumped version.
- Dscribe module is optional and not required.
- `backends.operations` renamed to `backends.oper`.  
- `atomistic.models.base`: added new `feature_preparation` function that
  prepares data in a way that vectorization is favored.
- `models.merger` the merger class now uses `torch.nn.ModuleList`
  instead of a plain list.
- `atomistic.models.neuralnetwork`:
    * `forward` is more efficient.
    * All data is scattered once to avoid problems with Dask messing up
      some tensors.
    * More steps forward to force training.
- `data.handler`:
    - New `get_largest_number_atoms` function for new vectorized forward
      improvements and avoid breaking compatibility.
    - New `target_keys` attribute that can be used to set any name for
      targets.

README.md

@@ -46,13 +46,13 @@ very nice.
 If you find this software useful, please use this bibtex to cite it:
 
 ```
-@article{El_Khatib2020,
-author = "Muammar El Khatib and Wibe de Jong",
-title = "{ML4Chem: A Machine Learning Package for Chemistry and Materials Science}",
-year = "2020",
-month = "3",
-url = "https://chemrxiv.org/articles/ML4Chem_A_Machine_Learning_Package_for_Chemistry_and_Materials_Science/11952516",
-doi = "10.26434/chemrxiv.11952516.v1"
+@article{Elkhatib2020ml4chem,
+    title={ML4Chem: A Machine Learning Package for Chemistry and Materials Science},
+    author={Muammar El Khatib and Wibe A de Jong},
+    year={2020},
+    eprint={2003.13388},
+    archivePrefix={arXiv},
+    primaryClass={physics.chem-ph}
 }
 ```
 

docs/environment.yml

@@ -12,3 +12,4 @@ dependencies:
         - sphinxcontrib-bibtex
         - sphinxcontrib-inlinesyntaxhighlight
         - m2r
+    - pip_install: True

ml4chem/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.9"
+__version__ = "0.1.0-dev"

ml4chem/atomistic/features/__init__.py

@@ -1,6 +1,13 @@
 from ml4chem.atomistic.features.autoencoders import LatentFeatures
 from ml4chem.atomistic.features.cartesian import Cartesian
-from ml4chem.atomistic.features.coulombmatrix import CoulombMatrix
 from ml4chem.atomistic.features.gaussian import Gaussian
 
-__all__ = ["LatentFeatures", "Cartesian", "CoulombMatrix", "Gaussian"]
+
+__all__ = ["LatentFeatures", "Cartesian", "Gaussian"]
+
+try:
+    from ml4chem.atomistic.features.coulombmatrix import CoulombMatrix
+
+    __all__.append("CoulombMatrix")
+except ModuleNotFoundError:
+    pass

ml4chem/atomistic/models/base.py

@@ -1,5 +1,8 @@
+import types
+import torch
 from ml4chem.atomistic import Potentials
 from abc import ABC, abstractmethod
+from collections import OrderedDict
 
 
 class DeepLearningModel(ABC):
@@ -56,3 +59,64 @@ def checkpoint_save(self, epoch, model, label=None, checkpoint=None, path=""):
             Potentials.save(model=model, label=label, path=path)
         elif epoch % checkpoint == 0:
             Potentials.save(model=model, label=label, path=path)
+
+    def feature_preparation(self, features, data):
+        """Feature preparation
+
+        This function takes features and rearrange the data to operate with a
+        DeepLearning class.
+
+        Parameters
+        ----------
+
+        """
+        purpose = "training"
+        data.get_largest_number_atoms(purpose)
+
+        rearrengements = []
+        conditions = []
+        if isinstance(features, (list, types.GeneratorType)):
+            for chunk in features:
+                chunk = OrderedDict(chunk)
+                rearrange = {
+                    symbol: [] for symbol in data.unique_element_symbols[purpose]
+                }
+
+                for _, values in chunk.items():
+                    image = {}
+                    for symbol, features_ in values:
+                        if symbol not in image.keys():
+                            image[symbol] = []
+                        image[symbol].append(features_)
+
+                    for symbol in data.unique_element_symbols[purpose]:
+                        tensors = image.get(symbol)
+
+                        if tensors == None:
+                            tensors = [torch.zeros(self.input_dimension)]
+
+                        tensors = torch.stack(tensors)
+
+                        tensor_size = tensors.size()[0]
+                        if tensor_size < data.largest_number_atoms[symbol]:
+
+                            diff = data.largest_number_atoms[symbol] - tensor_size
+                            expand = torch.zeros(diff, self.input_dimension)
+                            tensors = torch.cat([tensors, expand])
+
+                        rearrange[symbol].append(tensors)
+
+                rearrange = {
+                    symbol: torch.stack(tensors)
+                    for symbol, tensors in rearrange.items()
+                }
+
+                condition = {}
+
+                for symbol, tensors in rearrange.items():
+                    condition[symbol] = (tensors.sum(dim=2) != 0).float()
+
+                rearrengements.append(rearrange)
+                conditions.append(condition)
+
+        return rearrengements, conditions

ml4chem/atomistic/models/loss.py

@@ -13,7 +13,7 @@ def AtomicMSELoss(outputs, targets, atoms_per_image, uncertainty=None):
     outputs : tensor
         Outputs of the model.
     targets : tensor
-        Expected value of outputs.
+        Dictionary of tensors.
     atoms_per_image : tensor
         A tensor with the number of atoms per image. 
     uncertainty : tensor, optional

ml4chem/atomistic/models/merger.py

@@ -41,7 +41,7 @@ def name(cls):
 
     def __init__(self, models):
         super(ModelMerger, self).__init__()
-        self.models = models
+        self.models = torch.nn.ModuleList(models)
 
     def forward(self, X, models):
         """Forward propagation