Practical Examples

Complete Working Examples

Example 1: Loading and Exploring a Model

from dolo import yaml_import
import numpy as np

# Load model
model = yaml_import('consumption_savings_iid_egm.yaml')

# Explore structure
print("Model symbols:", model.symbols)
print("Equation blocks:", list(model.equations.keys()))
print("Parameters:", model.calibration.flat)

# Access compiled functions
transition_func = model.functions['transition']
arbitrage_func = model.functions['arbitrage']
expectation_func = model.functions['expectation']

Example 2: Evaluating Functions at a Single Point

# Set up a single evaluation point
y = np.array([0.0])      # No shock
w = np.array([2.0])      # Wealth = 2
c = np.array([0.9])      # Consumption = 0.9
mr = np.array([0.5])     # Marginal rate placeholder

# Parameters from calibration
params = np.array([
    model.calibration['β'],  # 0.96
    model.calibration['γ'],  # 4.0
    model.calibration['σ'],  # 0.1
    model.calibration['ρ'],  # 0.0
    model.calibration['r'],  # 1.0
])

# Evaluate auxiliary function: a = w - c
a = model.functions['auxiliary_direct_egm'](y, w, c, mr, params)
print(f"Assets (w={w[0]}, c={c[0]}): a = {a[0]}")
# Expected: a = 2.0 - 0.9 = 1.1 (approximately)

# Evaluate expectation: mr = β*(c[t+1])^(-γ)*r
c_next = np.array([1.0])
mr_value = model.functions['expectation'](
    y, w, c_next, mr, params  # Using c_next as future consumption
)
print(f"Marginal rate for c[t+1]={c_next[0]}: mr = {mr_value[0]}")

Example 3: Grid Evaluation

# Create a wealth grid
N = 100
w_grid = np.linspace(0.5, 5.0, N)

# Assume a consumption rule: c = 0.6 * w
c_grid = 0.6 * w_grid

# Other variables (constant across grid)
y_grid = np.zeros(N)
mr_grid = np.ones(N) * 0.5

# Evaluate assets at all grid points
a_grid = model.functions['auxiliary_direct_egm'](
    y_grid, w_grid, c_grid, mr_grid, params
)

print(f"Grid evaluation:")
print(f"  Wealth range: [{w_grid.min():.2f}, {w_grid.max():.2f}]")
print(f"  Consumption range: [{c_grid.min():.2f}, {c_grid.max():.2f}]")
print(f"  Assets range: [{a_grid.min():.2f}, {a_grid.max():.2f}]")

# Verify: a should equal w - c
assert np.allclose(a_grid, w_grid - c_grid, rtol=0.01)
print("✓ Verified: a = w - c")

Example 4: Solving for Equilibrium (Simple Example)

def find_steady_state_consumption(model, w_ss=1.0, tol=1e-6):
    """
    Find steady-state consumption for given wealth.
    In steady state: c[t] = c[t+1] = c_ss
    """
    params = np.array([
        model.calibration['β'],
        model.calibration['γ'],
        model.calibration['σ'],
        model.calibration['ρ'],
        model.calibration['r']
    ])

    # Initial guess
    c_ss = 0.5 * w_ss

    # Fixed other values
    y = np.array([0.0])
    w = np.array([w_ss])
    mr = np.array([0.5])

    for iter in range(100):
        c_current = np.array([c_ss])
        c_next = np.array([c_ss])  # Same in steady state

        # Evaluate Euler equation residual
        # In steady state: β*r = 1 (for log utility)
        residual = model.functions['arbitrage'](
            y, w, c_current,
            y, w, c_next,  # Same values in steady state
            params
        )

        if abs(residual[0]) < tol:
            print(f"Converged in {iter} iterations")
            break

        # Simple update (would use better solver in practice)
        c_ss = c_ss - 0.1 * residual[0]

    return c_ss

# Find steady state
c_steady = find_steady_state_consumption(model)
print(f"Steady-state consumption: {c_steady:.4f}")

Example 5: Time Iteration Pattern

