HANK_and_SAM_tutorial_utils.py

"""
Utilities for HANK-and-SAM-tutorial.ipynb

This module contains:
1. Calibration loading from JSON (generated by HA-Fiscal-HANK-SAM.py)
2. Plotting functions extracted from the tutorial notebook

When the main model (HA-Fiscal-HANK-SAM.py) is re-run, it regenerates
HA_Fiscal_calibration.json. The tutorial automatically picks up the new values.
"""

import json
import os
import numpy as np
import matplotlib.pyplot as plt


# ============================================================================
# Plot Configuration
# ============================================================================

# Color scheme for plots (consistent across all plotting functions)
PLOT_COLORS = {
    "green": "darkorange",  # Used for "Fixed Nominal Rate" lines
    "red": "red",            # Used for "Fixed Real Rate" lines
    "blue": "blue",          # Used for UI Extensions
    "stimulus": "green",     # Used for Stimulus Check in simple plots
}


# ============================================================================
# Calibration Loading
# ============================================================================

def load_calibration(calibration_path=None):
    """
    Load calibration parameters from JSON file.
    
    Parameters
    ----------
    calibration_path : str, optional
        Path to calibration JSON file. If None, looks in standard location.
    
    Returns
    -------
    dict
        Nested dictionary with calibration values organized by category:
        - labor_market: job_find, EU_prob, job_sep, N_ss, U_ss, ss_dstn
        - matching: alpha, phi_ss, v_ss, theta_ss, chi_ss, eta_ss
        - household: R, r_ss, C_ss, A_ss
        - bonds: delta, qb_ss, B_ss
        - wages_taxes: wage_ss, tau_ss, inc_ui_exhaust, UI
        - government: G_ss, Y_priv
        - firms: kappa, HC_ss, epsilon_p, MC_ss, Z_ss, Y_ss
        - inflation: pi_ss, varphi, kappa_p_ss
        - policy_defaults: phi_pi, phi_y, rho_r, phi_b, phi_w
        - computation: bigT
    """
    if calibration_path is None:
        # Try standard locations
        possible_paths = [
            'Code/HA-Models/FromPandemicCode/HA_Fiscal_calibration.json',
            '../Code/HA-Models/FromPandemicCode/HA_Fiscal_calibration.json',
            'HA_Fiscal_calibration.json',
        ]
        for path in possible_paths:
            if os.path.exists(path):
                calibration_path = path
                break
        else:
            raise FileNotFoundError(
                f"Could not find calibration file. Tried: {possible_paths}"
            )
    
    with open(calibration_path, 'r') as f:
        cal = json.load(f)
    
    # Convert ss_dstn back to numpy array
    if 'labor_market' in cal and 'ss_dstn' in cal['labor_market']:
        cal['labor_market']['ss_dstn'] = np.array(cal['labor_market']['ss_dstn'])
    
    return cal


def get_flat_calibration(calibration_path=None):
    """
    Load calibration and flatten into a single-level dictionary.
    
    Returns
    -------
    dict
        Flat dictionary with all calibration values
    """
    cal = load_calibration(calibration_path)
    flat = {}
    for category, values in cal.items():
        if isinstance(values, dict):
            flat.update(values)
    return flat


def load_steady_state(namespace):
    """
    Load all steady-state parameters into the caller's namespace.
    
    This silently injects all calibrated parameters (alpha, phi_ss, C_ss, etc.)
    so they can be used directly in the notebook without an explicit listing.
    
    Usage
    -----
    In notebook cell:
        load_steady_state(globals())
        # Now alpha, phi_ss, C_ss, etc. are all available
    
    Parameters
    ----------
    namespace : dict
        Pass globals() from the notebook to inject variables there
    """
    cal = load_calibration()
    flat = get_flat_calibration()
    
    # Inject all parameters into caller's namespace
    namespace.update(flat)
    
    # Also make full calibration dict available
    namespace['cal'] = cal
    
    # Derived quantities needed by notebook
    namespace['num_mrkv'] = len(cal['labor_market']['ss_dstn'])
    namespace['C_ss_sim'] = flat['C_ss']
    namespace['A_ss_sim'] = flat['A_ss']


