[DO NOT MERGE] Log root finding

ogcore/SS.py

@@ -1,6 +1,7 @@
 # imports
 from re import VERBOSE
 import numpy as np
+import time
 import scipy.optimize as opt
 from dask import delayed, compute
 import dask.multiprocessing
@@ -35,6 +36,19 @@
 ------------------------------------------------------------------------
 """
 
+def time_logger(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        pid = os.getpid()
+        print(f"SS {func.__name__} ZZZ: {pid} execution time: {end_time - start_time} seconds")
+        return result
+    return wrapper
+
+@time_logger
+def root_finder(*args, **kwargs):
+    return opt.root(*args, **kwargs)      
 
 def euler_equation_solver(guesses, *args):
     """
@@ -240,7 +254,7 @@ def inner_loop(outer_loop_vars, p, client):
             scattered_p,
         )
         lazy_values.append(
-            delayed(opt.root)(
+            delayed(root_finder)(
                 euler_equation_solver,
                 guesses * 0.9,
                 args=euler_params,
@@ -1214,7 +1228,7 @@ def run_SS(p, client=None):
                         + list(BQguess)
                         + [TRguess, factorguess]
                     )
-                sol = opt.root(
+                sol = root_finder(
                     SS_fsolve,
                     guesses,
                     args=ss_params_baseline,
@@ -1342,7 +1356,7 @@ def run_SS(p, client=None):
                     + list(BQguess)
                     + [TR_ss]
                 )
-            sol = opt.root(
+            sol = root_finder(
                 SS_fsolve,
                 guesses,
                 args=ss_params_reform,
@@ -1371,7 +1385,7 @@ def run_SS(p, client=None):
                     + list(BQguess)
                     + [TRguess]
                 )
-            sol = opt.root(
+            sol = root_finder(
                 SS_fsolve,
                 guesses,
                 args=ss_params_reform,

ogcore/TPI.py

@@ -1,6 +1,7 @@
 # imports
 import numpy as np
 import pickle
+import time
 import scipy.optimize as opt
 from dask import delayed, compute
 import dask.multiprocessing
@@ -25,6 +26,21 @@
 ENFORCE_SOLUTION_CHECKS = True
 
 
+def time_logger(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        pid = os.getpid()
+        print(f"TPI {func.__name__} ZZZ: {pid} execution time: {end_time - start_time} seconds")
+        return result
+    return wrapper
+
+@time_logger
+def root_finder(*args, **kwargs):
+    return opt.root(*args, **kwargs)      
+
+
 def get_initial_SS_values(p):
     """
     Get values of variables for the initial period and the steady state
@@ -373,7 +389,7 @@ def inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p):
     n_mat = np.zeros((p.T + p.S, p.S))
     euler_errors = np.zeros((p.T, 2 * p.S))
 
-    solutions = opt.root(
+    solutions = root_finder(
         firstdoughnutring,
         [guesses_b[0, -1], guesses_n[0, -1]],
         args=(
@@ -426,7 +442,7 @@ def inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p):
             [temp_mtry[i][j] for j in range(num_params)] for i in range(s + 2)
         ]
 
-        solutions = opt.root(
+        solutions = root_finder(
             twist_doughnut,
             list(b_guesses_to_use) + list(n_guesses_to_use),
             args=(
@@ -479,7 +495,7 @@ def inner_loop(guesses, outer_loop_vars, initial_values, ubi, j, ind, p):
             for s in range(p.S)
         ]
 
-        solutions = opt.root(
+        solutions = root_finder(
             twist_doughnut,
             list(b_guesses_to_use) + list(n_guesses_to_use),
             args=(