def simple_time_iteration(model, grid_size=50, max_iter=100, tol=1e-6):
    """
    Simplified time iteration example.
    """
    # Setup
    params = np.array([
        model.calibration['β'],
        model.calibration['γ'],
        model.calibration['σ'],
        model.calibration['ρ'],
        model.calibration['r']
    ])

    # State grid
    w_min, w_max = 0.1, 4.0
    w_grid = np.linspace(w_min, w_max, grid_size)

    # Initial policy: consume fixed fraction
    c_policy = 0.5 * w_grid

    # Fixed values
    y = np.zeros(grid_size)
    mr = np.zeros(grid_size)

    for iteration in range(max_iter):
        c_old = c_policy.copy()

        # For each grid point, solve for optimal consumption
        for i, w_i in enumerate(w_grid):
            # This is simplified - real implementation would:
            # 1. Compute expectation over future states
            # 2. Solve FOC for optimal c[i]
            # Here we just show the function evaluation pattern

            # Evaluate marginal value
            w_array = np.array([w_i])
            c_array = np.array([c_old[i]])

            # Get marginal rate (simplified)
            mr_i = model.functions['expectation'](
                np.array([0.0]), w_array, c_array,
                np.array([0.0]), params
            )

            # Update policy (simplified rule)
            c_policy[i] = min(c_old[i] * 1.01, w_i * 0.99)

        # Check convergence
        if np.max(np.abs(c_policy - c_old)) < tol:
            print(f"Converged in {iteration} iterations")
            break

    return w_grid, c_policy

# Run simplified time iteration
w_grid, c_policy = simple_time_iteration(model)
print(f"Policy function computed on {len(w_grid)} grid points")

Example 6: Working with Multiple State Variables

# Example model structure with multiple states
model_multi = {
    'symbols': {
        'states': ['k', 'h', 'z'],     # capital, human capital, productivity
        'controls': ['c', 'l', 'i'],   # consumption, labor, investment
        'exogenous': ['eps_z', 'eps_h'],
        'parameters': ['α', 'β', 'δ', 'η', 'ρ']
    }
}

# Grid over multiple dimensions
N = 30  # points per dimension
k_grid = np.linspace(1, 10, N)
h_grid = np.linspace(0.5, 5, N)
z_grid = np.linspace(-0.5, 0.5, N)

# Create mesh grid (all combinations)
K, H, Z = np.meshgrid(k_grid, h_grid, z_grid, indexing='ij')

# Flatten for function evaluation
states = np.column_stack([
    K.flatten(),
    H.flatten(),
    Z.flatten()
])  # Shape: (N^3, 3)

print(f"State space grid: {states.shape}")
print(f"  Total points: {states.shape[0]:,}")
print(f"  k range: [{states[:, 0].min():.2f}, {states[:, 0].max():.2f}]")
print(f"  h range: [{states[:, 1].min():.2f}, {states[:, 1].max():.2f}]")
print(f"  z range: [{states[:, 2].min():.2f}, {states[:, 2].max():.2f}]")

# Evaluate function at all grid points (vectorized)
# result = model.functions['some_function'](exo, states, controls, params)

Example 7: Monte Carlo Simulation

def simulate_paths(model, T=100, N=1000):
    """
    Simulate N paths for T periods.
    """
    params = np.array([
        model.calibration['β'],
        model.calibration['γ'],
        model.calibration['σ'],
        model.calibration['ρ'],
        model.calibration['r']
    ])

    # Initialize
    w = np.ones(N) * model.calibration['w']  # Initial wealth
    c = np.ones(N) * model.calibration['c']  # Initial consumption

    # Storage
    w_sim = np.zeros((T, N))
    c_sim = np.zeros((T, N))

    # Shock process (simplified IID)
    sigma_y = model.calibration['σ']
    y_shocks = np.random.normal(0, sigma_y, (T, N))

    for t in range(T):
        # Store current values
        w_sim[t] = w
        c_sim[t] = c

        # Simplified policy rule (would use solved policy in practice)
        c = 0.1 + 0.5 * w  # Linear rule for illustration

        # Ensure consumption within bounds
        c = np.minimum(c, w * 0.99)  # Can't consume more than wealth

        # Transition to next period
        if t > 0:
            y_prev = y_shocks[t-1]
        else:
            y_prev = np.zeros(N)

        # Other required inputs (placeholders)
        mr = np.zeros(N)

        # Next period wealth
        w_next = model.functions['transition'](
            y_prev, w, c, y_shocks[t], params
        )

        # Handle potential issues
        w_next = np.maximum(w_next, 0.01)  # Ensure positive

        # Update for next iteration
        w = w_next

    return w_sim, c_sim

# Run simulation
w_paths, c_paths = simulate_paths(model)

print(f"Simulation complete:")
print(f"  Periods: {w_paths.shape[0]}")
print(f"  Paths: {w_paths.shape[1]}")
print(f"  Mean wealth (final): {w_paths[-1].mean():.3f}")
print(f"  Std wealth (final): {w_paths[-1].std():.3f}")

Example 8: Debugging Function Calls