def build_steady_state_dict(SteadyStateDict, params):
    """
    Build the SteadyState_Dict for sequence_jacobian.
    
    This encapsulates the lengthy dictionary construction that would
    otherwise clutter the tutorial notebook.
    
    Parameters
    ----------
    SteadyStateDict : class
        The SteadyStateDict class from sequence_jacobian
    params : dict
        Dictionary of parameters (from globals() after load_steady_state)
    
    Returns
    -------
    SteadyStateDict
        The constructed steady state dictionary
    """
    p = params  # shorthand
    
    return SteadyStateDict({
        # Market clearing residuals (all zero in steady state)
        "asset_mkt": 0.0,
        "goods_mkt": 0.0,
        "arg_fisher_resid": 0.0,
        "lbp_resid": 0.0,
        "fiscal_resid": 0.0,
        "labor_evo_resid": 0.0,
        "taylor_resid": 0.0,
        "nkpc_resid": 0.0,
        
        # Unemployment by duration
        "U": (1 - p['N_ss']),
        "U1": p['ss_dstn'][1],
        "U2": p['ss_dstn'][2],
        "U3": p['ss_dstn'][3],
        "U4": p['ss_dstn'][4],
        "U5": p['ss_dstn'][5],
        
        # Firm variables
        "epsilon_p": p['epsilon_p'],
        "HC": p['MC_ss'] * p['Z_ss'],
        "MC": p['MC_ss'],
        "Z": p['Z_ss'],
        
        # Aggregate quantities
        "C": p['C_ss'],
        "Y": p['Y_ss'],
        "A": p['A_ss'],
        "N": p['N_ss'],
        "B": p['B_ss'],
        "G": p['G_ss'],
        
        # Interest rates
        "r": p['r_ss'],
        "r_ante": p['r_ss'],
        "i": p['r_ss'],
        
        # Labor market
        "phi": p['phi_ss'],
        "v": p['v_ss'],
        "ev": 0.0,
        "job_sep": p['job_sep'],
        "eta": p['eta_ss'],
        "chi": p['chi_ss'],
        "theta": p['theta_ss'],
        
        # Wages and prices
        "w": p['wage_ss'],
        "tau": p['tau_ss'],
        "pi": p['pi_ss'],
        "qb": p['qb_ss'],
        "varphi": p['varphi'],
        "kappa_p": p['kappa_p_ss'],
        
        # Policy parameters
        "phi_b": p['phi_b'],
        "phi_w": p['real_wage_rigidity'],  # Paper uses 0.837 (Gornemann et al. 2021)
        "rho_r": p['rho_r'],
        "phi_pi": p['phi_pi'],
        "phi_y": p['phi_y'],
        
        # Transfers and fiscal
        "UI": p['UI'],
        "transfers": 0.0,
        "UI_extend": 0.0,
        "deficit_T": -1,
        "UI_extension_cost": 0.0,
        "UI_rr": 0.0,
        "debt": p['qb_ss'] * p['B_ss'],
        "tax_cost": p['tau_ss'] * p['wage_ss'] * p['N_ss'],
        "lag": -1,
    })


# ============================================================================
# Plotting Functions
# ============================================================================

