From 10943b6fb4e6cebe600070f325c3257120adcaa6 Mon Sep 17 00:00:00 2001
From: Anna Grim <108307071+anna-grim@users.noreply.github.com>
Date: Fri, 19 Jan 2024 13:46:47 -0800
Subject: [PATCH] minor upds (#38)

Co-authored-by: anna-grim <anna.grim@alleninstitute.org>
---
 src/deep_neurographs/intake.py |  7 +------
 src/deep_neurographs/utils.py  | 26 ++++++++++++--------------
 2 files changed, 13 insertions(+), 20 deletions(-)

diff --git a/src/deep_neurographs/intake.py b/src/deep_neurographs/intake.py
index 72ca892..2ca0b8c 100644
--- a/src/deep_neurographs/intake.py
+++ b/src/deep_neurographs/intake.py
@@ -139,9 +139,7 @@ def build_neurograph_from_gcs_zips(
     print("Process swc files...")
     total_runtime, t0 = utils.init_timers()
     swc_dicts = download_gcs_zips(bucket_name, cloud_path, min_size)
-    t_stamp, unit_stamp = utils.time_writer(time() - total_runtime)
     t, unit = utils.time_writer(time() - t0)
-    print(f"\nTime Stamp: {round(t_stamp, 4)} {unit_stamp}")
     print(f"Module Runtime(): {round(t, 4)} {unit} \n")
 
     # Build neurograph
@@ -155,8 +153,6 @@ def build_neurograph_from_gcs_zips(
         smooth=smooth,
     )
     t, unit = utils.time_writer(time() - t0)
-    t_stamp, unit_stamp = utils.time_writer(time() - total_runtime)
-    print(f"Time Stamp: {round(t_stamp, 4)} {unit_stamp}")
     print(f"Module Runtime(): {round(t, 4)} {unit} \n")
 
     # Generate proposals
@@ -167,13 +163,12 @@ def build_neurograph_from_gcs_zips(
             search_radius, n_proposals_per_leaf=n_proposals_per_leaf
         )
         t, unit = utils.time_writer(time() - t0)
-        time_stamp, unit_stamp = utils.time_writer(time() - total_runtime)
         print("\n# proposals:", len(neurograph.mutable_edges))
-        print(f"Time Stamp: {round(t_stamp, 4)} {unit_stamp}")
         print(f"Module Runtime(): {round(t, 4)} {unit} \n")
 
     t, unit = utils.time_writer(time() - total_runtime)
     print(f"Total Runtime: {round(t, 4)} {unit}")
+    print(f"Memory Consumption: {round(utils.get_memory_usage(), 4)} GBs")
     return neurograph
 
 
diff --git a/src/deep_neurographs/utils.py b/src/deep_neurographs/utils.py
index 542f668..d6c444d 100644
--- a/src/deep_neurographs/utils.py
+++ b/src/deep_neurographs/utils.py
@@ -12,6 +12,7 @@
 import json
 import math
 import os
+import psutil
 import shutil
 from copy import deepcopy
 from io import BytesIO
@@ -508,16 +509,8 @@ def reformat_number(number):
     return f"{number:,}"
 
 
-def time_writer(t, unit="seconds"):
-    assert unit in ["seconds", "minutes", "hours"]
-    upd_unit = {"seconds": "minutes", "minutes": "hours"}
-    if t < 60 or unit == "hours":
-        return t, unit
-    else:
-        t /= 60
-        unit = upd_unit[unit]
-        t, unit = time_writer(t, unit=unit)
-    return t, unit
+def get_memory_usage():
+    return psutil.virtual_memory().used / 1e9
 
 
 def init_timers():
@@ -533,8 +526,13 @@ def progress_bar(current, total, bar_length=50, eta=None, runtime=None):
     print(f"\r{bar} {n_completed} | {eta} | {runtime}    ", end="", flush=True)
 
 
-def xor(a, b):
-    if (a and b) or (not a and not b):
-        return False
+def time_writer(t, unit="seconds"):
+    assert unit in ["seconds", "minutes", "hours"]
+    upd_unit = {"seconds": "minutes", "minutes": "hours"}
+    if t < 60 or unit == "hours":
+        return t, unit
     else:
-        return True
+        t /= 60
+        unit = upd_unit[unit]
+        t, unit = time_writer(t, unit=unit)
+    return t, unit