From 0c1c6d4bcdec6c1156115f9ef5e10cea94781fe2 Mon Sep 17 00:00:00 2001
From: Aaron Kanzer <aaronkanzer@Aarons-MacBook-Pro.local>
Date: Fri, 22 Nov 2024 14:13:52 -0500
Subject: [PATCH] move tswriter into helper function

---
 .../modalities/lsm/convert_tiff_to_zarr.py    | 43 ++++++++++++-------
 linc_convert/modalities/lsm/mosaic.py         | 15 ++++---
 2 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/linc_convert/modalities/lsm/convert_tiff_to_zarr.py b/linc_convert/modalities/lsm/convert_tiff_to_zarr.py
index 8a2d6eb..b9ce343 100644
--- a/linc_convert/modalities/lsm/convert_tiff_to_zarr.py
+++ b/linc_convert/modalities/lsm/convert_tiff_to_zarr.py
@@ -1,24 +1,33 @@
 import time
 import psutil
+from threading import Thread
 from linc_convert.modalities.lsm.mosaic import convert
 
 
+# Shared dictionary to store metrics
+metrics = {"cpu": [], "memory": []}
+
+
 def monitor_usage():
-    """Monitors CPU and memory usage."""
+    """Continuously monitors CPU and memory usage."""
     process = psutil.Process()
-    cpu_percent = process.cpu_percent(interval=0.1)  # Measure CPU over a small interval
-    memory_info = process.memory_info()
-    memory_mb = memory_info.rss / (1024 * 1024)  # Resident Set Size in MB
-    return cpu_percent, memory_mb
+    while True:
+        cpu_percent = process.cpu_percent(interval=0.1)  # Short interval for responsiveness
+        memory_info = process.memory_info()
+        memory_mb = memory_info.rss / (1024 * 1024)  # Convert to MB
+        metrics["cpu"].append(cpu_percent)
+        metrics["memory"].append(memory_mb)
+        time.sleep(0.1)  # Adjust monitoring frequency as needed
 
 
 def main():
-    # Start monitoring
-    start_time = time.time()
-    start_cpu_percent, start_memory_mb = monitor_usage()
+    # Start monitoring in a separate thread
+    monitor_thread = Thread(target=monitor_usage, daemon=True)
+    monitor_thread.start()
 
-    print(f"Initial CPU Usage: {start_cpu_percent}%")
-    print(f"Initial Memory Usage: {start_memory_mb:.2f} MB")
+    # Start the timer
+    start_time = time.time()
+    print("Starting execution...")
 
     # Execute the convert function
     convert(
@@ -28,14 +37,18 @@ def main():
         chunk=[32]
     )
 
-    # End monitoring
+    # End the timer
     end_time = time.time()
-    end_cpu_percent, end_memory_mb = monitor_usage()
-
     elapsed_time = end_time - start_time
+    print("Execution finished.")
+
+    # Collect metrics
+    avg_cpu = sum(metrics["cpu"]) / len(metrics["cpu"])
+    max_memory = max(metrics["memory"])
+
     print(f"\nExecution Time: {elapsed_time:.2f} seconds")
-    print(f"Final CPU Usage: {end_cpu_percent}%")
-    print(f"Final Memory Usage: {end_memory_mb:.2f} MB")
+    print(f"Average CPU Usage: {avg_cpu:.2f}%")
+    print(f"Peak Memory Usage: {max_memory:.2f} MB")
 
 
 if __name__ == "__main__":
diff --git a/linc_convert/modalities/lsm/mosaic.py b/linc_convert/modalities/lsm/mosaic.py
index 54743a1..90c0025 100644
--- a/linc_convert/modalities/lsm/mosaic.py
+++ b/linc_convert/modalities/lsm/mosaic.py
@@ -32,8 +32,12 @@
 mosaic = cyclopts.App(name="mosaic", help_format="markdown")
 lsm.command(mosaic)
 
-def write_plane_multiprocess(tswriter, subc, zstart, subz, ystart, yx_shape, dat):
+def write_plane_multiprocess(tswriter, fname, subc, zstart, subz, ystart, yx_shape):
     """Write a single plane of data into the Zarr file (multiprocessing)."""
+
+    with TiffFile(fname) as tiff:
+        dat = tiff.asarray()
+
     try:
         with ts.Transaction() as txn:
             tswriter.with_transaction(txn)[
@@ -295,12 +299,11 @@ def convert(
                 end="\r",
             )
 
-            dat = TiffFile(fname).asarray()
-            tasks.append((subc, zstart, subz, ystart, yx_shape, dat))
-
-    write_func = partial(write_plane_multiprocess, tswriter)
+            tasks.append((tswriter, fname, subc, zstart, subz, ystart, yx_shape))
+            
+    # write_func = partial(write_plane_multiprocess, tswriter)
     with Pool(processes=8) as pool:
-        pool.starmap(write_func, tasks)
+        pool.starmap(write_plane_multiprocess, tasks)
 
     print("")