def compute_all_multipliers(
    irfs_transfer, irfs_transfer_fixed_nominal_rate, irfs_transfer_fixed_real_rate,
    irfs_UI_extend, irfs_UI_extend_fixed_nominal_rate, irfs_UI_extension_fixed_real_rate,
    irfs_tau, irfs_tau_fixed_nominal_rate, irfs_tau_fixed_real_rate,
    NPV_func,
    horizon_length=20,
):
    """
    Compute fiscal multipliers for all policies under all monetary regimes.
    
    This function encapsulates the boilerplate array initialization and
    multiplier computation loop, keeping the tutorial notebook focused
    on the economics rather than bookkeeping.
    
    Parameters
    ----------
    irfs_* : dict
        IRF dictionaries from model.solve_impulse_linear() for each policy/regime
    NPV_func : callable
        Net present value function NPV(series, horizon)
    horizon_length : int
        Number of periods to compute multipliers for
    
    Returns
    -------
    dict
        Nested dictionary with multipliers organized by policy and regime:
        {
            'transfers': {'taylor': array, 'fixed_nominal': array, 'fixed_real': array},
            'UI_extend': {...},
            'tax_cut': {...}
        }
    """
    # Initialize result arrays
    results = {
        'transfers': {
            'taylor': np.zeros(horizon_length),
            'fixed_nominal': np.zeros(horizon_length),
            'fixed_real': np.zeros(horizon_length),
        },
        'UI_extend': {
            'taylor': np.zeros(horizon_length),
            'fixed_nominal': np.zeros(horizon_length),
            'fixed_real': np.zeros(horizon_length),
        },
        'tax_cut': {
            'taylor': np.zeros(horizon_length),
            'fixed_nominal': np.zeros(horizon_length),
            'fixed_real': np.zeros(horizon_length),
        },
    }
    
    # Compute multipliers at each horizon
    for i in range(horizon_length):
        # Standard Taylor rule
        results['transfers']['taylor'][i] = NPV_func(
            irfs_transfer["C"], i + 1
        ) / NPV_func(irfs_transfer["transfers"], 300)
        
        results['UI_extend']['taylor'][i] = NPV_func(
            irfs_UI_extend["C"], i + 1
        ) / NPV_func(irfs_UI_extend["UI_extension_cost"], 300)
        
        results['tax_cut']['taylor'][i] = -NPV_func(
            irfs_tau["C"], i + 1
        ) / NPV_func(irfs_tau["tax_cost"], 300)
        
        # Fixed nominal rate
        results['transfers']['fixed_nominal'][i] = NPV_func(
            irfs_transfer_fixed_nominal_rate["C"], i + 1
        ) / NPV_func(irfs_transfer_fixed_nominal_rate["transfers"], 300)
        
        results['UI_extend']['fixed_nominal'][i] = NPV_func(
            irfs_UI_extend_fixed_nominal_rate["C"], i + 1
        ) / NPV_func(irfs_UI_extend_fixed_nominal_rate["UI_extension_cost"], 300)
        
        results['tax_cut']['fixed_nominal'][i] = -NPV_func(
            irfs_tau_fixed_nominal_rate["C"], i + 1
        ) / NPV_func(irfs_tau_fixed_nominal_rate["tax_cost"], 300)
        
        # Fixed real rate
        results['transfers']['fixed_real'][i] = NPV_func(
            irfs_transfer_fixed_real_rate["C"], i + 1
        ) / NPV_func(irfs_transfer_fixed_real_rate["transfers"], 300)
        
        results['UI_extend']['fixed_real'][i] = NPV_func(
            irfs_UI_extension_fixed_real_rate["C"], i + 1
        ) / NPV_func(irfs_UI_extension_fixed_real_rate["UI_extension_cost"], 300)
        
        results['tax_cut']['fixed_real'][i] = -NPV_func(
            irfs_tau_fixed_real_rate["C"], i + 1
        ) / NPV_func(irfs_tau_fixed_real_rate["tax_cost"], 300)
    
    return results


def plot_multipliers_by_horizon(
    multipliers_transfers,
    multipliers_UI_extend,
    multipliers_tax_cut,
    horizon_length=None,
    xlim=(0.5, 12.5),
):
    """
    Plot fiscal multipliers for three policies across horizons (standard Taylor rule).
    
    Shows how cumulative multipliers evolve from short to long horizons.
    
    Parameters
    ----------
    multipliers_transfers : array-like
        Multiplier series for stimulus checks
    multipliers_UI_extend : array-like
        Multiplier series for UI extensions
    multipliers_tax_cut : array-like
        Multiplier series for tax cuts
    horizon_length : int, optional
        Number of periods to plot. If None, uses length of input arrays.
    xlim : tuple
        X-axis limits
    """
    if horizon_length is None:
        horizon_length = len(multipliers_transfers)
    
    x = np.arange(horizon_length) + 1
    
    plt.figure(figsize=(8, 5))
    plt.plot(x, multipliers_transfers, label="Stimulus Check", color="green", linewidth=2)
    plt.plot(x, multipliers_UI_extend, label="UI Extensions", color="blue", linewidth=2)
    plt.plot(x, multipliers_tax_cut, label="Tax Cut", color="red", linewidth=2)
    
    plt.legend(loc="lower right", fontsize=10)
    plt.ylabel("Consumption Multipliers", fontsize=11)
    plt.xlabel("Quarters", fontsize=11)
    plt.xlim(xlim)
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.show()


