diff --git a/ogcore/SS.py b/ogcore/SS.py index 61da4a400..8678d76c2 100644 --- a/ogcore/SS.py +++ b/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, diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 811459b8f..7996e5217 100644 --- a/ogcore/TPI.py +++ b/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=(