diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json index 2922d15f3..f34d6a596 100644 --- a/ogcore/default_parameters.json +++ b/ogcore/default_parameters.json @@ -4348,13 +4348,13 @@ "type": "float", "value": [ { - "value": 1e-09 + "value": 1e-03 } ], "validators": { "range": { "min": 1e-13, - "max": 0.001 + "max": 0.1 } } }, @@ -4384,13 +4384,13 @@ "type": "float", "value": [ { - "value": 1e-08 + "value": 1e-04 } ], "validators": { "range": { "min": 1e-13, - "max": 0.001 + "max": 0.1 } } }, diff --git a/ogcore/household.py b/ogcore/household.py index bc21eb140..ae5b0d9ee 100644 --- a/ogcore/household.py +++ b/ogcore/household.py @@ -10,6 +10,19 @@ import logging from ogcore import config +# Try to import Numba for JIT compilation (optional speedup) +try: + from numba import njit, prange + HAS_NUMBA = True +except ImportError: + HAS_NUMBA = False + # Fallback: decorator that does nothing + def njit(*args, **kwargs): + def decorator(func): + return func + return decorator + prange = range + """ ------------------------------------------------------------------------ Functions @@ -17,6 +30,36 @@ """ +# Numba-optimized core computation for marginal utility of consumption +@njit(cache=True) +def _marg_ut_cons_numba(c_flat, sigma, epsilon=0.003): + """ + Numba-optimized core computation for marginal utility of consumption. + + Args: + c_flat (1D array): flattened consumption array + sigma (float): coefficient of relative risk aversion + epsilon (float): threshold for constraint handling + + Returns: + MU_c (1D array): marginal utility values + """ + n = c_flat.shape[0] + MU_c = np.zeros(n) + b2 = (-sigma * (epsilon ** (-sigma - 1))) / 2 + b1 = (epsilon ** (-sigma)) - 2 * b2 * epsilon + + for i in prange(n): + if c_flat[i] < epsilon: + # Quadratic extrapolation for constrained values + MU_c[i] = 2 * b2 * c_flat[i] + b1 + else: + # Normal marginal utility + MU_c[i] = c_flat[i] ** (-sigma) + + return MU_c + + def marg_ut_cons(c, sigma): r""" Compute the marginal utility of consumption. @@ -34,105 +77,141 @@ def marg_ut_cons(c, sigma): """ if np.ndim(c) == 0: c = np.array([c]) - epsilon = 0.003 - cvec_cnstr = c < epsilon - MU_c = np.zeros(c.shape) - MU_c[~cvec_cnstr] = c[~cvec_cnstr] ** (-sigma) - b2 = (-sigma * (epsilon ** (-sigma - 1))) / 2 - b1 = (epsilon ** (-sigma)) - 2 * b2 * epsilon - MU_c[cvec_cnstr] = 2 * b2 * c[cvec_cnstr] + b1 - output = MU_c + + original_shape = c.shape + c_flat = c.ravel() + + # Use Numba-optimized function + MU_c_flat = _marg_ut_cons_numba(c_flat, sigma) + + # Reshape back to original shape + output = MU_c_flat.reshape(original_shape) output = np.squeeze(output) return output -def marg_ut_labor(n, chi_n, p): - r""" - Compute the marginal disutility of labor. - - .. math:: - MDU_{l} = \chi^n_{s}\biggl(\frac{b}{\tilde{l}}\biggr) - \biggl(\frac{n_{j,s,t}}{\tilde{l}}\biggr)^{\upsilon-1} - \Biggl[1-\biggl(\frac{n_{j,s,t}}{\tilde{l}}\biggr)^\upsilon - \Biggr]^{\frac{1-\upsilon}{\upsilon}} +# Numba-optimized core computation for marginal disutility of labor +@njit(cache=True) +def _marg_ut_labor_numba(nvec_flat, b_ellipse, ltilde, upsilon): + """ + Numba-optimized core computation for marginal disutility of labor. Args: - n (array_like): household labor supply - chi_n (array_like): utility weights on disutility of labor - p (OG-Core Specifications object): model parameters + nvec_flat (1D array): flattened labor supply array + b_ellipse (float): ellipse parameter + ltilde (float): time endowment + upsilon (float): Frisch elasticity parameter Returns: - output (array_like): marginal disutility of labor supply - + MDU_n (1D array): marginal disutility values """ - nvec = n - if np.ndim(nvec) == 0: - nvec = np.array([nvec]) + n = nvec_flat.shape[0] + MDU_n = np.zeros(n) eps_low = 0.000001 - eps_high = p.ltilde - 0.000001 - nvec_low = nvec < eps_low - nvec_high = nvec > eps_high - nvec_uncstr = np.logical_and(~nvec_low, ~nvec_high) - MDU_n = np.zeros(nvec.shape) - MDU_n[nvec_uncstr] = ( - (p.b_ellipse / p.ltilde) - * ((nvec[nvec_uncstr] / p.ltilde) ** (p.upsilon - 1)) - * ( - (1 - ((nvec[nvec_uncstr] / p.ltilde) ** p.upsilon)) - ** ((1 - p.upsilon) / p.upsilon) - ) - ) + eps_high = ltilde - 0.000001 + + # Pre-compute coefficients for low constraint b2 = ( 0.5 - * p.b_ellipse - * (p.ltilde ** (-p.upsilon)) - * (p.upsilon - 1) - * (eps_low ** (p.upsilon - 2)) + * b_ellipse + * (ltilde ** (-upsilon)) + * (upsilon - 1) + * (eps_low ** (upsilon - 2)) * ( - (1 - ((eps_low / p.ltilde) ** p.upsilon)) - ** ((1 - p.upsilon) / p.upsilon) + (1 - ((eps_low / ltilde) ** upsilon)) + ** ((1 - upsilon) / upsilon) ) * ( 1 - + ((eps_low / p.ltilde) ** p.upsilon) - * ((1 - ((eps_low / p.ltilde) ** p.upsilon)) ** (-1)) + + ((eps_low / ltilde) ** upsilon) + * ((1 - ((eps_low / ltilde) ** upsilon)) ** (-1)) ) ) - b1 = (p.b_ellipse / p.ltilde) * ( - (eps_low / p.ltilde) ** (p.upsilon - 1) + b1 = (b_ellipse / ltilde) * ( + (eps_low / ltilde) ** (upsilon - 1) ) * ( - (1 - ((eps_low / p.ltilde) ** p.upsilon)) - ** ((1 - p.upsilon) / p.upsilon) - ) - ( - 2 * b2 * eps_low - ) - MDU_n[nvec_low] = 2 * b2 * nvec[nvec_low] + b1 + (1 - ((eps_low / ltilde) ** upsilon)) + ** ((1 - upsilon) / upsilon) + ) - (2 * b2 * eps_low) + + # Pre-compute coefficients for high constraint d2 = ( 0.5 - * p.b_ellipse - * (p.ltilde ** (-p.upsilon)) - * (p.upsilon - 1) - * (eps_high ** (p.upsilon - 2)) + * b_ellipse + * (ltilde ** (-upsilon)) + * (upsilon - 1) + * (eps_high ** (upsilon - 2)) * ( - (1 - ((eps_high / p.ltilde) ** p.upsilon)) - ** ((1 - p.upsilon) / p.upsilon) + (1 - ((eps_high / ltilde) ** upsilon)) + ** ((1 - upsilon) / upsilon) ) * ( 1 - + ((eps_high / p.ltilde) ** p.upsilon) - * ((1 - ((eps_high / p.ltilde) ** p.upsilon)) ** (-1)) + + ((eps_high / ltilde) ** upsilon) + * ((1 - ((eps_high / ltilde) ** upsilon)) ** (-1)) ) ) - d1 = (p.b_ellipse / p.ltilde) * ( - (eps_high / p.ltilde) ** (p.upsilon - 1) + d1 = (b_ellipse / ltilde) * ( + (eps_high / ltilde) ** (upsilon - 1) ) * ( - (1 - ((eps_high / p.ltilde) ** p.upsilon)) - ** ((1 - p.upsilon) / p.upsilon) - ) - ( - 2 * d2 * eps_high + (1 - ((eps_high / ltilde) ** upsilon)) + ** ((1 - upsilon) / upsilon) + ) - (2 * d2 * eps_high) + + for i in prange(n): + nval = nvec_flat[i] + if nval < eps_low: + MDU_n[i] = 2 * b2 * nval + b1 + elif nval > eps_high: + MDU_n[i] = 2 * d2 * nval + d1 + else: + # Unconstrained case + MDU_n[i] = ( + (b_ellipse / ltilde) + * ((nval / ltilde) ** (upsilon - 1)) + * ( + (1 - ((nval / ltilde) ** upsilon)) + ** ((1 - upsilon) / upsilon) + ) + ) + + return MDU_n + + +def marg_ut_labor(n, chi_n, p): + r""" + Compute the marginal disutility of labor. + + .. math:: + MDU_{l} = \chi^n_{s}\biggl(\frac{b}{\tilde{l}}\biggr) + \biggl(\frac{n_{j,s,t}}{\tilde{l}}\biggr)^{\upsilon-1} + \Biggl[1-\biggl(\frac{n_{j,s,t}}{\tilde{l}}\biggr)^\upsilon + \Biggr]^{\frac{1-\upsilon}{\upsilon}} + + Args: + n (array_like): household labor supply + chi_n (array_like): utility weights on disutility of labor + p (OG-Core Specifications object): model parameters + + Returns: + output (array_like): marginal disutility of labor supply + + """ + nvec = n + if np.ndim(nvec) == 0: + nvec = np.array([nvec]) + + original_shape = nvec.shape + nvec_flat = nvec.ravel() + + # Use Numba-optimized function + MDU_n_flat = _marg_ut_labor_numba( + nvec_flat, p.b_ellipse, p.ltilde, p.upsilon ) - MDU_n[nvec_high] = 2 * d2 * nvec[nvec_high] + d1 + + # Reshape and apply chi_n weights + MDU_n = MDU_n_flat.reshape(original_shape) output = MDU_n * np.squeeze(chi_n) output = np.squeeze(output) return output diff --git a/ogcore/tax.py b/ogcore/tax.py index 4d0110463..ce2e0d040 100644 --- a/ogcore/tax.py +++ b/ogcore/tax.py @@ -9,6 +9,42 @@ from ogcore import utils, pensions from ogcore.txfunc import get_tax_rates +# Try to import Numba for JIT compilation +try: + from numba import njit, vectorize, float64 + HAS_NUMBA = True +except ImportError: + HAS_NUMBA = False + # Create dummy decorators if Numba not available + def njit(*args, **kwargs): + def decorator(func): + return func + if len(args) == 1 and callable(args[0]): + return args[0] + return decorator + def vectorize(*args, **kwargs): + def decorator(func): + return np.vectorize(func) + return decorator + float64 = None + + +# Numba-optimized wealth tax functions +if HAS_NUMBA: + @vectorize([float64(float64, float64, float64, float64)], nopython=True, cache=True) + def _etr_wealth_numba(b, h_wealth, m_wealth, p_wealth): + """Numba-optimized effective tax rate on wealth (scalar).""" + return (p_wealth * h_wealth * b) / (h_wealth * b + m_wealth) + + @vectorize([float64(float64, float64, float64, float64)], nopython=True, cache=True) + def _mtr_wealth_numba(b, h_wealth, m_wealth, p_wealth): + """Numba-optimized marginal tax rate on wealth (scalar).""" + etr = (p_wealth * h_wealth * b) / (h_wealth * b + m_wealth) + return etr * 2 - ((h_wealth**2 * p_wealth * b**2) / ((b * h_wealth + m_wealth) ** 2)) +else: + _etr_wealth_numba = None + _mtr_wealth_numba = None + """ ------------------------------------------------------------------------ Functions @@ -34,7 +70,15 @@ def ETR_wealth(b, h_wealth, m_wealth, p_wealth): tau_w (Numpy array): effective tax rate on wealth, size = SxJ """ - tau_w = (p_wealth * h_wealth * b) / (h_wealth * b + m_wealth) + if HAS_NUMBA and _etr_wealth_numba is not None: + # Ensure arrays are float64 for Numba + b_arr = np.asarray(b, dtype=np.float64) + h_arr = np.float64(h_wealth) if np.isscalar(h_wealth) else np.asarray(h_wealth, dtype=np.float64) + m_arr = np.float64(m_wealth) if np.isscalar(m_wealth) else np.asarray(m_wealth, dtype=np.float64) + p_arr = np.float64(p_wealth) if np.isscalar(p_wealth) else np.asarray(p_wealth, dtype=np.float64) + tau_w = _etr_wealth_numba(b_arr, h_arr, m_arr, p_arr) + else: + tau_w = (p_wealth * h_wealth * b) / (h_wealth * b + m_wealth) return tau_w @@ -57,9 +101,17 @@ def MTR_wealth(b, h_wealth, m_wealth, p_wealth): tau_prime (Numpy array): marginal tax rate on wealth, size = SxJ """ - tau_prime = ETR_wealth(b, h_wealth, m_wealth, p_wealth) * 2 - ( - (h_wealth**2 * p_wealth * b**2) / ((b * h_wealth + m_wealth) ** 2) - ) + if HAS_NUMBA and _mtr_wealth_numba is not None: + # Ensure arrays are float64 for Numba + b_arr = np.asarray(b, dtype=np.float64) + h_arr = np.float64(h_wealth) if np.isscalar(h_wealth) else np.asarray(h_wealth, dtype=np.float64) + m_arr = np.float64(m_wealth) if np.isscalar(m_wealth) else np.asarray(m_wealth, dtype=np.float64) + p_arr = np.float64(p_wealth) if np.isscalar(p_wealth) else np.asarray(p_wealth, dtype=np.float64) + tau_prime = _mtr_wealth_numba(b_arr, h_arr, m_arr, p_arr) + else: + tau_prime = ETR_wealth(b, h_wealth, m_wealth, p_wealth) * 2 - ( + (h_wealth**2 * p_wealth * b**2) / ((b * h_wealth + m_wealth) ** 2) + ) return tau_prime diff --git a/ogcore/txfunc.py b/ogcore/txfunc.py index a6df207fb..a8b235f5d 100644 --- a/ogcore/txfunc.py +++ b/ogcore/txfunc.py @@ -28,6 +28,119 @@ from matplotlib import cm import random +# Try to import Numba for JIT compilation +try: + from numba import njit, prange + HAS_NUMBA = True +except ImportError: + HAS_NUMBA = False + def njit(*args, **kwargs): + def decorator(func): + return func + if len(args) == 1 and callable(args[0]): + return args[0] + return decorator + prange = range + + +# Numba-optimized DEP tax rate calculation (for model runs, not estimation) +@njit(cache=True, parallel=True) +def _get_tax_rates_dep_numba(X, Y, A, B, C, D, max_x, max_y, share, + min_x, min_y, shift_x, shift_y, shift): + """ + Numba-optimized DEP tax rate calculation for model runs. + + This computes the effective tax rate using the DEP functional form + when for_estimation=False and analytical_mtrs=False. + """ + X2 = X * X + Y2 = Y * Y + + # tau_x = (max_x - min_x) * (A * X2 + B * X) / (A * X2 + B * X + 1) + min_x + # tau_y = (max_y - min_y) * (C * Y2 + D * Y) / (C * Y2 + D * Y + 1) + min_y + + n = X.shape[0] + txrates = np.empty(n, dtype=np.float64) + + for i in prange(n): + ax2_bx = A * X2[i] + B * X[i] + cy2_dy = C * Y2[i] + D * Y[i] + + tau_x = (max_x - min_x) * ax2_bx / (ax2_bx + 1.0) + min_x + tau_y = (max_y - min_y) * cy2_dy / (cy2_dy + 1.0) + min_y + + txrates[i] = ((tau_x + shift_x) ** share) * ((tau_y + shift_y) ** (1.0 - share)) + shift + + return txrates + + +@njit(cache=True, parallel=True) +def _get_tax_rates_dep_mtr_labor_numba(X, Y, A, B, C, D, max_x, max_y, share, + min_x, min_y, shift_x, shift_y, shift): + """ + Numba-optimized DEP marginal tax rate on labor income calculation. + + This computes the analytical MTR on labor when for_estimation=False. + """ + X2 = X * X + Y2 = Y * Y + income = X + Y + + n = X.shape[0] + txrates = np.empty(n, dtype=np.float64) + + for i in prange(n): + ax2_bx = A * X2[i] + B * X[i] + cy2_dy = C * Y2[i] + D * Y[i] + + tau_x = (max_x - min_x) * ax2_bx / (ax2_bx + 1.0) + min_x + tau_y = (max_y - min_y) * cy2_dy / (cy2_dy + 1.0) + min_y + + etr = ((tau_x + shift_x) ** share) * ((tau_y + shift_y) ** (1.0 - share)) + shift + + # d_etr for labor income (not capital) + d_etr = (share * ((tau_x + shift_x) ** (share - 1.0)) * + (max_x - min_x) * ((2.0 * A * X[i] + B) / ((ax2_bx + 1.0) ** 2)) * + ((tau_y + shift_y) ** (1.0 - share))) + + txrates[i] = d_etr * income[i] + etr + + return txrates + + +@njit(cache=True, parallel=True) +def _get_tax_rates_dep_mtr_capital_numba(X, Y, A, B, C, D, max_x, max_y, share, + min_x, min_y, shift_x, shift_y, shift): + """ + Numba-optimized DEP marginal tax rate on capital income calculation. + + This computes the analytical MTR on capital when for_estimation=False. + """ + X2 = X * X + Y2 = Y * Y + income = X + Y + + n = X.shape[0] + txrates = np.empty(n, dtype=np.float64) + + for i in prange(n): + ax2_bx = A * X2[i] + B * X[i] + cy2_dy = C * Y2[i] + D * Y[i] + + tau_x = (max_x - min_x) * ax2_bx / (ax2_bx + 1.0) + min_x + tau_y = (max_y - min_y) * cy2_dy / (cy2_dy + 1.0) + min_y + + etr = ((tau_x + shift_x) ** share) * ((tau_y + shift_y) ** (1.0 - share)) + shift + + # d_etr for capital income + d_etr = ((1.0 - share) * ((tau_y + shift_y) ** (-share)) * + (max_y - min_y) * ((2.0 * C * Y[i] + D) / ((cy2_dy + 1.0) ** 2)) * + ((tau_x + shift_x) ** share)) + + txrates[i] = d_etr * income[i] + etr + + return txrates + if not SHOW_RUNTIME: warnings.simplefilter("ignore", RuntimeWarning) @@ -171,46 +284,89 @@ def get_tax_rates( * ((tau_y + shift_y) ** (1 - share)) ) + shift else: + # Check if we can use Numba optimization (1D arrays with scalar params) + use_numba = (HAS_NUMBA and + np.ndim(X) == 1 and + np.ndim(A) == 0 and + X.size >= 10) # Only use Numba for larger arrays + if analytical_mtrs: - tau_x = (max_x - min_x) * (A * X2 + B * X) / ( - A * X2 + B * X + 1 - ) + min_x - tau_y = (max_y - min_y) * (C * Y2 + D * Y) / ( - C * Y2 + D * Y + 1 - ) + min_y - etr = ( - ((tau_x + shift_x) ** share) - * ((tau_y + shift_y) ** (1 - share)) - ) + shift - if mtr_capital: - d_etr = ( - (1 - share) - * ((tau_y + shift_y) ** (-share)) - * (max_y - min_y) - * ((2 * C * Y + D) / ((C * Y2 + D * Y + 1) ** 2)) - * ((tau_x + shift_x) ** share) - ) - txrates = d_etr * income + etr + if use_numba: + # Use Numba-optimized version + X_flat = np.ascontiguousarray(X.ravel(), dtype=np.float64) + Y_flat = np.ascontiguousarray(Y.ravel(), dtype=np.float64) + if mtr_capital: + txrates = _get_tax_rates_dep_mtr_capital_numba( + X_flat, Y_flat, + float(A), float(B), float(C), float(D), + float(max_x), float(max_y), float(share), + float(min_x), float(min_y), + float(shift_x), float(shift_y), float(shift) + ) + else: + txrates = _get_tax_rates_dep_mtr_labor_numba( + X_flat, Y_flat, + float(A), float(B), float(C), float(D), + float(max_x), float(max_y), float(share), + float(min_x), float(min_y), + float(shift_x), float(shift_y), float(shift) + ) + txrates = txrates.reshape(X.shape) else: - d_etr = ( - share - * ((tau_x + shift_x) ** (share - 1)) - * (max_x - min_x) - * ((2 * A * X + B) / ((A * X2 + B * X + 1) ** 2)) + # Original numpy version + tau_x = (max_x - min_x) * (A * X2 + B * X) / ( + A * X2 + B * X + 1 + ) + min_x + tau_y = (max_y - min_y) * (C * Y2 + D * Y) / ( + C * Y2 + D * Y + 1 + ) + min_y + etr = ( + ((tau_x + shift_x) ** share) * ((tau_y + shift_y) ** (1 - share)) - ) - txrates = d_etr * income + etr + ) + shift + if mtr_capital: + d_etr = ( + (1 - share) + * ((tau_y + shift_y) ** (-share)) + * (max_y - min_y) + * ((2 * C * Y + D) / ((C * Y2 + D * Y + 1) ** 2)) + * ((tau_x + shift_x) ** share) + ) + txrates = d_etr * income + etr + else: + d_etr = ( + share + * ((tau_x + shift_x) ** (share - 1)) + * (max_x - min_x) + * ((2 * A * X + B) / ((A * X2 + B * X + 1) ** 2)) + * ((tau_y + shift_y) ** (1 - share)) + ) + txrates = d_etr * income + etr else: - tau_x = ( - (max_x - min_x) * (A * X2 + B * X) / (A * X2 + B * X + 1) - ) + min_x - tau_y = ( - (max_y - min_y) * (C * Y2 + D * Y) / (C * Y2 + D * Y + 1) - ) + min_y - txrates = ( - ((tau_x + shift_x) ** share) - * ((tau_y + shift_y) ** (1 - share)) - ) + shift + if use_numba: + # Use Numba-optimized version for ETR + X_flat = np.ascontiguousarray(X.ravel(), dtype=np.float64) + Y_flat = np.ascontiguousarray(Y.ravel(), dtype=np.float64) + txrates = _get_tax_rates_dep_numba( + X_flat, Y_flat, + float(A), float(B), float(C), float(D), + float(max_x), float(max_y), float(share), + float(min_x), float(min_y), + float(shift_x), float(shift_y), float(shift) + ) + txrates = txrates.reshape(X.shape) + else: + # Original numpy version + tau_x = ( + (max_x - min_x) * (A * X2 + B * X) / (A * X2 + B * X + 1) + ) + min_x + tau_y = ( + (max_y - min_y) * (C * Y2 + D * Y) / (C * Y2 + D * Y + 1) + ) + min_y + txrates = ( + ((tau_x + shift_x) ** share) + * ((tau_y + shift_y) ** (1 - share)) + ) + shift elif tax_func_type == "DEP_totalinc": A, B, max_income, min_income, shift_income, shift = ( np.squeeze(params[..., 0]),