def plot_multipliers_three_experiments(
    multipliers_transfers,
    multipliers_transfers_fixed_nominal_rate,
    multipliers_transfers_fixed_real_rate,
    multipliers_UI_extend,
    multipliers_UI_extensions_fixed_nominal_rate,
    multipliers_UI_extensions_fixed_real_rate,
    multipliers_tax_cut,
    multipliers_tax_cut_fixed_nominal_rate,
    multipliers_tax_cut_fixed_real_rate,
    horizon_length=None,
):
    """
    Plot fiscal multipliers for three policies under three monetary regimes.
    
    Parameters
    ----------
    multipliers_*: array-like
        Multiplier time series for each policy/monetary regime combination
    horizon_length: int, optional
        Number of periods to plot. If None, uses length of input arrays.
    """
    if horizon_length is None:
        horizon_length = len(multipliers_transfers)
    
    green = PLOT_COLORS["green"]
    red = PLOT_COLORS["red"]

    Length = len(multipliers_transfers_fixed_nominal_rate) + 1
    fontsize = 10
    width = 2
    label_size = 8
    legend_size = 11
    ticksize = 8
    fig, axs = plt.subplots(1, 3, figsize=(12, 4))

    y_max1 = max(multipliers_transfers_fixed_nominal_rate) * 1.5
    y_max2 = max(multipliers_UI_extensions_fixed_real_rate) * 1.5
    y_max = max([y_max1, y_max2])
    for i in range(3):
        axs[i].set_ylim(-0.2, y_max)

    # Panel 1: Stimulus Check (labeled as UI Extension in original - keeping consistent)
    axs[1].plot(
        np.arange(horizon_length) + 1,
        multipliers_transfers,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[1].plot(
        np.arange(horizon_length) + 1,
        multipliers_transfers_fixed_nominal_rate,
        linewidth=width,
        label="Fixed Nominal Rate",
        linestyle="--",
        color=green,
    )
    axs[1].plot(
        np.arange(horizon_length) + 1,
        multipliers_transfers_fixed_real_rate,
        linewidth=width,
        label="Fixed Real ",
        linestyle=":",
        color=red,
    )
    axs[1].set_title("Stimulus Check", fontdict={"fontsize": fontsize})

    # Panel 0: UI Extension
    axs[0].plot(
        np.arange(horizon_length) + 1,
        multipliers_UI_extend,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[0].plot(
        np.arange(horizon_length) + 1,
        multipliers_UI_extensions_fixed_nominal_rate,
        linewidth=width,
        label="Fixed Nominal Rate",
        linestyle="--",
        color=green,
    )
    axs[0].plot(
        np.arange(horizon_length) + 1,
        multipliers_UI_extensions_fixed_real_rate,
        linewidth=width,
        label="Fixed Real ",
        linestyle=":",
        color=red,
    )
    axs[0].set_title("UI Extension", fontdict={"fontsize": fontsize})
    axs[0].legend(prop={"size": legend_size}, loc="upper left")

    # Panel 2: Tax Cut
    axs[2].plot(
        np.arange(horizon_length) + 1,
        multipliers_tax_cut,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[2].plot(
        np.arange(horizon_length) + 1,
        multipliers_tax_cut_fixed_nominal_rate,
        linewidth=width,
        label="Fixed Nominal Rate",
        linestyle="--",
        color=green,
    )
    axs[2].plot(
        np.arange(horizon_length) + 1,
        multipliers_tax_cut_fixed_real_rate,
        linewidth=width,
        label="Fixed Real ",
        linestyle=":",
        color=red,
    )
    axs[2].set_title("Tax Cut", fontdict={"fontsize": fontsize})

    for i in range(3):
        axs[i].plot(np.zeros(Length), "k")
        axs[i].tick_params(axis="both", labelsize=ticksize)
        axs[i].set_ylabel("Multipliers", fontsize=label_size)
        axs[i].set_xlabel("Quarters", fontsize=label_size)
        axs[i].locator_params(axis="both", nbins=7)
        axs[i].grid(alpha=0.3)

    fig.tight_layout()
    plt.show()


def plot_consumption_irfs_three_experiments(
    irf_UI1, irf_UI2, irf_UI3,
    irf_SC1, irf_SC2, irf_SC3,
    irf_TC1, irf_TC2, irf_TC3,
    C_ss,
):
    """
    Plot consumption IRFs for three policies under three monetary regimes.
    
    Parameters
    ----------
    irf_*: dict
        IRF dictionaries with 'C' key for consumption
    C_ss: float
        Steady state consumption for normalization
    """
    green = PLOT_COLORS["green"]
    red = PLOT_COLORS["red"]

    Length = 12
    fontsize = 10
    width = 2
    label_size = 8
    legend_size = 11
    ticksize = 8
    fig, axs = plt.subplots(1, 3, figsize=(12, 4))

    y_max1 = max(100 * irf_TC2["C"][:Length] / C_ss) * 1.05
    y_max2 = max(100 * irf_SC2["C"][:Length] / C_ss) * 1.05
    y_max = max([y_max1, y_max2])
    for i in range(3):
        axs[i].set_ylim(-0.2, y_max)

    # Panel 1: UI Extension
    axs[1].plot(
        100 * irf_UI1["C"][:Length] / C_ss,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[1].plot(
        100 * irf_UI2["C"][:Length] / C_ss,
        linewidth=width,
        label="Fixed Nominal Rate",
        linestyle="--",
        color=green,
    )
    axs[1].plot(
        100 * irf_UI3["C"][:Length] / C_ss,
        linewidth=width,
        label="Fixed Real ",
        linestyle=":",
        color=red,
    )
    axs[1].set_title("UI Extension", fontdict={"fontsize": fontsize})

    # Panel 0: Stimulus Check
    axs[0].plot(
        100 * irf_SC1["C"][:Length] / C_ss,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[0].plot(
        100 * irf_SC2["C"][:Length] / C_ss,
        linewidth=width,
        label="Fixed Nominal Rate",
        linestyle="--",
        color=green,
    )
    axs[0].plot(
        100 * irf_SC3["C"][:Length] / C_ss,
        linewidth=width,
        label="Fixed Real ",
        linestyle=":",
        color=red,
    )
    axs[0].set_title("Stimulus Check", fontdict={"fontsize": fontsize})
    axs[0].legend(prop={"size": legend_size})

    # Panel 2: Tax Cut
    axs[2].plot(
        100 * irf_TC1["C"][:Length] / C_ss,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[2].plot(
        100 * irf_TC2["C"][:Length] / C_ss,
        linewidth=width,
        label="Fixed Nominal Rate",
        linestyle="--",
        color=green,
    )
    axs[2].plot(
        100 * irf_TC3["C"][:Length] / C_ss,
        linewidth=width,
        label="Fixed Real ",
        linestyle=":",
        color=red,
    )
    axs[2].set_title("Tax Cut", fontdict={"fontsize": fontsize})

    for i in range(3):
        axs[i].plot(np.zeros(Length), "k")
        axs[i].tick_params(axis="both", labelsize=ticksize)
        axs[i].set_ylabel("% consumption deviation", fontsize=label_size)
        axs[i].set_xlabel("Quarters", fontsize=label_size)
        axs[i].locator_params(axis="both", nbins=7)
        axs[i].grid(alpha=0.3)

    fig.tight_layout()
    plt.show()


def plot_consumption_irfs_three(irf_SC1, irf_UI1, irf_TC1, C_ss):
    """
    Plot consumption IRFs for three policies under standard Taylor rule.
    
    Parameters
    ----------
    irf_*: dict
        IRF dictionaries with 'C' key
    C_ss: float
        Steady state consumption
    """
    Length = 12
    fontsize = 10
    width = 2
    label_size = 8
    legend_size = 8
    ticksize = 8
    fig, axs = plt.subplots(1, 3, figsize=(12, 4))

    y_max1 = max(100 * irf_TC1["C"][:Length] / C_ss) * 1.05
    y_max2 = max(100 * irf_SC1["C"][:Length] / C_ss) * 1.05
    y_max = max([y_max1, y_max2])
    for i in range(3):
        axs[i].set_ylim(-0.1, y_max)

    axs[1].plot(
        100 * irf_UI1["C"][:Length] / C_ss,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[1].set_title("UI Extension", fontdict={"fontsize": fontsize})

    axs[0].plot(
        100 * irf_SC1["C"][:Length] / C_ss,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[0].set_title("Stimulus Check", fontdict={"fontsize": fontsize})
    axs[0].legend(prop={"size": legend_size})

    axs[2].plot(
        100 * irf_TC1["C"][:Length] / C_ss,
        linewidth=width,
        label="Standard Taylor Rule",
    )
    axs[2].set_title("Tax Cut", fontdict={"fontsize": fontsize})

    for i in range(3):
        axs[i].plot(np.zeros(Length), "k")
        axs[i].tick_params(axis="both", labelsize=ticksize)
        axs[i].set_ylabel("% consumption deviation", fontsize=label_size)
        axs[i].set_xlabel("Quarters", fontsize=label_size)
        axs[i].locator_params(axis="both", nbins=7)
        axs[i].grid(alpha=0.3)
    
    fig.tight_layout()
    plt.show()


def plot_consumption_irf(irf1, irf2, irf3, C_ss, y_max, title="", legend=False):
    """
    Plot a single consumption IRF comparison across monetary regimes.
    
    Parameters
    ----------
    irf1, irf2, irf3: dict
        IRF dictionaries for Taylor rule, fixed nominal, fixed real
    C_ss: float
        Steady state consumption
    y_max: float
        Y-axis maximum
    title: str
        Plot title
    legend: bool
        Whether to show legend
    """
    green = PLOT_COLORS["green"]
    red = PLOT_COLORS["red"]

    Length = 12
    plt.figure(figsize=(4, 4))
    x_axis = np.arange(1, Length + 1)

    plt.plot(x_axis, 100 * irf1["C"][:Length] / C_ss, label="Active Taylor Rule")
    plt.plot(
        x_axis,
        100 * irf2["C"][:Length] / C_ss,
        label="Fixed Nominal Rate",
        linestyle="--",
        color=green,
    )
    plt.plot(
        x_axis,
        100 * irf3["C"][:Length] / C_ss,
        label="Fixed Real",
        linestyle=":",
        color=red,
    )

    plt.xticks(np.arange(min(x_axis), max(x_axis) + 1, 1.0))
    plt.xlabel("quarter")
    plt.ylim(0, y_max)
    if title:
        plt.title(title)
    if legend:
        plt.legend(loc="best")
        plt.ylabel("% consumption deviation")
    plt.show()


def plot_consumption_multipliers(
    multiplier1, multiplier2, multiplier3, y_max, title="", legend=False
):
    """
    Plot a single multiplier comparison across monetary regimes.
    
    Parameters
    ----------
    multiplier1, multiplier2, multiplier3: array-like
        Multiplier series for Taylor rule, fixed nominal, fixed real
    y_max: float
        Y-axis maximum
    title: str
        Plot title
    legend: bool
        Whether to show legend
    """
    green = PLOT_COLORS["green"]
    red = PLOT_COLORS["red"]

    Length = 12
    plt.figure(figsize=(4, 4))
    x_axis = np.arange(1, Length + 1)

    plt.plot(x_axis, multiplier1[0:Length], label="Active Taylor Rule")
    plt.plot(
        x_axis,
        multiplier2[0:Length],
        label="Fixed Nominal Rate",
        linestyle="--",
        color=green,
    )
    plt.plot(
        x_axis, multiplier3[0:Length], label="Fixed Real", linestyle=":", color=red
    )

    plt.xticks(np.arange(min(x_axis), max(x_axis) + 1, 1.0))
    plt.xlabel("quarter")
    plt.ylim(0, y_max)
    if title:
        plt.title(title)
    if legend:
        plt.legend(loc="best")
    plt.show()