Added missing descriptor interface

Author: RandomDefaultUser

mala/network/predictor.py

@@ -78,7 +78,7 @@ def predict_from_qeout(self, path_to_file, gather_ldos=False):
         )
 
     def predict_for_atoms(
-        self, atoms, gather_ldos=False, temperature=None, save_grads=False
+        self, atoms, gather_ldos=False, temperature=None, save_grads=False, pass_descriptors=None
     ):
         """
         Get predicted LDOS for an atomic configuration.
@@ -133,24 +133,29 @@ def predict_for_atoms(
         self.data.target_calculator.invalidate_target()
 
         # Calculate descriptors.
-        time_before = perf_counter()
-        snap_descriptors, local_size = (
-            self.data.descriptor_calculator.calculate_from_atoms(
-                atoms, self.data.grid_dimension
-            )
-        )
-        printout(
-            "Time for descriptor calculation: {:.8f}s".format(
-                perf_counter() - time_before
-            ),
-            min_verbosity=2,
-        )
 
+
+        if pass_descriptors is None:
+            time_before = perf_counter()
+            snap_descriptors, local_size = (
+                self.data.descriptor_calculator.calculate_from_atoms(
+                    atoms, self.data.grid_dimension
+                )
+            )
+            printout(
+                "Time for descriptor calculation: {:.8f}s".format(
+                    perf_counter() - time_before
+                ),
+                min_verbosity=2,
+            )
+            feature_length = self.data.descriptor_calculator.feature_size
+            descs_with_xyz = self.data.descriptor_calculator.descriptors_contain_xyz
+        else:
+            snap_descriptors, local_size, feature_length, descs_with_xyz = pass_descriptors
         # Provide info from current snapshot to target calculator.
         self.data.target_calculator.read_additional_calculation_data(
             [atoms, self.data.grid_dimension], "atoms+grid"
         )
-        feature_length = self.data.descriptor_calculator.feature_size
 
         # The actual calculation of the LDOS from the descriptors depends
         # on whether we run in parallel or serial. In the former case,
@@ -171,7 +176,7 @@ def predict_for_atoms(
                     return None
 
             else:
-                if self.data.descriptor_calculator.descriptors_contain_xyz:
+                if descs_with_xyz:
                     self.data.target_calculator.local_grid = snap_descriptors[
                         :, 0:3
                     ].copy()
@@ -196,7 +201,8 @@ def predict_for_atoms(
                 )
 
         if get_rank() == 0:
-            if self.data.descriptor_calculator.descriptors_contain_xyz:
+            #if self.data.descriptor_calculator.descriptors_contain_xyz:
+            if descs_with_xyz:
                 snap_descriptors = snap_descriptors[:, :, :, 3:]
                 feature_length -= 3
 
@@ -208,9 +214,8 @@ def predict_for_atoms(
             if save_grads is True:
                 self.input_data = snap_descriptors
                 self.input_data.requires_grad = True
-                return self._forward_snap_descriptors(
-                    snap_descriptors, save_torch_outputs=True
-                )
+                return self._forward_snap_descriptors(snap_descriptors,
+                                                      save_torch_outputs=True)
             else:
                 return self._forward_snap_descriptors(snap_descriptors)