def debug_function_call(model, func_name='auxiliary_direct_egm'):
    """
    Helper to understand what a function expects and returns.
    """
    print(f"Debugging function: {func_name}")
    print("=" * 50)

    # Get the function
    func = model.functions[func_name]

    # Get the equation
    if func_name in model.equations:
        eq = model.equations[func_name]
        print(f"Equation: {eq}")

    # Check function properties
    print(f"\nFunction properties:")
    print(f"  Type: {type(func)}")
    print(f"  Number of outputs: {func.n_output}")

    # Prepare test inputs
    params = np.array([
        model.calibration['β'],
        model.calibration['γ'],
        model.calibration['σ'],
        model.calibration['ρ'],
        model.calibration['r']
    ])

    # Single point test
    test_inputs = [
        np.array([0.0]),  # y
        np.array([2.0]),  # w
        np.array([0.9]),  # c
        np.array([0.5]),  # mr
        params
    ]

    print(f"\nTest inputs:")
    for i, inp in enumerate(test_inputs[:-1]):
        print(f"  Input {i}: {inp} (shape: {inp.shape})")
    print(f"  Params: {test_inputs[-1]} (shape: {test_inputs[-1].shape})")

    # Call function
    try:
        result = func(*test_inputs)
        print(f"\nResult: {result}")
        print(f"  Shape: {result.shape}")
        print(f"  Type: {result.dtype}")
    except Exception as e:
        print(f"\nError: {e}")

    # Try with vectorized inputs
    print("\n" + "=" * 50)
    print("Vectorized test (N=10):")

    N = 10
    test_inputs_vec = [
        np.zeros(N),                    # y
        np.linspace(1, 3, N),           # w
        np.linspace(0.5, 1.5, N),       # c
        np.ones(N) * 0.5,               # mr
        params
    ]

    try:
        result_vec = func(*test_inputs_vec)
        print(f"Result shape: {result_vec.shape}")
        print(f"First 5 values: {result_vec[:5]}")
    except Exception as e:
        print(f"Error: {e}")

# Debug a function
debug_function_call(model, 'expectation')

Common Patterns and Best Practices

Pattern 1: Extract Parameters Once

# Do this once
params = np.array([model.calibration[p] for p in model.symbols['parameters']])

# Use many times
for i in range(1000):
    result = func(..., params)

Pattern 2: Vectorize Over Grids

# Create grids
w_grid = np.linspace(w_min, w_max, N)
c_grid = policy_function(w_grid)  # Some policy

# Evaluate at all points simultaneously
results = model.functions['arbitrage'](
    y_grid, w_grid, c_grid, ..., params
)

Pattern 3: Handle Time Indices

# For time-dependent evaluations
for t in range(T):
    if t == 0:
        # Handle initial period
        y_prev = np.zeros(N)
        w_prev = w_init
    else:
        y_prev = y_sim[t-1]
        w_prev = w_sim[t-1]

    # Current period
    w_curr = model.functions['transition'](
        y_prev, w_prev, c_prev, y_sim[t], params
    )

Pattern 4: Error Handling

def safe_function_call(func, *args):
    """Wrapper with error handling."""
    try:
        # Ensure all inputs are arrays
        args_array = []
        for arg in args[:-1]:  # All except params
            if not isinstance(arg, np.ndarray):
                arg = np.array([arg])
            args_array.append(arg)
        args_array.append(args[-1])  # Add params

        # Check shapes match
        shapes = [a.shape[0] for a in args_array[:-1]]
        if len(set(shapes)) > 1:
            raise ValueError(f"Shape mismatch: {shapes}")

        return func(*args_array)

    except Exception as e:
        print(f"Function call failed: {e}")
        print(f"Input shapes: {[a.shape for a in args_array]}")
        raise

Tips for Performance

1. Preallocate arrays instead of growing them
2. Vectorize instead of looping when possible
3. Cache parameter arrays instead of recreating
4. Use views/slices instead of copies when possible
5. Profile your code to find bottlenecks

import time

# Timing example
N = 10000
inputs = [np.random.randn(N) for _ in range(4)]
inputs.append(params)

start = time.time()
result = model.functions['some_function'](*inputs)
elapsed = time.time() - start

print(f"Evaluated {N:,} points in {elapsed*1000:.2f} ms")
print(f"Rate: {N/elapsed:,.0f} points/second")

---

These examples demonstrate the core patterns for working with dolo's compiled functions. The key is understanding that everything is vectorized and all inputs must be arrays!

Christopher Carroll and Akshay Shanker — repository in development, please do not cite.