From c3489e76f64f771b49f7680f1406b19768b6f463 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Mon, 29 Jul 2024 12:00:26 +0200
Subject: [PATCH] Unifed timing calls

---
 mala/network/predictor.py | 16 +++++++++++++
 mala/targets/density.py   | 49 +++++++++++++++++++++++++--------------
 2 files changed, 47 insertions(+), 18 deletions(-)

diff --git a/mala/network/predictor.py b/mala/network/predictor.py
index 5a4a44588..847766ac7 100644
--- a/mala/network/predictor.py
+++ b/mala/network/predictor.py
@@ -1,5 +1,7 @@
 """Tester class for testing a network."""
 
+from time import perf_counter
+
 import numpy as np
 import torch
 
@@ -127,11 +129,18 @@ def predict_for_atoms(self, atoms, gather_ldos=False, temperature=None):
         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,
+        )
 
         # Provide info from current snapshot to target calculator.
         self.data.target_calculator.read_additional_calculation_data(
@@ -201,6 +210,7 @@ def _forward_snap_descriptors(
         # Ensure the Network is on the correct device.
         # This line is necessary because GPU acceleration may have been
         # activated AFTER loading a model.
+        time_before = perf_counter()
         self.network.to(self.network.params._configuration["device"])
 
         if local_data_size is None:
@@ -250,4 +260,10 @@ def _forward_snap_descriptors(
                 predicted_outputs
             )
         barrier()
+        printout(
+            "Time for network pass: {:.8f}s".format(
+                perf_counter() - time_before
+            ),
+            min_verbosity=2,
+        )
         return predicted_outputs
diff --git a/mala/targets/density.py b/mala/targets/density.py
index fab7913d7..b82bf9f2e 100644
--- a/mala/targets/density.py
+++ b/mala/targets/density.py
@@ -960,7 +960,7 @@ def __setup_total_energy_module(
 
         if Density.te_mutex is False:
             printout(
-                "MALA: Starting QuantumEspresso to get density-based"
+                "Starting QuantumEspresso to get density-based"
                 " energy contributions.",
                 min_verbosity=0,
             )
@@ -968,14 +968,18 @@ def __setup_total_energy_module(
             t0 = time.perf_counter()
             te.initialize(self.y_planes)
             barrier()
-            t1 = time.perf_counter()
-            printout("time used by total energy initialization: ", t1 - t0)
+            printout(
+                "Total energy module: Time used by total energy initialization: {:.8f}s".format(
+                    time.perf_counter() - t0
+                ),
+                min_verbosity=2,
+            )
 
             Density.te_mutex = True
-            printout("MALA: QuantumEspresso setup done.", min_verbosity=0)
+            printout("QuantumEspresso setup done.", min_verbosity=0)
         else:
             printout(
-                "MALA: QuantumEspresso is already running. Except for"
+                "QuantumEspresso is already running. Except for"
                 " the atomic positions, no new parameters will be used.",
                 min_verbosity=0,
             )
@@ -1087,10 +1091,10 @@ def __setup_total_energy_module(
                 )
             )
             barrier()
-            t1 = time.perf_counter()
             printout(
-                "time used by gaussian descriptors: ",
-                t1 - t0,
+                "Total energy module: Time used by gaussian descriptors: {:.8f}s".format(
+                    time.perf_counter() - t0
+                ),
                 min_verbosity=2,
             )
 
@@ -1119,10 +1123,10 @@ def __setup_total_energy_module(
                 )
             )
             barrier()
-            t1 = time.perf_counter()
             printout(
-                "time used by reference gaussian descriptors: ",
-                t1 - t0,
+                "Total energy module: Time used by reference gaussian descriptors: {:.8f}s".format(
+                    time.perf_counter() - t0
+                ),
                 min_verbosity=2,
             )
 
@@ -1149,9 +1153,12 @@ def __setup_total_energy_module(
             self._parameters_full.descriptors.use_atomic_density_energy_formula,
         )
         barrier()
-        t1 = time.perf_counter()
-        printout("time used by set_positions: ", t1 - t0, min_verbosity=2)
-
+        printout(
+            "Total energy module: Time used by set_positions: {:.8f}s".format(
+                time.perf_counter() - t0
+            ),
+            min_verbosity=2,
+        )
         barrier()
 
         if self._parameters_full.descriptors.use_atomic_density_energy_formula:
@@ -1191,9 +1198,11 @@ def __setup_total_energy_module(
                 1,
             )
             barrier()
-            t1 = time.perf_counter()
             printout(
-                "time used by set_positions_gauss: ", t1 - t0, min_verbosity=2
+                "Total energy module: Time used by set_positions_gauss: {:.8f}s".format(
+                    time.perf_counter() - t0
+                ),
+                min_verbosity=2,
             )
 
         # Now we can set the new density.
@@ -1201,8 +1210,12 @@ def __setup_total_energy_module(
         t0 = time.perf_counter()
         te.set_rho_of_r(density_for_qe, number_of_gridpoints, nr_spin_channels)
         barrier()
-        t1 = time.perf_counter()
-        printout("time used by set_rho_of_r: ", t1 - t0, min_verbosity=2)
+        printout(
+            "Total energy module: Time used by set_rho_of_r: {:.8f}s".format(
+                time.perf_counter() - t0
+            ),
+            min_verbosity=2,
+        )
 
         return atoms_Angstrom