├── docs └── index.md ├── tests ├── __init__.py ├── base │ ├── __init__.py │ ├── test_solved_block.py │ ├── test_jacobian_dict_block.py │ ├── test_steady_state.py │ ├── test_estimation.py │ ├── test_combined_block.py │ ├── test_remap.py │ ├── test_public_classes.py │ ├── test_simple_block.py │ ├── test_options.py │ ├── test_two_asset.py │ ├── test_jacobian.py │ ├── test_multiexog.py │ ├── test_transitional_dynamics.py │ ├── test_stage_block.py │ ├── test_workflow.py │ ├── test_het_support.py │ └── test_dchoice.py ├── performance │ └── __init__.py ├── robustness │ ├── __init__.py │ └── test_steady_state.py ├── utils │ ├── test_multidim.py │ ├── test_DAG.py │ ├── test_ordered_set.py │ └── test_function.py ├── conftest.py └── README.md ├── src └── sequence_jacobian │ ├── examples │ ├── __init__.py │ ├── rbc.py │ ├── hank.py │ ├── krusell_smith.py │ └── two_asset.py │ ├── blocks │ ├── __init__.py │ ├── auxiliary_blocks │ │ ├── __init__.py │ │ └── jacobiandict_block.py │ ├── support │ │ ├── __init__.py │ │ ├── parent.py │ │ ├── het_compiled.py │ │ ├── law_of_motion.py │ │ └── het_support.py │ ├── simple_block.py │ ├── solved_block.py │ └── combined_block.py │ ├── interpolate.py │ ├── misc.py │ ├── hetblocks │ ├── __init__.py │ ├── hh_sim.py │ ├── hh_labor.py │ └── hh_twoasset.py │ ├── grids.py │ ├── classes │ ├── __init__.py │ ├── steady_state_dict.py │ ├── result_dict.py │ └── impulse_dict.py │ ├── utilities │ ├── __init__.py │ ├── multidim.py │ ├── optimized_routines.py │ ├── differentiate.py │ ├── bijection.py │ ├── ordered_set.py │ ├── drawdag.py │ ├── solvers.py │ ├── misc.py │ ├── discretize.py │ ├── interpolate.py │ └── graph.py │ ├── __init__.py │ └── estimation.py ├── notebooks ├── dag │ └── RBC.png └── notebooks.zip ├── pyproject.toml ├── .gitignore ├── setup.cfg ├── .github └── workflows │ └── main.yml ├── LICENSE └── README.md /docs/index.md: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | """All tests""" -------------------------------------------------------------------------------- /src/sequence_jacobian/examples/__init__.py: -------------------------------------------------------------------------------- 1 | """Example models""" -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/__init__.py: -------------------------------------------------------------------------------- 1 | """Block-construction tools""" -------------------------------------------------------------------------------- /src/sequence_jacobian/interpolate.py: -------------------------------------------------------------------------------- 1 | from .utilities.interpolate import * 2 | -------------------------------------------------------------------------------- /tests/base/__init__.py: -------------------------------------------------------------------------------- 1 | """Tests for base-level functionality of the package""" -------------------------------------------------------------------------------- /tests/performance/__init__.py: -------------------------------------------------------------------------------- 1 | """Tests to check the performance of the code""" -------------------------------------------------------------------------------- /notebooks/dag/RBC.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shade-econ/sequence-jacobian/HEAD/notebooks/dag/RBC.png -------------------------------------------------------------------------------- /src/sequence_jacobian/misc.py: -------------------------------------------------------------------------------- 1 | # to be determined... 2 | from .utilities.optimized_routines import setmin 3 | -------------------------------------------------------------------------------- /notebooks/notebooks.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shade-econ/sequence-jacobian/HEAD/notebooks/notebooks.zip -------------------------------------------------------------------------------- /src/sequence_jacobian/hetblocks/__init__.py: -------------------------------------------------------------------------------- 1 | '''Heterogeneous agent blocks''' 2 | from . import hh_labor, hh_sim, hh_twoasset 3 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/auxiliary_blocks/__init__.py: -------------------------------------------------------------------------------- 1 | """Auxiliary Block types for building a coherent backend for Block handling""" 2 | -------------------------------------------------------------------------------- /src/sequence_jacobian/grids.py: -------------------------------------------------------------------------------- 1 | # ADD asset_grid in a minute! 2 | from .utilities.discretize import agrid, asset_grid, markov_rouwenhorst, markov_tauchen -------------------------------------------------------------------------------- /tests/robustness/__init__.py: -------------------------------------------------------------------------------- 1 | """Tests to check for code robustness, including error checking and attempts to use models with bad initializations.""" 2 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/support/__init__.py: -------------------------------------------------------------------------------- 1 | """Other classes and helpers to aid standard block functionality: .steady_state, .impulse_linear, .impulse_nonlinear, 2 | .jacobian""" 3 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = [ 3 | "setuptools>=42", 4 | "wheel", 5 | "numpy>=1.19.2", 6 | "scipy>=1.2", 7 | "numba>=0.49", 8 | ] 9 | build-backend = "setuptools.build_meta" -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | __pycache__/ 3 | notebooks/.ipynb_checkpoints/ 4 | .ipynb_checkpoints/ 5 | .DS_Store 6 | .vscode/ 7 | .pytest_cache/ 8 | 9 | *.zip 10 | !notebooks.zip 11 | 12 | scripts/ 13 | wiki/ 14 | *.egg-info/ 15 | 16 | dist/ 17 | -------------------------------------------------------------------------------- /src/sequence_jacobian/classes/__init__.py: -------------------------------------------------------------------------------- 1 | from .steady_state_dict import SteadyStateDict, UserProvidedSS 2 | from .impulse_dict import ImpulseDict 3 | from .jacobian_dict import JacobianDict, FactoredJacobianDict 4 | from .sparse_jacobians import IdentityMatrix, SimpleSparse 5 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/__init__.py: -------------------------------------------------------------------------------- 1 | """Utilities relating to: interpolation, forward step/transition, grids and Markov chains, solvers, sorting, etc.""" 2 | 3 | from . import (bijection, differentiate, discretize, drawdag, function, graph, interpolate, 4 | misc, multidim, optimized_routines, ordered_set, solvers) 5 | -------------------------------------------------------------------------------- /tests/utils/test_multidim.py: -------------------------------------------------------------------------------- 1 | from sequence_jacobian.utilities.multidim import outer 2 | import numpy as np 3 | 4 | def test_2d(): 5 | a = np.random.rand(10) 6 | b = np.random.rand(12) 7 | assert np.allclose(np.outer(a,b), outer([a,b])) 8 | 9 | def test_3d(): 10 | a = np.array([1., 2]) 11 | b = np.array([1., 7]) 12 | small = np.outer(a, b) 13 | 14 | c = np.array([2., 4]) 15 | product = np.empty((2,2,2)) 16 | product[..., 0] = 2*small 17 | product[..., 1] = 4*small 18 | 19 | assert np.array_equal(product, outer([a,b,c])) 20 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | """Fixtures used by tests.""" 2 | 3 | import pytest 4 | 5 | from sequence_jacobian.examples import rbc, krusell_smith, hank, two_asset 6 | 7 | 8 | @pytest.fixture(scope='session') 9 | def rbc_dag(): 10 | return rbc.dag() 11 | 12 | 13 | @pytest.fixture(scope='session') 14 | def krusell_smith_dag(): 15 | return krusell_smith.dag() 16 | 17 | 18 | @pytest.fixture(scope='session') 19 | def one_asset_hank_dag(): 20 | return hank.dag() 21 | 22 | 23 | @pytest.fixture(scope='session') 24 | def two_asset_hank_dag(): 25 | return two_asset.dag() 26 | 27 | 28 | @pytest.fixture(scope='session') 29 | def ks_remapped_dag(): 30 | return krusell_smith.remapped_dag() 31 | -------------------------------------------------------------------------------- /src/sequence_jacobian/classes/steady_state_dict.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | 3 | from .result_dict import ResultDict 4 | from ..utilities.misc import dict_diff 5 | from ..utilities.ordered_set import OrderedSet 6 | from ..utilities.bijection import Bijection 7 | 8 | import numpy as np 9 | 10 | from numbers import Real 11 | from typing import Any, Dict, Union 12 | Array = Any 13 | 14 | class SteadyStateDict(ResultDict): 15 | def difference(self, data_to_remove): 16 | return SteadyStateDict(dict_diff(self.toplevel, data_to_remove), deepcopy(self.internals)) 17 | 18 | def _vector_valued(self): 19 | return OrderedSet([k for k, v in self.toplevel.items() if np.size(v) > 1]) 20 | 21 | UserProvidedSS = Dict[str, Union[Real, Array]] 22 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = sequence-jacobian 3 | version = 1.0.0 4 | author = Sequence-Jacobian Team 5 | author_email = sequence.jacobian.team@gmail.com 6 | description = Sequence-Space Jacobian Methods for Solving and Estimating Heterogeneous Agent Models 7 | long_description = file: README.md 8 | long_description_content_type = text/markdown 9 | url = https://github.com/shade-econ/sequence-jacobian 10 | project_urls = 11 | Bug Tracker = https://github.com/shade-econ/sequence-jacobian/issues 12 | classifiers = 13 | Programming Language :: Python :: 3 14 | License :: OSI Approved :: MIT License 15 | Operating System :: OS Independent 16 | 17 | [options] 18 | package_dir = 19 | = src 20 | packages = find: 21 | python_requires = >=3.7 22 | install_requires = 23 | numpy >= 1.19.2 24 | scipy >= 1.2 25 | numba >= 0.49 26 | 27 | [options.packages.find] 28 | where = src -------------------------------------------------------------------------------- /tests/README.md: -------------------------------------------------------------------------------- 1 | # Sequence-Space Jacobian Test Suite 2 | 3 | This is a `pytest` test suite, containing unit tests for the Sequence-Space Jacobian package. 4 | Contained in the various directories in this test suite are major classes of tests: 5 | - `base`: The unit tests for the base functionality of the package. These should be run after every commit 6 | to the codebase. 7 | - `robustness`: The unit tests to ensure the utilization of various features of the package are robust to bad initial 8 | conditions or invalid user input. Some of these tests take a substantially longer time to execute and hence should be 9 | run subject to user discretion. 10 | - `performance`: The tests to gauge the performance (time-/allocation-wise) of the package to ensure no new bottlenecks 11 | are created and that code performance remains in line with previous iterations of the package. These tests can also be 12 | used for code optimization. -------------------------------------------------------------------------------- /.github/workflows/main.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies and run tests with a variety of Python versions 2 | # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python 3 | 4 | name: CI 5 | 6 | on: 7 | push: 8 | branches: [ "master" ] 9 | pull_request: 10 | branches: [ "master" ] 11 | 12 | jobs: 13 | build: 14 | runs-on: ubuntu-latest 15 | strategy: 16 | fail-fast: false 17 | matrix: 18 | python-version: ["3.9", "3.10", "3.11"] 19 | 20 | steps: 21 | - uses: actions/checkout@v4 22 | - name: Set up Python ${{ matrix.python-version }} 23 | uses: actions/setup-python@v5 24 | with: 25 | python-version: ${{ matrix.python-version }} 26 | - name: Install dependencies 27 | run: | 28 | python -m pip install --upgrade pip 29 | python -m pip install pytest setuptools wheel 30 | python -m pip install . 31 | - name: Test with pytest 32 | run: | 33 | pytest -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Sequence-Jacobian Team 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE.zoom 22 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/multidim.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def multiply_ith_dimension(Pi, i, X): 5 | """If Pi is a matrix, multiply Pi times the ith dimension of X and return""" 6 | X = X.swapaxes(0, i) 7 | shape = X.shape 8 | X = X.reshape((shape[0], -1)) 9 | 10 | # iterate forward using Pi 11 | X = Pi @ X 12 | 13 | # reverse steps 14 | X = X.reshape((Pi.shape[0], *shape[1:])) 15 | return X.swapaxes(0, i) 16 | 17 | 18 | def outer(pis): 19 | """Return n-dimensional outer product of list of n vectors""" 20 | pi = pis[0] 21 | for pi_i in pis[1:]: 22 | pi = np.kron(pi, pi_i) 23 | return pi.reshape(*(len(pi_i) for pi_i in pis)) 24 | 25 | 26 | def batch_multiply_ith_dimension(P, i, X): 27 | """If P is (D, X.shape) array, multiply P and X along ith dimension of X.""" 28 | # standardize arrays 29 | P = P.swapaxes(1, 1 + i) 30 | X = X.swapaxes(0, i) 31 | Pshape = P.shape 32 | P = P.reshape((*Pshape[:2], -1)) 33 | X = X.reshape((X.shape[0], -1)) 34 | 35 | # P[i, j, ...] @ X[j, ...] 36 | X = np.einsum('ijb,jb->ib', P, X) 37 | 38 | # original shape and order 39 | X = X.reshape(Pshape[0], *Pshape[2:]) 40 | return X.swapaxes(0, i) 41 | -------------------------------------------------------------------------------- /tests/base/test_solved_block.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from sequence_jacobian import simple, solved 3 | from sequence_jacobian.classes.steady_state_dict import SteadyStateDict 4 | from sequence_jacobian.classes.jacobian_dict import FactoredJacobianDict 5 | 6 | 7 | @simple 8 | def myblock(u, i): 9 | res = 0.5 * i(1) - u**2 - u(1) 10 | return res 11 | 12 | 13 | @solved(unknowns={'u': (-10.0, 10.0)}, targets=['res'], solver='brentq') 14 | def myblock_solved(u, i): 15 | res = 0.5 * i(1) - u**2 - u(1) 16 | return res 17 | 18 | def test_solved_block(): 19 | ss = SteadyStateDict({'u': 5, 'i': 10, 'res': 0.0}) 20 | 21 | # Compute jacobian of myblock_solved from scratch 22 | J1 = myblock_solved.jacobian(ss, inputs=['i'], T=20) 23 | 24 | # Compute jacobian of SolvedBlock using a pre-computed FactoredJacobian 25 | J_u = myblock.jacobian(ss, inputs=['u'], T=20) # square jac of underlying simple block 26 | J_factored = FactoredJacobianDict(J_u, T=20) 27 | J_i = myblock.jacobian(ss, inputs=['i'], T=20) # jac of underlying simple block wrt inputs that are NOT unknowns 28 | J2 = J_factored.compose(J_i) # obtain jac of unknown wrt to non-unknown inputs using factored jac 29 | 30 | assert np.allclose(J1['u']['i'], J2['u']['i']) 31 | 32 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/auxiliary_blocks/jacobiandict_block.py: -------------------------------------------------------------------------------- 1 | """A simple wrapper for JacobianDicts to be embedded in DAGs""" 2 | 3 | from ..block import Block 4 | from ...classes import ImpulseDict, JacobianDict 5 | 6 | class JacobianDictBlock(JacobianDict, Block): 7 | """A wrapper for nested dicts/JacobianDicts passed directly into DAGs to ensure method compatibility""" 8 | def __init__(self, nesteddict, outputs=None, inputs=None, name=None): 9 | super().__init__(nesteddict, outputs=outputs, inputs=inputs, name=name) 10 | Block.__init__(self) 11 | 12 | def __repr__(self): 13 | return f"" 14 | 15 | def _impulse_linear(self, ss, inputs, outputs, Js): 16 | return ImpulseDict(self.jacobian(ss, list(inputs.keys()), outputs, inputs.T, Js).apply(inputs)) 17 | 18 | def _jacobian(self, ss, inputs, outputs, T): 19 | if not inputs <= self.inputs: 20 | raise KeyError(f'Asking JacobianDictBlock for {inputs - self.inputs}, which are among its inputs {self.inputs}') 21 | if not outputs <= self.outputs: 22 | raise KeyError(f'Asking JacobianDictBlock for {outputs - self.outputs}, which are among its outputs {self.outputs}') 23 | return self[outputs, inputs] 24 | -------------------------------------------------------------------------------- /tests/base/test_jacobian_dict_block.py: -------------------------------------------------------------------------------- 1 | """Test JacobianDictBlock functionality""" 2 | 3 | import numpy as np 4 | 5 | from sequence_jacobian import combine 6 | from sequence_jacobian.examples import rbc 7 | from sequence_jacobian.blocks.auxiliary_blocks.jacobiandict_block import JacobianDictBlock 8 | from sequence_jacobian import SteadyStateDict 9 | 10 | 11 | def test_jacobian_dict_block_impulses(rbc_dag): 12 | rbc_model, ss, unknowns, _, exogenous = rbc_dag 13 | 14 | T = 10 15 | J_pe = rbc_model.jacobian(ss, inputs=unknowns + exogenous, T=10) 16 | J_block = JacobianDictBlock(J_pe) 17 | 18 | J_block_Z = J_block.jacobian(SteadyStateDict({}), ["Z"]) 19 | for o in J_block_Z.outputs: 20 | assert np.all(J_block[o].get("Z") == J_block_Z[o].get("Z")) 21 | 22 | dZ = 0.8 ** np.arange(T) 23 | 24 | dO1 = J_block @ {"Z": dZ} 25 | dO2 = J_block_Z @ {"Z": dZ} 26 | 27 | for k in J_block: 28 | assert np.all(dO1[k] == dO2[k]) 29 | 30 | 31 | def test_jacobian_dict_block_combine(rbc_dag): 32 | _, ss, _, _, exogenous = rbc_dag 33 | 34 | J_firm = rbc.firm.jacobian(ss, inputs=exogenous) 35 | blocks_w_jdict = [rbc.household, J_firm, rbc.mkt_clearing] 36 | cblock_w_jdict = combine(blocks_w_jdict) 37 | 38 | # Using `combine` converts JacobianDicts to JacobianDictBlocks 39 | assert isinstance(cblock_w_jdict.blocks[0], JacobianDictBlock) 40 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/optimized_routines.py: -------------------------------------------------------------------------------- 1 | """Njitted routines to speed up some steps in backward iteration or aggregation""" 2 | 3 | import numpy as np 4 | from numba import njit 5 | 6 | 7 | @njit 8 | def setmin(x, xmin): 9 | """Set 2-dimensional array x where each row is ascending equal to equal to max(x, xmin).""" 10 | ni, nj = x.shape 11 | for i in range(ni): 12 | for j in range(nj): 13 | if x[i, j] < xmin: 14 | x[i, j] = xmin 15 | else: 16 | break 17 | 18 | 19 | @njit 20 | def within_tolerance(x1, x2, tol): 21 | """Efficiently test max(abs(x1-x2)) <= tol for arrays of same dimensions x1, x2.""" 22 | y1 = x1.ravel() 23 | y2 = x2.ravel() 24 | 25 | for i in range(y1.shape[0]): 26 | if np.abs(y1[i] - y2[i]) > tol: 27 | return False 28 | return True 29 | 30 | 31 | @njit 32 | def fast_aggregate(X, Y): 33 | """If X has dims (T, ...) and Y has dims (T, ...), do dot product for each T to get length-T vector. 34 | 35 | Identical to np.sum(X*Y, axis=(1,...,X.ndim-1)) but avoids costly creation of intermediates, 36 | useful for speeding up aggregation in td by factor of 4 to 5.""" 37 | T = X.shape[0] 38 | Xnew = X.reshape(T, -1) 39 | Ynew = Y.reshape(T, -1) 40 | Z = np.empty(T) 41 | for t in range(T): 42 | Z[t] = Xnew[t, :] @ Ynew[t, :] 43 | return Z 44 | -------------------------------------------------------------------------------- /src/sequence_jacobian/hetblocks/hh_sim.py: -------------------------------------------------------------------------------- 1 | '''Standard Incomplete Market model''' 2 | 3 | import numpy as np 4 | 5 | from ..blocks.het_block import het 6 | from .. import interpolate, misc, grids 7 | 8 | '''Core HetBlock''' 9 | 10 | def hh_init(a_grid, y, r, eis): 11 | coh = (1 + r) * a_grid[np.newaxis, :] + y[:, np.newaxis] 12 | Va = (1 + r) * (0.1 * coh) ** (-1 / eis) 13 | return Va 14 | 15 | 16 | @het(exogenous='Pi', policy='a', backward='Va', backward_init=hh_init) 17 | def hh(Va_p, a_grid, y, r, beta, eis): 18 | uc_nextgrid = beta * Va_p 19 | c_nextgrid = uc_nextgrid ** (-eis) 20 | coh = (1 + r) * a_grid[np.newaxis, :] + y[:, np.newaxis] 21 | a = interpolate.interpolate_y(c_nextgrid + a_grid, coh, a_grid) 22 | misc.setmin(a, a_grid[0]) 23 | c = coh - a 24 | Va = (1 + r) * c ** (-1 / eis) 25 | return Va, a, c 26 | 27 | 28 | '''Extended HetBlock with grid and income process inputs added, and example calibration''' 29 | 30 | def make_grids(rho_e, sd_e, n_e, min_a, max_a, n_a): 31 | e_grid, _, Pi = grids.markov_rouwenhorst(rho_e, sd_e, n_e) 32 | a_grid = grids.asset_grid(min_a, max_a, n_a) 33 | return e_grid, Pi, a_grid 34 | 35 | 36 | def income(w, e_grid): 37 | y = w * e_grid 38 | return y 39 | 40 | 41 | hh_extended = hh.add_hetinputs([income, make_grids]) 42 | 43 | 44 | def example_calibration(): 45 | return dict(min_a=0, max_a=1000, rho_e=0.975, sd_e=0.7, n_a=200, n_e=7, 46 | w=1, r=0.01/4, beta=1-0.08/4, eis=1) 47 | -------------------------------------------------------------------------------- /src/sequence_jacobian/examples/rbc.py: -------------------------------------------------------------------------------- 1 | from sequence_jacobian import simple, create_model 2 | 3 | 4 | '''Part 1: Blocks''' 5 | 6 | @simple 7 | def firm(K, L, Z, alpha, delta): 8 | r = alpha * Z * (K(-1) / L) ** (alpha-1) - delta 9 | w = (1 - alpha) * Z * (K(-1) / L) ** alpha 10 | Y = Z * K(-1) ** alpha * L ** (1 - alpha) 11 | return r, w, Y 12 | 13 | 14 | @simple 15 | def household(K, L, w, eis, frisch, vphi, delta): 16 | C = (w / vphi / L ** (1 / frisch)) ** eis 17 | I = K - (1 - delta) * K(-1) 18 | return C, I 19 | 20 | 21 | @simple 22 | def mkt_clearing(r, C, Y, I, K, L, w, eis, beta): 23 | goods_mkt = Y - C - I 24 | euler = C ** (-1 / eis) - beta * (1 + r(+1)) * C(+1) ** (-1 / eis) 25 | walras = C + K - (1 + r) * K(-1) - w * L 26 | return goods_mkt, euler, walras 27 | 28 | 29 | '''Part 2: Assembling the model''' 30 | 31 | def dag(): 32 | # Combine blocks 33 | blocks = [household, firm, mkt_clearing] 34 | rbc_model = create_model(blocks, name="RBC") 35 | 36 | # Steady state 37 | calibration = {'eis': 1., 'frisch': 1., 'delta': 0.025, 'alpha': 0.11, 'L': 1.} 38 | unknowns_ss = {'vphi': 0.92, 'beta': 1 / (1 + 0.01), 'K': 2., 'Z': 1.} 39 | targets_ss = {'goods_mkt': 0., 'r': 0.01, 'euler': 0., 'Y': 1.} 40 | ss = rbc_model.solve_steady_state(calibration, unknowns_ss, targets_ss, solver='hybr') 41 | 42 | # Transitional dynamics 43 | unknowns = ['K', 'L'] 44 | targets = ['goods_mkt', 'euler'] 45 | exogenous = ['Z'] 46 | 47 | return rbc_model, ss, unknowns, targets, exogenous 48 | -------------------------------------------------------------------------------- /tests/base/test_steady_state.py: -------------------------------------------------------------------------------- 1 | """Test all models' steady state computations""" 2 | 3 | import numpy as np 4 | 5 | from sequence_jacobian.examples import rbc, krusell_smith, hank, two_asset 6 | 7 | 8 | # def test_rbc_steady_state(rbc_dag): 9 | # _, ss, *_ = rbc_dag 10 | # ss_ref = rbc.rbc_ss() 11 | # assert set(ss.keys()) == set(ss_ref.keys()) 12 | # for k in ss.keys(): 13 | # assert np.all(np.isclose(ss[k], ss_ref[k])) 14 | 15 | 16 | # def test_ks_steady_state(krusell_smith_dag): 17 | # _, ss, *_ = krusell_smith_dag 18 | # ss_ref = krusell_smith.ks_ss(nS=2, nA=10, amax=200) 19 | # assert set(ss.keys()) == set(ss_ref.keys()) 20 | # for k in ss.keys(): 21 | # assert np.all(np.isclose(ss[k], ss_ref[k])) 22 | 23 | 24 | # def test_hank_steady_state(one_asset_hank_dag): 25 | # _, ss, *_ = one_asset_hank_dag 26 | # ss_ref = hank.hank_ss(nS=2, nA=10, amax=150) 27 | # assert set(ss.keys()) == set(ss_ref.keys()) 28 | # for k in ss.keys(): 29 | # assert np.all(np.isclose(ss[k], ss_ref[k])) 30 | 31 | 32 | # def test_two_asset_steady_state(two_asset_hank_dag): 33 | # _, ss, *_ = two_asset_hank_dag 34 | # ss_ref = two_asset.two_asset_ss(nZ=3, nB=10, nA=16, nK=4, verbose=False) 35 | # assert set(ss.keys()) == set(ss_ref.keys()) 36 | # for k in ss.keys(): 37 | # assert np.all(np.isclose(ss[k], ss_ref[k])) 38 | 39 | 40 | # def test_remap_steady_state(ks_remapped_dag): 41 | # _, _, _, _, ss = ks_remapped_dag 42 | # assert ss['beta_impatient'] < ss['beta_patient'] 43 | # assert ss['A_impatient'] < ss['A_patient'] 44 | -------------------------------------------------------------------------------- /tests/base/test_estimation.py: -------------------------------------------------------------------------------- 1 | """Test all models' estimation calculations""" 2 | '' 3 | import pytest 4 | import numpy as np 5 | 6 | from sequence_jacobian import estimation 7 | 8 | 9 | # See test_determinacy.py for the to-do describing this suppression 10 | @pytest.mark.filterwarnings("ignore:.*cannot be safely interpreted as an integer.*:DeprecationWarning") 11 | def test_krusell_smith_estimation(krusell_smith_dag): 12 | _, ss, ks_model, unknowns, targets, exogenous = krusell_smith_dag 13 | 14 | np.random.seed(41234) 15 | T = 50 16 | G = ks_model.solve_jacobian(ss, unknowns, targets, exogenous, T=T) 17 | 18 | # Step 1: Stacked impulse responses 19 | rho = 0.9 20 | sigma_persist = 0.1 21 | sigma_trans = 0.2 22 | 23 | dZ1 = rho**(np.arange(T)) 24 | dY1, dC1, dK1 = G['Y']['Z'] @ dZ1, G['C']['Z'] @ dZ1, G['K']['Z'] @ dZ1 25 | dX1 = np.stack([dZ1, dY1, dC1, dK1], axis=1) 26 | 27 | dZ2 = np.arange(T) == 0 28 | dY2, dC2, dK2 = G['Y']['Z'] @ dZ2, G['C']['Z'] @ dZ2, G['K']['Z'] @ dZ2 29 | dX2 = np.stack([dZ2, dY2, dC2, dK2], axis=1) 30 | 31 | dX = np.stack([dX1, dX2], axis=2) 32 | 33 | # Step 2: Obtain covariance at all leads and lags 34 | sigmas = np.array([sigma_persist, sigma_trans]) 35 | Sigma = estimation.all_covariances(dX, sigmas) 36 | 37 | # Step 3: Log-likelihood calculation 38 | # random 100 observations 39 | Y = np.random.randn(100, 4) 40 | 41 | # 0.05 measurement error in each variable 42 | sigma_measurement = np.full(4, 0.05) 43 | 44 | # calculate log-likelihood 45 | ll = estimation.log_likelihood(Y, Sigma, sigma_measurement) 46 | assert np.isclose(ll, -59921.410111251025) -------------------------------------------------------------------------------- /tests/base/test_combined_block.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sequence_jacobian as sj 3 | 4 | 5 | def test_jacobian_accumulation(): 6 | 7 | # Define two blocks. Notice: Second one does not use output from the first! 8 | @sj.solved(unknowns={'p': (-10, 1000)}, targets=['valuation'] , solver="brentq") 9 | def equity(r1, p, Y): 10 | valuation = Y + p(+1) / (1 + r1) - p 11 | return valuation 12 | 13 | @sj.simple 14 | def mkt_clearing(r0, r1, A0, A1, Y, B): 15 | asset_mkt_0 = A0 + (r0 + Y - 0.5*r1) - B 16 | asset_mkt_1 = A1 + (r1 + Y - 0.5*r0) - B 17 | return asset_mkt_0, asset_mkt_1 18 | 19 | both_blocks = sj.create_model([equity, mkt_clearing]) 20 | only_second = sj.create_model([mkt_clearing]) 21 | 22 | calibration = {'B': 0, 'Y': 0, 'r0': 0.01/4, 'r1': 0.01/4, 'A0': 1, 'A1': 1} 23 | ss_both = both_blocks.steady_state(calibration) 24 | ss_second = only_second.steady_state(calibration) 25 | 26 | # Second block alone gives us Jacobian without issues. 27 | 28 | unknowns_td = ['Y', 'r1'] 29 | targets_td = ['asset_mkt_0', 'asset_mkt_1'] 30 | T = 300 31 | shock = {'r0': 0.95**np.arange(T)} 32 | irf = only_second.solve_impulse_linear(ss_second, unknowns_td, targets_td, shock) 33 | G = only_second.solve_jacobian(ss_second, unknowns_td, targets_td, ['r0'], T=T) 34 | 35 | # Both blocks give us trouble. Even though solve_impulse_linear runs through... 36 | 37 | unknowns_td = ['Y', 'r1'] 38 | targets_td = ['asset_mkt_0', 'asset_mkt_1'] 39 | T = 300 40 | shock = {'r0': 0.95**np.arange(T)} 41 | irf = both_blocks.solve_impulse_linear(ss_both, unknowns_td, targets_td, shock) 42 | G = both_blocks.solve_jacobian(ss_both, unknowns_td, targets_td, ['r0'], T=T) -------------------------------------------------------------------------------- /tests/base/test_remap.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from sequence_jacobian import simple, solved, combine 3 | 4 | 5 | @simple 6 | def matching(theta, ell, kappa): 7 | f = theta / (1 + theta ** ell) ** (1 / ell) 8 | qfill = f / theta 9 | hiring_cost = kappa / qfill 10 | return f, qfill, hiring_cost 11 | 12 | 13 | @solved(unknowns={'h': (0, 1)}, targets=['jc_res']) 14 | def job_creation(h, w, beta, s, hiring_cost): 15 | jc_res = h - w + beta * (1 - s(+1)) * hiring_cost(+1) - hiring_cost 16 | return jc_res 17 | 18 | 19 | @solved(unknowns={'N': (0.5, 1)}, targets=['N_lom']) 20 | def labor_lom(h, w, N, s, qfill, f, theta, hiring_cost): 21 | N_lom = (1 - s * (1 - f)) * N(-1) + f * (1 - N(-1)) - N 22 | U = 1 - N 23 | v = theta * U 24 | vacancy_cost = hiring_cost * qfill * v 25 | Div_labor = (h - w) * N - vacancy_cost 26 | return N_lom, U, v, vacancy_cost, Div_labor 27 | 28 | 29 | @simple 30 | def dmp_aggregate(U_men, U_women, Div_labor_men, Div_labor_women, vacancy_cost_men, vacancy_cost_women): 31 | U = (U_men + U_women) / 2 32 | Div_labor = (Div_labor_men + Div_labor_women) / 2 33 | vacancy_cost = (vacancy_cost_men + vacancy_cost_women) / 2 34 | return U, Div_labor, vacancy_cost 35 | 36 | 37 | def test_remap_combined_block(): 38 | dmp = combine([matching, job_creation, labor_lom], name='DMP') 39 | dmp_men = dmp.rename(suffix='_men') 40 | dmp_women = dmp.rename(suffix='_women') 41 | 42 | # remap some inputs and all outputs 43 | to_remap = ['theta', 'ell', 's'] + list(dmp_men.outputs) 44 | dmp_men = dmp_men.remap({k: k + '_men' for k in to_remap}) 45 | dmp_women = dmp_women.remap({k: k + '_women' for k in to_remap}) 46 | 47 | # combine remapped blocks 48 | dmp_all = combine([dmp_men, dmp_women, dmp_aggregate], name='dmp_all') 49 | -------------------------------------------------------------------------------- /tests/base/test_public_classes.py: -------------------------------------------------------------------------------- 1 | """Test public-facing classes""" 2 | 3 | import numpy as np 4 | import pytest 5 | 6 | from sequence_jacobian import het 7 | from sequence_jacobian.classes.steady_state_dict import SteadyStateDict 8 | from sequence_jacobian.classes.impulse_dict import ImpulseDict 9 | from sequence_jacobian.utilities.bijection import Bijection 10 | 11 | def test_impulsedict(krusell_smith_dag): 12 | _, ss, ks_model, unknowns, targets, _ = krusell_smith_dag 13 | T = 200 14 | 15 | # Linearized impulse responses as deviations 16 | ir_lin = ks_model.solve_impulse_linear(ss, unknowns, targets, inputs={'Z': 0.01 * 0.5**np.arange(T)}, outputs=['C', 'K', 'r']) 17 | 18 | # Get method 19 | assert isinstance(ir_lin, ImpulseDict) 20 | assert isinstance(ir_lin[['C']], ImpulseDict) 21 | assert isinstance(ir_lin['C'], np.ndarray) 22 | 23 | # Merge method 24 | temp = ir_lin[['C', 'K']] | ir_lin[['r']] 25 | assert list(temp.keys()) == ['C', 'K', 'r'] 26 | 27 | # SS and scalar multiplication 28 | dC1 = 100 * ir_lin['C'] / ss['C'] 29 | dC2 = 100 * ir_lin[['C']] / ss 30 | assert np.allclose(dC1, dC2['C']) 31 | 32 | 33 | def test_bijection(): 34 | # generate and invert 35 | mymap = Bijection({'a': 'a1', 'b': 'b1'}) 36 | mymapinv = mymap.inv 37 | assert mymap['a'] == 'a1' and mymap['b'] == 'b1' 38 | assert mymapinv['a1'] == 'a' and mymapinv['b1'] == 'b' 39 | 40 | # duplicate keys rejected 41 | with pytest.raises(ValueError): 42 | Bijection({'a': 'a1', 'b': 'a1'}) 43 | 44 | # composition with another bijection (flows backwards) 45 | mymap2 = Bijection({'a1': 'a2'}) 46 | assert (mymap2 @ mymap)['a'] == 'a2' 47 | 48 | # composition with SteadyStateDict 49 | ss = SteadyStateDict({'a': 2.0, 'b': 1.0}) 50 | ss_remapped = ss @ mymap 51 | assert isinstance(ss_remapped, SteadyStateDict) 52 | assert ss_remapped['a1'] == ss['a'] and ss_remapped['b1'] == ss['b'] 53 | -------------------------------------------------------------------------------- /tests/utils/test_DAG.py: -------------------------------------------------------------------------------- 1 | from sequence_jacobian.utilities.graph import DAG 2 | from sequence_jacobian.utilities.ordered_set import OrderedSet 3 | from sequence_jacobian import simple, combine 4 | import pytest 5 | 6 | class Block: 7 | def __init__(self, inputs, outputs): 8 | self.inputs = OrderedSet(inputs) 9 | self.outputs = OrderedSet(outputs) 10 | 11 | 12 | test_dag = DAG([Block(inputs=['a', 'b', 'z'], outputs=['c', 'd']), 13 | Block(inputs=['a', 'e'], outputs=['b']), 14 | Block(inputs = ['d'], outputs=['f'])]) 15 | 16 | 17 | def test_dag_constructor(): 18 | # the blocks should be ordered 1, 0, 2 19 | assert list(test_dag.blocks[0].inputs) == ['a', 'e'] 20 | assert list(test_dag.blocks[1].inputs) == ['a', 'b', 'z'] 21 | assert list(test_dag.blocks[2].inputs) == ['d'] 22 | 23 | assert set(test_dag.inmap['a']) == {0, 1} 24 | assert set(test_dag.inmap['b']) == {1} 25 | 26 | assert test_dag.outmap['c'] == 1 27 | assert test_dag.outmap['f'] == 2 28 | assert test_dag.outmap['d'] == 1 29 | 30 | assert set(test_dag.adj[0]) == {1} 31 | assert set(test_dag.adj[1]) == {2} 32 | assert set(test_dag.revadj[2]) == {1} 33 | assert set(test_dag.revadj[1]) == {0} 34 | 35 | 36 | def test_visited(): 37 | test_dag.visit_from_outputs(['f']) == OrderedSet([0, 1, 2]) 38 | test_dag.visit_from_outputs(['b']) == OrderedSet([0]) 39 | test_dag.visit_from_outputs(['d']) == OrderedSet([0, 1]) 40 | 41 | test_dag.visit_from_inputs(['e']) == OrderedSet([0, 1, 2]) 42 | test_dag.visit_from_inputs(['z']) == OrderedSet([1, 2]) 43 | 44 | 45 | def test_find_cycle(): 46 | @simple 47 | def f1(x): 48 | y = x 49 | return y 50 | 51 | @simple 52 | def f2(y, theta): 53 | z = y 54 | return z 55 | 56 | @simple 57 | def f3(x, z): 58 | w = x * z 59 | return w 60 | 61 | @simple 62 | def f4(z): 63 | x = z 64 | return x 65 | 66 | with pytest.raises(Exception) as exception: 67 | a = combine([f1, f2, f3, f4]) 68 | assert "Topological sort failed: cyclic dependency f1 -> f4 -> f2 -> f1" in str(exception.value) 69 | -------------------------------------------------------------------------------- /tests/utils/test_ordered_set.py: -------------------------------------------------------------------------------- 1 | from sequence_jacobian.utilities.ordered_set import OrderedSet 2 | 3 | def test_ordered_set(): 4 | # order matters 5 | assert OrderedSet([1,2,3]) != OrderedSet([3,2,1]) 6 | 7 | # first insertion determines order 8 | assert OrderedSet([5,1,6,5]) == OrderedSet([5,1,6]) 9 | 10 | # union preserves first and second order 11 | assert (OrderedSet([6,1,3]) | OrderedSet([3,1,7,9])) == OrderedSet([6,1,3,7,9]) 12 | 13 | # intersection preserves first order 14 | assert (OrderedSet([6,1,3]) & OrderedSet([3,1,7])) == OrderedSet([1,3]) 15 | 16 | # difference works 17 | assert (OrderedSet([6,1,3,2]) - OrderedSet([3,1,7])) == OrderedSet([6,2]) 18 | 19 | # symmetric difference: first then second 20 | assert (OrderedSet([6,1,3,8]) ^ OrderedSet([3,1,7,9])) == OrderedSet([6,8,7,9]) 21 | 22 | # in-place versions of these 23 | s = OrderedSet([6,1,3]) 24 | s2 = s 25 | s2 |= OrderedSet([3,1,7,9]) 26 | assert s == OrderedSet([6,1,3,7,9]) 27 | 28 | s = OrderedSet([6,1,3]) 29 | s2 = s 30 | s2 &= OrderedSet([3,1,7]) 31 | assert s == OrderedSet([1,3]) 32 | 33 | s = OrderedSet([6,1,3,2]) 34 | s2 = s 35 | s2 -= OrderedSet([3,1,7]) 36 | assert s == OrderedSet([6,2]) 37 | 38 | s = OrderedSet([6,1,3,8]) 39 | s2 = s 40 | s2 ^= OrderedSet([3,1,7,9]) 41 | assert s == OrderedSet([6,8,7,9]) 42 | 43 | # comparisons (order not used for these) 44 | assert OrderedSet([4,3,2,1]) <= OrderedSet([1,2,3,4]) 45 | assert not (OrderedSet([4,3,2,1]) < OrderedSet([1,2,3,4])) 46 | assert OrderedSet([3,2,1]) < OrderedSet([1,2,3,4]) 47 | 48 | # allow second argument (but ONLY second argument) to be any iterable, not just ordered set 49 | # we use the order from the iterable... 50 | assert (OrderedSet([6,1,3]) | [3,1,7,9]) == OrderedSet([6,1,3,7,9]) 51 | assert (OrderedSet([6,1,3]) & [3,1,7]) == OrderedSet([1,3]) 52 | assert (OrderedSet([6,1,3,2]) - [3,1,7]) == OrderedSet([6,2]) 53 | assert (OrderedSet([6,1,3,8]) ^ [3,1,7,9]) == OrderedSet([6,8,7,9]) 54 | 55 | 56 | def test_ordered_set_dict_from(): 57 | assert OrderedSet(['a','b','c']).dict_from([1, 2, 3]) == {'a': 1, 'b': 2, 'c': 3} -------------------------------------------------------------------------------- /tests/robustness/test_steady_state.py: -------------------------------------------------------------------------------- 1 | """Tests for steady_state with worse initial guesses, making use of the constrained solution functionality""" 2 | 3 | import pytest 4 | import numpy as np 5 | 6 | 7 | # Filter out warnings when the solver is trying to search in bad regions 8 | @pytest.mark.filterwarnings("ignore:.*invalid value encountered in.*:RuntimeWarning") 9 | def test_hank_steady_state_w_bad_init_guesses_and_bounds(one_asset_hank_dag): 10 | dag_ss, ss, dag, *_ = one_asset_hank_dag 11 | 12 | calibration = {"r": 0.005, "rstar": 0.005, "eis": 0.5, "frisch": 0.5, "B": 5.6, "mu": 1.2, 13 | "rho_s": 0.966, "sigma_s": 0.5, "kappa": 0.1, "phi": 1.5, "Y": 1, "Z": 1, "L": 1, 14 | "pi": 0, "nS": 2, "amax": 150, "nA": 10} 15 | unknowns_ss = {"beta": (0.95, 0.97, 0.999 / (1 + 0.005)), "vphi": (0.001, 1.0, 10.)} 16 | targets_ss = {"asset_mkt": 0, "labor_mkt": 0} 17 | cali = dag_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, solver="hybr", 18 | constrained_kwargs={"boundary_epsilon": 5e-3, "penalty_scale": 100}) 19 | ss_ref = dag.steady_state(cali) 20 | 21 | for k in ss.keys(): 22 | assert np.all(np.isclose(ss[k], ss_ref[k])) 23 | 24 | 25 | @pytest.mark.filterwarnings("ignore:.*invalid value encountered in.*:RuntimeWarning") 26 | def test_two_asset_steady_state_w_bad_init_guesses_and_bounds(two_asset_hank_dag): 27 | dag_ss, ss, dag, *_ = two_asset_hank_dag 28 | 29 | # Steady State 30 | calibration = {"Y": 1., "r": 0.0125, "rstar": 0.0125, "tot_wealth": 14, "delta": 0.02, 31 | "kappap": 0.1, "muw": 1.1, 'N': 1.0, 'K': 10., 'pi': 0.0, 32 | "Bh": 1.04, "Bg": 2.8, "G": 0.2, "eis": 0.5, "frisch": 1, "chi0": 0.25, "chi2": 2, 33 | "epsI": 4, "omega": 0.005, "kappaw": 0.1, "phi": 1.5, "nZ": 3, "nB": 10, "nA": 16, 34 | "nK": 4, "bmax": 50, "amax": 4000, "kmax": 1, "rho_z": 0.966, "sigma_z": 0.92} 35 | unknowns_ss = {"beta": 0.976, "chi1": 6.5} 36 | targets_ss = {"asset_mkt": 0., "B": "Bh"} 37 | cali = dag_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, 38 | solver="broyden_custom") 39 | ss_ref = dag.steady_state(cali) 40 | for k in ss.keys(): 41 | assert np.all(np.isclose(ss[k], ss_ref[k])) 42 | -------------------------------------------------------------------------------- /src/sequence_jacobian/__init__.py: -------------------------------------------------------------------------------- 1 | """Public-facing objects.""" 2 | 3 | from . import estimation, utilities, grids, interpolate, misc, hetblocks 4 | 5 | from .blocks.simple_block import simple 6 | from .blocks.het_block import het 7 | from .blocks.solved_block import solved 8 | from .blocks.combined_block import combine, create_model 9 | from .blocks.support.simple_displacement import apply_function 10 | from .classes.steady_state_dict import SteadyStateDict 11 | from .classes.impulse_dict import ImpulseDict 12 | from .classes.jacobian_dict import JacobianDict 13 | from .utilities.drawdag import drawdag 14 | 15 | # Ensure warning uniformity across package 16 | import warnings 17 | 18 | # Force warnings.warn() to omit the source code line in the message 19 | formatwarning_orig = warnings.formatwarning 20 | warnings.formatwarning = lambda message, category, filename, lineno, line=None: \ 21 | formatwarning_orig(message, category, filename, lineno, line='') 22 | 23 | # deprecation of old ways for calling things 24 | def agrid(*args, **kwargs): 25 | warnings.warn("The function 'agrid' is deprecated and will be removed in a subsequent version.\n" 26 | "Please call sj.grids.asset_grid(amin, amax, n) instead.") 27 | return utilities.discretize.agrid(*args, **kwargs) 28 | 29 | def markov_rouwenhorst(*args, **kwargs): 30 | warnings.warn("Calling sj.markov_rouwenhorst() is deprecated and will be disallowed in a subsequent version.\n" 31 | "Please call sj.grids.markov_rouwenhorst() instead.") 32 | return grids.markov_rouwenhorst(*args, **kwargs) 33 | 34 | def markov_tauchen(*args, **kwargs): 35 | warnings.warn("Calling sj.markov_tauchen() is deprecated and will be disallowed in a subsequent version.\n" 36 | "Please call sj.grids.markov_tauchen() instead.") 37 | return grids.markov_tauchen(*args, **kwargs) 38 | 39 | def interpolate_y(*args, **kwargs): 40 | warnings.warn("Calling sj.interpolate_y() is deprecated and will be disallowed in a subsequent version.\n" 41 | "Please call sj.interpolate.interpolate_y() instead.") 42 | return interpolate.interpolate_y(*args, **kwargs) 43 | 44 | def setmin(*args, **kwargs): 45 | warnings.warn("Calling sj.setmin() is deprecated and will be disallowed in a subsequent version.\n" 46 | "Please call sj.misc.setmin() instead.") 47 | misc.setmin(*args, **kwargs) 48 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/differentiate.py: -------------------------------------------------------------------------------- 1 | """Numerical differentiation""" 2 | 3 | from .misc import make_tuple 4 | 5 | 6 | def numerical_diff(func, ssinputs_dict, shock_dict, h=1E-4, y_ss_list=None): 7 | """Differentiate function numerically via forward difference, i.e. calculate 8 | 9 | f'(xss)*shock = (f(xss + h*shock) - f(xss))/h 10 | 11 | for small h. (Variable names inspired by application of differentiating around ss.) 12 | 13 | Parameters 14 | ---------- 15 | func : function, 'f' to be differentiated 16 | ssinputs_dict : dict, values in 'xss' around which to differentiate 17 | shock_dict : dict, values in 'shock' for which we're taking derivative 18 | (keys in shock_dict are weak subset of keys in ssinputs_dict) 19 | h : [optional] scalar, scaling of forward difference 'h' 20 | y_ss_list : [optional] list, value of y=f(xss) if we already have it 21 | 22 | Returns 23 | ---------- 24 | dy_list : list, output f'(xss)*shock of numerical differentiation 25 | """ 26 | # compute ss output if not supplied 27 | if y_ss_list is None: 28 | y_ss_list = make_tuple(func(**ssinputs_dict)) 29 | 30 | # response to small shock 31 | shocked_inputs = {**ssinputs_dict, **{k: ssinputs_dict[k] + h * shock for k, shock in shock_dict.items()}} 32 | y_list = make_tuple(func(**shocked_inputs)) 33 | 34 | # scale responses back up, dividing by h 35 | dy_list = [(y - y_ss) / h for y, y_ss in zip(y_list, y_ss_list)] 36 | 37 | return dy_list 38 | 39 | 40 | def numerical_diff_symmetric(func, ssinputs_dict, shock_dict, h=1E-4): 41 | """Same as numerical_diff, but differentiate numerically using central (symmetric) difference, i.e. 42 | 43 | f'(xss)*shock = (f(xss + h*shock) - f(xss - h*shock))/(2*h) 44 | """ 45 | 46 | # response to small shock in each direction 47 | shocked_inputs_up = {**ssinputs_dict, **{k: ssinputs_dict[k] + h * shock for k, shock in shock_dict.items()}} 48 | y_up_list = make_tuple(func(**shocked_inputs_up)) 49 | 50 | shocked_inputs_down = {**ssinputs_dict, **{k: ssinputs_dict[k] - h * shock for k, shock in shock_dict.items()}} 51 | y_down_list = make_tuple(func(**shocked_inputs_down)) 52 | 53 | # scale responses back up, dividing by h 54 | dy_list = [(y_up - y_down) / (2*h) for y_up, y_down in zip(y_up_list, y_down_list)] 55 | 56 | return dy_list 57 | -------------------------------------------------------------------------------- /tests/utils/test_function.py: -------------------------------------------------------------------------------- 1 | from sequence_jacobian.utilities.ordered_set import OrderedSet 2 | from sequence_jacobian.utilities.function import (DifferentiableExtendedFunction, ExtendedFunction, 3 | CombinedExtendedFunction, metadata) 4 | import numpy as np 5 | 6 | def f1(a, b, c): 7 | k = a + 1 8 | l = b - c 9 | return k, l 10 | 11 | def f2(b): 12 | k = b + 4 13 | return k 14 | 15 | 16 | def test_metadata(): 17 | assert metadata(f1) == ('f1', OrderedSet(['a', 'b', 'c']), OrderedSet(['k', 'l'])) 18 | assert metadata(f2) == ('f2', OrderedSet(['b']), OrderedSet(['k'])) 19 | 20 | 21 | def test_extended_function(): 22 | inputs = {'a': 1, 'b': 2, 'c': 3} 23 | assert ExtendedFunction(f1)(inputs) == {'k': 2, 'l': -1} 24 | assert ExtendedFunction(f2)(inputs) == {'k': 6} 25 | 26 | 27 | def f3(a, b): 28 | c = a*b - 5*a 29 | d = 3*b**2 30 | return c, d 31 | 32 | 33 | def test_differentiable_extended_function(): 34 | extf3 = ExtendedFunction(f3) 35 | 36 | ss1 = {'a': 1, 'b': 2} 37 | inputs1 = {'a': 0.5} 38 | 39 | diff = extf3.differentiable(ss1).diff(inputs1) 40 | assert np.isclose(diff['c'], -1.5) 41 | assert np.isclose(diff['d'], 0) 42 | 43 | 44 | def f4(a, c, e): 45 | f = a / c + a * e - c 46 | return f 47 | 48 | 49 | def test_differentiable_combined_extended_function(): 50 | # swapping in combined extended function to see if it works! 51 | fs = CombinedExtendedFunction([f3, f4]) 52 | 53 | ss1 = {'a': 1, 'b': 2, 'e': 4} 54 | ss1.update(fs(ss1)) 55 | 56 | inputs1 = {'a': 0.5, 'e': 1} 57 | 58 | diff = fs.differentiable(ss1).diff(inputs1) 59 | assert np.isclose(diff['c'], -1.5) 60 | assert np.isclose(diff['d'], 0) 61 | assert np.isclose(diff['f'], 4.5) 62 | 63 | # test narrowing down outputs 64 | diff = fs.differentiable(ss1).diff(inputs1, outputs=['c','d']) 65 | assert np.isclose(diff['c'], -1.5) 66 | assert np.isclose(diff['d'], 0) 67 | assert list(diff) == ['c', 'd'] 68 | 69 | # if no shocks to first function, hide first function 70 | inputs2 = {'e': -2} 71 | diff = fs.differentiable(ss1).diff2(inputs2) 72 | assert list(diff) == ['f'] 73 | assert np.isclose(diff['f'], -2) 74 | 75 | # if we ask for output from first function but no inputs shocked, shouldn't be there! 76 | diff = fs.differentiable(ss1).diff(inputs2, outputs=['c', 'f']) 77 | assert list(diff) == ['f'] 78 | assert np.isclose(diff['f'], -2) 79 | 80 | -------------------------------------------------------------------------------- /tests/base/test_simple_block.py: -------------------------------------------------------------------------------- 1 | """Test SimpleBlock functionality""" 2 | import copy 3 | 4 | import numpy as np 5 | import pytest 6 | 7 | from sequence_jacobian import simple 8 | from sequence_jacobian.classes.steady_state_dict import SteadyStateDict 9 | 10 | 11 | @simple 12 | def F(K, L, Z, alpha): 13 | Y = Z * K(-1)**alpha * L**(1-alpha) 14 | FK = alpha * Y / K 15 | FL = (1-alpha) * Y / L 16 | return Y, FK, FL 17 | 18 | 19 | @simple 20 | def investment(Q, K, r, N, mc, Z, delta, epsI, alpha): 21 | inv = (K/K(-1) - 1) / (delta * epsI) + 1 - Q 22 | val = alpha * Z(+1) * (N(+1) / K) ** (1-alpha) * mc(+1) - (K(+1)/K - 23 | (1-delta) + (K(+1)/K - 1)**2 / (2*delta*epsI)) + K(+1)/K*Q(+1) - (1 + r(+1))*Q 24 | return inv, val 25 | 26 | 27 | @simple 28 | def taylor(r, pi, phi): 29 | i = r.ss + phi * (pi - pi.ss) 30 | return i 31 | 32 | 33 | @pytest.mark.parametrize("block,ss", [(F, SteadyStateDict({"K": 1, "L": 1, "Z": 1, "alpha": 0.5})), 34 | (investment, SteadyStateDict({"Q": 1, "K": 1, "r": 0.05, "N": 1, "mc": 1, 35 | "Z": 1, "delta": 0.05, "epsI": 2, "alpha": 0.5})), 36 | (taylor, SteadyStateDict({"r": 0.05, "pi": 0.01, "phi": 1.5}))]) 37 | def test_block_consistency(block, ss): 38 | """Make sure ss, td, and jac methods are all consistent with each other. 39 | Requires that all inputs of simple block allow calculating Jacobians""" 40 | # get ss output 41 | ss_results = block.steady_state(ss) 42 | 43 | # now if we put in constant inputs, td should give us the same! 44 | td_results = block.impulse_nonlinear(ss_results, {k: np.zeros(20) for k in ss.keys()}) 45 | for v in td_results.values(): 46 | assert np.all(v == 0) 47 | 48 | # now get the Jacobian 49 | J = block.jacobian(ss, inputs=block.inputs) 50 | 51 | # now perturb the steady state by small random vectors 52 | # and verify that the second-order numerical derivative implied by .td 53 | # is equivalent to what we get from jac 54 | 55 | h = 1E-5 56 | all_shocks = {i: np.random.rand(10) for i in block.inputs} 57 | td_up = block.impulse_nonlinear(ss_results, {i: h*shock for i, shock in all_shocks.items()}) 58 | td_dn = block.impulse_nonlinear(ss_results, {i: -h*shock for i, shock in all_shocks.items()}) 59 | 60 | linear_impulses = {o: (td_up[o] - td_dn[o])/(2*h) for o in block.outputs} 61 | linear_impulses_from_jac = {o: sum(J[o][i] @ all_shocks[i] for i in all_shocks if i in J[o]) for o in block.outputs} 62 | 63 | for o in linear_impulses: 64 | assert np.all(np.abs(linear_impulses[o] - linear_impulses_from_jac[o]) < 1E-5) 65 | -------------------------------------------------------------------------------- /tests/base/test_options.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pytest 3 | from sequence_jacobian.examples import krusell_smith 4 | 5 | def test_jacobian_h(krusell_smith_dag): 6 | _, ss, dag, *_ = krusell_smith_dag 7 | hh = dag['hh'] 8 | 9 | lowacc = hh.jacobian(ss, inputs=['r'], outputs=['C'], T=10, h=0.05) 10 | midacc = hh.jacobian(ss, inputs=['r'], outputs=['C'], T=10, h=1E-3) 11 | usual = hh.jacobian(ss, inputs=['r'], outputs=['C'], T=10, h=1E-4) 12 | nooption = hh.jacobian(ss, inputs=['r'], outputs=['C'], T=10) 13 | 14 | assert np.array_equal(usual['C','r'], nooption['C','r']) 15 | assert np.linalg.norm(usual['C','r'] - midacc['C','r']) < np.linalg.norm(usual['C','r'] - lowacc['C','r']) 16 | 17 | midacc_alt = hh.jacobian(ss, inputs=['r'], outputs=['C'], T=10, options={'hh': {'h': 1E-3}}) 18 | assert np.array_equal(midacc['C', 'r'], midacc_alt['C', 'r']) 19 | 20 | lowacc = dag.jacobian(ss, inputs=['K'], outputs=['C'], T=10, options={'hh': {'h': 0.05}}) 21 | midacc = dag.jacobian(ss, inputs=['K'], outputs=['C'], T=10, options={'hh': {'h': 1E-3}}) 22 | usual = dag.jacobian(ss, inputs=['K'], outputs=['C'], T=10, options={'hh': {'h': 1E-4}}) 23 | 24 | assert np.linalg.norm(usual['C','K'] - midacc['C','K']) < np.linalg.norm(usual['C','K'] - lowacc['C','K']) 25 | 26 | 27 | def test_jacobian_steady_state(krusell_smith_dag): 28 | dag = krusell_smith_dag[2] 29 | calibration = {"eis": 1, "delta": 0.025, "alpha": 0.11, "rho": 0.966, "sigma": 0.5, 30 | "L": 1.0, "nS": 2, "nA": 10, "amax": 200, "r": 0.01, 'beta': 0.96, 31 | "Z": 0.85, "K": 3.} 32 | 33 | pytest.raises(ValueError, dag.steady_state, calibration, options={'hh': {'backward_maxit': 10}}) 34 | 35 | ss1 = dag.steady_state(calibration) 36 | ss2 = dag.steady_state(calibration, options={'hh': {'backward_maxit': 100000}}) 37 | assert ss1['A'] == ss2['A'] 38 | 39 | 40 | def test_steady_state_solution(krusell_smith_dag): 41 | dag_ss, ss, *_ = krusell_smith_dag 42 | 43 | calibration = {'eis': 1.0, 'delta': 0.025, 'alpha': 0.11, 'rho': 0.966, 'sigma': 0.5, 44 | 'Y': 1.0, 'L': 1.0, 'nS': 2, 'nA': 10, 'amax': 200, 'r': 0.01} 45 | unknowns_ss = {'beta': (0.98 / 1.01, 0.999 / 1.01)} 46 | targets_ss = {'asset_mkt': 0.} 47 | 48 | 49 | # less accurate solution 50 | ss2 = dag_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, solver="brentq", 51 | ttol=1E-2, ctol=1E-2) 52 | 53 | assert not np.isclose(ss['asset_mkt'], ss2['asset_mkt']) 54 | 55 | # different solution method (Newton needs other inputs) 56 | with pytest.raises(ValueError): 57 | ss3 = dag_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, 58 | solver="newton") 59 | -------------------------------------------------------------------------------- /src/sequence_jacobian/classes/result_dict.py: -------------------------------------------------------------------------------- 1 | import copy 2 | 3 | from ..utilities.bijection import Bijection 4 | 5 | class ResultDict: 6 | def __init__(self, data, internals=None): 7 | if isinstance(data, ResultDict): 8 | if internals is not None: 9 | raise ValueError(f'Supplying {type(self).__name__} and also internals to constructor not allowed') 10 | self.toplevel = data.toplevel.copy() 11 | self.internals = data.internals.copy() 12 | else: 13 | self.toplevel: dict = data.copy() 14 | self.internals: dict = {} if internals is None else internals.copy() 15 | 16 | def __repr__(self): 17 | if self.internals: 18 | return f"<{type(self).__name__}: {list(self.toplevel.keys())}, internals={list(self.internals.keys())}>" 19 | else: 20 | return f"<{type(self).__name__}: {list(self.toplevel.keys())}>" 21 | 22 | def __iter__(self): 23 | return iter(self.toplevel) 24 | 25 | def __getitem__(self, k, **kwargs): 26 | if isinstance(k, str): 27 | return self.toplevel[k] 28 | elif isinstance(k, tuple): 29 | raise TypeError(f'Key {k} to {type(self).__name__} cannot be tuple') 30 | else: 31 | try: 32 | return type(self)({ki: self.toplevel[ki] for ki in k}, **kwargs) 33 | except TypeError: 34 | raise TypeError(f'Key {k} to {type(self).__name__} needs to be a string or an iterable (list, set, etc) of strings') 35 | 36 | def __setitem__(self, k, v): 37 | self.toplevel[k] = v 38 | 39 | def __matmul__(self, x): 40 | # remap keys in toplevel 41 | if isinstance(x, Bijection): 42 | new = copy.deepcopy(self) 43 | new.toplevel = x @ self.toplevel 44 | return new 45 | else: 46 | return NotImplemented 47 | 48 | def __rmatmul__(self, x): 49 | return self.__matmul__(x) 50 | 51 | def __len__(self): 52 | return len(self.toplevel) 53 | 54 | def __or__(self, other): 55 | if not isinstance(other, type(self)): 56 | raise ValueError(f'Trying to merge a {type(self).__name__} with a {type(other).__name__}.') 57 | merged = self.copy() 58 | merged.update(other) 59 | return merged 60 | 61 | def keys(self): 62 | return self.toplevel.keys() 63 | 64 | def values(self): 65 | return self.toplevel.values() 66 | 67 | def items(self): 68 | return self.toplevel.items() 69 | 70 | def update(self, rdict): 71 | if isinstance(rdict, ResultDict): 72 | self.toplevel.update(rdict.toplevel) 73 | self.internals.update(rdict.internals) 74 | else: 75 | self.toplevel.update(dict(rdict)) 76 | 77 | def copy(self): 78 | return type(self)(self) 79 | -------------------------------------------------------------------------------- /src/sequence_jacobian/hetblocks/hh_labor.py: -------------------------------------------------------------------------------- 1 | '''Standard Incomplete Market model with Endogenous Labor Supply''' 2 | 3 | import numpy as np 4 | from numba import vectorize, njit 5 | 6 | from ..blocks.het_block import het 7 | from .. import interpolate 8 | 9 | 10 | def hh_init(a_grid, we, r, eis, T): 11 | fininc = (1 + r) * a_grid + T[:, np.newaxis] - a_grid[0] 12 | coh = (1 + r) * a_grid[np.newaxis, :] + we[:, np.newaxis] + T[:, np.newaxis] 13 | Va = (1 + r) * (0.1 * coh) ** (-1 / eis) 14 | return fininc, Va 15 | 16 | 17 | @het(exogenous='Pi', policy='a', backward='Va', backward_init=hh_init) 18 | def hh(Va_p, a_grid, we, T, r, beta, eis, frisch, vphi): 19 | '''Single backward step via EGM.''' 20 | uc_nextgrid = beta * Va_p 21 | c_nextgrid, n_nextgrid = cn(uc_nextgrid, we[:, np.newaxis], eis, frisch, vphi) 22 | 23 | lhs = c_nextgrid - we[:, np.newaxis] * n_nextgrid + a_grid[np.newaxis, :] - T[:, np.newaxis] 24 | rhs = (1 + r) * a_grid 25 | c = interpolate.interpolate_y(lhs, rhs, c_nextgrid) 26 | n = interpolate.interpolate_y(lhs, rhs, n_nextgrid) 27 | 28 | a = rhs + we[:, np.newaxis] * n + T[:, np.newaxis] - c 29 | iconst = np.nonzero(a < a_grid[0]) 30 | a[iconst] = a_grid[0] 31 | 32 | if iconst[0].size != 0 and iconst[1].size != 0: 33 | c[iconst], n[iconst] = solve_cn(we[iconst[0]], 34 | rhs[iconst[1]] + T[iconst[0]] - a_grid[0], 35 | eis, frisch, vphi, Va_p[iconst]) 36 | 37 | Va = (1 + r) * c ** (-1 / eis) 38 | 39 | return Va, a, c, n 40 | 41 | 42 | '''Supporting functions for HA block''' 43 | 44 | @njit 45 | def cn(uc, w, eis, frisch, vphi): 46 | """Return optimal c, n as function of u'(c) given parameters""" 47 | return uc ** (-eis), (w * uc / vphi) ** frisch 48 | 49 | 50 | def solve_cn(w, T, eis, frisch, vphi, uc_seed): 51 | uc = solve_uc(w, T, eis, frisch, vphi, uc_seed) 52 | return cn(uc, w, eis, frisch, vphi) 53 | 54 | 55 | @vectorize 56 | def solve_uc(w, T, eis, frisch, vphi, uc_seed): 57 | """Solve for optimal uc given in log uc space. 58 | 59 | max_{c, n} c**(1-1/eis) + vphi*n**(1+1/frisch) s.t. c = w*n + T 60 | """ 61 | log_uc = np.log(uc_seed) 62 | for i in range(30): 63 | ne, ne_p = netexp(log_uc, w, T, eis, frisch, vphi) 64 | if abs(ne) < 1E-11: 65 | break 66 | else: 67 | log_uc -= ne / ne_p 68 | else: 69 | raise ValueError("Cannot solve constrained household's problem: No convergence after 30 iterations!") 70 | 71 | return np.exp(log_uc) 72 | 73 | 74 | @njit 75 | def netexp(log_uc, w, T, eis, frisch, vphi): 76 | """Return net expenditure as a function of log uc and its derivative.""" 77 | c, n = cn(np.exp(log_uc), w, eis, frisch, vphi) 78 | ne = c - w * n - T 79 | 80 | # c and n have elasticities of -eis and frisch wrt log u'(c) 81 | c_loguc = -eis * c 82 | n_loguc = frisch * n 83 | netexp_loguc = c_loguc - w * n_loguc 84 | 85 | return ne, netexp_loguc 86 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/support/parent.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | 3 | class Parent: 4 | # see tests in test_parent_block.py 5 | 6 | def __init__(self, blocks, name=None): 7 | # dict from names to immediate kid blocks themselves 8 | # dict from descendants to the names of kid blocks through which to access them 9 | # "descendants" of a block include itself 10 | if not hasattr(self, 'name') and name is not None: 11 | self.name = name 12 | 13 | kids = {} 14 | descendants = {} 15 | 16 | for block in blocks: 17 | kids[block.name] = block 18 | 19 | if isinstance(block, Parent): 20 | for k in block.descendants: 21 | if k in descendants: 22 | raise ValueError(f'Overlapping block name {k}') 23 | descendants[k] = block.name 24 | else: 25 | descendants[block.name] = block.name 26 | 27 | # add yourself to descendants too! but you don't belong to any kid... 28 | if self.name in descendants: 29 | raise ValueError(f'Overlapping block name {self.name}') 30 | descendants[self.name] = None 31 | 32 | self.kids = kids 33 | self.descendants = descendants 34 | 35 | def __getitem__(self, k): 36 | if k == self.name: 37 | return self 38 | elif k in self.kids: 39 | return self.kids[k] 40 | else: 41 | return self.kids[self.descendants[k]][k] 42 | 43 | def select(self, d, kid): 44 | """If d is a dict with block names as keys and kid is a kid, select only the entries in d that are descendants of kid""" 45 | return {k: v for k, v in d.items() if k in self.kids[kid].descendants} 46 | 47 | def path(self, k, reverse=True): 48 | if k not in self.descendants: 49 | raise KeyError(f'Cannot get path to {k} because it is not a descendant of current block') 50 | 51 | if k != self.name: 52 | kid = self.kids[self.descendants[k]] 53 | if isinstance(kid, Parent): 54 | p = kid.path(k, reverse=False) 55 | else: 56 | p = [k] 57 | else: 58 | p = [] 59 | p.append(self.name) 60 | 61 | if reverse: 62 | return list(reversed(p)) 63 | else: 64 | return p 65 | 66 | def get_attribute(self, k, attr): 67 | """Gets attribute attr from descendant k, respecting any remapping 68 | along the way (requires that attr is list, dict, set)""" 69 | if k == self.name: 70 | inner = getattr(self, attr) 71 | else: 72 | kid = self.kids[self.descendants[k]] 73 | if isinstance(kid, Parent): 74 | inner = kid.get_attribute(k, attr) 75 | else: 76 | inner = getattr(kid, attr) 77 | if hasattr(kid, 'M'): 78 | inner = kid.M @ inner 79 | 80 | if hasattr(self, 'M'): 81 | return self.M @ inner 82 | else: 83 | return inner 84 | -------------------------------------------------------------------------------- /tests/base/test_two_asset.py: -------------------------------------------------------------------------------- 1 | """Test the two asset HANK steady state computation""" 2 | 3 | import numpy as np 4 | 5 | from sequence_jacobian.hetblocks import hh_twoasset as hh 6 | from sequence_jacobian import utilities as utils 7 | 8 | 9 | def test_hank_ss(): 10 | A, B, UCE = hank_ss_singlerun() 11 | assert np.isclose(A, 12.526539492650361) 12 | assert np.isclose(B, 1.0840860793350566) 13 | assert np.isclose(UCE, 4.5102870939550055) 14 | 15 | 16 | def hank_ss_singlerun(beta=0.976, r=0.0125, tot_wealth=14, K=10, delta=0.02, Bg=2.8, G=0.2, 17 | eis=0.5, chi0=0.25, chi1=6.5, chi2=2, omega=0.005, nZ=3, nB=50, 18 | nA=70, nK=50, bmax=50, amax=4000, kmax=1, rho_z=0.966, sigma_z=0.92): 19 | """Mostly cribbed from two_asset.hank_ss(), but just does backward iteration to get 20 | a partial equilibrium household steady state given parameters, not solving for equilibrium. 21 | Convenient for testing.""" 22 | 23 | # set up grid 24 | b_grid = utils.discretize.agrid(amax=bmax, n=nB) 25 | a_grid = utils.discretize.agrid(amax=amax, n=nA) 26 | k_grid = utils.discretize.agrid(amax=kmax, n=nK)[::-1].copy() 27 | e_grid, _, Pi = utils.discretize.markov_rouwenhorst(rho=rho_z, sigma=sigma_z, N=nZ) 28 | 29 | # solve analytically what we can 30 | mc = 1 - r * (tot_wealth - Bg - K) 31 | alpha = (r + delta) * K / mc 32 | w = (1 - alpha) * mc 33 | tax = (r * Bg + G) / w 34 | ra = r 35 | rb = r - omega 36 | z_grid = (1 - tax) * w * e_grid 37 | 38 | # figure out initializer 39 | calibration = {'Pi': Pi, 'a_grid': a_grid, 'b_grid': b_grid, 'e_grid': e_grid, 40 | 'z_grid': z_grid, 'k_grid': k_grid, 'beta': beta, 'N': 1.0, 41 | 'tax': tax, 'w': w, 'eis': eis, 'rb': rb, 'ra': ra, 42 | 'chi0': chi0, 'chi1': chi1, 'chi2': chi2} 43 | 44 | out = hh.hh.steady_state(calibration) 45 | 46 | return out['A'], out['B'], out['UCE'] 47 | 48 | 49 | def test_Psi(): 50 | np.random.seed(41234) 51 | chi0, chi1, chi2 = 0.25, 6.5, 2.3 52 | ra = 0.05 53 | 54 | a = np.random.rand(50) + 1 55 | ap = np.random.rand(50) + 1 56 | 57 | oPsi, oPsi1, oPsi2 = hh.get_Psi_and_deriv(ap, a, ra, chi0, chi1, chi2) 58 | 59 | Psi = Psi_correct(ap, a, ra, chi0, chi1, chi2) 60 | assert np.allclose(oPsi, Psi) 61 | 62 | # compare two-sided numerical derivative to our analytical one 63 | # numerical doesn't work well at kink of "abs" function, so this would fail 64 | # for some seeds if chi2 was less than 2 65 | Psi1 = (Psi_correct(ap+1E-4, a, ra, chi0, chi1, chi2) - 66 | Psi_correct(ap-1E-4, a, ra, chi0, chi1, chi2)) / 2E-4 67 | assert np.allclose(oPsi1, Psi1) 68 | 69 | Psi2 = (Psi_correct(ap, a+1E-4, ra, chi0, chi1, chi2) - 70 | Psi_correct(ap, a-1E-4, ra, chi0, chi1, chi2)) / 2E-4 71 | assert np.allclose(oPsi2, Psi2) 72 | 73 | 74 | def Psi_correct(ap, a, ra, chi0, chi1, chi2): 75 | """Original Psi function that we know is correct, once denominator has power 76 | chi2-1 rather than 1 (error in original code)""" 77 | return chi1 / chi2 * np.abs((ap - (1 + ra) * a)) ** chi2 / ((1 + ra) * a + chi0) ** (chi2 - 1) -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/bijection.py: -------------------------------------------------------------------------------- 1 | from .ordered_set import OrderedSet 2 | 3 | class Bijection: 4 | def __init__(self, map): 5 | # identity always implicit, remove if there explicitly 6 | self.map = {k: v for k, v in map.items() if k != v} 7 | invmap = {} 8 | for k, v in map.items(): 9 | if v in invmap: 10 | raise ValueError(f'Duplicate value {v}, for keys {invmap[v]} and {k}') 11 | invmap[v] = k 12 | self.invmap = invmap 13 | 14 | @property 15 | def inv(self): 16 | invmap = Bijection.__new__(Bijection) # better way to do this? 17 | invmap.map = self.invmap 18 | invmap.invmap = self.map 19 | return invmap 20 | 21 | def __repr__(self): 22 | return f'Bijection({repr(self.map)})' 23 | 24 | def __getitem__(self, k): 25 | return self.map.get(k, k) 26 | 27 | def __matmul__(self, x): 28 | if x is None: 29 | return None 30 | elif isinstance(x, str) or isinstance(x, int): 31 | return self[x] 32 | elif isinstance(x, Bijection): 33 | # compose self: v -> u with x: w -> v 34 | # assume everything missing in either is the identity 35 | M = {} 36 | for v, u in self.map.items(): 37 | w = x.invmap.get(v, v) 38 | M[w] = u 39 | for w, v in x.map.items(): 40 | if v not in self.map: 41 | M[w] = v 42 | return Bijection(M) 43 | elif isinstance(x, dict): 44 | d = {} 45 | for k, v in x.items(): 46 | if k in self.map: 47 | d[self.map[k]] = v 48 | elif k not in d: 49 | # don't overwrite if we already mapped to this 50 | # effectively this prioritizes the remapped names over others 51 | d[k] = v 52 | return d 53 | elif isinstance(x, list): 54 | return [self[k] for k in x] 55 | elif isinstance(x, set): 56 | return {self[k] for k in x} 57 | elif isinstance(x, tuple): 58 | return tuple(self[k] for k in x) 59 | elif isinstance(x, OrderedSet): 60 | return OrderedSet([self[k] for k in x]) 61 | else: 62 | return NotImplemented 63 | 64 | def __rmatmul__(self, x): 65 | if isinstance(x, str): 66 | return self[x] 67 | elif isinstance(x, dict): 68 | d = {} 69 | for k, v in x.items(): 70 | if k in self.map: 71 | d[self.map[k]] = v 72 | elif k not in d: 73 | # don't overwrite if we already mapped to this 74 | # effectively this prioritizes the remapped names over others 75 | d[k] = v 76 | return d 77 | elif isinstance(x, list): 78 | return [self[k] for k in x] 79 | elif isinstance(x, set): 80 | return {self[k] for k in x} 81 | elif isinstance(x, tuple): 82 | return tuple(self[k] for k in x) 83 | else: 84 | return NotImplemented 85 | 86 | def __bool__(self): 87 | return bool(self.map) -------------------------------------------------------------------------------- /src/sequence_jacobian/estimation.py: -------------------------------------------------------------------------------- 1 | """Functions for calculating the log likelihood of a model from its impulse responses""" 2 | 3 | import numpy as np 4 | import scipy.linalg as linalg 5 | from numba import njit 6 | 7 | '''Part 1: compute covariances at all lags and log likelihood''' 8 | 9 | 10 | def all_covariances(M, sigmas): 11 | """Use Fast Fourier Transform to compute covariance function between O vars up to T-1 lags. 12 | 13 | See equation (108) in appendix B.5 of paper for details. 14 | 15 | Parameters 16 | ---------- 17 | M : array (T*O*Z), stacked impulse responses of nO variables to nZ shocks (MA(T-1) representation) 18 | sigmas : array (Z), standard deviations of shocks 19 | 20 | Returns 21 | ---------- 22 | Sigma : array (T*O*O), covariance function between O variables for 0, ..., T-1 lags 23 | """ 24 | T = M.shape[0] 25 | dft = np.fft.rfftn(M, s=(2 * T - 2,), axes=(0,)) 26 | total = (dft.conjugate() * sigmas**2) @ dft.swapaxes(1, 2) 27 | return np.fft.irfftn(total, s=(2 * T - 2,), axes=(0,))[:T] 28 | 29 | 30 | def log_likelihood(Y, Sigma, sigma_measurement=None): 31 | """Given second moments, compute log-likelihood of data Y. 32 | 33 | Parameters 34 | ---------- 35 | Y : array (Tobs*O) 36 | stacked data for O observables over Tobs periods 37 | Sigma : array (T*O*O) 38 | covariance between observables in model for 0, ... , T lags (e.g. from all_covariances) 39 | sigma_measurement : [optional] array (O) 40 | std of measurement error for each observable, assumed zero if not provided 41 | 42 | Returns 43 | ---------- 44 | L : scalar, log-likelihood 45 | """ 46 | Tobs, nO = Y.shape 47 | if sigma_measurement is None: 48 | sigma_measurement = np.zeros(nO) 49 | V = build_full_covariance_matrix(Sigma, sigma_measurement, Tobs) 50 | y = Y.ravel() 51 | return log_likelihood_formula(y, V) 52 | 53 | 54 | '''Part 2: helper functions''' 55 | 56 | 57 | def log_likelihood_formula(y, V): 58 | """Implements multivariate normal log-likelihood formula using Cholesky with data vector y and variance V. 59 | Calculates -log det(V)/2 - y'V^(-1)y/2 60 | """ 61 | V_factored = linalg.cho_factor(V) 62 | quadratic_form = np.dot(y, linalg.cho_solve(V_factored, y)) 63 | log_determinant = 2*np.sum(np.log(np.diag(V_factored[0]))) 64 | return -(log_determinant + quadratic_form) / 2 65 | 66 | 67 | @njit 68 | def build_full_covariance_matrix(Sigma, sigma_measurement, Tobs): 69 | """Takes in T*O*O array Sigma with covariances at each lag t, 70 | assembles them into (Tobs*O)*(Tobs*O) matrix of covariances, including measurement errors. 71 | """ 72 | T, O, O = Sigma.shape 73 | V = np.empty((Tobs, O, Tobs, O)) 74 | for t1 in range(Tobs): 75 | for t2 in range(Tobs): 76 | if abs(t1-t2) >= T: 77 | V[t1, :, t2, :] = np.zeros((O, O)) 78 | else: 79 | if t1 < t2: 80 | V[t1, : , t2, :] = Sigma[t2-t1, :, :] 81 | elif t1 > t2: 82 | V[t1, : , t2, :] = Sigma[t1-t2, :, :].T 83 | else: 84 | # want exactly symmetric 85 | V[t1, :, t2, :] = (np.diag(sigma_measurement**2) + (Sigma[0, :, :]+Sigma[0, :, :].T)/2) 86 | return V.reshape((Tobs*O, Tobs*O)) 87 | -------------------------------------------------------------------------------- /src/sequence_jacobian/examples/hank.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from sequence_jacobian import grids, simple, create_model, hetblocks 4 | hh = hetblocks.hh_labor.hh 5 | 6 | 7 | '''Part 1: Blocks''' 8 | 9 | 10 | @simple 11 | def firm(Y, w, Z, pi, mu, kappa): 12 | L = Y / Z 13 | Div = Y - w * L - mu/(mu-1)/(2*kappa) * (1+pi).apply(np.log)**2 * Y 14 | return L, Div 15 | 16 | 17 | @simple 18 | def monetary(pi, rstar, phi): 19 | r = (1 + rstar(-1) + phi * pi(-1)) / (1 + pi) - 1 20 | return r 21 | 22 | 23 | @simple 24 | def nkpc(pi, w, Z, Y, r, mu, kappa): 25 | nkpc_res = kappa * (w / Z - 1 / mu) + Y(+1) / Y * (1 + pi(+1)).apply(np.log) / (1 + r(+1))\ 26 | - (1 + pi).apply(np.log) 27 | return nkpc_res 28 | 29 | 30 | @simple 31 | def fiscal(r, B): 32 | Tax = r * B 33 | return Tax 34 | 35 | 36 | @simple 37 | def mkt_clearing(A, NE, C, L, Y, B, pi, mu, kappa): 38 | asset_mkt = A - B 39 | labor_mkt = NE - L 40 | goods_mkt = Y - C - mu/(mu-1)/(2*kappa) * (1+pi).apply(np.log)**2 * Y 41 | return asset_mkt, labor_mkt, goods_mkt 42 | 43 | 44 | @simple 45 | def nkpc_ss(Z, mu): 46 | '''Solve (w) to hit targets for (nkpc_res)''' 47 | w = Z / mu 48 | return w 49 | 50 | 51 | '''Part 2: Embed HA block''' 52 | 53 | 54 | def make_grids(rho_s, sigma_s, nS, amax, nA): 55 | e_grid, pi_e, Pi = grids.markov_rouwenhorst(rho=rho_s, sigma=sigma_s, N=nS) 56 | a_grid = grids.agrid(amax=amax, n=nA) 57 | return e_grid, pi_e, Pi, a_grid 58 | 59 | 60 | def transfers(pi_e, Div, Tax, e_grid): 61 | # hardwired incidence rules are proportional to skill; scale does not matter 62 | tax_rule, div_rule = e_grid, e_grid 63 | div = Div / np.sum(pi_e * div_rule) * div_rule 64 | tax = Tax / np.sum(pi_e * tax_rule) * tax_rule 65 | T = div - tax 66 | return T 67 | 68 | 69 | def wages(w, e_grid): 70 | we = w * e_grid 71 | return we 72 | 73 | 74 | def labor_supply(n, e_grid): 75 | ne = e_grid[:, np.newaxis] * n 76 | return ne 77 | 78 | 79 | '''Part 3: DAG''' 80 | 81 | def dag(): 82 | # Combine blocks 83 | household = hh.add_hetinputs([transfers, wages, make_grids]) 84 | household = household.add_hetoutputs([labor_supply]) 85 | blocks = [household, firm, monetary, fiscal, mkt_clearing, nkpc] 86 | blocks_ss = [household, firm, monetary, fiscal, mkt_clearing, nkpc_ss] 87 | hank_model = create_model(blocks, name="One-Asset HANK") 88 | hank_model_ss = create_model(blocks_ss, name="One-Asset HANK SS") 89 | 90 | # Steady state 91 | calibration = {'r': 0.005, 'rstar': 0.005, 'eis': 0.5, 'frisch': 0.5, 'B': 5.6, 92 | 'mu': 1.2, 'rho_s': 0.966, 'sigma_s': 0.5, 'kappa': 0.1, 'phi': 1.5, 93 | 'Y': 1., 'Z': 1., 'pi': 0., 'nS': 2, 'amax': 150, 'nA': 10} 94 | unknowns_ss = {'beta': 0.986, 'vphi': 0.8} 95 | targets_ss = {'asset_mkt': 0., 'NE': 1.} 96 | cali = hank_model_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, 97 | solver='broyden_custom') 98 | ss = hank_model.steady_state(cali) 99 | 100 | # Transitional dynamics 101 | unknowns = ['w', 'Y', 'pi'] 102 | targets = ['asset_mkt', 'goods_mkt', 'nkpc_res'] 103 | exogenous = ['rstar', 'Z'] 104 | 105 | return hank_model_ss, ss, hank_model, unknowns, targets, exogenous 106 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/support/het_compiled.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numba import njit 3 | 4 | @njit 5 | def forward_policy_1d(D, x_i, x_pi): 6 | nZ, nX = D.shape 7 | Dnew = np.zeros_like(D) 8 | for iz in range(nZ): 9 | for ix in range(nX): 10 | i = x_i[iz, ix] 11 | pi = x_pi[iz, ix] 12 | d = D[iz, ix] 13 | 14 | Dnew[iz, i] += d * pi 15 | Dnew[iz, i+1] += d * (1 - pi) 16 | 17 | return Dnew 18 | 19 | 20 | @njit 21 | def expectation_policy_1d(X, x_i, x_pi): 22 | nZ, nX = X.shape 23 | Xnew = np.zeros_like(X) 24 | for iz in range(nZ): 25 | for ix in range(nX): 26 | i = x_i[iz, ix] 27 | pi = x_pi[iz, ix] 28 | Xnew[iz, ix] = pi * X[iz, i] + (1-pi) * X[iz, i+1] 29 | return Xnew 30 | 31 | 32 | @njit 33 | def forward_policy_shock_1d(Dss, x_i_ss, x_pi_shock): 34 | """forward_step_1d linearized wrt x_pi""" 35 | nZ, nX = Dss.shape 36 | Dshock = np.zeros_like(Dss) 37 | for iz in range(nZ): 38 | for ix in range(nX): 39 | i = x_i_ss[iz, ix] 40 | dshock = x_pi_shock[iz, ix] * Dss[iz, ix] 41 | Dshock[iz, i] += dshock 42 | Dshock[iz, i + 1] -= dshock 43 | 44 | return Dshock 45 | 46 | 47 | @njit 48 | def forward_policy_2d(D, x_i, y_i, x_pi, y_pi): 49 | nZ, nX, nY = D.shape 50 | Dnew = np.zeros_like(D) 51 | for iz in range(nZ): 52 | for ix in range(nX): 53 | for iy in range(nY): 54 | ixp = x_i[iz, ix, iy] 55 | iyp = y_i[iz, ix, iy] 56 | beta = x_pi[iz, ix, iy] 57 | alpha = y_pi[iz, ix, iy] 58 | 59 | Dnew[iz, ixp, iyp] += alpha * beta * D[iz, ix, iy] 60 | Dnew[iz, ixp+1, iyp] += alpha * (1 - beta) * D[iz, ix, iy] 61 | Dnew[iz, ixp, iyp+1] += (1 - alpha) * beta * D[iz, ix, iy] 62 | Dnew[iz, ixp+1, iyp+1] += (1 - alpha) * (1 - beta) * D[iz, ix, iy] 63 | return Dnew 64 | 65 | 66 | @njit 67 | def expectation_policy_2d(X, x_i, y_i, x_pi, y_pi): 68 | nZ, nX, nY = X.shape 69 | Xnew = np.empty_like(X) 70 | for iz in range(nZ): 71 | for ix in range(nX): 72 | for iy in range(nY): 73 | ixp = x_i[iz, ix, iy] 74 | iyp = y_i[iz, ix, iy] 75 | alpha = x_pi[iz, ix, iy] 76 | beta = y_pi[iz, ix, iy] 77 | 78 | Xnew[iz, ix, iy] = (alpha * beta * X[iz, ixp, iyp] + alpha * (1-beta) * X[iz, ixp, iyp+1] + 79 | (1-alpha) * beta * X[iz, ixp+1, iyp] + 80 | (1-alpha) * (1-beta) * X[iz, ixp+1, iyp+1]) 81 | return Xnew 82 | 83 | 84 | @njit 85 | def forward_policy_shock_2d(Dss, x_i_ss, y_i_ss, x_pi_ss, y_pi_ss, x_pi_shock, y_pi_shock): 86 | """Endogenous update part of forward_step_shock_2d""" 87 | nZ, nX, nY = Dss.shape 88 | Dshock = np.zeros_like(Dss) 89 | for iz in range(nZ): 90 | for ix in range(nX): 91 | for iy in range(nY): 92 | ixp = x_i_ss[iz, ix, iy] 93 | iyp = y_i_ss[iz, ix, iy] 94 | alpha = x_pi_ss[iz, ix, iy] 95 | beta = y_pi_ss[iz, ix, iy] 96 | 97 | dalpha = x_pi_shock[iz, ix, iy] * Dss[iz, ix, iy] 98 | dbeta = y_pi_shock[iz, ix, iy] * Dss[iz, ix, iy] 99 | 100 | Dshock[iz, ixp, iyp] += dalpha * beta + alpha * dbeta 101 | Dshock[iz, ixp+1, iyp] += dbeta * (1-alpha) - beta * dalpha 102 | Dshock[iz, ixp, iyp+1] += dalpha * (1-beta) - alpha * dbeta 103 | Dshock[iz, ixp+1, iyp+1] -= dalpha * (1-beta) + dbeta * (1-alpha) 104 | return Dshock 105 | -------------------------------------------------------------------------------- /tests/base/test_jacobian.py: -------------------------------------------------------------------------------- 1 | """Test all models' Jacobian calculations""" 2 | 3 | import numpy as np 4 | 5 | def test_ks_jac(krusell_smith_dag): 6 | _, ss, ks_model, unknowns, targets, exogenous = krusell_smith_dag 7 | household, firm = ks_model['hh'], ks_model['firm'] 8 | T = 10 9 | 10 | # Automatically calculate the general equilibrium Jacobian 11 | G2 = ks_model.solve_jacobian(ss, unknowns, targets, exogenous, T=T) 12 | 13 | # Manually calculate the general equilibrium Jacobian 14 | J_firm = firm.jacobian(ss, inputs=['K', 'Z']) 15 | J_ha = household.jacobian(ss, T=T, inputs=['r', 'w']) 16 | J_curlyK_K = J_ha['A']['r'] @ J_firm['r']['K'] + J_ha['A']['w'] @ J_firm['w']['K'] 17 | J_curlyK_Z = J_ha['A']['r'] @ J_firm['r']['Z'] + J_ha['A']['w'] @ J_firm['w']['Z'] 18 | J_curlyK = {'curlyK': {'K': J_curlyK_K, 'Z': J_curlyK_Z}} 19 | 20 | H_K = J_curlyK['curlyK']['K'] - np.eye(T) 21 | H_Z = J_curlyK['curlyK']['Z'] 22 | 23 | G = {'K': -np.linalg.solve(H_K, H_Z)} # H_K^(-1)H_Z 24 | G['r'] = J_firm['r']['Z'] + J_firm['r']['K'] @ G['K'] 25 | G['w'] = J_firm['w']['Z'] + J_firm['w']['K'] @ G['K'] 26 | G['Y'] = J_firm['Y']['Z'] + J_firm['Y']['K'] @ G['K'] 27 | G['C'] = J_ha['C']['r'] @ G['r'] + J_ha['C']['w'] @ G['w'] 28 | 29 | for o in G: 30 | assert np.allclose(G2[o]['Z'], G[o]) 31 | 32 | 33 | # TODO: decide whether to get rid of this or revise it with manual solve_jacobian stuff 34 | # def test_hank_jac(one_asset_hank_dag): 35 | # hank_model, exogenous, unknowns, targets, ss = one_asset_hank_dag 36 | # T = 10 37 | 38 | # # Automatically calculate the general equilibrium Jacobian 39 | # G2 = hank_model.solve_jacobian(ss, unknowns, targets, exogenous, T=T) 40 | 41 | # # Manually calculate the general equilibrium Jacobian 42 | # curlyJs, required = curlyJ_sorted(hank_model.blocks, unknowns + exogenous, ss, T) 43 | # J_curlyH_U = forward_accumulate(curlyJs, unknowns, targets, required) 44 | # J_curlyH_Z = forward_accumulate(curlyJs, exogenous, targets, required) 45 | # H_U = J_curlyH_U[targets, unknowns].pack(T) 46 | # H_Z = J_curlyH_Z[targets, exogenous].pack(T) 47 | # G_U = JacobianDict.unpack(-np.linalg.solve(H_U, H_Z), unknowns, exogenous, T) 48 | # curlyJs = [G_U] + curlyJs 49 | # outputs = set().union(*(curlyJ.outputs for curlyJ in curlyJs)) - set(targets) 50 | # G = forward_accumulate(curlyJs, exogenous, outputs, required | set(unknowns)) 51 | 52 | # for o in G: 53 | # for i in G[o]: 54 | # assert np.allclose(G[o][i], G2[o][i]) 55 | 56 | 57 | 58 | def test_fake_news_v_direct_method(one_asset_hank_dag): 59 | hank_model, ss, *_ = one_asset_hank_dag 60 | 61 | household = hank_model['hh'] 62 | T = 40 63 | exogenous = ['r'] 64 | output_list = household.non_backward_outputs 65 | h = 1E-4 66 | 67 | Js = household.jacobian(ss, exogenous, T=T) 68 | Js_direct = {o.upper(): {i: np.empty((T, T)) for i in exogenous} for o in output_list} 69 | 70 | # run td once without any shocks to get paths to subtract against 71 | # (better than subtracting by ss since ss not exact) 72 | # monotonic=True lets us know there is monotonicity of policy rule, makes TD run faster 73 | # .impulse_nonlinear requires at least one input 'shock', so we put in steady-state w 74 | td_noshock = household.impulse_nonlinear(ss, {'w': np.zeros(T)}) 75 | 76 | for i in exogenous: 77 | # simulate with respect to a shock at each date up to T 78 | for t in range(T): 79 | td_out = household.impulse_nonlinear(ss, {i: h * (np.arange(T) == t)}) 80 | 81 | # store results as column t of J[o][i] for each outcome o 82 | for o in output_list: 83 | Js_direct[o.upper()][i][:, t] = (td_out[o.upper()] - td_noshock[o.upper()]) / h 84 | 85 | assert np.linalg.norm(Js['C']['r'] - Js_direct['C']['r'], np.inf) < 3e-4 86 | -------------------------------------------------------------------------------- /src/sequence_jacobian/examples/krusell_smith.py: -------------------------------------------------------------------------------- 1 | from sequence_jacobian import grids, simple, create_model, hetblocks 2 | 3 | hh = hetblocks.hh_sim.hh 4 | 5 | '''Part 1: Blocks''' 6 | 7 | @simple 8 | def firm(K, L, Z, alpha, delta): 9 | r = alpha * Z * (K(-1) / L) ** (alpha-1) - delta 10 | w = (1 - alpha) * Z * (K(-1) / L) ** alpha 11 | Y = Z * K(-1) ** alpha * L ** (1 - alpha) 12 | return r, w, Y 13 | 14 | 15 | @simple 16 | def mkt_clearing(K, A, Y, C, delta): 17 | asset_mkt = A - K 18 | I = K - (1 - delta) * K(-1) 19 | goods_mkt = Y - C - I 20 | return asset_mkt, goods_mkt, I 21 | 22 | 23 | @simple 24 | def firm_ss(r, Y, L, delta, alpha): 25 | '''Solve for (Z, K) given targets for (Y, r).''' 26 | rk = r + delta 27 | K = alpha * Y / rk 28 | Z = Y / K ** alpha / L ** (1 - alpha) 29 | w = (1 - alpha) * Z * (K / L) ** alpha 30 | return K, Z, w 31 | 32 | 33 | '''Part 2: Embed HA block''' 34 | 35 | def make_grids(rho, sigma, nS, amax, nA): 36 | e_grid, _, Pi = grids.markov_rouwenhorst(rho=rho, sigma=sigma, N=nS) 37 | a_grid = grids.agrid(amax=amax, n=nA) 38 | return e_grid, Pi, a_grid 39 | 40 | 41 | def income(w, e_grid): 42 | y = w * e_grid 43 | return y 44 | 45 | 46 | '''Part 3: DAG''' 47 | 48 | def dag(): 49 | # Combine blocks 50 | household = hh.add_hetinputs([income, make_grids]) 51 | ks_model = create_model([household, firm, mkt_clearing], name="Krusell-Smith") 52 | ks_model_ss = create_model([household, firm_ss, mkt_clearing], name="Krusell-Smith SS") 53 | 54 | # Steady state 55 | calibration = {'eis': 1.0, 'delta': 0.025, 'alpha': 0.11, 'rho': 0.966, 'sigma': 0.5, 56 | 'Y': 1.0, 'L': 1.0, 'nS': 2, 'nA': 10, 'amax': 200, 'r': 0.01} 57 | unknowns_ss = {'beta': (0.98 / 1.01, 0.999 / 1.01)} 58 | targets_ss = {'asset_mkt': 0.} 59 | ss = ks_model_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, solver='brentq') 60 | 61 | # Transitional dynamics 62 | inputs = ['Z'] 63 | unknowns = ['K'] 64 | targets = ['asset_mkt'] 65 | 66 | return ks_model_ss, ss, ks_model, unknowns, targets, inputs 67 | 68 | 69 | '''Part 3: Permanent beta heterogeneity''' 70 | 71 | @simple 72 | def aggregate(A_patient, A_impatient, C_patient, C_impatient, mass_patient): 73 | C = mass_patient * C_patient + (1 - mass_patient) * C_impatient 74 | A = mass_patient * A_patient + (1 - mass_patient) * A_impatient 75 | return C, A 76 | 77 | 78 | def remapped_dag(): 79 | # Create 2 versions of the household block using `remap` 80 | household = hh.household.add_hetinputs([income, make_grids]) 81 | to_map = ['beta', *household.outputs] 82 | hh_patient = household.remap({k: k + '_patient' for k in to_map}).rename('hh_patient') 83 | hh_impatient = household.remap({k: k + '_impatient' for k in to_map}).rename('hh_impatient') 84 | blocks = [hh_patient, hh_impatient, firm, mkt_clearing, aggregate] 85 | blocks_ss = [hh_patient, hh_impatient, firm_ss, mkt_clearing, aggregate] 86 | ks_remapped = create_model(blocks, name='KS-beta-het') 87 | ks_remapped_ss = create_model(blocks_ss, name='KS-beta-het') 88 | 89 | # Steady State 90 | calibration = {'eis': 1., 'delta': 0.025, 'alpha': 0.3, 'rho': 0.966, 'sigma': 0.5, 'Y': 1.0, 'L': 1.0, 91 | 'nS': 3, 'nA': 100, 'amax': 1000, 'beta_impatient': 0.985, 'mass_patient': 0.5} 92 | unknowns_ss = {'beta_patient': (0.98 / 1.01, 0.999 / 1.01)} 93 | targets_ss = {'asset_mkt': 0.} 94 | ss = ks_remapped_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, solver='brentq') 95 | 96 | # Transitional Dynamics/Jacobian Calculation 97 | unknowns = ['K'] 98 | targets = ['asset_mkt'] 99 | exogenous = ['Z'] 100 | 101 | return ks_remapped_ss, ss, ks_remapped, unknowns, targets, ss, exogenous 102 | -------------------------------------------------------------------------------- /tests/base/test_multiexog.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sequence_jacobian as sj 3 | from sequence_jacobian import het, simple, combine 4 | 5 | 6 | def household_init(a_grid, y, r, sigma): 7 | c = np.maximum(1e-8, y[..., np.newaxis] + np.maximum(r, 0.04) * a_grid) 8 | Va = (1 + r) * (c ** (-sigma)) 9 | return Va 10 | 11 | 12 | def search_frictions(f, s): 13 | Pi_e = np.vstack(([1 - s, s], [f, 1 - f])) 14 | return Pi_e 15 | 16 | 17 | def labor_income(z, w, b): 18 | y = np.vstack((w * z, b * w * z)) 19 | return y 20 | 21 | 22 | @simple 23 | def income_state_vars(rho_z, sd_z, nZ): 24 | z, _, Pi_z = sj.utilities.discretize.markov_rouwenhorst(rho=rho_z, sigma=sd_z, N=nZ) 25 | return z, Pi_z 26 | 27 | @simple 28 | def asset_state_vars(amin, amax, nA): 29 | a_grid = sj.utilities.discretize.agrid(amin=amin, amax=amax, n=nA) 30 | return a_grid 31 | 32 | 33 | @het(exogenous=['Pi_e', 'Pi_z'], policy='a', backward='Va', backward_init=household_init) 34 | def household_multidim(Va_p, a_grid, y, r, beta, sigma): 35 | c_nextgrid = (beta * Va_p) ** (-1 / sigma) 36 | coh = (1 + r) * a_grid + y[..., np.newaxis] 37 | a = sj.utilities.interpolate.interpolate_y(c_nextgrid + a_grid, coh, a_grid) # (x, xq, y) 38 | a = np.maximum(a, a_grid[0]) 39 | c = coh - a 40 | uc = c ** (-sigma) 41 | Va = (1 + r) * uc 42 | 43 | return Va, a, c 44 | 45 | @het(exogenous='Pi', policy='a', backward='Va', backward_init=household_init) 46 | def household_onedim(Va_p, a_grid, y, r, beta, sigma): 47 | c_nextgrid = (beta * Va_p) ** (-1 / sigma) 48 | coh = (1 + r) * a_grid[np.newaxis, :] + y[:, np.newaxis] 49 | a = sj.utilities.interpolate.interpolate_y(c_nextgrid + a_grid, coh, a_grid) # (x, xq, y) 50 | sj.utilities.optimized_routines.setmin(a, a_grid[0]) 51 | c = coh - a 52 | uc = c ** (-sigma) 53 | Va = (1 + r) * uc 54 | 55 | return Va, a, c 56 | 57 | def test_equivalence(): 58 | calibration = dict(beta=0.95, r=0.01, sigma=2, a_grid = sj.utilities.discretize.agrid(1000, 50)) 59 | 60 | e1, _, Pi1 = sj.utilities.discretize.markov_rouwenhorst(rho=0.7, sigma=0.7, N=3) 61 | e2, _, Pi2 = sj.utilities.discretize.markov_rouwenhorst(rho=0.3, sigma=0.5, N=3) 62 | e_multidim = np.outer(e1, e2) 63 | 64 | e_onedim = np.kron(e1, e2) 65 | Pi = np.kron(Pi1, Pi2) 66 | 67 | ss_multidim = household_multidim.steady_state({**calibration, 'y': e_multidim, 'Pi_e': Pi1, 'Pi_z': Pi2}) 68 | ss_onedim = household_onedim.steady_state({**calibration, 'y': e_onedim, 'Pi': Pi}) 69 | 70 | assert np.isclose(ss_multidim['A'], ss_onedim['A']) and np.isclose(ss_multidim['C'], ss_onedim['C']) 71 | 72 | D_onedim = ss_onedim.internals['household_onedim']['D'] 73 | D_multidim = ss_multidim.internals['household_multidim']['D'] 74 | 75 | assert np.allclose(D_onedim, D_multidim.reshape(*D_onedim.shape)) 76 | 77 | J_multidim = household_multidim.jacobian(ss_multidim, inputs = ['r'], outputs=['A'], T=10) 78 | J_onedim = household_onedim.jacobian(ss_onedim, inputs = ['r'], outputs=['A'], T=10) 79 | 80 | assert np.allclose(J_multidim['A','r'], J_onedim['A','r']) 81 | 82 | 83 | def test_pishock(): 84 | calibration = dict(beta=0.95, r=0.01, sigma=2., f=0.4, s=0.1, w=1., b=0.5, 85 | rho_z=0.9, sd_z=0.5, nZ=3, amin=0., amax=1000, nA=50) 86 | 87 | household = household_multidim.add_hetinputs([search_frictions, labor_income]) 88 | hh = combine([household, income_state_vars, asset_state_vars]) 89 | 90 | ss = hh.steady_state(calibration) 91 | 92 | J = hh.jacobian(ss, inputs=['f', 's', 'r'], outputs=['C'], T=10) 93 | 94 | assert np.max(np.triu(J['C']['r'], 1)) <= 0 # low C before hike in r 95 | assert np.min(np.tril(J['C']['r'])) >= 0 # high C after hike in r 96 | 97 | assert np.all(J['C']['f'] > 0) # high f increases C everywhere 98 | assert np.all(J['C']['s'] < 0) # high s decreases C everywhere 99 | 100 | shock = 0.8**np.arange(10) 101 | C_up = hh.impulse_nonlinear(ss, {'f': 1E-4*shock})['C'] 102 | C_dn = hh.impulse_nonlinear(ss, {'f': -1E-4*shock})['C'] 103 | dC = (C_up - C_dn)/2E-4 104 | assert np.allclose(dC, J['C', 'f'] @ shock, atol=2E-6) 105 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Sequence-Space Jacobian (SSJ) 2 | 3 | [![CI](https://github.com/shade-econ/sequence-jacobian/actions/workflows/main.yml/badge.svg)](https://github.com/shade-econ/sequence-jacobian/actions/workflows/main.yml) 4 | 5 | SSJ is a toolkit for analyzing dynamic macroeconomic models with (or without) rich microeconomic heterogeneity. 6 | 7 | The conceptual framework is based on our paper Adrien Auclert, Bence Bardóczy, Matthew Rognlie, Ludwig Straub (2021), [Using the Sequence-Space Jacobian to Solve and Estimate Heterogeneous-Agent Models](https://doi.org/10.3982/ECTA17434), Econometrica 89(5), pp. 2375–2408 [[ungated copy]](http://mattrognlie.com/sequence_space_jacobian.pdf). 8 | 9 | ## Requirements and installation 10 | 11 | SSJ runs on Python 3.7 or newer, and requires Python's core numerical libraries (NumPy, SciPy, Numba). We recommend that you first install the latest [Anaconda](https://www.anaconda.com/distribution/) distribution. This includes all of the packages and tools that you will need to run our code. 12 | 13 | To install SSJ, open a terminal and type 14 | ``` 15 | pip install sequence-jacobian 16 | ``` 17 | 18 | *Optional package*: There is an optional interface for plotting the directed acyclic graph (DAG) representation of models, which requires [Graphviz for Python](https://github.com/xflr6/graphviz#graphviz). With Anaconda, you can install this by typing `conda install -c conda-forge python-graphviz`. 19 | 20 | ## Using SSJ: introductory notebooks 21 | 22 | To learn how to use the toolkit, it's best to work through our introductory Jupyter notebooks, which show how SSJ can be used to represent and solve various models. We recommend working through the notebooks in the order listed below. [Click here to download all notebooks as a zip](https://github.com/shade-econ/sequence-jacobian/raw/master/notebooks/notebooks.zip). 23 | 24 | - [RBC](https://github.com/shade-econ/sequence-jacobian/blob/master/notebooks/rbc.ipynb) 25 | - represent macro models as collections of blocks (DAG) 26 | - write SimpleBlocks and CombinedBlocks 27 | - compute linearized and non-linear (perfect-foresight) impulse responses 28 | - [Krusell-Smith](https://github.com/shade-econ/sequence-jacobian/blob/master/notebooks/krusell_smith.ipynb) 29 | - write HetBlocks to represent heterogeneous agents 30 | - construct general-equilibrium Jacobians manually 31 | - compute the log-likelihood of the model given time-series data 32 | - [One-asset HANK](https://github.com/shade-econ/sequence-jacobian/blob/master/notebooks/hank.ipynb) 33 | - adapt an off-the-shelf HetBlock to any macro environment using helper functions 34 | - see a more advanced example of calibration 35 | - [Two-asset HANK](https://github.com/shade-econ/sequence-jacobian/blob/master/notebooks/two_asset.ipynb) 36 | - write SolvedBlocks to represent implicit aggregate equilibrium conditions 37 | - re-use saved Jacobians 38 | - fine tune options of block methods 39 | - [Labor search](https://github.com/shade-econ/sequence-jacobian/blob/master/notebooks/labor_search.ipynb) 40 | - example with multiple exogenous states 41 | - shocks to transition matrix of exogenous states 42 | 43 | ## Resources 44 | 45 | If you'd like to learn more about Python, its numerical libraries, and Jupyter notebooks, the [introductory lectures at QuantEcon](https://python-programming.quantecon.org/intro.html) are a terrific place to start. More advanced tutorials for numerical Python include the [SciPy Lecture Notes](http://scipy-lectures.org/intro/language/python_language.html) and the [Python Data Science Handbook](https://jakevdp.github.io/PythonDataScienceHandbook/). There are many other good options as well: thanks to Python's popularity, nearly limitless answers are available via Google, Stack Overflow, and YouTube. 46 | 47 | If you have questions or issues specific to this package, consider posting them on our [GitHub issue tracker](https://github.com/shade-econ/sequence-jacobian/issues). 48 | 49 | For those who used our pre-1.0 toolkit, which had a number of differences relative post-1.0, you can go back to our [early toolkit page](https://github.com/shade-econ/sequence-jacobian/tree/bcca2eff6041abc77d0a777e6c64f9ac6ff44305) if needed. 50 | 51 | ## Team 52 | 53 | The current development team for SSJ is 54 | 55 | - Bence Bardóczy 56 | - Michael Cai 57 | - Matthew Rognlie 58 | 59 | with contributions also from Adrien Auclert, Martin Souchier, and Ludwig Straub. -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/ordered_set.py: -------------------------------------------------------------------------------- 1 | from typing import Iterable 2 | 3 | class OrderedSet: 4 | """Ordered set implemented as dict (where key insertion order is preserved) mapping all to None. 5 | 6 | Operations on multiple ordered sets (e.g. union) order all members of first argument first, then 7 | second argument. If a member is in both, order is as early as possible. 8 | 9 | See test_misc_support.test_ordered_set() for examples.""" 10 | 11 | def __init__(self, members: Iterable = []): 12 | self.d = {k: None for k in members} 13 | 14 | def dict_from(self, s): 15 | return dict(zip(self, s)) 16 | 17 | def __iter__(self): 18 | return iter(self.d) 19 | 20 | def __reversed__(self): 21 | return OrderedSet(list(self)[::-1]) 22 | 23 | def __repr__(self): 24 | return f"OrderedSet({list(self)})" 25 | 26 | def __str__(self): 27 | return str(list(self.d)) 28 | 29 | def __contains__(self, k): 30 | return k in self.d 31 | 32 | def __len__(self): 33 | return len(self.d) 34 | 35 | def __getitem__(self, i): 36 | return list(self.d)[i] 37 | 38 | def add(self, x): 39 | self.d[x] = None 40 | 41 | def difference(self, s): 42 | return OrderedSet(k for k in self if k not in s) 43 | 44 | def difference_update(self, s): 45 | self.d = self.difference(s).d 46 | return self 47 | 48 | def discard(self, k): 49 | self.d.pop(k, None) 50 | 51 | def intersection(self, s): 52 | return OrderedSet(k for k in self if k in s) 53 | 54 | def intersection_update(self, s): 55 | self.d = self.intersection(s).d 56 | return self 57 | 58 | def isdisjoint(self, s): 59 | return len(self.intersection(s)) == 0 60 | 61 | def issubset(self, s): 62 | return len(self.difference(s)) == 0 63 | 64 | def issuperset(self, s): 65 | return len(self.intersection(s)) == len(s) 66 | 67 | def remove(self, k): 68 | self.d.pop(k) 69 | 70 | def symmetric_difference(self, s): 71 | diff = self.difference(s) 72 | for k in s: 73 | if k not in self: 74 | diff.add(k) 75 | return diff 76 | 77 | def symmetric_difference_update(self, s): 78 | self.d = self.symmetric_difference(s).d 79 | return self 80 | 81 | def union(self, s): 82 | return self.copy().update(s) 83 | 84 | def update(self, s): 85 | for k in s: 86 | self.add(k) 87 | return self 88 | 89 | def copy(self): 90 | return OrderedSet(self) 91 | 92 | def __eq__(self, s): 93 | if isinstance(s, OrderedSet): 94 | return list(self) == list(s) 95 | else: 96 | return False 97 | 98 | def __le__(self, s): 99 | return self.issubset(s) 100 | 101 | def __lt__(self, s): 102 | return self.issubset(s) and (len(self) != len(s)) 103 | 104 | def __ge__(self, s): 105 | return self.issuperset(s) 106 | 107 | def __gt__(self, s): 108 | return self.issuperset(s) and (len(self) != len(s)) 109 | 110 | def __or__(self, s): 111 | return self.union(s) 112 | 113 | def __ior__(self, s): 114 | return self.update(s) 115 | 116 | def __ror__(self, s): 117 | return self.union(s) 118 | 119 | def __and__(self, s): 120 | return self.intersection(s) 121 | 122 | def __iand__(self, s): 123 | return self.intersection_update(s) 124 | 125 | def __rand__(self, s): 126 | return self.intersection(s) 127 | 128 | def __sub__(self, s): 129 | return self.difference(s) 130 | 131 | def __isub__(self, s): 132 | return self.difference_update(s) 133 | 134 | def __rsub__(self, s): 135 | return OrderedSet(s).difference(self) 136 | 137 | def __xor__(self, s): 138 | return self.symmetric_difference(s) 139 | 140 | def __ixor__(self, s): 141 | return self.symmetric_difference_update(s) 142 | 143 | def __rxor__(self, s): 144 | return OrderedSet(s).symmetric_difference(self) 145 | 146 | """Compatibility methods, regular use not advised""" 147 | 148 | def pop(self): 149 | k = self.top() 150 | del self.d[k] 151 | return k 152 | 153 | def top(self): 154 | return list(self.d)[-1] 155 | 156 | def index(self, k): 157 | return list(self.d).index(k) 158 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/simple_block.py: -------------------------------------------------------------------------------- 1 | """Class definition of a simple block""" 2 | 3 | import numpy as np 4 | from copy import deepcopy 5 | 6 | from .support.simple_displacement import ignore, Displace, AccumulatedDerivative 7 | from .block import Block 8 | from ..classes import SteadyStateDict, ImpulseDict, JacobianDict, SimpleSparse 9 | from ..utilities import misc 10 | from ..utilities.function import ExtendedFunction 11 | 12 | '''Part 1: SimpleBlock class and @simple decorator to generate it''' 13 | 14 | 15 | def simple(f): 16 | return SimpleBlock(f) 17 | 18 | 19 | class SimpleBlock(Block): 20 | """Generated from simple block written in Dynare-ish style and decorated with @simple, e.g. 21 | 22 | @simple 23 | def production(Z, K, L, alpha): 24 | Y = Z * K(-1) ** alpha * L ** (1 - alpha) 25 | return Y 26 | 27 | which is a SimpleBlock that takes in Z, K, L, and alpha, all of which can be either constants 28 | or series, and implements a Cobb-Douglas production function, noting that for production today 29 | we use the capital K(-1) determined yesterday. 30 | 31 | Key methods are .ss, .td, and .jac, like HetBlock. 32 | """ 33 | 34 | def __init__(self, f): 35 | super().__init__() 36 | self.f = ExtendedFunction(f) 37 | self.name = self.f.name 38 | self.inputs = self.f.inputs 39 | self.outputs = self.f.outputs 40 | 41 | def __repr__(self): 42 | return f"" 43 | 44 | def _steady_state(self, ss): 45 | outputs = self.f.wrapped_call(ss, preprocess=ignore, postprocess=misc.numeric_primitive) 46 | return SteadyStateDict({**ss, **outputs}) 47 | 48 | def _impulse_nonlinear(self, ss, inputs, outputs, ss_initial): 49 | if ss_initial is None: 50 | ss_initial = ss 51 | ss_initial_flag = False 52 | else: 53 | ss_initial_flag = True 54 | 55 | input_args = {} 56 | for k, v in inputs.items(): 57 | if np.isscalar(v): 58 | raise ValueError(f'Keyword argument {k}={v} is scalar, should be time path.') 59 | input_args[k] = Displace(v + ss[k], ss[k], ss_initial[k], k) 60 | 61 | for k in self.inputs: 62 | if k not in input_args: 63 | if not ss_initial_flag or (ss_initial_flag and np.array_equal(ss_initial[k], ss[k])): 64 | input_args[k] = ignore(ss[k]) 65 | else: 66 | input_args[k] = Displace(np.full(inputs.T, ss[k]), ss[k], ss_initial[k], k) 67 | 68 | return ImpulseDict(make_impulse_uniform_length(self.f(input_args)))[outputs] - ss 69 | 70 | def _impulse_linear(self, ss, inputs, outputs, Js): 71 | return ImpulseDict(self.jacobian(ss, list(inputs.keys()), outputs, inputs.T, Js).apply(inputs)) 72 | 73 | def _jacobian(self, ss, inputs, outputs, T): 74 | invertedJ = {i: {} for i in inputs} 75 | 76 | # Loop over all inputs/shocks which we want to differentiate with respect to 77 | for i in inputs: 78 | invertedJ[i] = self.compute_single_shock_J(ss, i) 79 | 80 | # Because we computed the Jacobian of all outputs with respect to each shock (invertedJ[i][o]), 81 | # we need to loop back through to have J[o][i] to map for a given output `o`, shock `i`, 82 | # the Jacobian curlyJ^{o,i}. 83 | J = {o: {} for o in outputs} 84 | for o in outputs: 85 | for i in inputs: 86 | # drop zeros from JacobianDict 87 | if invertedJ[i][o] and not invertedJ[i][o].iszero: 88 | J[o][i] = invertedJ[i][o] 89 | 90 | return JacobianDict(J, outputs, inputs, self.name, T) 91 | 92 | def compute_single_shock_J(self, ss, i): 93 | input_args = {i: ignore(ss[i]) for i in self.inputs} 94 | input_args[i] = AccumulatedDerivative(f_value=ss[i]) 95 | 96 | J = {o: {} for o in self.outputs} 97 | for o_name, o in self.f(input_args).items(): 98 | if isinstance(o, AccumulatedDerivative): 99 | J[o_name] = SimpleSparse(o.elements) 100 | 101 | return J 102 | 103 | 104 | # TODO: move this to impulse.py? 105 | def make_impulse_uniform_length(out): 106 | T = np.max([np.size(v) for v in out.values()]) 107 | return {k: (np.full(T, misc.numeric_primitive(v)) if np.isscalar(v) else misc.numeric_primitive(v)) 108 | for k, v in out.items()} 109 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/drawdag.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | from sequence_jacobian.blocks.solved_block import SolvedBlock 3 | from sequence_jacobian.blocks.het_block import HetBlock 4 | 5 | """ 6 | Adrien's DAG Graph routine, updated for SSJ v1.0 7 | 8 | Requires installing graphviz package and executables 9 | https://www.graphviz.org/ 10 | 11 | On a mac this can be done as follows: 12 | 1) Download macports at: 13 | https://www.macports.org/install.php 14 | 2) On the command line, install graphviz with macports by typing 15 | sudo port install graphviz 16 | 17 | """ 18 | 19 | try: 20 | from graphviz import Digraph 21 | from IPython.display import display 22 | 23 | def drawdag(model, exogenous=[], unknowns=[], targets=[], leftright=False, save=False, savepath=None): 24 | ''' 25 | Routine that draws DAG 26 | :param model: combined block to be represented as dag 27 | :param exogenous: (optional) exogenous variables, to be represented on DAG 28 | :param unknowns: (optional) unknown variables, to be represented on DAG 29 | :param unknowns: (optional) target variables, to be represented on DAG 30 | :bool leftright: if True, plots dag from left to right instead of top to bottom 31 | :return: none 32 | ''' 33 | # Start DAG 34 | dot = Digraph(comment='Model DAG') 35 | 36 | # Make it left-to-right if asked 37 | if leftright: 38 | dot.attr(rankdir='LR', ratio='compress', center='true') 39 | else: 40 | dot.attr(ratio='auto', center='true') 41 | 42 | # add initial nodes (one for exogenous, one for unknowns) provided those are not empty lists 43 | if exogenous: 44 | dot.node('exog', 'exogenous', shape='box') 45 | if unknowns: 46 | dot.node('unknowns', 'unknowns', shape='box') 47 | if targets: 48 | dot.node('targets', 'targets', shape='diamond') 49 | 50 | # add nodes sequentially in order 51 | for i, b in enumerate(model.blocks): 52 | if isinstance(b, HetBlock): 53 | dot.node(str(i), b.name + ' [HA, ' + str(i) + ']') 54 | elif isinstance(b, SolvedBlock) : 55 | dot.node(str(i), b.name + ' [solved,' + str(i) + ']') 56 | else: 57 | dot.node(str(i), b.name + ' [' + str(i) + ']') 58 | 59 | # nodes from exogenous to i (figure out if needed and draw) 60 | if exogenous: 61 | edgelabel = b.inputs & set(exogenous) 62 | if len(edgelabel) != 0: 63 | edgelabel_list = list(edgelabel) 64 | edgelabel_str = ', '.join(str(e) for e in edgelabel_list) 65 | dot.edge('exog', str(i), label=str(edgelabel_str)) 66 | 67 | # nodes from unknowns to i (figure out if needed, then draw) 68 | if unknowns: 69 | edgelabel = b.inputs & set(unknowns) 70 | if len(edgelabel) != 0: 71 | edgelabel_list = list(edgelabel) 72 | edgelabel_str = ', '.join(str(e) for e in edgelabel_list) 73 | dot.edge('unknowns', str(i), label=str(edgelabel_str)) 74 | 75 | # nodes from i to final targets 76 | for target in targets: 77 | if target in b.outputs: 78 | dot.edge(str(i), 'targets', label=target) 79 | 80 | # nodes from any interior block to i 81 | for j in model.revadj[i]: 82 | # figure out inputs of i that are also outputs of j 83 | edgelabel = b.inputs & model.blocks[j].outputs 84 | edgelabel_list = list(edgelabel) 85 | edgelabel_str = ', '.join(str(e) for e in edgelabel_list) 86 | 87 | # draw edge from j to i 88 | dot.edge(str(j), str(i), label=str(edgelabel_str)) 89 | 90 | if save: 91 | if savepath is None: 92 | savepath = 'dag/' + model.name 93 | dot.render(savepath, format='png', cleanup=True) 94 | display(dot) 95 | 96 | except ImportError: 97 | def drawdag(*args, **kwargs): 98 | warnings.warn("\nAttempted to use `drawdag` when the package `graphviz` has not yet been installed. \n" 99 | "DAG visualization tools, i.e. drawdag, will not produce any figures unless this dependency has been installed. \n" 100 | "If you want to install, try typing 'conda install -c conda-forge python-graphviz' at the terminal,\n" 101 | "or see README for more instructions. Once installed, re-load sequence-jacobian to produce DAG figures.") 102 | -------------------------------------------------------------------------------- /src/sequence_jacobian/classes/impulse_dict.py: -------------------------------------------------------------------------------- 1 | """ImpulseDict class for manipulating impulse responses.""" 2 | 3 | import numpy as np 4 | 5 | from .result_dict import ResultDict 6 | 7 | from ..utilities.ordered_set import OrderedSet 8 | from ..utilities.bijection import Bijection 9 | from .steady_state_dict import SteadyStateDict 10 | 11 | class ImpulseDict(ResultDict): 12 | def __init__(self, data, internals=None, T=None): 13 | if isinstance(data, ImpulseDict): 14 | if internals is not None or T is not None: 15 | raise ValueError('Supplying ImpulseDict and also internal or T to constructor not allowed') 16 | super().__init__(data) 17 | self.T = data.T 18 | else: 19 | if not isinstance(data, dict): 20 | raise ValueError('ImpulseDicts are initialized with a `dict` of top-level impulse responses.') 21 | super().__init__(data, internals) 22 | self.T = (T if T is not None else self.infer_length()) 23 | 24 | def __getitem__(self, k): 25 | return super().__getitem__(k, T=self.T) 26 | 27 | def __add__(self, other): 28 | return self.binary_operation(other, lambda a, b: a + b) 29 | 30 | def __radd__(self, other): 31 | return self.__add__(other) 32 | 33 | def __sub__(self, other): 34 | return self.binary_operation(other, lambda a, b: a - b) 35 | 36 | def __rsub__(self, other): 37 | return self.binary_operation(other, lambda a, b: b - a) 38 | 39 | def __mul__(self, other): 40 | return self.binary_operation(other, lambda a, b: a * b) 41 | 42 | def __rmul__(self, other): 43 | return self.__mul__(other) 44 | 45 | def __truediv__(self, other): 46 | return self.binary_operation(other, lambda a, b: a / b) 47 | 48 | def __rtruediv__(self, other): 49 | return self.binary_operation(other, lambda a, b: b / a) 50 | 51 | def __neg__(self): 52 | return self.unary_operation(lambda a: -a) 53 | 54 | def __pos__(self): 55 | return self 56 | 57 | def __abs__(self): 58 | return self.unary_operation(lambda a: abs(a)) 59 | 60 | def binary_operation(self, other, op): 61 | if isinstance(other, (SteadyStateDict, ImpulseDict)): 62 | toplevel = {k: op(v, other[k]) for k, v in self.toplevel.items()} 63 | internals = {} 64 | for b in self.internals: 65 | other_internals = other.internals[b] 66 | internals[b] = {k: op(v, other_internals[k]) for k, v in self.internals[b].items()} 67 | return ImpulseDict(toplevel, internals, self.T) 68 | elif isinstance(other, (float, int)): 69 | toplevel = {k: op(v, other) for k, v in self.toplevel.items()} 70 | internals = {} 71 | for b in self.internals: 72 | internals[b] = {k: op(v, other) for k, v in self.internals[b].items()} 73 | return ImpulseDict(toplevel, internals, self.T) 74 | else: 75 | return NotImplementedError(f'Can only perform operations with ImpulseDicts and other ImpulseDicts, SteadyStateDicts, or numbers, not {type(other).__name__}') 76 | 77 | def unary_operation(self, op): 78 | toplevel = {k: op(v) for k, v in self.toplevel.items()} 79 | internals = {} 80 | for b in self.internals: 81 | internals[b] = {k: op(v) for k, v in self.internals[b].items()} 82 | return ImpulseDict(toplevel, internals, self.T) 83 | 84 | def pack(self): 85 | T = self.T 86 | bigv = np.empty(T*len(self.toplevel)) 87 | for i, v in enumerate(self.toplevel.values()): 88 | bigv[i*T:(i+1)*T] = v 89 | return bigv 90 | 91 | @staticmethod 92 | def unpack(bigv, outputs, T): 93 | impulse = {} 94 | for i, o in enumerate(outputs): 95 | impulse[o] = bigv[i*T:(i+1)*T] 96 | return ImpulseDict(impulse, T=T) 97 | 98 | def infer_length(self): 99 | lengths = [len(v) for v in self.toplevel.values()] 100 | length = max(lengths) 101 | if length != min(lengths): 102 | raise ValueError(f'Building ImpulseDict with inconsistent lengths {max(lengths)} and {min(lengths)}') 103 | return length 104 | 105 | def get(self, k): 106 | """Like __getitem__ but with default of zero impulse""" 107 | if isinstance(k, str): 108 | return self.toplevel.get(k, np.zeros(self.T)) 109 | elif isinstance(k, tuple): 110 | raise TypeError(f'Key {k} to {type(self).__name__} cannot be tuple') 111 | else: 112 | try: 113 | return type(self)({ki: self.toplevel.get(ki, np.zeros(self.T)) for ki in k}, T=self.T) 114 | except TypeError: 115 | raise TypeError(f'Key {k} to {type(self).__name__} needs to be a string or an iterable (list, set, etc) of strings') 116 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/solved_block.py: -------------------------------------------------------------------------------- 1 | from .block import Block 2 | from .simple_block import simple 3 | from .support.parent import Parent 4 | from ..classes import FactoredJacobianDict 5 | from ..utilities.ordered_set import OrderedSet 6 | 7 | 8 | def solved(unknowns, targets, solver=None, solver_kwargs={}, name=""): 9 | """Convenience @solved(unknowns=..., targets=...) decorator on a single SimpleBlock""" 10 | # call as decorator, return function of function 11 | def singleton_solved_block(f): 12 | return SolvedBlock(simple(f).rename(f.__name__ + '_inner'), f.__name__, unknowns, targets, solver=solver, solver_kwargs=solver_kwargs) 13 | return singleton_solved_block 14 | 15 | 16 | class SolvedBlock(Block, Parent): 17 | """SolvedBlocks are mini SHADE models embedded as blocks inside larger SHADE models. 18 | 19 | When creating them, we need to provide the basic ingredients of a SHADE model: the list of 20 | blocks comprising the model, the list on unknowns, and the list of targets. 21 | 22 | When we use .jac to ask for the Jacobian of a SolvedBlock, we are really solving for the 'G' 23 | matrices of the mini SHADE models, which then become the 'curlyJ' Jacobians of the block. 24 | 25 | Similarly, when we use .td to evaluate a SolvedBlock on a path, we are really solving for the 26 | nonlinear transition path such that all internal targets of the mini SHADE model are zero. 27 | """ 28 | 29 | def __init__(self, block: Block, name, unknowns, targets, solver=None, solver_kwargs={}): 30 | super().__init__() 31 | 32 | # since we dispatch to solve methods, same set of options 33 | self.impulse_nonlinear_options = self.solve_impulse_nonlinear_options 34 | self.steady_state_options = self.solve_steady_state_options 35 | 36 | self.block = block 37 | self.name = name 38 | self.unknowns = unknowns 39 | self.targets = targets 40 | self.solver = solver 41 | self.solver_kwargs = solver_kwargs 42 | 43 | Parent.__init__(self, [self.block]) 44 | 45 | # validate unknowns and targets 46 | if not len(unknowns) == len(targets): 47 | raise ValueError(f'Unknowns {set(unknowns)} and targets {set(targets)} different sizes in SolvedBlock {name}') 48 | if not set(unknowns) <= block.inputs: 49 | raise ValueError(f'Unknowns has element {set(unknowns) - block.inputs} not in inputs in SolvedBlock {name}') 50 | if not set(targets) <= block.outputs: 51 | raise ValueError(f'Targets has element {set(targets) - block.outputs} not in outputs in SolvedBlock {name}') 52 | 53 | # what are overall outputs and inputs? 54 | self.outputs = block.outputs | set(unknowns) 55 | self.inputs = block.inputs - set(unknowns) 56 | 57 | def __repr__(self): 58 | return f"" 59 | 60 | def _steady_state(self, calibration, dissolve, options, **kwargs): 61 | if self.name in dissolve: 62 | kwargs['solver'] = "solved" 63 | unknowns = {k: v for k, v in calibration.items() if k in self.unknowns} 64 | else: 65 | unknowns = self.unknowns 66 | if 'solver' not in kwargs: 67 | # TODO: replace this with default option 68 | kwargs['solver'] = self.solver 69 | 70 | return self.block.solve_steady_state(calibration, unknowns, self.targets, options, **kwargs) 71 | 72 | def _impulse_nonlinear(self, ss, inputs, outputs, internals, Js, options, ss_initial, **kwargs): 73 | return self.block.solve_impulse_nonlinear(ss, OrderedSet(self.unknowns), OrderedSet(self.targets), 74 | inputs, outputs, internals, Js, options, self._get_H_U_factored(Js), ss_initial, **kwargs) 75 | 76 | def _impulse_linear(self, ss, inputs, outputs, Js, options): 77 | return self.block.solve_impulse_linear(ss, OrderedSet(self.unknowns), OrderedSet(self.targets), 78 | inputs, outputs, Js, options, self._get_H_U_factored(Js)) 79 | 80 | def _jacobian(self, ss, inputs, outputs, T, Js, options): 81 | return self.block.solve_jacobian(ss, OrderedSet(self.unknowns), OrderedSet(self.targets), 82 | inputs, outputs, T, Js, options, self._get_H_U_factored(Js))[outputs] 83 | 84 | def _partial_jacobians(self, ss, inputs, outputs, T, Js, options): 85 | # call it on the child first 86 | inner_Js = self.block.partial_jacobians(ss, (OrderedSet(self.unknowns) | inputs), 87 | (OrderedSet(self.targets) | outputs - self.unknowns.keys()), T, Js, options) 88 | 89 | # with these inner Js, also compute H_U and factorize 90 | H_U = self.block.jacobian(ss, OrderedSet(self.unknowns), OrderedSet(self.targets), T, inner_Js, options) 91 | H_U_factored = FactoredJacobianDict(H_U, T) 92 | 93 | return {**inner_Js, self.name: H_U_factored} 94 | 95 | def _get_H_U_factored(self, Js): 96 | if self.name in Js and isinstance(Js[self.name], FactoredJacobianDict): 97 | return Js[self.name] 98 | else: 99 | return None 100 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/solvers.py: -------------------------------------------------------------------------------- 1 | """Simple nonlinear solvers""" 2 | 3 | import numpy as np 4 | import warnings 5 | 6 | 7 | def newton_solver(f, x0, y0=None, tol=1E-9, maxcount=100, backtrack_c=0.5, verbose=True): 8 | """Simple line search solver for root x satisfying f(x)=0 using Newton direction. 9 | 10 | Backtracks if input invalid or improvement is not at least half the predicted improvement. 11 | 12 | Parameters 13 | ---------- 14 | f : function, to solve for f(x)=0, input and output are arrays of same length 15 | x0 : array (n), initial guess for x 16 | y0 : [optional] array (n), y0=f(x0), if already known 17 | tol : [optional] scalar, solver exits successfully when |f(x)| < tol 18 | maxcount : [optional] int, maximum number of Newton steps 19 | backtrack_c : [optional] scalar, fraction to backtrack if step unsuccessful, i.e. 20 | if we tried step from x to x+dx, now try x+backtrack_c*dx 21 | 22 | Returns 23 | ---------- 24 | x : array (n), (approximate) root of f(x)=0 25 | y : array (n), y=f(x), satisfies |y| < tol 26 | """ 27 | 28 | x, y = x0, y0 29 | if y is None: 30 | y = f(x) 31 | 32 | for count in range(maxcount): 33 | if verbose: 34 | printit(count, x, y) 35 | 36 | if np.max(np.abs(y)) < tol: 37 | return x, y 38 | 39 | J = obtain_J(f, x, y) 40 | dx = np.linalg.solve(J, -y) 41 | 42 | # backtrack at most 29 times 43 | for bcount in range(30): 44 | try: 45 | ynew = f(x + dx) 46 | except ValueError: 47 | if verbose: 48 | print('backtracking\n') 49 | dx *= backtrack_c 50 | else: 51 | predicted_improvement = -np.sum((J @ dx) * y) * ((1 - 1 / 2 ** bcount) + 1) / 2 52 | actual_improvement = (np.sum(y ** 2) - np.sum(ynew ** 2)) / 2 53 | if actual_improvement < predicted_improvement / 2: 54 | if verbose: 55 | print('backtracking\n') 56 | dx *= backtrack_c 57 | else: 58 | y = ynew 59 | x += dx 60 | break 61 | else: 62 | raise ValueError('Too many backtracks, maybe bad initial guess?') 63 | else: 64 | raise ValueError(f'No convergence after {maxcount} iterations') 65 | 66 | 67 | def broyden_solver(f, x0, y0=None, tol=1E-9, maxcount=100, backtrack_c=0.5, verbose=True): 68 | """Similar to newton_solver, but solves f(x)=0 using approximate rather than exact Newton direction, 69 | obtaining approximate Jacobian J=f'(x) from Broyden updating (starting from exact Newton at f'(x0)). 70 | 71 | Backtracks only if error raised by evaluation of f, since improvement criterion no longer guaranteed 72 | to work for any amount of backtracking if Jacobian not exact. 73 | """ 74 | 75 | x, y = x0, y0 76 | if y is None: 77 | y = f(x) 78 | 79 | for count in range(maxcount): 80 | if verbose: 81 | printit(count, x, y) 82 | 83 | if np.max(np.abs(y)) < tol: 84 | return x, y 85 | 86 | # initialize J with Newton! 87 | if count == 0: 88 | J = obtain_J(f, x, y) 89 | 90 | if len(x) == len(y): 91 | dx = np.linalg.solve(J, -y) 92 | elif len(x) < len(y): 93 | warnings.warn(f"Dimension of x, {len(x)} is less than dimension of y, {len(y)}." 94 | f" Using least-squares criterion to solve for approximate root.") 95 | dx = np.linalg.lstsq(J, -y, rcond=None)[0] 96 | else: 97 | raise ValueError(f"Dimension of x, {len(x)} is greater than dimension of y, {len(y)}." 98 | f" Cannot solve underdetermined system.") 99 | 100 | # backtrack at most 29 times 101 | for bcount in range(30): 102 | # note: can't test for improvement with Broyden because maybe 103 | # the function doesn't improve locally in this direction, since 104 | # J isn't the exact Jacobian 105 | try: 106 | ynew = f(x + dx) 107 | except ValueError: 108 | if verbose: 109 | print('backtracking\n') 110 | dx *= backtrack_c 111 | else: 112 | J = broyden_update(J, dx, ynew - y) 113 | y = ynew 114 | x += dx 115 | break 116 | else: 117 | raise ValueError('Too many backtracks, maybe bad initial guess?') 118 | else: 119 | raise ValueError(f'No convergence after {maxcount} iterations') 120 | 121 | 122 | def obtain_J(f, x, y, h=1E-5): 123 | """Finds Jacobian f'(x) around y=f(x)""" 124 | nx = x.shape[0] 125 | ny = y.shape[0] 126 | J = np.empty((ny, nx)) 127 | 128 | for i in range(nx): 129 | dx = h * (np.arange(nx) == i) 130 | J[:, i] = (f(x + dx) - y) / h 131 | return J 132 | 133 | 134 | def broyden_update(J, dx, dy): 135 | """Returns Broyden update to approximate Jacobian J, given that last change in inputs to function 136 | was dx and led to output change of dy.""" 137 | return J + np.outer(((dy - J @ dx) / np.linalg.norm(dx) ** 2), dx) 138 | 139 | 140 | def printit(it, x, y, **kwargs): 141 | """Convenience printing function for verbose iterations""" 142 | print(f'On iteration {it}') 143 | print(('x = %.3f' + ',%.3f' * (len(x) - 1)) % tuple(x)) 144 | print(('y = %.3f' + ',%.3f' * (len(y) - 1)) % tuple(y)) 145 | for kw, val in kwargs.items(): 146 | print(f'{kw} = {val:.3f}') 147 | print('\n') 148 | -------------------------------------------------------------------------------- /tests/base/test_transitional_dynamics.py: -------------------------------------------------------------------------------- 1 | """Test all models' non-linear transitional dynamics computations""" 2 | 3 | import numpy as np 4 | 5 | from sequence_jacobian import combine 6 | from sequence_jacobian.examples import two_asset 7 | from sequence_jacobian.hetblocks import hh_twoasset as hh 8 | 9 | 10 | # TODO: Figure out a more robust way to check similarity of the linear and non-linear solution. 11 | # As of now just checking that the tolerance for difference (by infinity norm) is below a manually checked threshold 12 | def test_rbc_td(rbc_dag): 13 | rbc_model, ss, unknowns, targets, exogenous = rbc_dag 14 | 15 | T, impact, rho, news = 30, 0.01, 0.8, 10 16 | G = rbc_model.solve_jacobian(ss, unknowns, targets, exogenous, T=T) 17 | 18 | dZ = np.empty((T, 2)) 19 | dZ[:, 0] = impact * ss['Z'] * rho**np.arange(T) 20 | dZ[:, 1] = np.concatenate((np.zeros(news), dZ[:-news, 0])) 21 | dC = 100 * G['C']['Z'] @ dZ / ss['C'] 22 | 23 | td_nonlin = rbc_model.solve_impulse_nonlinear(ss, unknowns, targets, inputs={"Z": dZ[:, 0]}, outputs=['C']) 24 | td_nonlin_news = rbc_model.solve_impulse_nonlinear(ss, unknowns, targets, inputs={"Z": dZ[:, 1]}, outputs=['C']) 25 | 26 | dC_nonlin = 100 * td_nonlin['C'] / ss['C'] 27 | dC_nonlin_news = 100 * td_nonlin_news['C'] / ss['C'] 28 | 29 | assert np.linalg.norm(dC[:, 0] - dC_nonlin, np.inf) < 3e-2 30 | assert np.linalg.norm(dC[:, 1] - dC_nonlin_news, np.inf) < 7e-2 31 | 32 | 33 | def test_ks_td(krusell_smith_dag): 34 | _, ss, ks_model, unknowns, targets, exogenous = krusell_smith_dag 35 | 36 | T = 30 37 | G = ks_model.solve_jacobian(ss, unknowns, targets, exogenous, T=T) 38 | 39 | for shock_size, tol in [(0.01, 7e-3), (0.1, 0.6)]: 40 | dZ = shock_size * 0.8 ** np.arange(T) 41 | 42 | td_nonlin = ks_model.solve_impulse_nonlinear(ss, unknowns, targets, {"Z": dZ}) 43 | dr_nonlin = 10000 * td_nonlin['r'] 44 | dr_lin = 10000 * G['r']['Z'] @ dZ 45 | 46 | assert np.linalg.norm(dr_nonlin - dr_lin, np.inf) < tol 47 | 48 | 49 | def test_hank_td(one_asset_hank_dag): 50 | _, ss, hank_model, unknowns, targets, exogenous = one_asset_hank_dag 51 | 52 | T = 30 53 | household = hank_model['hh'] 54 | J_ha = household.jacobian(ss=ss, T=T, inputs=['Div', 'Tax', 'r', 'w']) 55 | G = hank_model.solve_jacobian(ss, unknowns, targets, exogenous, T=T, Js={'hh': J_ha}) 56 | 57 | rho_r, sig_r = 0.61, -0.01/4 58 | drstar = sig_r * rho_r ** (np.arange(T)) 59 | 60 | td_nonlin = hank_model.solve_impulse_nonlinear(ss, unknowns, targets, {"rstar": drstar}, Js={'hh': J_ha}) 61 | 62 | dC_nonlin = 100 * td_nonlin['C'] / ss['C'] 63 | dC_lin = 100 * G['C']['rstar'] @ drstar / ss['C'] 64 | 65 | assert np.linalg.norm(dC_nonlin - dC_lin, np.inf) < 3e-3 66 | 67 | 68 | # TODO: needs to compute Jacobian of hetoutput `Chi` 69 | def test_two_asset_td(two_asset_hank_dag): 70 | _, ss, two_asset_model, unknowns, targets, exogenous = two_asset_hank_dag 71 | 72 | T = 30 73 | household = two_asset_model['hh'] 74 | J_ha = household.jacobian(ss=ss, T=T, inputs=['N', 'r', 'ra', 'rb', 'tax', 'w']) 75 | G = two_asset_model.solve_jacobian(ss, unknowns, targets, exogenous, T=T, Js={'hh': J_ha}) 76 | 77 | for shock_size, tol in [(0.1, 3e-4), (1, 2e-2)]: 78 | drstar = shock_size * -0.0025 * 0.6 ** np.arange(T) 79 | 80 | td_nonlin = two_asset_model.solve_impulse_nonlinear(ss, unknowns, targets, {"rstar": drstar}, 81 | Js={'hh': J_ha}) 82 | 83 | dY_nonlin = 100 * td_nonlin['Y'] 84 | dY_lin = 100 * G['Y']['rstar'] @ drstar 85 | 86 | assert np.linalg.norm(dY_nonlin - dY_lin, np.inf) < tol 87 | 88 | 89 | def test_two_asset_solved_v_simple_td(two_asset_hank_dag): 90 | _, ss, two_asset_model, unknowns, targets, exogenous = two_asset_hank_dag 91 | 92 | household = hh.hh.add_hetinputs([two_asset.income, two_asset.make_grids]) 93 | blocks_simple = [household, two_asset.pricing, two_asset.arbitrage, 94 | two_asset.labor, two_asset.investment, two_asset.dividend, 95 | two_asset.taylor, two_asset.fiscal, two_asset.share_value, 96 | two_asset.finance, two_asset.wage, two_asset.union, 97 | two_asset.mkt_clearing] 98 | two_asset_model_simple = combine(blocks_simple, name="Two-Asset HANK w/ SimpleBlocks") 99 | unknowns_simple = ["r", "w", "Y", "pi", "p", "Q", "K"] 100 | targets_simple = ["asset_mkt", "fisher", "wnkpc", "nkpc", "equity", "inv", "val"] 101 | 102 | T = 30 103 | household = two_asset_model['hh'] 104 | J_ha = household.jacobian(ss=ss, T=T, inputs=['N', 'r', 'ra', 'rb', 'tax', 'w']) 105 | G = two_asset_model.solve_jacobian(ss, unknowns, targets, exogenous, T=T, Js={'hh': J_ha}) 106 | G_simple = two_asset_model_simple.solve_jacobian(ss, unknowns_simple, targets_simple, exogenous, T=T, 107 | Js={'hh': J_ha}) 108 | 109 | drstar = -0.0025 * 0.6 ** np.arange(T) 110 | 111 | dY = 100 * G['Y']['rstar'] @ drstar 112 | td_nonlin = two_asset_model.solve_impulse_nonlinear(ss, unknowns, targets, {"rstar": drstar}, 113 | Js={'hh': J_ha}) 114 | dY_nonlin = 100 * (td_nonlin['Y'] - 1) 115 | 116 | dY_simple = 100 * G_simple['Y']['rstar'] @ drstar 117 | td_nonlin_simple = two_asset_model_simple.solve_impulse_nonlinear(ss, 118 | unknowns_simple, targets_simple, 119 | {"rstar": drstar}, Js={'hh': J_ha}) 120 | 121 | dY_nonlin_simple = 100 * (td_nonlin_simple['Y'] - 1) 122 | 123 | assert np.linalg.norm(dY_nonlin - dY_nonlin_simple, np.inf) < 2e-7 124 | assert np.linalg.norm(dY - dY_simple, np.inf) < 0.02 125 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/combined_block.py: -------------------------------------------------------------------------------- 1 | """CombinedBlock class and the combine function to generate it""" 2 | 3 | from .block import Block 4 | from .auxiliary_blocks.jacobiandict_block import JacobianDictBlock 5 | from .support.parent import Parent 6 | from ..classes import ImpulseDict, JacobianDict 7 | from ..utilities.graph import DAG, find_intermediate_inputs 8 | 9 | 10 | def combine(blocks, name="", model_alias=False): 11 | return CombinedBlock(blocks, name=name, model_alias=model_alias) 12 | 13 | 14 | # Useful functional alias 15 | def create_model(blocks, **kwargs): 16 | return combine(blocks, model_alias=True, **kwargs) 17 | 18 | 19 | class CombinedBlock(Block, Parent, DAG): 20 | """A combined `Block` object comprised of several `Block` objects, which topologically sorts them and provides 21 | a set of partial and general equilibrium methods for evaluating their steady state, computes impulse responses, 22 | and calculates Jacobians along the DAG""" 23 | # To users: Do *not* manually change the attributes via assignment. Instantiating a 24 | # CombinedBlock has some automated features that are inferred from initial instantiation but not from 25 | # re-assignment of attributes post-instantiation. 26 | def __init__(self, blocks, name="", model_alias=False, sorted_indices=None, intermediate_inputs=None): 27 | super().__init__() 28 | 29 | blocks_unsorted = [b if isinstance(b, Block) else JacobianDictBlock(b) for b in blocks] 30 | DAG.__init__(self, blocks_unsorted) 31 | 32 | # TODO: deprecate this, use DAG methods instead 33 | self._required = find_intermediate_inputs(blocks) if intermediate_inputs is None else intermediate_inputs 34 | 35 | if not name: 36 | self.name = f"{self.blocks[0].name}_to_{self.blocks[-1].name}_combined" 37 | else: 38 | self.name = name 39 | 40 | # now that it has a name, do Parent initialization 41 | Parent.__init__(self, blocks) 42 | 43 | # If the create_model() is used instead of combine(), we will have __repr__ show this object as a 'Model' 44 | self._model_alias = model_alias 45 | 46 | def __repr__(self): 47 | if self._model_alias: 48 | return f"" 49 | else: 50 | return f"" 51 | 52 | def _steady_state(self, calibration, dissolve, **kwargs): 53 | """Evaluate a partial equilibrium steady state of the CombinedBlock given a `calibration`""" 54 | 55 | ss = calibration.copy() 56 | for block in self.blocks: 57 | # TODO: make this inner_dissolve better, clumsy way to dispatch dissolve only to correct children 58 | inner_dissolve = [k for k in dissolve if self.descendants[k] == block.name] 59 | outputs = block.steady_state(ss, dissolve=inner_dissolve, **kwargs) 60 | ss.update(outputs) 61 | 62 | return ss 63 | 64 | def _impulse_nonlinear(self, ss, inputs, outputs, internals, Js, options, ss_initial): 65 | original_outputs = outputs 66 | outputs = (outputs | self._required) - ss._vector_valued() 67 | 68 | impulses = inputs.copy() 69 | for block in self.blocks: 70 | input_args = {k: v for k, v in impulses.items() if k in block.inputs} 71 | 72 | if input_args or ss_initial is not None: 73 | # If this block is actually perturbed, or we start from different initial ss 74 | # TODO: be more selective about ss_initial here - did any inputs change that matter for this one block? 75 | impulses.update(block.impulse_nonlinear(ss, input_args, outputs & block.outputs, internals, Js, options, ss_initial)) 76 | 77 | return ImpulseDict({k: impulses.toplevel[k] for k in original_outputs if k in impulses.toplevel}, impulses.internals, impulses.T) 78 | 79 | def _impulse_linear(self, ss, inputs, outputs, Js, options): 80 | original_outputs = outputs 81 | outputs = (outputs | self._required) - ss._vector_valued() 82 | 83 | impulses = inputs.copy() 84 | for block in self.blocks: 85 | input_args = {k: v for k, v in impulses.items() if k in block.inputs} 86 | 87 | if input_args: # If this block is actually perturbed 88 | impulses.update(block.impulse_linear(ss, input_args, outputs & block.outputs, Js, options)) 89 | 90 | return ImpulseDict({k: impulses.toplevel[k] for k in original_outputs if k in impulses.toplevel}, T=impulses.T) 91 | 92 | def _partial_jacobians(self, ss, inputs, outputs, T, Js, options): 93 | vector_valued = ss._vector_valued() 94 | inputs = (inputs | self._required) - vector_valued 95 | outputs = (outputs | self._required) - vector_valued 96 | 97 | curlyJs = {} 98 | for block in self.blocks: 99 | curlyJ = block.partial_jacobians(ss, inputs & block.inputs, outputs & block.outputs, T, Js, options) 100 | curlyJs.update(curlyJ) 101 | 102 | return curlyJs 103 | 104 | def _jacobian(self, ss, inputs, outputs, T, Js, options): 105 | Js = self._partial_jacobians(ss, inputs, outputs, T, Js, options) 106 | 107 | original_outputs = outputs 108 | total_Js = JacobianDict.identity(inputs) 109 | 110 | # TODO: horrible, redoing work from partial_jacobians, also need more efficient sifting of intermediates! 111 | vector_valued = ss._vector_valued() 112 | inputs = (inputs | self._required) - vector_valued 113 | outputs = (outputs | self._required) - vector_valued 114 | for block in self.blocks: 115 | if (inputs & block.inputs) and (outputs & block.outputs): 116 | J = block.jacobian(ss, inputs & block.inputs, outputs & block.outputs, T, Js, options) 117 | total_Js.update(J @ total_Js) 118 | 119 | return total_Js[original_outputs & total_Js.outputs, :] 120 | 121 | # Useful type aliases 122 | Model = CombinedBlock 123 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/misc.py: -------------------------------------------------------------------------------- 1 | """Assorted other utilities""" 2 | 3 | import numpy as np 4 | import scipy.linalg 5 | from numba import njit, guvectorize 6 | 7 | 8 | def make_tuple(x): 9 | """If not tuple or list, make into tuple with one element. 10 | 11 | Wrapping with this allows user to write, e.g.: 12 | "return r" rather than "return (r,)" 13 | "policy='a'" rather than "policy=('a',)" 14 | """ 15 | return (x,) if not (isinstance(x, tuple) or isinstance(x, list)) else x 16 | 17 | 18 | def numeric_primitive(instance): 19 | # If it is already a primitive, just return it 20 | if type(instance) in {int, float}: 21 | return instance 22 | elif isinstance(instance, np.ndarray): 23 | if np.issubdtype(instance.dtype, np.number): 24 | return np.array(instance) 25 | else: 26 | raise ValueError(f"The tuple/list argument provided to numeric_primitive has dtype: {instance.dtype}," 27 | f" which is not a valid numeric type.") 28 | elif type(instance) in {tuple, list}: 29 | instance_array = np.asarray(instance) 30 | if np.issubdtype(instance_array.dtype, np.number): 31 | return type(instance)(instance_array) 32 | else: 33 | raise ValueError(f"The tuple/list argument provided to numeric_primitive has dtype: {instance_array.dtype}," 34 | f" which is not a valid numeric type.") 35 | else: 36 | return instance.real if np.isscalar(instance) else instance.base 37 | 38 | 39 | def demean(x): 40 | return x - x.sum()/x.size 41 | 42 | 43 | # simpler aliases for LU factorization and solution 44 | def factor(X): 45 | return scipy.linalg.lu_factor(X) 46 | 47 | 48 | def factored_solve(Z, y): 49 | return scipy.linalg.lu_solve(Z, y) 50 | 51 | 52 | # The below functions are used in steady_state 53 | def unprime(s): 54 | """Given a variable's name as a `str`, check if the variable is a prime, i.e. has "_p" at the end. 55 | If so, return the unprimed version, if not return itself.""" 56 | if s[-2:] == "_p": 57 | return s[:-2] 58 | else: 59 | return s 60 | 61 | 62 | def uncapitalize(s): 63 | return s[0].lower() + s[1:] 64 | 65 | 66 | def list_diff(l1, l2): 67 | """Returns the list that is the "set difference" between l1 and l2 (based on element values)""" 68 | o_list = [] 69 | for k in set(l1) - set(l2): 70 | o_list.append(k) 71 | return o_list 72 | 73 | 74 | def dict_diff(d1, d2): 75 | """Returns the dictionary that is the "set difference" between d1 and d2 (based on keys, not key-value pairs) 76 | E.g. d1 = {"a": 1, "b": 2}, d2 = {"b": 5}, then dict_diff(d1, d2) = {"a": 1} 77 | """ 78 | o_dict = {} 79 | for k in set(d1.keys()) - set(d2.keys()): 80 | o_dict[k] = d1[k] 81 | return o_dict 82 | 83 | 84 | def smart_set(data): 85 | # We want set to construct a single-element set for strings, i.e. ignoring the .iter method of strings 86 | if isinstance(data, str): 87 | return {data} 88 | else: 89 | return set(data) 90 | 91 | 92 | def smart_zip(keys, values): 93 | """For handling the case where keys and values may be scalars""" 94 | if isinstance(values, float): 95 | return zip(keys, [values]) 96 | else: 97 | return zip(keys, values) 98 | 99 | 100 | def smart_zeros(n): 101 | """Return either the float 0. or a np.ndarray of length 0 depending on whether n > 1""" 102 | if n > 1: 103 | return np.zeros(n) 104 | else: 105 | return 0. 106 | 107 | 108 | '''Tools for taste shocks used in discrete choice problems''' 109 | 110 | 111 | def logit(V, scale): 112 | """Logit choice probability of choosing along 0th axis""" 113 | Vnorm = V - V.max(axis=0) 114 | Vexp = np.exp(Vnorm / scale) 115 | P = Vexp / Vexp.sum(axis=0) 116 | return P 117 | 118 | 119 | def logsum(V, scale): 120 | """Logsum formula along 0th axis""" 121 | const = V.max(axis=0) 122 | Vnorm = V - const 123 | EV = const + scale * np.log(np.exp(Vnorm / scale).sum(axis=0)) 124 | return EV 125 | 126 | 127 | def logit_choice(V, scale): 128 | """Logit choice probabilities and logsum along 0th axis""" 129 | const = V.max(axis=0) 130 | Vnorm = V - const 131 | Vexp = np.exp(Vnorm / scale) 132 | Vexpsum = Vexp.sum(axis=0) 133 | 134 | P = Vexp / Vexpsum 135 | EV = const + scale * np.log(Vexpsum) 136 | return P, EV 137 | 138 | 139 | @guvectorize(['void(float64[:], uint32[:], uint32[:])'], '(nA) -> (),()', nopython=True) 140 | def nonconcave(Va, ilower, iupper): 141 | """ 142 | Let V(..., a) be the value function associated with a non-convex dynamic program. `Va` is its derivative with respect to the **single** continuous state variable `a`. 143 | 144 | Find ilower and iupper such that {a_{ilower + 1}, ..., a_{iupper - 1}} is the region where V is non-concave. 145 | 146 | Reference: Fella (2014): A generalized endogenous grid method for non-smooth and non-concave problems 147 | """ 148 | nA = Va.shape[-1] 149 | vmin = np.inf 150 | vmax = -np.inf 151 | # Find vmin & vmax 152 | for ia in range(nA - 1): 153 | if Va[ia + 1] > Va[ia]: 154 | vmin_temp = Va[ia] 155 | vmax_temp = Va[ia + 1] 156 | if vmin_temp < vmin: 157 | vmin = vmin_temp 158 | if vmax_temp > vmax: 159 | vmax = vmax_temp 160 | 161 | # Find ilower 162 | if vmax == -np.inf: 163 | ilower_ = nA 164 | else: 165 | ia = nA 166 | while ia > 0: 167 | if Va[ia] > vmax: 168 | break 169 | ia -= 1 170 | ilower_ = ia 171 | 172 | # Find iupper 173 | if vmin == np.inf: 174 | iupper_ = 0 175 | else: 176 | ia = 0 177 | while ia < nA: 178 | if Va[ia] < vmin: 179 | break 180 | ia += 1 181 | iupper_ = ia 182 | 183 | ilower[:] = ilower_ 184 | iupper[:] = iupper_ 185 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/discretize.py: -------------------------------------------------------------------------------- 1 | """Grids and Markov chains""" 2 | 3 | import numpy as np 4 | from scipy.stats import norm 5 | 6 | 7 | def asset_grid(amin, amax, n): 8 | # find maximum ubar of uniform grid corresponding to desired maximum amax of asset grid 9 | ubar = np.log(1 + np.log(1 + amax - amin)) 10 | 11 | # make uniform grid 12 | u_grid = np.linspace(0, ubar, n) 13 | 14 | # double-exponentiate uniform grid and add amin to get grid from amin to amax 15 | return amin + np.exp(np.exp(u_grid) - 1) - 1 16 | 17 | 18 | def agrid(amax, n, amin=0): 19 | """Create grid between amin and amax that is equidistant in logs.""" 20 | pivot = np.abs(amin) + 0.25 21 | a_grid = np.geomspace(amin + pivot, amax + pivot, n) - pivot 22 | a_grid[0] = amin # make sure *exactly* equal to amin 23 | return a_grid 24 | 25 | 26 | # TODO: Temporarily include the old way of constructing grids from ikc_old for comparability of results 27 | def agrid_old(amax, N, amin=0, frac=1/25): 28 | """crappy discretization method we've been using, generates N point 29 | log-spaced grid between bmin and bmax, choosing pivot such that 'frac' of 30 | total log space between log(1+amin) and log(1+amax) beneath it""" 31 | apivot = (1+amin)**(1-frac)*(1+amax)**frac - 1 32 | a = np.geomspace(amin+apivot,amax+apivot,N) - apivot 33 | a[0] = amin 34 | return a 35 | 36 | 37 | def nonlinspace(amax, n, phi, amin=0): 38 | """Create grid between amin and amax. phi=1 is equidistant, phi>1 dense near amin. Extra flexibility may be useful in non-convex problems in which policy functions have nonlinear (even non-monotonic) sections far from the borrowing limit.""" 39 | a_grid = np.zeros(n) 40 | a_grid[0] = amin 41 | for i in range(1, n): 42 | a_grid[i] = a_grid[i-1] + (amax - a_grid[i-1]) / (n-i)**phi 43 | return a_grid 44 | 45 | 46 | def stationary(Pi, pi_seed=None, tol=1E-11, maxit=10_000): 47 | """Find invariant distribution of a Markov chain by iteration.""" 48 | if pi_seed is None: 49 | pi = np.ones(Pi.shape[0]) / Pi.shape[0] 50 | else: 51 | pi = pi_seed 52 | 53 | for it in range(maxit): 54 | pi_new = pi @ Pi 55 | if np.max(np.abs(pi_new - pi)) < tol: 56 | break 57 | pi = pi_new 58 | else: 59 | raise ValueError(f'No convergence after {maxit} forward iterations!') 60 | pi = pi_new 61 | 62 | return pi 63 | 64 | 65 | def mean(x, pi): 66 | """Mean of discretized random variable with support x and probability mass function pi.""" 67 | return np.sum(pi * x) 68 | 69 | 70 | def variance(x, pi): 71 | """Variance of discretized random variable with support x and probability mass function pi.""" 72 | return np.sum(pi * (x - np.sum(pi * x)) ** 2) 73 | 74 | 75 | def std(x, pi): 76 | """Standard deviation of discretized random variable with support x and probability mass function pi.""" 77 | return np.sqrt(variance(x, pi)) 78 | 79 | 80 | def cov(x, y, pi): 81 | """Covariance of two discretized random variables with supports x and y common probability mass function pi.""" 82 | return np.sum(pi * (x - mean(x, pi)) * (y - mean(y, pi))) 83 | 84 | 85 | def corr(x, y, pi): 86 | """Correlation of two discretized random variables with supports x and y common probability mass function pi.""" 87 | return cov(x, y, pi) / (std(x, pi) * std(y, pi)) 88 | 89 | 90 | def markov_tauchen(rho, sigma, N=7, m=3, normalize=True): 91 | """Tauchen method discretizing AR(1) s_t = rho*s_(t-1) + eps_t. 92 | 93 | Parameters 94 | ---------- 95 | rho : scalar, persistence 96 | sigma : scalar, unconditional sd of s_t 97 | N : int, number of states in discretized Markov process 98 | m : scalar, discretized s goes from approx -m*sigma to m*sigma 99 | 100 | Returns 101 | ---------- 102 | y : array (N), states proportional to exp(s) s.t. E[y] = 1 103 | pi : array (N), stationary distribution of discretized process 104 | Pi : array (N*N), Markov matrix for discretized process 105 | """ 106 | 107 | # make normalized grid, start with cross-sectional sd of 1 108 | s = np.linspace(-m, m, N) 109 | ds = s[1] - s[0] 110 | sd_innov = np.sqrt(1 - rho ** 2) 111 | 112 | # standard Tauchen method to generate Pi given N and m 113 | Pi = np.empty((N, N)) 114 | Pi[:, 0] = norm.cdf(s[0] - rho * s + ds / 2, scale=sd_innov) 115 | Pi[:, -1] = 1 - norm.cdf(s[-1] - rho * s - ds / 2, scale=sd_innov) 116 | for j in range(1, N - 1): 117 | Pi[:, j] = (norm.cdf(s[j] - rho * s + ds / 2, scale=sd_innov) - 118 | norm.cdf(s[j] - rho * s - ds / 2, scale=sd_innov)) 119 | 120 | # invariant distribution and scaling 121 | pi = stationary(Pi) 122 | s *= (sigma / np.sqrt(variance(s, pi))) 123 | if normalize: 124 | y = np.exp(s) / np.sum(pi * np.exp(s)) 125 | else: 126 | y = s 127 | 128 | return y, pi, Pi 129 | 130 | 131 | def markov_rouwenhorst(rho, sigma, N=7): 132 | """Rouwenhorst method analog to markov_tauchen""" 133 | 134 | # Explicitly typecast N as an integer, since when the grid constructor functions 135 | # (e.g. the function that makes a_grid) are implemented as blocks, they interpret the integer-valued calibration 136 | # as a float. 137 | N = int(N) 138 | 139 | # parametrize Rouwenhorst for n=2 140 | p = (1 + rho) / 2 141 | Pi = np.array([[p, 1 - p], [1 - p, p]]) 142 | 143 | # implement recursion to build from n=3 to n=N 144 | for n in range(3, N + 1): 145 | P1, P2, P3, P4 = (np.zeros((n, n)) for _ in range(4)) 146 | P1[:-1, :-1] = p * Pi 147 | P2[:-1, 1:] = (1 - p) * Pi 148 | P3[1:, :-1] = (1 - p) * Pi 149 | P4[1:, 1:] = p * Pi 150 | Pi = P1 + P2 + P3 + P4 151 | Pi[1:-1] /= 2 152 | 153 | # invariant distribution and scaling 154 | pi = stationary(Pi) 155 | s = np.linspace(-1, 1, N) 156 | s *= (sigma / np.sqrt(variance(s, pi))) 157 | y = np.exp(s) / np.sum(pi * np.exp(s)) 158 | 159 | return y, pi, Pi 160 | -------------------------------------------------------------------------------- /tests/base/test_stage_block.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from sequence_jacobian.blocks.stage_block import StageBlock 4 | from sequence_jacobian.hetblocks.hh_sim import hh, hh_init 5 | from sequence_jacobian.blocks.support.stages import Continuous1D, ExogenousMaker 6 | from sequence_jacobian import interpolate, grids, misc, combine 7 | from sequence_jacobian.classes import ImpulseDict 8 | 9 | def make_grids(rho_e, sd_e, nE, amin, amax, nA): 10 | e_grid, e_dist, Pi_ss = grids.markov_rouwenhorst(rho=rho_e, sigma=sd_e, N=nE) 11 | a_grid = grids.agrid(amin=amin, amax=amax, n=nA) 12 | return e_grid, e_dist, Pi_ss, a_grid 13 | 14 | def alter_Pi(Pi_ss, shift): 15 | Pi = Pi_ss.copy() 16 | Pi[:, 0] -= shift 17 | Pi[:, -1] += shift 18 | return Pi 19 | 20 | def income(atw, N, e_grid, transfer): 21 | y = atw * N * e_grid + transfer 22 | return y 23 | 24 | # copy original household hetblock but get rid of _p on Va 25 | def household_new(Va, a_grid, y, r, beta, eis): 26 | uc_nextgrid = beta * Va 27 | c_nextgrid = uc_nextgrid ** (-eis) 28 | coh = (1 + r) * a_grid[np.newaxis, :] + y[:, np.newaxis] 29 | a = interpolate.interpolate_y(c_nextgrid + a_grid, coh, a_grid) 30 | misc.setmin(a, a_grid[0]) 31 | c = coh - a 32 | Va = (1 + r) * c ** (-1 / eis) 33 | return Va, a, c 34 | 35 | def marginal_utility(c, eis): 36 | uc = c ** (-1 / eis) 37 | return uc 38 | 39 | #het_stage = Continuous1D(backward='Va', policy='a', f=household_new, name='stage1') 40 | het_stage = Continuous1D(backward='Va', policy='a', f=household_new, name='stage1', hetoutputs=[marginal_utility]) 41 | hh2 = StageBlock([ExogenousMaker('Pi', 0, 'stage0'), het_stage], name='hh', 42 | backward_init=hh_init, hetinputs=(make_grids, income, alter_Pi)) 43 | 44 | def test_equivalence(): 45 | hh1 = hh.add_hetinputs([make_grids, income, alter_Pi]).add_hetoutputs([marginal_utility]) 46 | calibration = {'r': 0.004, 'eis': 0.5, 'rho_e': 0.91, 'sd_e': 0.92, 'nE': 3, 47 | 'amin': 0.0, 'amax': 200, 'nA': 100, 'transfer': 0.143, 'N': 1, 48 | 'atw': 1, 'beta': 0.97, 'shift': 0} 49 | ss1 = hh1.steady_state(calibration) 50 | ss2 = hh2.steady_state(calibration) 51 | 52 | # test steady-state equivalence 53 | assert np.isclose(ss1['A'], ss2['A']) 54 | assert np.isclose(ss1['C'], ss2['C']) 55 | assert np.allclose(ss1.internals['hh']['Dbeg'], ss2.internals['hh']['stage0']['D']) 56 | assert np.allclose(ss1.internals['hh']['a'], ss2.internals['hh']['stage1']['a']) 57 | assert np.allclose(ss1.internals['hh']['c'], ss2.internals['hh']['stage1']['c']) 58 | assert np.allclose(ss1.internals['hh']['Va'], ss2.internals['hh']['stage0']['Va']) 59 | 60 | # find Jacobians... 61 | inputs = ['r', 'atw', 'shift'] 62 | outputs = ['A', 'C', 'UC'] 63 | T = 200 64 | J1 = hh1.jacobian(ss1, inputs, outputs, T) 65 | J2 = hh2.jacobian(ss2, inputs, outputs, T) 66 | 67 | # test Jacobian equivalence 68 | for i in inputs: 69 | for o in outputs: 70 | if o == 'UC': 71 | # not sure why numerical differences somewhat larger here? 72 | assert np.max(np.abs(J1[o, i] - J2[o, i])) < 2E-4 73 | else: 74 | assert np.allclose(J1[o, i], J2[o, i]) 75 | 76 | # impulse linear 77 | shock = ImpulseDict({'r': 0.5 ** np.arange(20)}) 78 | td_lin1 = hh1.impulse_linear(ss1, shock, outputs=['C', 'UC']) 79 | td_lin2 = hh2.impulse_linear(ss2, shock, outputs=['C', 'UC']) 80 | assert np.allclose(td_lin1['C'], td_lin2['C']) 81 | assert np.max(np.abs(td_lin1['UC'] - td_lin2['UC'])) < 2E-4 82 | 83 | # impulse nonlinear 84 | td_nonlin1 = hh1.impulse_nonlinear(ss1, shock * 1E-4, outputs=['C', 'UC']) 85 | td_nonlin2 = hh2.impulse_nonlinear(ss2, shock * 1E-4, outputs=['C', 'UC']) 86 | assert np.allclose(td_nonlin1['C'], td_nonlin2['C']) 87 | assert np.allclose(td_nonlin1['UC'], td_nonlin2['UC']) 88 | 89 | 90 | def test_remap(): 91 | # hetblock 92 | hh1 = hh.add_hetinputs([make_grids, income, alter_Pi]) 93 | hh1_men = hh1.remap({k: k + '_men' for k in hh1.outputs | ['sd_e']}).rename('men') 94 | hh1_women = hh1.remap({k: k + '_women' for k in hh1.outputs | ['sd_e']}).rename('women') 95 | hh1_all = combine([hh1_men, hh1_women]) 96 | 97 | # stageblock 98 | hh2_men = hh2.remap({k: k + '_men' for k in hh2.outputs| ['sd_e']}).rename('men') 99 | hh2_women = hh2.remap({k: k + '_women' for k in hh2.outputs | ['sd_e']}).rename('women') 100 | hh2_all = combine([hh2_men, hh2_women]) 101 | 102 | # steady state 103 | calibration = {'sd_e_men': 0.92, 'sd_e_women': 0.82, 104 | 'r': 0.004, 'eis': 0.5, 'rho_e': 0.91, 'nE': 3, 105 | 'amin': 0.0, 'amax': 200, 'nA': 100, 'transfer': 0.143, 'N': 1, 106 | 'atw': 1, 'beta': 0.97, 'shift': 0} 107 | 108 | ss1 = hh1_all.steady_state(calibration) 109 | ss2 = hh2_all.steady_state(calibration) 110 | 111 | # test steady-state equivalence 112 | assert np.isclose(ss1['A_men'], ss2['A_men']) 113 | assert np.isclose(ss1['C_women'], ss2['C_women']) 114 | 115 | # find Jacobians... 116 | inputs = ['r', 'atw', 'shift'] 117 | outputs = ['A_men', 'A_women'] 118 | T = 100 119 | J1 = hh1_all.jacobian(ss1, inputs, outputs, T) 120 | J2 = hh2_all.jacobian(ss2, inputs, outputs, T) 121 | 122 | # test Jacobian equivalence 123 | for i in inputs: 124 | for o in outputs: 125 | assert np.allclose(J1[o, i], J2[o, i]) 126 | 127 | # impulse linear 128 | shock = ImpulseDict({'r': 0.5 ** np.arange(20)}) 129 | td_lin1 = hh1_all.impulse_linear(ss1, shock, outputs=['C_men', 'C_women']) 130 | td_lin2 = hh2_all.impulse_linear(ss2, shock, outputs=['C_men', 'C_women']) 131 | assert np.allclose(td_lin1['C_women'], td_lin2['C_women']) 132 | 133 | # impulse nonlinear 134 | td_nonlin1 = hh1_all.impulse_nonlinear(ss1, shock * 1E-4, outputs=['C_men']) 135 | td_nonlin2 = hh2_all.impulse_nonlinear(ss2, shock * 1E-4, outputs=['C_men']) 136 | assert np.allclose(td_nonlin1['C_men'], td_nonlin2['C_men']) 137 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/support/law_of_motion.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from . import het_compiled 3 | from ...utilities.interpolate import interpolate_coord_robust, interpolate_coord 4 | from ...utilities.multidim import batch_multiply_ith_dimension, multiply_ith_dimension 5 | from typing import Optional, Sequence, Any, List, Tuple, Union 6 | import copy 7 | 8 | class LawOfMotion: 9 | """Abstract class representing a matrix that operates on state space. 10 | Rather than giant Ns*Ns matrix (even if sparse), some other representation 11 | almost always desirable; such representations are subclasses of this.""" 12 | 13 | def __matmul__(self, X): 14 | pass 15 | 16 | @property 17 | def T(self): 18 | pass 19 | 20 | 21 | def lottery_1d(a, a_grid, monotonic=False): 22 | if not monotonic: 23 | return PolicyLottery1D(*interpolate_coord_robust(a_grid, a), a_grid) 24 | else: 25 | return PolicyLottery1D(*interpolate_coord(a_grid, a), a_grid) 26 | 27 | 28 | class PolicyLottery1D(LawOfMotion): 29 | # TODO: always operates on final dimension, make more general! 30 | def __init__(self, i, pi, grid, forward=True): 31 | # flatten non-policy dimensions into one because that's what methods accept 32 | self.i = i.reshape((-1,) + grid.shape) 33 | self.flatshape = self.i.shape 34 | 35 | self.pi = pi.reshape(self.flatshape) 36 | 37 | # but store original shape so we can convert all outputs to it 38 | self.shape = i.shape 39 | self.grid = grid 40 | 41 | # also store shape of the endogenous grid itself 42 | self.endog_shape = self.shape[-1:] 43 | 44 | self.forward = forward 45 | 46 | @property 47 | def T(self): 48 | newself = copy.copy(self) 49 | newself.forward = not self.forward 50 | return newself 51 | 52 | def __matmul__(self, X): 53 | if self.forward: 54 | return het_compiled.forward_policy_1d(X.reshape(self.flatshape), self.i, self.pi).reshape(self.shape) 55 | else: 56 | return het_compiled.expectation_policy_1d(X.reshape(self.flatshape), self.i, self.pi).reshape(self.shape) 57 | 58 | 59 | class ShockedPolicyLottery1D(PolicyLottery1D): 60 | def __matmul__(self, X): 61 | if self.forward: 62 | return het_compiled.forward_policy_shock_1d(X.reshape(self.flatshape), self.i, self.pi).reshape(self.shape) 63 | else: 64 | raise NotImplementedError 65 | 66 | 67 | def lottery_2d(a, b, a_grid, b_grid, monotonic=False): 68 | if not monotonic: 69 | return PolicyLottery2D(*interpolate_coord_robust(a_grid, a), 70 | *interpolate_coord_robust(b_grid, b), a_grid, b_grid) 71 | if monotonic: 72 | # right now we have no monotonic 2D examples, so this shouldn't be called 73 | return PolicyLottery2D(*interpolate_coord(a_grid, a), 74 | *interpolate_coord(b_grid, b), a_grid, b_grid) 75 | 76 | 77 | class PolicyLottery2D(LawOfMotion): 78 | def __init__(self, i1, pi1, i2, pi2, grid1, grid2, forward=True): 79 | # flatten non-policy dimensions into one because that's what methods accept 80 | self.i1 = i1.reshape((-1,) + grid1.shape + grid2.shape) 81 | self.flatshape = self.i1.shape 82 | 83 | self.i2 = i2.reshape(self.flatshape) 84 | self.pi1 = pi1.reshape(self.flatshape) 85 | self.pi2 = pi2.reshape(self.flatshape) 86 | 87 | # but store original shape so we can convert all outputs to it 88 | self.shape = i1.shape 89 | self.grid1 = grid1 90 | self.grid2 = grid2 91 | 92 | # also store shape of the endogenous grid itself 93 | self.endog_shape = self.shape[-2:] 94 | 95 | self.forward = forward 96 | 97 | @property 98 | def T(self): 99 | newself = copy.copy(self) 100 | newself.forward = not self.forward 101 | return newself 102 | 103 | def __matmul__(self, X): 104 | if self.forward: 105 | return het_compiled.forward_policy_2d(X.reshape(self.flatshape), self.i1, self.i2, 106 | self.pi1, self.pi2).reshape(self.shape) 107 | else: 108 | return het_compiled.expectation_policy_2d(X.reshape(self.flatshape), self.i1, self.i2, 109 | self.pi1, self.pi2).reshape(self.shape) 110 | 111 | 112 | class ShockedPolicyLottery2D(PolicyLottery2D): 113 | def __matmul__(self, X): 114 | if self.forward: 115 | return het_compiled.forward_policy_shock_2d(X.reshape(self.flatshape), self.i, self.pi).reshape(self.shape) 116 | else: 117 | raise NotImplementedError 118 | 119 | 120 | class Markov(LawOfMotion): 121 | def __init__(self, Pi, i): 122 | self.Pi = Pi 123 | self.i = i 124 | 125 | @property 126 | def T(self): 127 | newself = copy.copy(self) 128 | newself.Pi = newself.Pi.T 129 | if isinstance(newself.Pi, np.ndarray): 130 | # optimizing: copy to get right order in memory 131 | newself.Pi = newself.Pi.copy() 132 | return newself 133 | 134 | def __matmul__(self, X): 135 | return multiply_ith_dimension(self.Pi, self.i, X) 136 | 137 | 138 | class DiscreteChoice(LawOfMotion): 139 | def __init__(self, P, i): 140 | self.P = P # choice prob P(d|...s_i...), 0 for unavailable choices 141 | self.i = i # dimension of state space that will be updated 142 | 143 | # cache "transposed" version of this, since we'll always need both! 144 | self.forward = True 145 | self.P_T = P.swapaxes(0, 1+self.i).copy() 146 | 147 | @property 148 | def T(self): 149 | newself = copy.copy(self) 150 | newself.forward = not self.forward 151 | return newself 152 | 153 | def __matmul__(self, X): 154 | if self.forward: 155 | return batch_multiply_ith_dimension(self.P, self.i, X) 156 | else: 157 | return batch_multiply_ith_dimension(self.P_T, self.i, X) 158 | -------------------------------------------------------------------------------- /tests/base/test_workflow.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from sequence_jacobian import simple, solved, create_model, markov_rouwenhorst, agrid 3 | from sequence_jacobian.classes.impulse_dict import ImpulseDict 4 | from sequence_jacobian.hetblocks.hh_sim import hh 5 | 6 | 7 | '''Part 1: Household block''' 8 | 9 | def make_grids(rho_e, sd_e, nE, amin, amax, nA): 10 | e_grid, e_dist, Pi = markov_rouwenhorst(rho=rho_e, sigma=sd_e, N=nE) 11 | a_grid = agrid(amin=amin, amax=amax, n=nA) 12 | return e_grid, e_dist, Pi, a_grid 13 | 14 | 15 | def income(atw, N, e_grid, transfer): 16 | y = atw * N * e_grid + transfer 17 | return y 18 | 19 | 20 | def get_mpcs(c, a, a_grid, r): 21 | mpcs_ = np.empty_like(c) 22 | post_return = (1 + r) * a_grid 23 | mpcs_[:, 1:-1] = (c[:, 2:] - c[:, 0:-2]) / (post_return[2:] - post_return[:-2]) 24 | mpcs_[:, 0] = (c[:, 1] - c[:, 0]) / (post_return[1] - post_return[0]) 25 | mpcs_[:, -1] = (c[:, -1] - c[:, -2]) / (post_return[-1] - post_return[-2]) 26 | mpcs_[a == a_grid[0]] = 1 27 | return mpcs_ 28 | 29 | 30 | def mpcs(c, a, a_grid, r): 31 | mpc = get_mpcs(c, a, a_grid, r) 32 | return mpc 33 | 34 | 35 | def weighted_uc(c, e_grid, eis): 36 | uce = c ** (-1 / eis) * e_grid[:, np.newaxis] 37 | return uce 38 | 39 | 40 | '''Part 2: rest of the model''' 41 | 42 | @solved(unknowns={'C': 1.0, 'A': 1.0}, targets=['euler', 'budget_constraint'], solver='broyden_custom') 43 | def household_ra(C, A, r, atw, N, transfer, beta, eis): 44 | euler = beta * (1 + r(1)) * C(1) ** (-1 / eis) - C ** (-1 / eis) 45 | budget_constraint = (1 + r) * A(-1) + atw * N + transfer - C - A 46 | UCE = C ** (-1 / eis) 47 | return euler, budget_constraint, UCE 48 | 49 | 50 | @simple 51 | def firm(N, Z): 52 | Y = Z * N 53 | w = Z 54 | return Y, w 55 | 56 | 57 | @simple 58 | def union(UCE, tau, w, N, pi, muw, kappaw, nu, vphi, beta): 59 | wnkpc = kappaw * N * (vphi * N ** nu - (1 - tau) * w * UCE / muw) + \ 60 | beta * (1 + pi(+1)).apply(np.log) - (1 + pi).apply(np.log) 61 | return wnkpc 62 | 63 | 64 | @solved(unknowns={'B': (0.0, 10.0)}, targets=['B_rule'], solver='brentq') 65 | def fiscal(B, G, r, w, N, transfer, rho_B): 66 | B_rule = B.ss + rho_B * (B(-1) - B.ss + G - G.ss) - B 67 | rev = (1 + r) * B(-1) + G + transfer - B # revenue to be raised 68 | tau = rev / (w * N) 69 | atw = (1 - tau) * w 70 | return B_rule, rev, tau, atw 71 | 72 | 73 | # Use this to test zero impulse once we have it 74 | # @simple 75 | # def real_bonds(r): 76 | # rb = r 77 | # return rb 78 | 79 | 80 | @simple 81 | def mkt_clearing(A, B, C, G, Y): 82 | asset_mkt = A - B 83 | goods_mkt = C + G - Y 84 | return asset_mkt, goods_mkt 85 | 86 | 87 | '''Part 3: Helper blocks''' 88 | 89 | @simple 90 | def household_ra_ss(r, B, tau, w, N, transfer, eis): 91 | beta = 1 / (1 + r) 92 | A = B 93 | C = r * A + (1 - tau) * w * N + transfer 94 | UCE = C ** (-1 / eis) 95 | return beta, A, C, UCE 96 | 97 | 98 | @simple 99 | def union_ss(atw, UCE, muw, N, nu, kappaw, beta, pi): 100 | vphi = atw * UCE / (muw * N ** nu) 101 | wnkpc = kappaw * N * (vphi * N ** nu - atw * UCE / muw) + \ 102 | beta * (1 + pi(+1)).apply(np.log) - (1 + pi).apply(np.log) 103 | return wnkpc, vphi 104 | 105 | 106 | '''Tests''' 107 | 108 | def test_all(): 109 | # Assemble HA block (want to test nesting) 110 | household_ha = hh.add_hetinputs([make_grids, income]) 111 | household_ha = household_ha.add_hetoutputs([mpcs, weighted_uc]).rename('household_ha') 112 | 113 | # Assemble DAG (for transition dynamics) 114 | dag = {} 115 | common_blocks = [firm, union, fiscal, mkt_clearing] 116 | dag['ha'] = create_model([household_ha] + common_blocks, name='HANK') 117 | dag['ra'] = create_model([household_ra] + common_blocks, name='RANK') 118 | unknowns = ['N', 'pi'] 119 | targets = ['asset_mkt', 'wnkpc'] 120 | 121 | # Solve steady state 122 | calibration = {'N': 1.0, 'Z': 1.0, 'r': 0.005, 'pi': 0.0, 'eis': 0.5, 'nu': 0.5, 123 | 'rho_e': 0.91, 'sd_e': 0.92, 'nE': 3, 'amin': 0.0, 'amax': 200, 124 | 'nA': 100, 'kappaw': 0.1, 'muw': 1.2, 'transfer': 0.143, 'rho_B': 0.9} 125 | 126 | ss = {} 127 | # Constructing ss-dag manually works just fine 128 | dag_ss = {} 129 | dag_ss['ha'] = create_model([household_ha, union_ss, firm, fiscal, mkt_clearing]) 130 | ss['ha'] = dag_ss['ha'].solve_steady_state(calibration, dissolve=['fiscal'], solver='hybr', 131 | unknowns={'beta': 0.96, 'B': 3.0, 'G': 0.2}, 132 | targets={'asset_mkt': 0.0, 'MPC': 0.25, 'tau': 0.334}) 133 | assert np.isclose(ss['ha']['goods_mkt'], 0.0) 134 | assert np.isclose(ss['ha']['asset_mkt'], 0.0) 135 | assert np.isclose(ss['ha']['wnkpc'], 0.0) 136 | 137 | dag_ss['ra'] = create_model([household_ra_ss, union_ss, firm, fiscal, mkt_clearing]) 138 | ss['ra'] = dag_ss['ra'].steady_state(ss['ha'], dissolve=['fiscal']) 139 | assert np.isclose(ss['ra']['goods_mkt'], 0.0) 140 | assert np.isclose(ss['ra']['asset_mkt'], 0.0) 141 | assert np.isclose(ss['ra']['wnkpc'], 0.0) 142 | 143 | # Precompute HA Jacobian 144 | Js = {'ra': {}, 'ha': {}} 145 | Js['ha']['household_ha'] = household_ha.jacobian(ss['ha'], 146 | inputs=['N', 'atw', 'r', 'transfer'], outputs=['C', 'A', 'UCE'], T=300) 147 | 148 | # Linear impulse responses from Jacobian vs directly 149 | shock = ImpulseDict({'G': 0.9 ** np.arange(300)}) 150 | G, td_lin1, td_lin2 = dict(), dict(), dict() 151 | for k in ['ra', 'ha']: 152 | G[k] = dag[k].solve_jacobian(ss[k], unknowns, targets, inputs=['G'], T=300, Js=Js[k]) 153 | td_lin1[k] = G[k] @ shock 154 | td_lin2[k] = dag[k].solve_impulse_linear(ss[k], unknowns, targets, shock, Js=Js[k]) 155 | assert all(np.allclose(td_lin1[k][i], td_lin2[k][i]) for i in td_lin1[k]) 156 | 157 | # Nonlinear vs linear impulses (sneak in test of ss_initial here too) 158 | td_nonlin = dag['ha'].solve_impulse_nonlinear(ss['ha'], unknowns, targets, inputs=shock*1E-2, 159 | Js=Js, internals=['household_ha'], ss_initial=ss['ha']) 160 | assert np.max(np.abs(td_nonlin['goods_mkt'])) < 1E-8 161 | 162 | # See if D change matches up with aggregate assets 163 | td_nonlin_lvl = td_nonlin + ss['ha'] 164 | td_A = np.sum(td_nonlin_lvl.internals['household_ha']['a'] * td_nonlin_lvl.internals['household_ha']['D'], axis=(1, 2)) 165 | assert np.allclose(td_A - ss['ha']['A'], td_nonlin['A']) 166 | -------------------------------------------------------------------------------- /src/sequence_jacobian/examples/two_asset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from sequence_jacobian import simple, solved, combine, create_model, grids, hetblocks 4 | hh = hetblocks.hh_twoasset.hh 5 | 6 | 7 | '''Part 1: Blocks''' 8 | 9 | @simple 10 | def pricing(pi, mc, r, Y, kappap, mup): 11 | nkpc = kappap * (mc - 1 / mup) + Y(+1) / Y * (1 + pi(+1)).apply(np.log) \ 12 | / (1 + r(+1)) - (1 + pi).apply(np.log) 13 | return nkpc 14 | 15 | 16 | @simple 17 | def arbitrage(div, p, r): 18 | equity = div(+1) + p(+1) - p * (1 + r(+1)) 19 | return equity 20 | 21 | 22 | @simple 23 | def labor(Y, w, K, Z, alpha): 24 | N = (Y / Z / K(-1) ** alpha) ** (1 / (1 - alpha)) 25 | mc = w * N / (1 - alpha) / Y 26 | return N, mc 27 | 28 | 29 | @simple 30 | def investment(Q, K, r, N, mc, Z, delta, epsI, alpha): 31 | inv = (K / K(-1) - 1) / (delta * epsI) + 1 - Q 32 | val = alpha * Z(+1) * (N(+1) / K) ** (1 - alpha) * mc(+1) -\ 33 | (K(+1) / K - (1 - delta) + (K(+1) / K - 1) ** 2 / (2 * delta * epsI)) +\ 34 | K(+1) / K * Q(+1) - (1 + r(+1)) * Q 35 | return inv, val 36 | 37 | 38 | @simple 39 | def dividend(Y, w, N, K, pi, mup, kappap, delta, epsI): 40 | psip = mup / (mup - 1) / 2 / kappap * (1 + pi).apply(np.log) ** 2 * Y 41 | k_adjust = K(-1) * (K / K(-1) - 1) ** 2 / (2 * delta * epsI) 42 | I = K - (1 - delta) * K(-1) + k_adjust 43 | div = Y - w * N - I - psip 44 | return psip, I, div 45 | 46 | 47 | @simple 48 | def taylor(rstar, pi, phi): 49 | i = rstar + phi * pi 50 | return i 51 | 52 | 53 | @simple 54 | def fiscal(r, w, N, G, Bg): 55 | tax = (r * Bg + G) / w / N 56 | return tax 57 | 58 | 59 | @simple 60 | def finance(i, p, pi, r, div, omega, pshare): 61 | rb = r - omega 62 | ra = pshare(-1) * (div + p) / p(-1) + (1 - pshare(-1)) * (1 + r) - 1 63 | fisher = 1 + i(-1) - (1 + r) * (1 + pi) 64 | return rb, ra, fisher 65 | 66 | 67 | @simple 68 | def wage(pi, w): 69 | piw = (1 + pi) * w / w(-1) - 1 70 | return piw 71 | 72 | 73 | @simple 74 | def union(piw, N, tax, w, UCE, kappaw, muw, vphi, frisch, beta): 75 | wnkpc = kappaw * (vphi * N ** (1 + 1 / frisch) - (1 - tax) * w * N * UCE / muw) + beta * \ 76 | (1 + piw(+1)).apply(np.log) - (1 + piw).apply(np.log) 77 | return wnkpc 78 | 79 | 80 | @simple 81 | def mkt_clearing(p, A, B, Bg, C, I, G, CHI, psip, omega, Y): 82 | wealth = A + B 83 | asset_mkt = p + Bg - wealth 84 | goods_mkt = C + I + G + CHI + psip + omega * B - Y 85 | return asset_mkt, wealth, goods_mkt 86 | 87 | 88 | @simple 89 | def share_value(p, tot_wealth, Bh): 90 | pshare = p / (tot_wealth - Bh) 91 | return pshare 92 | 93 | 94 | @solved(unknowns={'pi': (-0.1, 0.1)}, targets=['nkpc'], solver="brentq") 95 | def pricing_solved(pi, mc, r, Y, kappap, mup): 96 | nkpc = kappap * (mc - 1 / mup) + Y(+1) / Y * (1 + pi(+1)).apply(np.log) / \ 97 | (1 + r(+1)) - (1 + pi).apply(np.log) 98 | return nkpc 99 | 100 | 101 | @solved(unknowns={'p': (5, 15)}, targets=['equity'], solver="brentq") 102 | def arbitrage_solved(div, p, r): 103 | equity = div(+1) + p(+1) - p * (1 + r(+1)) 104 | return equity 105 | 106 | 107 | @simple 108 | def partial_ss(Y, N, K, r, tot_wealth, Bg, delta): 109 | """Solves for (mup, alpha, Z, w) to hit (tot_wealth, Y, K, pi).""" 110 | # 1. Solve for markup to hit total wealth 111 | p = tot_wealth - Bg 112 | mc = 1 - r * (p - K) / Y 113 | mup = 1 / mc 114 | 115 | # 2. Solve for capital share to hit K 116 | alpha = (r + delta) * K / Y / mc 117 | 118 | # 3. Solve for TFP to hit Y 119 | Z = Y * K ** (-alpha) * N ** (alpha - 1) 120 | 121 | # 4. Solve for w such that piw = 0 122 | w = mc * (1 - alpha) * Y / N 123 | 124 | return p, mc, mup, alpha, Z, w 125 | 126 | 127 | @simple 128 | def union_ss(tax, w, UCE, N, muw, frisch): 129 | """Solves for (vphi) to hit (wnkpc).""" 130 | vphi = (1 - tax) * w * UCE / muw / N ** (1 + 1 / frisch) 131 | wnkpc = vphi * N ** (1 + 1 / frisch) - (1 - tax) * w * UCE / muw 132 | return vphi, wnkpc 133 | 134 | 135 | '''Part 2: Embed HA block''' 136 | 137 | def make_grids(bmax, amax, kmax, nB, nA, nK, nZ, rho_z, sigma_z): 138 | b_grid = grids.agrid(amax=bmax, n=nB) 139 | a_grid = grids.agrid(amax=amax, n=nA) 140 | k_grid = grids.agrid(amax=kmax, n=nK)[::-1].copy() 141 | e_grid, _, Pi = grids.markov_rouwenhorst(rho=rho_z, sigma=sigma_z, N=nZ) 142 | return b_grid, a_grid, k_grid, e_grid, Pi 143 | 144 | 145 | def income(e_grid, tax, w, N): 146 | z_grid = (1 - tax) * w * N * e_grid 147 | return z_grid 148 | 149 | 150 | '''Part 3: DAG''' 151 | 152 | def dag(): 153 | # Combine Blocks 154 | household = hh.add_hetinputs([income, make_grids]) 155 | production = combine([labor, investment]) 156 | production_solved = production.solved(unknowns={'Q': 1., 'K': 10.}, 157 | targets=['inv', 'val'], solver='broyden_custom') 158 | blocks = [household, pricing_solved, arbitrage_solved, production_solved, 159 | dividend, taylor, fiscal, share_value, finance, wage, union, mkt_clearing] 160 | two_asset_model = create_model(blocks, name='Two-Asset HANK') 161 | 162 | # Steadt state DAG 163 | blocks_ss = [household, partial_ss, 164 | dividend, taylor, fiscal, share_value, finance, union_ss, mkt_clearing] 165 | two_asset_model_ss = create_model(blocks_ss, name='Two-Asset HANK SS') 166 | 167 | # Steady State 168 | calibration = {'Y': 1., 'N': 1.0, 'K': 10., 'r': 0.0125, 'rstar': 0.0125, 'tot_wealth': 14, 169 | 'delta': 0.02, 'pi': 0., 170 | 'kappap': 0.1, 'muw': 1.1, 'Bh': 1.04, 'Bg': 2.8, 'G': 0.2, 'eis': 0.5, 171 | 'frisch': 1, 'chi0': 0.25, 'chi2': 2, 'epsI': 4, 'omega': 0.005, 172 | 'kappaw': 0.1, 'phi': 1.5, 'nZ': 3, 'nB': 10, 'nA': 16, 'nK': 4, 173 | 'bmax': 50, 'amax': 4000, 'kmax': 1, 'rho_z': 0.966, 'sigma_z': 0.92} 174 | unknowns_ss = {'beta': 0.976, 'chi1': 6.5} 175 | targets_ss = {'asset_mkt': 0., 'B': 'Bh'} 176 | cali = two_asset_model_ss.solve_steady_state(calibration, unknowns_ss, targets_ss, solver='broyden_custom') 177 | ss = two_asset_model.steady_state(cali) 178 | 179 | # Transitional Dynamics/Jacobian Calculation 180 | unknowns = ['r', 'w', 'Y'] 181 | targets = ['asset_mkt', 'fisher', 'wnkpc'] 182 | exogenous = ['rstar', 'Z', 'G'] 183 | 184 | return two_asset_model_ss, ss, two_asset_model, unknowns, targets, exogenous 185 | -------------------------------------------------------------------------------- /tests/base/test_het_support.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from sequence_jacobian.blocks.support.het_support import (Transition, 3 | PolicyLottery1D, PolicyLottery2D, Markov, CombinedTransition, 4 | lottery_1d, lottery_2d) 5 | from sequence_jacobian.utilities.multidim import batch_multiply_ith_dimension 6 | 7 | 8 | def test_combined_markov(): 9 | shape = (5, 6, 7) 10 | np.random.seed(12345) 11 | 12 | for _ in range(10): 13 | D = np.random.rand(*shape) 14 | Pis = [np.random.rand(s, s) for s in shape[:2]] 15 | markovs = [Markov(Pi, i) for i, Pi in enumerate(Pis)] 16 | combined = CombinedTransition(markovs) 17 | 18 | Dout = combined.expectation(D) 19 | Dout_forward = combined.forward(D) 20 | 21 | D_kron = D.reshape((-1, D.shape[2])) 22 | Pi_kron = np.kron(Pis[0], Pis[1]) 23 | Dout2 = (Pi_kron @ D_kron).reshape(Dout.shape) 24 | Dout2_forward = (Pi_kron.T @ D_kron).reshape(Dout.shape) 25 | 26 | assert np.allclose(Dout, Dout2) 27 | assert np.allclose(Dout_forward, Dout2_forward) 28 | 29 | 30 | def test_many_markov_shock(): 31 | shape = (5, 6, 7) 32 | np.random.seed(12345) 33 | 34 | for _ in range(10): 35 | D = np.random.rand(*shape) 36 | Pis = [np.random.rand(s, s) for s in shape[:2]] 37 | dPis = [np.random.rand(s, s) for s in shape[:2]] 38 | 39 | h = 1E-4 40 | Dout_up = CombinedTransition([Markov(Pi + h*dPi, i) for i, (Pi, dPi) in enumerate(zip(Pis, dPis))]).forward(D) 41 | Dout_dn = CombinedTransition([Markov(Pi - h*dPi, i) for i, (Pi, dPi) in enumerate(zip(Pis, dPis))]).forward(D) 42 | Dder = (Dout_up - Dout_dn) / (2*h) 43 | 44 | Dder2 = CombinedTransition([Markov(Pi, i) for i, Pi in enumerate(Pis)]).forward_shockable(D).forward_shock(dPis) 45 | 46 | assert np.allclose(Dder, Dder2) 47 | 48 | 49 | def test_policy_shock(): 50 | shape = (3, 4, 30) 51 | grid = np.geomspace(0.5, 10, shape[-1]) 52 | np.random.seed(98765) 53 | 54 | a = (np.full(shape[0], 0.01)[:, np.newaxis, np.newaxis] 55 | + np.linspace(0, 1, shape[1])[:, np.newaxis] 56 | + 0.001*grid**2 + 0.9*grid + 0.5) 57 | 58 | for _ in range(10): 59 | D = np.random.rand(*shape) 60 | 61 | da = np.random.rand(*shape) 62 | h = 1E-5 63 | Dout_up = lottery_1d(a + h*da, grid).forward(D) 64 | Dout_dn = lottery_1d(a - h*da, grid).forward(D) 65 | Dder = (Dout_up - Dout_dn) / (2*h) 66 | 67 | Dder2 = lottery_1d(a, grid).forward_shockable(D).forward_shock(da) 68 | 69 | assert np.allclose(Dder, Dder2, atol=1E-4) 70 | 71 | 72 | def test_law_of_motion_shock(): 73 | # shock everything in the law of motion, and see if it works! 74 | shape = (3, 4, 30) 75 | grid = np.geomspace(0.5, 10, shape[-1]) 76 | np.random.seed(98765) 77 | 78 | a = (np.full(shape[0], 0.01)[:, np.newaxis, np.newaxis] 79 | + np.linspace(0, 1, shape[1])[:, np.newaxis] 80 | + 0.001*grid**2 + 0.9*grid + 0.5) 81 | 82 | for _ in range(10): 83 | D = np.random.rand(*shape) 84 | Pis = [np.random.rand(s, s) for s in shape[:2]] 85 | 86 | da = np.random.rand(*shape) 87 | dPis = [np.random.rand(s, s) for s in shape[:2]] 88 | 89 | h = 1E-5 90 | policy_up = lottery_1d(a + h*da, grid) 91 | policy_dn = lottery_1d(a - h*da, grid) 92 | markovs_up = [Markov(Pi + h*dPi, i) for i, (Pi, dPi) in enumerate(zip(Pis, dPis))] 93 | markovs_dn =[Markov(Pi - h*dPi, i) for i, (Pi, dPi) in enumerate(zip(Pis, dPis))] 94 | Dout_up = CombinedTransition([policy_up, *markovs_up]).forward(D) 95 | Dout_dn = CombinedTransition([policy_dn, *markovs_dn]).forward(D) 96 | Dder = (Dout_up - Dout_dn) / (2*h) 97 | 98 | markovs = [Markov(Pi, i) for i, Pi, in enumerate(Pis)] 99 | Dder2 = CombinedTransition([lottery_1d(a, grid), *markovs]).forward_shockable(D).forward_shock([da, *dPis]) 100 | 101 | assert np.allclose(Dder, Dder2, atol=1E-4) 102 | 103 | 104 | def test_2d_policy_shock(): 105 | shape = (3, 4, 20, 30) 106 | a_grid = np.geomspace(0.5, 10, shape[-2]) 107 | b_grid = np.geomspace(0.2, 8, shape[-1]) 108 | np.random.seed(98765) 109 | 110 | a = (0.001*a_grid**2 + 0.9*a_grid + 0.5)[:, np.newaxis] 111 | b = (-0.001*b_grid**2 + 0.9*b_grid + 0.5) 112 | 113 | a = np.broadcast_to(a, shape) 114 | b = np.broadcast_to(b, shape) 115 | 116 | for _ in range(10): 117 | D = np.random.rand(*shape) 118 | Pis = [np.random.rand(s, s) for s in shape[:2]] 119 | 120 | da = np.random.rand(*shape) 121 | db = np.random.rand(*shape) 122 | dPis = [np.random.rand(s, s) for s in shape[:2]] 123 | 124 | h = 1E-5 125 | 126 | policy_up = lottery_2d(a + h*da, b + h*db, a_grid, b_grid) 127 | policy_dn = lottery_2d(a - h*da, b - h*db, a_grid, b_grid) 128 | markovs_up = [Markov(Pi + h*dPi, i) for i, (Pi, dPi) in enumerate(zip(Pis, dPis))] 129 | markovs_dn = [Markov(Pi - h*dPi, i) for i, (Pi, dPi) in enumerate(zip(Pis, dPis))] 130 | Dout_up = CombinedTransition([policy_up, *markovs_up]).forward(D) 131 | Dout_dn = CombinedTransition([policy_dn, *markovs_dn]).forward(D) 132 | Dder = (Dout_up - Dout_dn) / (2*h) 133 | 134 | policy = lottery_2d(a, b, a_grid, b_grid) 135 | 136 | markovs = [Markov(Pi, i) for i, Pi, in enumerate(Pis)] 137 | Dder2 = CombinedTransition([policy, *markovs]).forward_shockable(D).forward_shock([[da, db], *dPis]) 138 | 139 | assert np.allclose(Dder, Dder2, atol=1E-4) 140 | 141 | 142 | def test_forward_expectations_symmetry(): 143 | # given a random law of motion, should be identical to iterate forward on distribution, 144 | # then aggregate, or take expectations backward on outcome, then aggregate 145 | shape = (3, 4, 30) 146 | grid = np.geomspace(0.5, 10, shape[-1]) 147 | np.random.seed(1423) 148 | 149 | a = (np.full(shape[0], 0.01)[:, np.newaxis, np.newaxis] 150 | + np.linspace(0, 1, shape[1])[:, np.newaxis] 151 | + 0.001*grid**2 + 0.9*grid + 0.5) 152 | 153 | for _ in range(10): 154 | D = np.random.rand(*shape) 155 | X = np.random.rand(*shape) 156 | Pis = [np.random.rand(s, s) for s in shape[:2]] 157 | 158 | markovs = [Markov(Pi, i) for i, Pi, in enumerate(Pis)] 159 | lom = CombinedTransition([lottery_1d(a, grid), *markovs]) 160 | 161 | Dforward = D 162 | for _ in range(30): 163 | Dforward = lom.forward(Dforward) 164 | outcome = np.vdot(Dforward, X) 165 | 166 | Xbackward = X 167 | for _ in range(30): 168 | Xbackward = lom.expectation(Xbackward) 169 | outcome2 = np.vdot(D, Xbackward) 170 | 171 | assert np.isclose(outcome, outcome2) 172 | 173 | 174 | def test_einsum(): 175 | D = np.random.rand(2, 5, 10) 176 | P = np.random.rand(3, 2, 5, 10) 177 | Dnew = np.einsum('xij,dxij->dij', D, P) 178 | assert Dnew[0, 1, 1] == np.sum(P[0, :, 1, 1] * D[:, 1, 1]) 179 | 180 | # can I generalize this? reshape and then einsum 181 | Dnew2 = batch_multiply_ith_dimension(P, 0, D) 182 | 183 | assert (Dnew == Dnew2).all() -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/interpolate.py: -------------------------------------------------------------------------------- 1 | """Efficient linear interpolation exploiting monotonicity. 2 | 3 | Interpolates increasing query points xq against increasing data points x. 4 | 5 | - interpolate_y: (x, xq, y) -> yq 6 | get interpolated values of yq at xq 7 | 8 | - interpolate_coord: (x, xq) -> (xqi, xqpi) 9 | get representation xqi, xqpi of xq interpolated against x 10 | xq = xqpi * x[xqi] + (1-xqpi) * x[xqi+1] 11 | 12 | - apply_coord: (xqi, xqpi, y) -> yq 13 | use representation xqi, xqpi to get yq at xq 14 | yq = xqpi * y[xqi] + (1-xqpi) * y[xqi+1] 15 | 16 | Composing interpolate_coord and apply_coord gives interpolate_y. 17 | 18 | All three functions are written for vectors but can be broadcast to other dimensions 19 | since we use Numba's guvectorize decorator. In these cases, interpolation is always 20 | done on the final dimension. 21 | """ 22 | 23 | import numpy as np 24 | from numba import njit, guvectorize 25 | 26 | 27 | @guvectorize(['void(float64[:], float64[:], float64[:], float64[:])'], '(n),(nq),(n)->(nq)') 28 | def interpolate_y(x, xq, y, yq): 29 | """Efficient linear interpolation exploiting monotonicity. 30 | 31 | Complexity O(n+nq), so most efficient when x and xq have comparable number of points. 32 | Extrapolates linearly when xq out of domain of x. 33 | 34 | Parameters 35 | ---------- 36 | x : array (n), ascending data points 37 | xq : array (nq), ascending query points 38 | y : array (n), data points 39 | 40 | Returns 41 | ---------- 42 | yq : array (nq), interpolated points 43 | """ 44 | nxq, nx = xq.shape[0], x.shape[0] 45 | 46 | xi = 0 47 | x_low = x[0] 48 | x_high = x[1] 49 | for xqi_cur in range(nxq): 50 | xq_cur = xq[xqi_cur] 51 | while xi < nx - 2: 52 | if x_high >= xq_cur: 53 | break 54 | xi += 1 55 | x_low = x_high 56 | x_high = x[xi + 1] 57 | 58 | xqpi_cur = (x_high - xq_cur) / (x_high - x_low) 59 | yq[xqi_cur] = xqpi_cur * y[xi] + (1 - xqpi_cur) * y[xi + 1] 60 | 61 | 62 | @guvectorize(['void(float64[:], float64[:], uint32[:], float64[:])'], '(n),(nq)->(nq),(nq)') 63 | def interpolate_coord(x, xq, xqi, xqpi): 64 | """Get representation xqi, xqpi of xq interpolated against x: 65 | xq = xqpi * x[xqi] + (1-xqpi) * x[xqi+1] 66 | 67 | Parameters 68 | ---------- 69 | x : array (n), ascending data points 70 | xq : array (nq), ascending query points 71 | 72 | Returns 73 | ---------- 74 | xqi : array (nq), indices of lower bracketing gridpoints 75 | xqpi : array (nq), weights on lower bracketing gridpoints 76 | """ 77 | nxq, nx = xq.shape[0], x.shape[0] 78 | 79 | xi = 0 80 | x_low = x[0] 81 | x_high = x[1] 82 | for xqi_cur in range(nxq): 83 | xq_cur = xq[xqi_cur] 84 | while xi < nx - 2: 85 | if x_high >= xq_cur: 86 | break 87 | xi += 1 88 | x_low = x_high 89 | x_high = x[xi + 1] 90 | 91 | xqpi[xqi_cur] = (x_high - xq_cur) / (x_high - x_low) 92 | xqi[xqi_cur] = xi 93 | 94 | 95 | @guvectorize(['void(int64[:], float64[:], float64[:], float64[:])', 96 | 'void(uint32[:], float64[:], float64[:], float64[:])'], '(nq),(nq),(n)->(nq)') 97 | def apply_coord(x_i, x_pi, y, yq): 98 | """Use representation xqi, xqpi to get yq at xq: 99 | yq = xqpi * y[xqi] + (1-xqpi) * y[xqi+1] 100 | 101 | Parameters 102 | ---------- 103 | xqi : array (nq), indices of lower bracketing gridpoints 104 | xqpi : array (nq), weights on lower bracketing gridpoints 105 | y : array (n), data points 106 | 107 | Returns 108 | ---------- 109 | yq : array (nq), interpolated points 110 | """ 111 | nq = x_i.shape[0] 112 | for iq in range(nq): 113 | y_low = y[x_i[iq]] 114 | y_high = y[x_i[iq]+1] 115 | yq[iq] = x_pi[iq]*y_low + (1-x_pi[iq])*y_high 116 | 117 | 118 | '''Part 2: More robust linear interpolation that does not require monotonicity in query points. 119 | 120 | Intended for general use in interpolating policy rules that we cannot be sure are monotonic. 121 | Only get xqi, xqpi representation, for case where x is one-dimensional, in this application. 122 | ''' 123 | 124 | 125 | def interpolate_coord_robust(x, xq, check_increasing=False): 126 | """Linear interpolation exploiting monotonicity only in data x, not in query points xq. 127 | Simple binary search, less efficient but more robust. 128 | xq = xqpi * x[xqi] + (1-xqpi) * x[xqi+1] 129 | 130 | Main application intended to be universally-valid interpolation of policy rules. 131 | Dimension k is optional. 132 | 133 | Parameters 134 | ---------- 135 | x : array (n), ascending data points 136 | xq : array (k, nq), query points (in any order) 137 | 138 | Returns 139 | ---------- 140 | xqi : array (k, nq), indices of lower bracketing gridpoints 141 | xqpi : array (k, nq), weights on lower bracketing gridpoints 142 | """ 143 | if x.ndim != 1: 144 | raise ValueError('Data input to interpolate_coord_robust must have exactly one dimension') 145 | 146 | if check_increasing and np.any(x[:-1] >= x[1:]): 147 | raise ValueError('Data input to interpolate_coord_robust must be strictly increasing') 148 | 149 | if xq.ndim == 1: 150 | return interpolate_coord_robust_vector(x, xq) 151 | else: 152 | i, pi = interpolate_coord_robust_vector(x, xq.ravel()) 153 | return i.reshape(xq.shape), pi.reshape(xq.shape) 154 | 155 | 156 | @njit 157 | def interpolate_coord_robust_vector(x, xq): 158 | """Does interpolate_coord_robust where xq must be a vector, more general function is wrapper""" 159 | 160 | n = len(x) 161 | nq = len(xq) 162 | xqi = np.empty(nq, dtype=np.uint32) 163 | xqpi = np.empty(nq) 164 | 165 | for iq in range(nq): 166 | if xq[iq] < x[0]: 167 | ilow = 0 168 | elif xq[iq] > x[-2]: 169 | ilow = n-2 170 | else: 171 | # start binary search 172 | # should end with ilow and ihigh exactly 1 apart, bracketing variable 173 | ihigh = n-1 174 | ilow = 0 175 | while ihigh - ilow > 1: 176 | imid = (ihigh + ilow) // 2 177 | if xq[iq] > x[imid]: 178 | ilow = imid 179 | else: 180 | ihigh = imid 181 | 182 | xqi[iq] = ilow 183 | xqpi[iq] = (x[ilow+1] - xq[iq]) / (x[ilow+1] - x[ilow]) 184 | 185 | return xqi, xqpi 186 | 187 | 188 | '''Used in discrete choice problems''' 189 | 190 | @njit 191 | def interpolate_coord_njit(x, xq): 192 | nxq, nx = xq.shape[0], x.shape[0] 193 | xqi = np.empty(nxq, dtype=np.uint32) 194 | xqpi = np.empty(nxq) 195 | 196 | xi = 0 197 | x_low = x[0] 198 | x_high = x[1] 199 | for xqi_cur in range(nxq): 200 | xq_cur = xq[xqi_cur] 201 | while xi < nx - 2: 202 | if x_high >= xq_cur: 203 | break 204 | xi += 1 205 | x_low = x_high 206 | x_high = x[xi + 1] 207 | 208 | xqpi[xqi_cur] = (x_high - xq_cur) / (x_high - x_low) 209 | xqi[xqi_cur] = xi 210 | 211 | return xqi, xqpi 212 | 213 | 214 | @njit 215 | def apply_coord_njit(x_i, x_pi, y): 216 | nq = x_i.shape[0] 217 | yq = np.empty(nq) 218 | 219 | for iq in range(nq): 220 | y_low = y[x_i[iq]] 221 | y_high = y[x_i[iq]+1] 222 | yq[iq] = x_pi[iq]*y_low + (1-x_pi[iq])*y_high 223 | 224 | return yq 225 | 226 | 227 | @njit 228 | def interpolate_point(x, x0, x1, y0, y1): 229 | y = y0 + (x - x0) * (y1 - y0) / (x1 - x0) 230 | return y 231 | -------------------------------------------------------------------------------- /src/sequence_jacobian/hetblocks/hh_twoasset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numba import guvectorize 3 | 4 | from ..blocks.het_block import het 5 | from .. import interpolate 6 | 7 | 8 | def hh_init(b_grid, a_grid, z_grid, eis): 9 | Va = (0.6 + 1.1 * b_grid[:, np.newaxis] + a_grid) ** (-1 / eis) * np.ones((z_grid.shape[0], 1, 1)) 10 | Vb = (0.5 + b_grid[:, np.newaxis] + 1.2 * a_grid) ** (-1 / eis) * np.ones((z_grid.shape[0], 1, 1)) 11 | return Va, Vb 12 | 13 | 14 | def adjustment_costs(a, a_grid, ra, chi0, chi1, chi2): 15 | chi = get_Psi_and_deriv(a, a_grid, ra, chi0, chi1, chi2)[0] 16 | return chi 17 | 18 | 19 | def marginal_cost_grid(a_grid, ra, chi0, chi1, chi2): 20 | # precompute Psi1(a', a) on grid of (a', a) for steps 3 and 5 21 | Psi1 = get_Psi_and_deriv(a_grid[:, np.newaxis], 22 | a_grid[np.newaxis, :], ra, chi0, chi1, chi2)[1] 23 | return Psi1 24 | 25 | 26 | # policy and bacward order as in grid! 27 | @het(exogenous='Pi', policy=['b', 'a'], backward=['Vb', 'Va'], 28 | hetinputs=[marginal_cost_grid], hetoutputs=[adjustment_costs], backward_init=hh_init) 29 | def hh(Va_p, Vb_p, a_grid, b_grid, z_grid, e_grid, k_grid, beta, eis, rb, ra, chi0, chi1, chi2, Psi1): 30 | # === STEP 2: Wb(z, b', a') and Wa(z, b', a') === 31 | # (take discounted expectation of tomorrow's value function) 32 | Wb = beta * Vb_p 33 | Wa = beta * Va_p 34 | W_ratio = Wa / Wb 35 | 36 | # === STEP 3: a'(z, b', a) for UNCONSTRAINED === 37 | 38 | # for each (z, b', a), linearly interpolate to find a' between gridpoints 39 | # satisfying optimality condition W_ratio == 1+Psi1 40 | i, pi = lhs_equals_rhs_interpolate(W_ratio, 1 + Psi1) 41 | 42 | # use same interpolation to get Wb and then c 43 | a_endo_unc = interpolate.apply_coord(i, pi, a_grid) 44 | c_endo_unc = interpolate.apply_coord(i, pi, Wb) ** (-eis) 45 | 46 | # === STEP 4: b'(z, b, a), a'(z, b, a) for UNCONSTRAINED === 47 | 48 | # solve out budget constraint to get b(z, b', a) 49 | b_endo = (c_endo_unc + a_endo_unc + addouter(-z_grid, b_grid, -(1 + ra) * a_grid) 50 | + get_Psi_and_deriv(a_endo_unc, a_grid, ra, chi0, chi1, chi2)[0]) / (1 + rb) 51 | 52 | # interpolate this b' -> b mapping to get b -> b', so we have b'(z, b, a) 53 | # and also use interpolation to get a'(z, b, a) 54 | # (note utils.interpolate.interpolate_coord and utils.interpolate.apply_coord work on last axis, 55 | # so we need to swap 'b' to the last axis, then back when done) 56 | i, pi = interpolate.interpolate_coord(b_endo.swapaxes(1, 2), b_grid) 57 | a_unc = interpolate.apply_coord(i, pi, a_endo_unc.swapaxes(1, 2)).swapaxes(1, 2) 58 | b_unc = interpolate.apply_coord(i, pi, b_grid).swapaxes(1, 2) 59 | 60 | # === STEP 5: a'(z, kappa, a) for CONSTRAINED === 61 | 62 | # for each (z, kappa, a), linearly interpolate to find a' between gridpoints 63 | # satisfying optimality condition W_ratio/(1+kappa) == 1+Psi1, assuming b'=0 64 | lhs_con = W_ratio[:, 0:1, :] / (1 + k_grid[np.newaxis, :, np.newaxis]) 65 | i, pi = lhs_equals_rhs_interpolate(lhs_con, 1 + Psi1) 66 | 67 | # use same interpolation to get Wb and then c 68 | a_endo_con = interpolate.apply_coord(i, pi, a_grid) 69 | c_endo_con = ((1 + k_grid[np.newaxis, :, np.newaxis]) ** (-eis) 70 | * interpolate.apply_coord(i, pi, Wb[:, 0:1, :]) ** (-eis)) 71 | 72 | # === STEP 6: a'(z, b, a) for CONSTRAINED === 73 | 74 | # solve out budget constraint to get b(z, kappa, a), enforcing b'=0 75 | b_endo = (c_endo_con + a_endo_con 76 | + addouter(-z_grid, np.full(len(k_grid), b_grid[0]), -(1 + ra) * a_grid) 77 | + get_Psi_and_deriv(a_endo_con, a_grid, ra, chi0, chi1, chi2)[0]) / (1 + rb) 78 | 79 | # interpolate this kappa -> b mapping to get b -> kappa 80 | # then use the interpolated kappa to get a', so we have a'(z, b, a) 81 | # (utils.interpolate.interpolate_y does this in one swoop, but since it works on last 82 | # axis, we need to swap kappa to last axis, and then b back to middle when done) 83 | a_con = interpolate.interpolate_y(b_endo.swapaxes(1, 2), b_grid, 84 | a_endo_con.swapaxes(1, 2)).swapaxes(1, 2) 85 | 86 | # === STEP 7: obtain policy functions and update derivatives of value function === 87 | 88 | # combine unconstrained solution and constrained solution, choosing latter 89 | # when unconstrained goes below minimum b 90 | a, b = a_unc.copy(), b_unc.copy() 91 | b[b <= b_grid[0]] = b_grid[0] 92 | a[b <= b_grid[0]] = a_con[b <= b_grid[0]] 93 | 94 | # calculate adjustment cost and its derivative 95 | Psi, _, Psi2 = get_Psi_and_deriv(a, a_grid, ra, chi0, chi1, chi2) 96 | 97 | # solve out budget constraint to get consumption and marginal utility 98 | c = addouter(z_grid, (1 + rb) * b_grid, (1 + ra) * a_grid) - Psi - a - b 99 | uc = c ** (-1 / eis) 100 | uce = e_grid[:, np.newaxis, np.newaxis] * uc 101 | 102 | # update derivatives of value function using envelope conditions 103 | Va = (1 + ra - Psi2) * uc 104 | Vb = (1 + rb) * uc 105 | 106 | return Va, Vb, a, b, c, uce 107 | 108 | 109 | '''Supporting functions for HA block''' 110 | 111 | def get_Psi_and_deriv(ap, a, ra, chi0, chi1, chi2): 112 | """Adjustment cost Psi(ap, a) and its derivatives with respect to 113 | first argument (ap) and second argument (a)""" 114 | a_with_return = (1 + ra) * a 115 | a_change = ap - a_with_return 116 | abs_a_change = np.abs(a_change) 117 | sign_change = np.sign(a_change) 118 | 119 | adj_denominator = a_with_return + chi0 120 | core_factor = (abs_a_change / adj_denominator) ** (chi2 - 1) 121 | 122 | Psi = chi1 / chi2 * abs_a_change * core_factor 123 | Psi1 = chi1 * sign_change * core_factor 124 | Psi2 = -(1 + ra) * (Psi1 + (chi2 - 1) * Psi / adj_denominator) 125 | return Psi, Psi1, Psi2 126 | 127 | 128 | def matrix_times_first_dim(A, X): 129 | """Take matrix A times vector X[:, i1, i2, i3, ... , in] separately 130 | for each i1, i2, i3, ..., in. Same output as A @ X if X is 1D or 2D""" 131 | # flatten all dimensions of X except first, then multiply, then restore shape 132 | return (A @ X.reshape(X.shape[0], -1)).reshape(X.shape) 133 | 134 | 135 | def addouter(z, b, a): 136 | """Take outer sum of three arguments: result[i, j, k] = z[i] + b[j] + a[k]""" 137 | return z[:, np.newaxis, np.newaxis] + b[:, np.newaxis] + a 138 | 139 | 140 | @guvectorize(['void(float64[:], float64[:,:], uint32[:], float64[:])'], '(ni),(ni,nj)->(nj),(nj)') 141 | def lhs_equals_rhs_interpolate(lhs, rhs, iout, piout): 142 | """ 143 | Given lhs (i) and rhs (i,j), for each j, find the i such that 144 | 145 | lhs[i] > rhs[i,j] and lhs[i+1] < rhs[i+1,j] 146 | 147 | i.e. where given j, lhs == rhs in between i and i+1. 148 | 149 | Also return the pi such that 150 | 151 | pi*(lhs[i] - rhs[i,j]) + (1-pi)*(lhs[i+1] - rhs[i+1,j]) == 0 152 | 153 | i.e. such that the point at pi*i + (1-pi)*(i+1) satisfies lhs == rhs by linear interpolation. 154 | 155 | If lhs[0] < rhs[0,j] already, just return u=0 and pi=1. 156 | 157 | ***IMPORTANT: Assumes that solution i is monotonically increasing in j 158 | and that lhs - rhs is monotonically decreasing in i.*** 159 | """ 160 | 161 | ni, nj = rhs.shape 162 | assert len(lhs) == ni 163 | 164 | i = 0 165 | for j in range(nj): 166 | while True: 167 | if lhs[i] < rhs[i, j]: 168 | break 169 | elif i < nj - 1: 170 | i += 1 171 | else: 172 | break 173 | 174 | if i == 0: 175 | iout[j] = 0 176 | piout[j] = 1 177 | else: 178 | iout[j] = i - 1 179 | err_upper = rhs[i, j] - lhs[i] 180 | err_lower = rhs[i - 1, j] - lhs[i - 1] 181 | piout[j] = err_upper / (err_upper - err_lower) 182 | -------------------------------------------------------------------------------- /src/sequence_jacobian/utilities/graph.py: -------------------------------------------------------------------------------- 1 | """Topological sort and related code""" 2 | from .ordered_set import OrderedSet 3 | from .bijection import Bijection 4 | 5 | class DAG: 6 | """Represents "blocks" that each have inputs and outputs, where output-input relationships between 7 | blocks form a DAG. Fundamental DAG object intended to underlie CombinedBlock and CombinedExtendedFunction. 8 | 9 | Initialized with list of blocks, which are then topologically sorted""" 10 | 11 | def __init__(self, blocks): 12 | inmap = get_input_map(blocks) 13 | outmap = get_output_map(blocks) 14 | adj = get_block_adjacency_list(blocks, inmap) 15 | revadj = get_block_reverse_adjacency_list(blocks, outmap) 16 | topsort = topological_sort(adj, revadj, names=[getattr(block, 'name', '[NO BLOCK NAME]') for block in blocks]) 17 | 18 | M = Bijection({t: i for i, t in enumerate(topsort)}) 19 | 20 | self.blocks = [blocks[t] for t in topsort] 21 | self.inmap = {k: M @ v for k, v in inmap.items()} 22 | self.outmap = {k: M @ v for k, v in outmap.items()} 23 | self.adj = [M @ adj[t] for t in topsort] 24 | self.revadj = [M @ revadj[t] for t in topsort] 25 | 26 | self.inputs = OrderedSet(k for k in inmap if k not in outmap) 27 | self.outputs = OrderedSet(outmap) 28 | 29 | def visit_from_inputs(self, inputs): 30 | """Which block numbers are ultimately dependencies of 'inputs'?""" 31 | inputs = inputs & self.inputs 32 | visited = OrderedSet() 33 | for n, (block, parentset) in enumerate(zip(self.blocks, self.revadj)): 34 | # first see if block has its input directly changed 35 | for i in inputs: 36 | if i in block.inputs: 37 | visited.add(n) 38 | break 39 | else: 40 | if not parentset.isdisjoint(visited): 41 | visited.add(n) 42 | 43 | return visited 44 | 45 | def visit_from_outputs(self, outputs): 46 | """Which block numbers are 'outputs' ultimately dependent on?""" 47 | outputs = outputs & self.outputs 48 | visited = OrderedSet() 49 | for n in reversed(range(len(self.blocks))): 50 | block = self.blocks[n] 51 | childset = self.adj[n] 52 | 53 | # first see if block has its output directly used 54 | for o in outputs: 55 | if o in block.outputs: 56 | visited.add(n) 57 | break 58 | else: 59 | if not childset.isdisjoint(visited): 60 | visited.add(n) 61 | 62 | return reversed(visited) 63 | 64 | 65 | def topological_sort(adj, revadj, names=None): 66 | """Given directed graph pointing from each node to the nodes it depends on, topologically sort nodes""" 67 | # get complete set version of dep, and its reversal, and build initial stack of nodes with no dependencies 68 | revdep = adj 69 | dep = [s.copy() for s in revadj] 70 | nodeps = [n for n, depset in enumerate(dep) if not depset] 71 | topsorted = [] 72 | 73 | # Kahn's algorithm: find something with no dependency, delete its edges and update 74 | while nodeps: 75 | n = nodeps.pop() 76 | topsorted.append(n) 77 | for n2 in revdep[n]: 78 | dep[n2].remove(n) 79 | if not dep[n2]: 80 | nodeps.append(n2) 81 | 82 | # should be done: topsorted should be topologically sorted with same # of elements as original graphs! 83 | if len(topsorted) != len(dep): 84 | cycle_ints = find_cycle(dep, set(range(len(dep))) - set(topsorted)) 85 | assert cycle_ints is not None, 'topological sort failed but no cycle, THIS SHOULD NEVER EVER HAPPEN' 86 | cycle = [names[i] for i in cycle_ints] if names else cycle_ints 87 | raise Exception(f'Topological sort failed: cyclic dependency {" -> ".join([str(n) for n in cycle])}') 88 | 89 | return topsorted 90 | 91 | 92 | def get_input_map(blocks: list): 93 | """inmap[i] gives set of block numbers where i is an input""" 94 | inmap = dict() 95 | for num, block in enumerate(blocks): 96 | for i in block.inputs: 97 | inset = inmap.setdefault(i, OrderedSet()) 98 | inset.add(num) 99 | 100 | return inmap 101 | 102 | 103 | def get_output_map(blocks: list): 104 | """outmap[o] gives unique block number where o is an output""" 105 | outmap = dict() 106 | for num, block in enumerate(blocks): 107 | for o in block.outputs: 108 | if o in outmap: 109 | raise ValueError(f'{o} is output twice') 110 | outmap[o] = num 111 | 112 | return outmap 113 | 114 | 115 | def get_block_adjacency_list(blocks, inmap): 116 | """adj[n] for block number n gives set of block numbers which this block points to""" 117 | adj = [] 118 | for block in blocks: 119 | current_adj = OrderedSet() 120 | for o in block.outputs: 121 | # for each output, if that output is used as an input by some blocks, add those blocks to adj 122 | if o in inmap: 123 | current_adj |= inmap[o] 124 | adj.append(current_adj) 125 | return adj 126 | 127 | 128 | def get_block_reverse_adjacency_list(blocks, outmap): 129 | """revadj[n] for block number n gives set of block numbers that point to this block""" 130 | revadj = [] 131 | for block in blocks: 132 | current_revadj = OrderedSet() 133 | for i in block.inputs: 134 | if i in outmap: 135 | current_revadj.add(outmap[i]) 136 | revadj.append(current_revadj) 137 | return revadj 138 | 139 | 140 | def find_intermediate_inputs(blocks): 141 | # TODO: should be deprecated 142 | """Find outputs of the blocks in blocks that are inputs to other blocks in blocks. 143 | This is useful to ensure that all of the relevant curlyJ Jacobians (of all inputs to all outputs) are computed. 144 | """ 145 | required = OrderedSet() 146 | outmap = get_output_map(blocks) 147 | for num, block in enumerate(blocks): 148 | if hasattr(block, 'inputs'): 149 | inputs = block.inputs 150 | else: 151 | inputs = OrderedSet(i for o in block for i in block[o]) 152 | for i in inputs: 153 | if i in outmap: 154 | required.add(i) 155 | return required 156 | 157 | 158 | def find_cycle(dep, onlyset): 159 | """Return list giving cycle if there is one, otherwise None""" 160 | 161 | # supposed to look only within 'onlyset', so filter out everything else 162 | # awkward holdover: 'dep' is transformed here into a dict with integer keys 163 | dep = {k: (dep[k] & onlyset) for k in range(len(dep)) if k in onlyset} 164 | 165 | tovisit = set(dep.keys()) 166 | stack = OrderedSet() 167 | while tovisit or stack: 168 | if stack: 169 | # if stack has something, still need to proceed with DFS 170 | n = stack.top() 171 | if dep[n]: 172 | # if there are any dependencies left, let's look at them 173 | n2 = dep[n].pop() 174 | if n2 in stack: 175 | # we have a cycle, since this is already in our stack 176 | i2loc = stack.index(n2) 177 | return stack[i2loc:] + [stack[i2loc]] 178 | else: 179 | # no cycle, visit this node only if we haven't already visited it 180 | if n2 in tovisit: 181 | tovisit.remove(n2) 182 | stack.add(n2) 183 | else: 184 | # if no dependencies left, then we're done with this node, so let's forget about it 185 | stack.pop(n) 186 | else: 187 | # nothing left on stack, let's start the DFS from something new 188 | n = tovisit.pop() 189 | stack.add(n) 190 | 191 | # if we never find a cycle, we're done 192 | return None 193 | -------------------------------------------------------------------------------- /tests/base/test_dchoice.py: -------------------------------------------------------------------------------- 1 | ''' 2 | SIM model with labor force participation choice 3 | - state space: (s, x, e, a) 4 | - s is employment 5 | - 0: employed, 1: unemployed, 2: out of labor force 6 | - x is matching 7 | - 0: matched, 1: unmatched 8 | - e is labor productivity 9 | - a is assets 10 | ''' 11 | import numpy as np 12 | from numba import njit 13 | 14 | from sequence_jacobian.blocks.stage_block import StageBlock 15 | from sequence_jacobian.blocks.support.stages import Continuous1D, ExogenousMaker, LogitChoice 16 | from sequence_jacobian import markov_rouwenhorst, agrid 17 | from sequence_jacobian.classes.impulse_dict import ImpulseDict 18 | from sequence_jacobian.utilities.misc import nonconcave 19 | from sequence_jacobian.utilities.interpolate import interpolate_coord_njit, apply_coord_njit, interpolate_point 20 | 21 | 22 | '''Setup: utility function, hetinputs, initializer''' 23 | 24 | 25 | @njit 26 | def util(c, eis): 27 | if eis == 1: 28 | u = np.log(c) 29 | else: 30 | u = c ** (1 - 1 / eis) / (1 - 1 / eis) 31 | return u 32 | 33 | 34 | def make_grids(rho_e, sd_e, nE, amin, amax, nA): 35 | e_grid, e_dist, Pi_e = markov_rouwenhorst(rho=rho_e, sigma=sd_e, N=nE) 36 | a_grid = agrid(amin=amin, amax=amax, n=nA) 37 | return e_grid, e_dist, Pi_e, a_grid 38 | 39 | 40 | def labor_income(a_grid, e_grid, atw, b, s, f, r): 41 | y = e_grid[np.newaxis, :] * np.array([atw, b, b])[:, np.newaxis] 42 | coh = (1 + r) * a_grid[np.newaxis, np.newaxis, :] + y[..., np.newaxis] 43 | Pi_s = np.array([[1 - s, s], [f, (1 - f)], [0, 1]]) 44 | return y, coh, Pi_s 45 | 46 | 47 | def backward_init(coh, a_grid, eis): 48 | V = util(0.1 * coh, eis) / 0.01 49 | Va = np.empty_like(V) 50 | Va[..., 1:-1] = (V[..., 2:] - V[..., :-2]) / (a_grid[2:] - a_grid[:-2]) 51 | Va[..., 0] = (V[..., 1] - V[..., 0]) / (a_grid[1] - a_grid[0]) 52 | Va[..., -1] = (V[..., -1] - V[..., -2]) / (a_grid[-1] - a_grid[-2]) 53 | return V, Va 54 | 55 | 56 | '''Consumption-savings stage: : (s, e, a) -> (s, e, a')''' 57 | 58 | 59 | def consav(V, Va, a_grid, coh, y, r, beta, eis): 60 | """DC-EGM algorithm""" 61 | # EGM step 62 | W = beta * V 63 | uc_endo= beta * Va 64 | c_endo= uc_endo** (-eis) 65 | a_endo= (c_endo+ a_grid[np.newaxis, np.newaxis, :] - y[:, :, np.newaxis]) / (1 + r) 66 | 67 | # upper envelope step 68 | V, c = upper_envelope(Va, W, a_endo, c_endo, coh, a_grid, eis) 69 | 70 | # update Va, report asset policy 71 | uc = c ** (-1 / eis) 72 | Va = (1 + r) * uc 73 | a = coh - c 74 | 75 | return V, Va, a, c 76 | 77 | 78 | def upper_envelope(Va, W, a_endo, c_endo, coh, a_grid, *args): 79 | # identify bounds of nonconcave region 80 | ilower, iupper = nonconcave(Va) 81 | 82 | # upper envelope 83 | shape = W.shape 84 | W = W.reshape((-1, shape[-1])) 85 | a_endo = a_endo.reshape((-1, shape[-1])) 86 | c_endo = c_endo.reshape((-1, shape[-1])) 87 | coh = coh.reshape((-1, shape[-1])) 88 | ilower = ilower.reshape(-1) 89 | iupper = iupper.reshape(-1) 90 | V, c = upper_envelope_core(ilower, iupper, W, a_endo, c_endo, coh, a_grid, *args) 91 | 92 | return V.reshape(shape), c.reshape(shape) 93 | 94 | 95 | @njit 96 | def upper_envelope_core(ilower, iupper, W, a_endo, c_endo, coh, a_grid, *args): 97 | """Interpolate value function and consumption to exogenous grid.""" 98 | nB, nA = W.shape 99 | c = np.zeros_like(W) 100 | V = -np.inf * np.ones_like(W) 101 | 102 | for ib in range(nB): 103 | ilower_cur = ilower[ib] 104 | iupper_cur = iupper[ib] 105 | 106 | # Below nonconcave region: exploit monotonicity 107 | if ilower_cur > 0: 108 | ai, api = interpolate_coord_njit(a_endo[ib, :ilower_cur], a_grid[:ilower_cur]) 109 | c0 = apply_coord_njit(ai, api, c_endo[ib, :ilower_cur]) 110 | W0 = apply_coord_njit(ai, api, W[ib, :ilower_cur]) 111 | c[ib, :ilower_cur] = c0 112 | V[ib, :ilower_cur] = util(c0, *args) + W0 113 | 114 | # Nonconcave region: check everything 115 | for ia in range(ilower_cur, iupper_cur): 116 | acur = a_grid[ia] 117 | for ja in range(nA - 1): 118 | ap_low = a_endo[ib, ja] 119 | ap_high = a_endo[ib, ja + 1] 120 | 121 | interp = (ap_low <= acur <= ap_high) or (ap_low >= acur >= ap_high) 122 | extrap = (ja == nA - 2) and (acur > a_endo[ib, nA - 1]) 123 | 124 | if interp or extrap: 125 | c0 = interpolate_point(acur, ap_low, ap_high, c_endo[ib, ja], c_endo[ib, ja+1]) 126 | W0 = interpolate_point(acur, ap_low, ap_high, W[ib, ja], W[ib, ja + 1]) 127 | V0 = util(c0, *args) + W0 128 | 129 | if V0 > V[ib, ia]: 130 | V[ib, ia] = V0 131 | c[ib, ia] = c0 132 | 133 | # Above nonconcave region: exploit monotonicity 134 | if iupper_cur > 0: 135 | ai, api = interpolate_coord_njit(a_endo[ib, iupper_cur:], a_grid[iupper_cur:]) 136 | c0 = apply_coord_njit(ai, api, c_endo[ib, iupper_cur:]) 137 | W0 = apply_coord_njit(ai, api, W[ib, iupper_cur:]) 138 | c[ib, iupper_cur:] = c0 139 | V[ib, iupper_cur:] = util(c0, *args) + W0 140 | 141 | # Enforce borrowing constraint 142 | ia = 0 143 | while ia < nA and a_grid[ia] <= a_endo[ib, 0]: 144 | c[ib, ia] = coh[ib, ia] 145 | V[ib, ia] = util(c[ib, ia], *args) + W[ib, 0] 146 | ia += 1 147 | 148 | return V, c 149 | 150 | 151 | '''Logit choice stage: (x, z, a) -> (s, z, a)''' 152 | 153 | 154 | def participation(V, vphi, chi): 155 | '''adjustments to flow utility associated with x -> s choice, implements constraints on discrete choice''' 156 | flow_u = np.zeros((3, 2,) + V.shape[-2:]) # (s, x, z, a) 157 | flow_u[0, ...] = -vphi # employed 158 | flow_u[1, ...] = -chi # unemployed 159 | flow_u[0, 1, ...] = -np.inf # unmatched -> employed 160 | return flow_u 161 | 162 | 163 | '''Put stages together''' 164 | 165 | consav_stage = Continuous1D(backward=['Va', 'V'], policy='a', f=consav, name='consav') 166 | labsup_stage = LogitChoice(value='V', backward='Va', index=0, 167 | taste_shock_scale='taste_shock', 168 | f=participation, name='dchoice') 169 | search_stage = ExogenousMaker(markov_name='Pi_s', index=0, name='search_shock') 170 | prod_stage = ExogenousMaker(markov_name='Pi_e', index=1, name='prod_shock') 171 | 172 | hh = StageBlock([prod_stage, search_stage, labsup_stage, consav_stage], 173 | backward_init=backward_init, hetinputs=[make_grids, labor_income], name='household') 174 | 175 | def test_runs(): 176 | calibration = {'taste_shock': 0.01, 'r': 0.005, 'beta': 0.97, 'eis': 0.5, 177 | 'vphi': 0.3, 'chi': 0.3, 'rho_e': 0.95, 'sd_e': 0.5, 'nE': 7, 'amin': .0, 'amax': 200.0, 'nA': 200, 'atw': 1.0, 'b': 0.5, 's': 0.1, 'f': 0.4} 178 | 179 | ss1 = hh.steady_state(calibration) 180 | ss2 = hh.steady_state({**calibration, 181 | 'V': 0.9*ss1.internals['household']['consav']['V'], 182 | 'Va': 0.9*ss1.internals['household']['consav']['Va']}) 183 | 184 | # test steady-state equivalence (from different starting point) 185 | assert np.isclose(ss1['A'], ss2['A']) 186 | assert np.isclose(ss1['C'], ss2['C']) 187 | assert np.allclose(ss1.internals['household']['consav']['D'], ss2.internals['household']['consav']['D']) 188 | assert np.allclose(ss1.internals['household']['consav']['a'], ss2.internals['household']['consav']['a']) 189 | assert np.allclose(ss1.internals['household']['consav']['c'], ss2.internals['household']['consav']['c']) 190 | 191 | inputs = ['r', 'atw', 'f'] 192 | outputs = ['A', 'C'] 193 | T = 50 194 | J = hh.jacobian(ss1, inputs, outputs, T) 195 | 196 | # impulse responses 197 | shock = ImpulseDict({'f': 0.5 ** np.arange(50)}) 198 | td_lin = hh.impulse_linear(ss1, shock, outputs=['C']) 199 | td_nonlin = hh.impulse_nonlinear(ss1, shock * 1E-4, outputs=['C']) 200 | td_ghost = hh.impulse_nonlinear(ss1, shock * 0.0, outputs=['C']) 201 | td_nonlin = td_nonlin - td_ghost 202 | assert np.allclose(td_lin['C'], td_nonlin['C'] / 1E-4, atol=1E-5) 203 | -------------------------------------------------------------------------------- /src/sequence_jacobian/blocks/support/het_support.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from . import het_compiled 3 | from ...utilities.discretize import stationary as general_stationary 4 | from ...utilities.interpolate import interpolate_coord_robust, interpolate_coord 5 | from ...utilities.multidim import batch_multiply_ith_dimension, multiply_ith_dimension 6 | from ...utilities.misc import logsum 7 | from typing import Optional, Sequence, Any, List, Tuple, Union 8 | 9 | class Transition: 10 | """Abstract class for PolicyLottery or ManyMarkov, i.e. some part of state-space transition""" 11 | def forward(self, D): 12 | pass 13 | 14 | def expectation(self, X): 15 | pass 16 | 17 | def forward_shockable(self, Dss): 18 | pass 19 | 20 | def expectation_shockable(self, Xss): 21 | raise NotImplementedError(f'Shockable expectation not implemented for {type(self)}') 22 | 23 | 24 | class ForwardShockableTransition(Transition): 25 | """Abstract class extending Transition, allowing us to find effect of shock to transition rule 26 | on one-period-ahead distribution. This functionality isn't included in the regular Transition 27 | because it requires knowledge of the incoming ("steady-state") distribution and also sometimes 28 | some precomputation. 29 | 30 | One crucial thing here is the order of shock arguments in shocks. Also, is None is the default 31 | argument for a shock, we allow that shock to be None. We always allow shocks in lists to be None.""" 32 | 33 | def forward_shock(self, shocks): 34 | pass 35 | 36 | 37 | class ExpectationShockableTransition(Transition): 38 | def expectation_shock(self, shocks): 39 | pass 40 | 41 | 42 | 43 | def lottery_1d(a, a_grid, monotonic=False): 44 | if not monotonic: 45 | return PolicyLottery1D(*interpolate_coord_robust(a_grid, a), a_grid) 46 | else: 47 | return PolicyLottery1D(*interpolate_coord(a_grid, a), a_grid) 48 | 49 | 50 | class PolicyLottery1D(Transition): 51 | # TODO: always operates on final dimension, highly non-generic in that sense 52 | def __init__(self, i, pi, grid): 53 | # flatten non-policy dimensions into one because that's what methods accept 54 | self.i = i.reshape((-1,) + grid.shape) 55 | self.flatshape = self.i.shape 56 | 57 | self.pi = pi.reshape(self.flatshape) 58 | 59 | # but store original shape so we can convert all outputs to it 60 | self.shape = i.shape 61 | self.grid = grid 62 | 63 | # also store shape of the endogenous grid itself 64 | self.endog_shape = self.shape[-1:] 65 | 66 | def forward(self, D): 67 | return het_compiled.forward_policy_1d(D.reshape(self.flatshape), self.i, self.pi).reshape(self.shape) 68 | 69 | def expectation(self, X): 70 | return het_compiled.expectation_policy_1d(X.reshape(self.flatshape), self.i, self.pi).reshape(self.shape) 71 | 72 | def forward_shockable(self, Dss): 73 | return ForwardShockablePolicyLottery1D(self.i.reshape(self.shape), self.pi.reshape(self.shape), 74 | self.grid, Dss) 75 | 76 | 77 | class ForwardShockablePolicyLottery1D(PolicyLottery1D, ForwardShockableTransition): 78 | def __init__(self, i, pi, grid, Dss): 79 | super().__init__(i, pi, grid) 80 | self.Dss = Dss.reshape(self.flatshape) 81 | self.space = grid[self.i+1] - grid[self.i] 82 | 83 | def forward_shock(self, da): 84 | pi_shock = - da.reshape(self.flatshape) / self.space 85 | return het_compiled.forward_policy_shock_1d(self.Dss, self.i, pi_shock).reshape(self.shape) 86 | 87 | 88 | def lottery_2d(a, b, a_grid, b_grid, monotonic=False): 89 | if not monotonic: 90 | return PolicyLottery2D(*interpolate_coord_robust(a_grid, a), 91 | *interpolate_coord_robust(b_grid, b), a_grid, b_grid) 92 | if monotonic: 93 | # right now we have no monotonic 2D examples, so this shouldn't be called 94 | return PolicyLottery2D(*interpolate_coord(a_grid, a), 95 | *interpolate_coord(b_grid, b), a_grid, b_grid) 96 | 97 | 98 | class PolicyLottery2D(Transition): 99 | def __init__(self, i1, pi1, i2, pi2, grid1, grid2): 100 | # flatten non-policy dimensions into one because that's what methods accept 101 | self.i1 = i1.reshape((-1,) + grid1.shape + grid2.shape) 102 | self.flatshape = self.i1.shape 103 | 104 | self.i2 = i2.reshape(self.flatshape) 105 | self.pi1 = pi1.reshape(self.flatshape) 106 | self.pi2 = pi2.reshape(self.flatshape) 107 | 108 | # but store original shape so we can convert all outputs to it 109 | self.shape = i1.shape 110 | self.grid1 = grid1 111 | self.grid2 = grid2 112 | 113 | # also store shape of the endogenous grid itself 114 | self.endog_shape = self.shape[-2:] 115 | 116 | def forward(self, D): 117 | return het_compiled.forward_policy_2d(D.reshape(self.flatshape), self.i1, self.i2, 118 | self.pi1, self.pi2).reshape(self.shape) 119 | 120 | def expectation(self, X): 121 | return het_compiled.expectation_policy_2d(X.reshape(self.flatshape), self.i1, self.i2, 122 | self.pi1, self.pi2).reshape(self.shape) 123 | 124 | def forward_shockable(self, Dss): 125 | return ForwardShockablePolicyLottery2D(self.i1.reshape(self.shape), self.pi1.reshape(self.shape), 126 | self.i2.reshape(self.shape), self.pi2.reshape(self.shape), 127 | self.grid1, self.grid2, Dss) 128 | 129 | 130 | class ForwardShockablePolicyLottery2D(PolicyLottery2D, ForwardShockableTransition): 131 | def __init__(self, i1, pi1, i2, pi2, grid1, grid2, Dss): 132 | super().__init__(i1, pi1, i2, pi2, grid1, grid2) 133 | self.Dss = Dss.reshape(self.flatshape) 134 | self.space1 = grid1[self.i1+1] - grid1[self.i1] 135 | self.space2 = grid2[self.i2+1] - grid2[self.i2] 136 | 137 | def forward_shock(self, da): 138 | da1, da2 = da 139 | pi_shock1 = -da1.reshape(self.flatshape) / self.space1 140 | pi_shock2 = -da2.reshape(self.flatshape) / self.space2 141 | 142 | return het_compiled.forward_policy_shock_2d(self.Dss, self.i1, self.i2, self.pi1, self.pi2, 143 | pi_shock1, pi_shock2).reshape(self.shape) 144 | 145 | 146 | class Markov(Transition): 147 | def __init__(self, Pi, i): 148 | self.Pi = Pi 149 | self.Pi_T = self.Pi.T 150 | if isinstance(self.Pi_T, np.ndarray): 151 | # optimization: copy to get right order in memory 152 | self.Pi_T = self.Pi_T.copy() 153 | self.i = i 154 | 155 | def forward(self, D): 156 | return multiply_ith_dimension(self.Pi_T, self.i, D) 157 | 158 | def expectation(self, X): 159 | return multiply_ith_dimension(self.Pi, self.i, X) 160 | 161 | def forward_shockable(self, Dss): 162 | return ForwardShockableMarkov(self.Pi, self.i, Dss) 163 | 164 | def expectation_shockable(self, Xss): 165 | return ExpectationShockableMarkov(self.Pi, self.i, Xss) 166 | 167 | def stationary(self, pi_seed, tol=1E-11, maxit=10_000): 168 | return general_stationary(self.Pi, pi_seed, tol, maxit) 169 | 170 | 171 | class ForwardShockableMarkov(Markov, ForwardShockableTransition): 172 | def __init__(self, Pi, i, Dss): 173 | super().__init__(Pi, i) 174 | self.Dss = Dss 175 | 176 | def forward_shock(self, dPi): 177 | return multiply_ith_dimension(dPi.T, self.i, self.Dss) 178 | 179 | 180 | class ExpectationShockableMarkov(Markov, ExpectationShockableTransition): 181 | def __init__(self, Pi, i, Xss): 182 | super().__init__(Pi, i) 183 | self.Xss = Xss 184 | 185 | def expectation_shock(self, dPi): 186 | return multiply_ith_dimension(dPi, self.i, self.Xss) 187 | 188 | 189 | class CombinedTransition(Transition): 190 | def __init__(self, stages: Sequence[Transition]): 191 | self.stages = stages 192 | 193 | def forward(self, D): 194 | for stage in self.stages: 195 | D = stage.forward(D) 196 | return D 197 | 198 | def expectation(self, X): 199 | for stage in reversed(self.stages): 200 | X = stage.expectation(X) 201 | return X 202 | 203 | def forward_shockable(self, Dss): 204 | shockable_stages = [] 205 | for stage in self.stages: 206 | shockable_stages.append(stage.forward_shockable(Dss)) 207 | Dss = stage.forward(Dss) 208 | 209 | return ForwardShockableCombinedTransition(shockable_stages) 210 | 211 | def expectation_shockable(self, Xss): 212 | shockable_stages = [] 213 | for stage in reversed(self.stages): 214 | shockable_stages.append(stage.expectation_shockable(Xss)) 215 | Xss = stage.expectation(Xss) 216 | 217 | return ExpectationShockableCombinedTransition(list(reversed(shockable_stages))) 218 | 219 | def __getitem__(self, i): 220 | return self.stages[i] 221 | 222 | 223 | Shock = Any 224 | ListTupleShocks = Union[List[Shock], Tuple[Shock]] 225 | 226 | class ForwardShockableCombinedTransition(CombinedTransition, ForwardShockableTransition): 227 | def __init__(self, stages: Sequence[ForwardShockableTransition]): 228 | self.stages = stages 229 | self.Dss = stages[0].Dss 230 | 231 | def forward_shock(self, shocks: Optional[Sequence[Optional[Union[Shock, ListTupleShocks]]]]): 232 | if shocks is None: 233 | return None 234 | 235 | # each entry of shocks is either a sequence (list or tuple) 236 | dD = None 237 | 238 | for stage, shock in zip(self.stages, shocks): 239 | if shock is not None: 240 | dD_shock = stage.forward_shock(shock) 241 | else: 242 | dD_shock = None 243 | 244 | if dD is not None: 245 | dD = stage.forward(dD) 246 | 247 | if shock is not None: 248 | dD += dD_shock 249 | else: 250 | dD = dD_shock 251 | 252 | return dD 253 | 254 | 255 | class ExpectationShockableCombinedTransition(CombinedTransition, ExpectationShockableTransition): 256 | def __init__(self, stages: Sequence[ExpectationShockableTransition]): 257 | self.stages = stages 258 | self.Xss = stages[-1].Xss 259 | 260 | def expectation_shock(self, shocks: Sequence[Optional[Union[Shock, ListTupleShocks]]]): 261 | dX = None 262 | 263 | for stage, shock in zip(reversed(self.stages), reversed(shocks)): 264 | if shock is not None: 265 | dX_shock = stage.expectation_shock(shock) 266 | else: 267 | dX_shock = None 268 | 269 | if dX is not None: 270 | dX = stage.expectation(dX) 271 | 272 | if shock is not None: 273 | dX += dX_shock 274 | else: 275 | dX = dX_shock 276 | 277 | return dX 278 | --------------------------------------------------------------------------------