From 17aae47ca5c2b3992ef9720cd23d87b8619b1dae Mon Sep 17 00:00:00 2001
From: avalluvan <62253557+avalluvan@users.noreply.github.com>
Date: Mon, 16 Dec 2024 12:02:35 -0800
Subject: [PATCH] Calculate parallel and master_node flags in
 RLparallelscript.py

---
 cosipy/image_deconvolution/RLparallelscript.py   |  9 ++++++++-
 cosipy/image_deconvolution/RichardsonLucy.py     |  6 +++---
 .../image_deconvolution/image_deconvolution.py   | 16 +++-------------
 3 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/cosipy/image_deconvolution/RLparallelscript.py b/cosipy/image_deconvolution/RLparallelscript.py
index 14d24d2c..561a6f38 100644
--- a/cosipy/image_deconvolution/RLparallelscript.py
+++ b/cosipy/image_deconvolution/RLparallelscript.py
@@ -48,8 +48,15 @@ def main():
     parameter_filepath = DATA_DIR / 'imagedeconvolution_parfile_gal_511keV.yml'
     image_deconvolution.read_parameterfile(parameter_filepath)
 
+    parallel_computation = True
+    if comm.Get_rank() == MASTER:
+        master_node = True
+    else:
+        master_node = False
+
     # Initialize model
-    image_deconvolution.initialize(comm=comm)
+    image_deconvolution.initialize(parallel_computation = parallel_computation,
+                                   master_node = master_node)
 
     # Execute deconvolution
     image_deconvolution.run_deconvolution()
diff --git a/cosipy/image_deconvolution/RichardsonLucy.py b/cosipy/image_deconvolution/RichardsonLucy.py
index 3a128710..c790339e 100644
--- a/cosipy/image_deconvolution/RichardsonLucy.py
+++ b/cosipy/image_deconvolution/RichardsonLucy.py
@@ -262,9 +262,9 @@ def finalization(self):
                 for this_result in self.results:
                     iteration_count = this_result["iteration"]
 
-                    this_result["model"].write(f"{self.save_results_directory}/model_itr{iteration_count}_2.hdf5", overwrite = True)
-                    this_result["delta_model"].write(f"{self.save_results_directory}/delta_model_itr{iteration_count}_2.hdf5", overwrite = True)
-                    this_result["processed_delta_model"].write(f"{self.save_results_directory}/processed_delta_model_itr{iteration_count}_2.hdf5", overwrite = True)
+                    this_result["model"].write(f"{self.save_results_directory}/model_itr{iteration_count}.hdf5", overwrite = True)
+                    this_result["delta_model"].write(f"{self.save_results_directory}/delta_model_itr{iteration_count}.hdf5", overwrite = True)
+                    this_result["processed_delta_model"].write(f"{self.save_results_directory}/processed_delta_model_itr{iteration_count}.hdf5", overwrite = True)
 
                 #fits
                 primary_hdu = fits.PrimaryHDU()
diff --git a/cosipy/image_deconvolution/image_deconvolution.py b/cosipy/image_deconvolution/image_deconvolution.py
index ff1360fe..804cbd00 100644
--- a/cosipy/image_deconvolution/image_deconvolution.py
+++ b/cosipy/image_deconvolution/image_deconvolution.py
@@ -11,9 +11,6 @@
 from .RichardsonLucy import RichardsonLucy
 from .RichardsonLucySimple import RichardsonLucySimple
 
-# MPI Master node
-MASTER = 0
-
 class ImageDeconvolution:
     """
     A class to reconstruct all-sky images from COSI data based on image deconvolution methods.
@@ -121,21 +118,14 @@ def results(self):
         """
         return self._deconvolution.results
 
-    def initialize(self, comm = None):      # comm is required if the user intends to use parallel features
+    def initialize(self, parallel_computation = False, master_node = True):
         """
         Initialize an initial model and an image deconvolution algorithm.
         It is mandatory to execute this method before running the image deconvolution.
         """
 
-        if comm is None:
-            self.parallel_computation = False
-            self.master_node = True
-        elif comm.Get_rank() == MASTER:
-            self.parallel_computation = True
-            self.master_node = True
-        else:
-            self.parallel_computation = True
-            self.master_node = False
+        self.parallel_computation = parallel_computation
+        self.master_node = master_node
 
         logger.info("#### Initialization Starts ####")