From 2e9a03ed1c39b84186a8751a8221c670638a4f02 Mon Sep 17 00:00:00 2001
From: Roden Luo <roden.luo@gmail.com>
Date: Fri, 18 Oct 2024 10:41:17 +0300
Subject: [PATCH] better log for disk mode: input, params, and time

---
 src/DiffAtomComp.py |  3 ---
 src/tool.py         | 39 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/DiffAtomComp.py b/src/DiffAtomComp.py
index 5470b39..96041e6 100644
--- a/src/DiffAtomComp.py
+++ b/src/DiffAtomComp.py
@@ -986,9 +986,6 @@ def diff_atom_comp(target_vol_path: str,
     log_idx = 0
     os.makedirs(out_dir, exist_ok=out_dir_exist_ok)
 
-    with open(f"{out_dir}/log.log", "a") as log_file:
-        log_file.write(f"Wall clock time: {datetime.now()}\n")
-
     # Create the optimizer with different learning rates
     optimizer = torch.optim.Adam([
         {'params': [e_shifts], 'lr': target_size.mean() * 0.01},
diff --git a/src/tool.py b/src/tool.py
index b8d93cd..02caee5 100644
--- a/src/tool.py
+++ b/src/tool.py
@@ -1559,16 +1559,37 @@ def run_button_clicked(self):
                                       "and then launch DiffFit again to run the fitting in Disk mode.")
             return
 
+        _out_dir_exist_ok = self.settings.out_dir_exist_ok
+        _out_dir = self.settings.output_directory
+        os.makedirs(_out_dir, exist_ok=_out_dir_exist_ok)
+        with open(f"{_out_dir}/log.log", "a") as log_file:
+            log_file.write(f"=======\n"
+                           f"Wall clock time: {datetime.now()}\n"
+                           f"-------\n"
+                           f"Disk mode\n"
+                           f"Target Volume: {self.settings.target_vol_path}\n"
+                           f"Structures Folder: {self.settings.structures_directory}\n"
+                           f"Sim-map Folder: {self.settings.structures_sim_map_dir}\n"
+                           f"Target Surface Threshold: {self.settings.target_surface_threshold}\n"
+                           f"-------\n"
+                           f"# shifts: {self.settings.N_shifts}\n"
+                           f"# quaternions: {self.settings.N_quaternions}\n"
+                           f"Gaussian mode: {self.Gaussian_mode}\n"
+                           f"Conv. loops: {self.settings.conv_loops}\n"
+                           f"Conv. kernel sizes: {self.settings.conv_kernel_sizes}\n"
+                           f"Conv. weights: {self.settings.conv_weights}\n"
+                           f"-------\n")
+
+
         self.disable_spheres_clicked()
         disk_fit_timer_start = datetime.now()
 
         print("Running the computation...")
 
-        #target_vol_path = "D:\\GIT\\DiffFit\dev_data\input\domain_fit_demo_3domains\density2.mrc"
-        #output_folder = "D:\\GIT\\DiffFit\dev_data\output"
+        with open(f"{_out_dir}/log.log", "a") as log_file:
+            log_file.write(f"DiffFit optimization starts: {disk_fit_timer_start}\n")
 
         timer_start = datetime.now()
-
         mol_centers, e_sqd_log = diff_atom_comp(
             target_vol_path=self.settings.target_vol_path,
             target_surface_threshold=self.settings.target_surface_threshold,
@@ -1593,6 +1614,10 @@ def run_button_clicked(self):
         timer_stop = datetime.now()
         print(f"\nDiffFit optimization time elapsed: {timer_stop - timer_start}\n\n")
 
+        with open(f"{_out_dir}/log.log", "a") as log_file:
+            log_file.write(f"-------\n"
+                           f"DiffFit optimization time elapsed: {timer_stop - timer_start}\n")
+
         # copy the directories
         self.target_vol.setText(self.settings.target_vol_path)     
         self.structures_folder.setText(self.settings.structures_directory)
@@ -1607,6 +1632,14 @@ def run_button_clicked(self):
         timer_stop = datetime.now()
         print(f"\nDiffFit total time elapsed: {timer_stop - disk_fit_timer_start}\n\n")
 
+        metric_json = self.return_cluster_metric_json(0)
+        with open(f"{_out_dir}/log.log", "a") as log_file:
+            log_file.write(f"DiffFit total time elapsed: {timer_stop - disk_fit_timer_start}\n"
+                           f"-------\n"
+                           f"DiffFit top fit metric:\n"
+                           f"{metric_json}\n"
+                           f"=======\n\n")
+
     def load_button_clicked(self):
         if self.fit_input_mode == "interactive":
             if self.interactive_fit_result_ready: