├── notebooks_source ├── build.sh ├── clean.sh ├── equilibrium.md ├── equilib_nd.md ├── opt_invest.md ├── european_option.md ├── jax_intro.md └── demo.md ├── slides ├── gpu.jpg ├── main.pdf ├── ppf.pdf ├── dual_core.png ├── ppf_plus.pdf ├── tradeoff.pdf ├── tradeoff2.pdf ├── tradeoff3.pdf ├── tradeoff4.pdf ├── sloan_logo.png ├── python_vs_rest.png ├── processor_clock.png ├── tradeoff.fig ├── tradeoff2.fig ├── tradeoff3.fig ├── tradeoff4.fig └── main.tex ├── qe-logo-large.png ├── notebooks ├── launch.sh ├── equilibrium.ipynb ├── equilib_nd.ipynb ├── opt_invest.ipynb └── european_option.ipynb ├── README.md └── .gitignore /notebooks_source/build.sh: -------------------------------------------------------------------------------- 1 | jupytext --to notebook *.md 2 | mv *.ipynb ../notebooks 3 | -------------------------------------------------------------------------------- /slides/gpu.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/gpu.jpg -------------------------------------------------------------------------------- /slides/main.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/main.pdf -------------------------------------------------------------------------------- /slides/ppf.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/ppf.pdf -------------------------------------------------------------------------------- /qe-logo-large.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/qe-logo-large.png -------------------------------------------------------------------------------- /notebooks_source/clean.sh: -------------------------------------------------------------------------------- 1 | rm ../notebooks/*.ipynb 2 | rm ../notebooks/*.f90 3 | rm ../notebooks/*.py 4 | -------------------------------------------------------------------------------- /slides/dual_core.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/dual_core.png -------------------------------------------------------------------------------- /slides/ppf_plus.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/ppf_plus.pdf -------------------------------------------------------------------------------- /slides/tradeoff.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/tradeoff.pdf -------------------------------------------------------------------------------- /slides/tradeoff2.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/tradeoff2.pdf -------------------------------------------------------------------------------- /slides/tradeoff3.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/tradeoff3.pdf -------------------------------------------------------------------------------- /slides/tradeoff4.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/tradeoff4.pdf -------------------------------------------------------------------------------- /slides/sloan_logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/sloan_logo.png -------------------------------------------------------------------------------- /slides/python_vs_rest.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/python_vs_rest.png -------------------------------------------------------------------------------- /slides/processor_clock.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/QuantEcon/keio_comp_econ_2023/master/slides/processor_clock.png -------------------------------------------------------------------------------- /notebooks/launch.sh: -------------------------------------------------------------------------------- 1 | jupytext --to notebook ../notebooks_source/*.md 2 | mv ../notebooks_source/*.ipynb . 3 | jupyter-notebook --no-browser --port=8080 4 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Computational Methods for Quantitative Economics 2 | 3 | ## Keio University 4 | 5 | 6 | ![](qe-logo-large.png) 7 | 8 | ## Schedule 9 | 10 | * Date: 27th April 2023 11 | * Timing: 16:30-18:00 12 | * Instructor: [John Stachurski](https://johnstachurski.net/) 13 | 14 | ## Abstract 15 | 16 | Computational methods and scientific computing are becoming increasingly central to research, analysis and policy work in economics, finance and social science. 17 | 18 | This is a workshop for students and policy makers who are interested in learning about the evolution of modern scientific computing tools and how they can be applied to economic problems. The main language used in the workshop will be Python, although other languages will also be discussed. Topics will include simulation, vectorization, JIT compilers, and parallelization. 19 | 20 | ### Resources 21 | 22 | 23 | * [QuantEcon lectures](https://lectures.quantecon.org/) 24 | * [Anaconda Python](https://www.anaconda.com/distribution/) 25 | -------------------------------------------------------------------------------- /slides/tradeoff.fig: -------------------------------------------------------------------------------- 1 | #FIG 3.2 Produced by xfig version 3.2.5c 2 | Landscape 3 | Center 4 | Metric 5 | A4 6 | 100.00 7 | Single 8 | -2 9 | 1200 2 10 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3960 3375 101 101 3859 3375 4061 3375 11 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 4500 4050 101 101 4399 4050 4601 4050 12 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 540 1215 101 101 439 1215 641 1215 13 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 11745 8145 101 101 11644 8145 11846 8145 14 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 7965 3555 101 101 7864 3555 8066 3555 15 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3600 2880 101 101 3499 2880 3701 2880 16 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3825 2430 101 101 3724 2430 3926 2430 17 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 9585 6480 101 101 9484 6480 9686 6480 18 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 8550 5985 101 101 8449 5985 8651 5985 19 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 8325 5670 101 101 8224 5670 8426 5670 20 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 6615 5445 101 101 6514 5445 6716 5445 21 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 5490 4770 101 101 5389 4770 5591 4770 22 | 2 1 0 2 0 7 50 -1 -1 0.000 0 0 -1 1 0 2 23 | 0 0 1.00 60.00 120.00 24 | 2745 7515 2745 1935 25 | 2 1 0 2 0 7 50 -1 -1 0.000 0 0 -1 1 0 2 26 | 0 0 1.00 60.00 120.00 27 | 2115 6795 11025 6795 28 | 4 0 0 50 -1 0 20 0.0000 4 165 1170 9315 7515 Machine speed\001 29 | 4 0 0 50 -1 16 18 0.0000 4 165 720 9720 5715 assembly\001 30 | 4 0 0 50 -1 16 18 0.0000 4 135 630 7335 6345 Fortran\001 31 | 4 0 0 50 -1 16 18 0.0000 4 165 360 3285 3735 Ruby\001 32 | 4 0 7 50 -1 16 18 0.0000 4 135 450 8235 3195 Julia\001 33 | 4 0 0 50 -1 16 18 0.0000 4 165 540 4005 2205 Python\001 34 | 4 0 0 50 -1 16 18 0.0000 4 135 90 8280 5400 C\001 35 | 4 0 0 50 -1 0 20 0.0000 4 165 1080 810 2070 Productivity\001 36 | -------------------------------------------------------------------------------- /slides/tradeoff2.fig: -------------------------------------------------------------------------------- 1 | #FIG 3.2 Produced by xfig version 3.2.7 2 | Landscape 3 | Center 4 | Metric 5 | A4 6 | 100.00 7 | Single 8 | -2 9 | 1200 2 10 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3960 3375 101 101 3859 3375 4061 3375 11 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 4500 4050 101 101 4399 4050 4601 4050 12 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 540 1215 101 101 439 1215 641 1215 13 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 11745 8145 101 101 11644 8145 11846 8145 14 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 7965 3555 101 101 7864 3555 8066 3555 15 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3600 2880 101 101 3499 2880 3701 2880 16 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3825 2430 101 101 3724 2430 3926 2430 17 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 9585 6480 101 101 9484 6480 9686 6480 18 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 8550 5985 101 101 8449 5985 8651 5985 19 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 8325 5670 101 101 8224 5670 8426 5670 20 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 6615 5445 101 101 6514 5445 6716 5445 21 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 5490 4770 101 101 5389 4770 5591 4770 22 | 1 4 0 2 0 8 50 -1 20 0.000 1 0.0000 8145 2835 101 101 8044 2835 8246 2835 23 | 2 1 0 2 0 7 50 -1 -1 0.000 0 0 -1 1 0 2 24 | 0 0 1.00 60.00 120.00 25 | 2745 7515 2745 1935 26 | 2 1 0 2 0 7 50 -1 -1 0.000 0 0 -1 1 0 2 27 | 0 0 1.00 60.00 120.00 28 | 2115 6795 11025 6795 29 | 4 0 0 50 -1 0 20 0.0000 4 165 1170 9315 7515 Machine speed\001 30 | 4 0 0 50 -1 16 18 0.0000 4 165 720 9720 5715 assembly\001 31 | 4 0 0 50 -1 16 18 0.0000 4 135 630 7335 6345 Fortran\001 32 | 4 0 0 50 -1 16 18 0.0000 4 165 360 3285 3735 Ruby\001 33 | 4 0 0 50 -1 16 18 0.0000 4 165 540 4005 2205 Python\001 34 | 4 0 0 50 -1 16 18 0.0000 4 135 90 8280 5400 C\001 35 | 4 0 0 50 -1 0 20 0.0000 4 165 1080 810 2070 Productivity\001 36 | 4 0 7 50 -1 16 18 0.0000 4 135 450 8055 3195 Julia\001 37 | 4 0 0 50 -1 16 18 0.0000 4 135 540 6750 3150 MATLAB\001 38 | -------------------------------------------------------------------------------- /slides/tradeoff3.fig: -------------------------------------------------------------------------------- 1 | #FIG 3.2 Produced by xfig version 3.2.8 2 | Landscape 3 | Center 4 | Metric 5 | A4 6 | 100.00 7 | Single 8 | -2 9 | 1200 2 10 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3960 3375 101 101 3859 3375 4061 3375 11 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 4500 4050 101 101 4399 4050 4601 4050 12 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 540 1215 101 101 439 1215 641 1215 13 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 11745 8145 101 101 11644 8145 11846 8145 14 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 7965 3555 101 101 7864 3555 8066 3555 15 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3600 2880 101 101 3499 2880 3701 2880 16 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3825 2430 101 101 3724 2430 3926 2430 17 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 9585 6480 101 101 9484 6480 9686 6480 18 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 8550 5985 101 101 8449 5985 8651 5985 19 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 8325 5670 101 101 8224 5670 8426 5670 20 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 6615 5445 101 101 6514 5445 6716 5445 21 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 5490 4770 101 101 5389 4770 5591 4770 22 | 1 4 0 2 0 8 50 -1 20 0.000 1 0.0000 8145 2835 101 101 8044 2835 8246 2835 23 | 1 4 0 2 0 8 50 -1 20 0.000 1 0.0000 9090 2970 101 101 8989 2970 9191 2970 24 | 2 1 0 2 0 7 50 -1 -1 0.000 0 0 -1 1 0 2 25 | 0 0 1.00 60.00 120.00 26 | 2745 7515 2745 1935 27 | 2 1 0 2 0 7 50 -1 -1 0.000 0 0 -1 1 0 2 28 | 0 0 1.00 60.00 120.00 29 | 2115 6795 11025 6795 30 | 4 0 0 50 -1 0 20 0.0000 4 165 1170 9315 7515 Machine speed\001 31 | 4 0 0 50 -1 16 18 0.0000 4 165 720 9720 5715 assembly\001 32 | 4 0 0 50 -1 16 18 0.0000 4 135 630 7335 6345 Fortran\001 33 | 4 0 0 50 -1 16 18 0.0000 4 165 360 3285 3735 Ruby\001 34 | 4 0 0 50 -1 16 18 0.0000 4 165 540 4005 2205 Python\001 35 | 4 0 0 50 -1 16 18 0.0000 4 135 90 8280 5400 C\001 36 | 4 0 0 50 -1 0 20 0.0000 4 165 1080 810 2070 Productivity\001 37 | 4 0 7 50 -1 16 18 0.0000 4 135 450 8055 3195 Julia\001 38 | 4 0 0 50 -1 16 18 0.0000 4 135 540 6750 3150 MATLAB\001 39 | 4 0 0 50 -1 16 18 0.0000 4 135 450 9000 3465 Julia\001 40 | -------------------------------------------------------------------------------- /slides/tradeoff4.fig: -------------------------------------------------------------------------------- 1 | #FIG 3.2 Produced by xfig version 3.2.8 2 | Landscape 3 | Center 4 | Metric 5 | A4 6 | 100.00 7 | Single 8 | -2 9 | 1200 2 10 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3960 3375 101 101 3859 3375 4061 3375 11 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 4500 4050 101 101 4399 4050 4601 4050 12 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 540 1215 101 101 439 1215 641 1215 13 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 11745 8145 101 101 11644 8145 11846 8145 14 | 1 4 0 2 7 7 50 -1 20 0.000 1 0.0000 7965 3555 101 101 7864 3555 8066 3555 15 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3600 2880 101 101 3499 2880 3701 2880 16 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 3825 2430 101 101 3724 2430 3926 2430 17 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 9585 6480 101 101 9484 6480 9686 6480 18 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 8550 5985 101 101 8449 5985 8651 5985 19 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 8325 5670 101 101 8224 5670 8426 5670 20 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 6615 5445 101 101 6514 5445 6716 5445 21 | 1 4 0 2 0 13 50 -1 20 0.000 1 0.0000 5490 4770 101 101 5389 4770 5591 4770 22 | 1 4 0 2 0 8 50 -1 20 0.000 1 0.0000 8145 2835 101 101 8044 2835 8246 2835 23 | 1 4 0 2 0 8 50 -1 20 0.000 1 0.0000 9225 2385 101 101 9124 2385 9326 2385 24 | 1 4 0 2 0 8 50 -1 20 0.000 1 0.0000 9090 2970 101 101 8989 2970 9191 2970 25 | 2 1 0 2 0 7 50 -1 -1 0.000 0 0 -1 1 0 2 26 | 0 0 1.00 60.00 120.00 27 | 2745 7515 2745 1935 28 | 2 1 0 2 0 7 50 -1 -1 0.000 0 0 -1 1 0 2 29 | 0 0 1.00 60.00 120.00 30 | 2115 6795 11025 6795 31 | 2 1 0 1 0 0 50 -1 -1 0.000 0 0 -1 1 0 2 32 | 1 1 1.00 60.00 120.00 33 | 4500 2475 7560 2475 34 | 4 0 0 50 -1 0 20 0.0000 4 165 1170 9315 7515 Machine speed\001 35 | 4 0 0 50 -1 16 18 0.0000 4 165 720 9720 5715 assembly\001 36 | 4 0 0 50 -1 16 18 0.0000 4 135 630 7335 6345 Fortran\001 37 | 4 0 0 50 -1 16 18 0.0000 4 165 360 3285 3735 Ruby\001 38 | 4 0 0 50 -1 16 18 0.0000 4 165 540 4005 2205 Python\001 39 | 4 0 0 50 -1 16 18 0.0000 4 135 90 8280 5400 C\001 40 | 4 0 0 50 -1 0 20 0.0000 4 165 1080 810 2070 Productivity\001 41 | 4 0 7 50 -1 16 18 0.0000 4 135 450 8055 3195 Julia\001 42 | 4 0 0 50 -1 16 18 0.0000 4 135 540 6750 3150 MATLAB\001 43 | 4 0 0 50 -1 16 18 0.0000 4 135 450 9000 3465 Julia\001 44 | 4 0 0 50 -1 16 18 0.0000 4 165 2520 7065 2025 Python + NumPy + Numba + JAX\001 45 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *__pycache__* 2 | *ipynb_checkpoints* 3 | 4 | *_minted* 5 | 6 | #/notebooks/ 7 | 8 | *.DStore 9 | 10 | ## Core latex/pdflatex auxiliary files: 11 | *.aux 12 | *.lof 13 | *.log 14 | *.lot 15 | *.fls 16 | *.out 17 | *.toc 18 | *.fmt 19 | *.fot 20 | *.cb 21 | *.cb2 22 | .*.lb 23 | *.ain 24 | 25 | ## Intermediate documents: 26 | *.dvi 27 | *.xdv 28 | *-converted-to.* 29 | # these rules might exclude image files for figures etc. 30 | # *.ps 31 | # *.eps 32 | # *.pdf 33 | 34 | ## Generated if empty string is given at "Please type another file name for output:" 35 | .pdf 36 | 37 | ## Bibliography auxiliary files (bibtex/biblatex/biber): 38 | *.bbl 39 | *.bcf 40 | *.blg 41 | *-blx.aux 42 | *-blx.bib 43 | *.run.xml 44 | 45 | ## Build tool auxiliary files: 46 | *.fdb_latexmk 47 | *.synctex 48 | *.synctex(busy) 49 | *.synctex.gz 50 | *.synctex.gz(busy) 51 | *.pdfsync 52 | 53 | ## Build tool directories for auxiliary files 54 | # latexrun 55 | latex.out/ 56 | 57 | ## Auxiliary and intermediate files from other packages: 58 | # algorithms 59 | *.alg 60 | *.loa 61 | 62 | # achemso 63 | acs-*.bib 64 | 65 | # amsthm 66 | *.thm 67 | 68 | # beamer 69 | *.nav 70 | *.pre 71 | *.snm 72 | *.vrb 73 | 74 | # changes 75 | *.soc 76 | 77 | # comment 78 | *.cut 79 | 80 | # cprotect 81 | *.cpt 82 | 83 | # elsarticle (documentclass of Elsevier journals) 84 | *.spl 85 | 86 | # endnotes 87 | *.ent 88 | 89 | # fixme 90 | *.lox 91 | 92 | # feynmf/feynmp 93 | *.mf 94 | *.mp 95 | *.t[1-9] 96 | *.t[1-9][0-9] 97 | *.tfm 98 | 99 | #(r)(e)ledmac/(r)(e)ledpar 100 | *.end 101 | *.?end 102 | *.[1-9] 103 | *.[1-9][0-9] 104 | *.[1-9][0-9][0-9] 105 | *.[1-9]R 106 | *.[1-9][0-9]R 107 | *.[1-9][0-9][0-9]R 108 | *.eledsec[1-9] 109 | *.eledsec[1-9]R 110 | *.eledsec[1-9][0-9] 111 | *.eledsec[1-9][0-9]R 112 | *.eledsec[1-9][0-9][0-9] 113 | *.eledsec[1-9][0-9][0-9]R 114 | 115 | # glossaries 116 | *.acn 117 | *.acr 118 | *.glg 119 | *.glo 120 | *.gls 121 | *.glsdefs 122 | *.lzo 123 | *.lzs 124 | 125 | # uncomment this for glossaries-extra (will ignore makeindex's style files!) 126 | # *.ist 127 | 128 | # gnuplottex 129 | *-gnuplottex-* 130 | 131 | # gregoriotex 132 | *.gaux 133 | *.gtex 134 | 135 | # htlatex 136 | *.4ct 137 | *.4tc 138 | *.idv 139 | *.lg 140 | *.trc 141 | *.xref 142 | 143 | # hyperref 144 | *.brf 145 | 146 | # knitr 147 | *-concordance.tex 148 | # TODO Comment the next line if you want to keep your tikz graphics files 149 | *.tikz 150 | *-tikzDictionary 151 | 152 | # listings 153 | *.lol 154 | 155 | # luatexja-ruby 156 | *.ltjruby 157 | 158 | # makeidx 159 | *.idx 160 | *.ilg 161 | *.ind 162 | 163 | # minitoc 164 | *.maf 165 | *.mlf 166 | *.mlt 167 | *.mtc[0-9]* 168 | *.slf[0-9]* 169 | *.slt[0-9]* 170 | *.stc[0-9]* 171 | 172 | # minted 173 | _minted* 174 | *.pyg 175 | 176 | # morewrites 177 | *.mw 178 | 179 | # nomencl 180 | *.nlg 181 | *.nlo 182 | *.nls 183 | 184 | # pax 185 | *.pax 186 | 187 | # pdfpcnotes 188 | *.pdfpc 189 | 190 | # sagetex 191 | *.sagetex.sage 192 | *.sagetex.py 193 | *.sagetex.scmd 194 | 195 | # scrwfile 196 | *.wrt 197 | 198 | # sympy 199 | *.sout 200 | *.sympy 201 | sympy-plots-for-*.tex/ 202 | 203 | # pdfcomment 204 | *.upa 205 | *.upb 206 | 207 | # pythontex 208 | *.pytxcode 209 | pythontex-files-*/ 210 | 211 | # tcolorbox 212 | *.listing 213 | 214 | # thmtools 215 | *.loe 216 | 217 | # TikZ & PGF 218 | *.dpth 219 | *.md5 220 | *.auxlock 221 | 222 | # todonotes 223 | *.tdo 224 | 225 | # vhistory 226 | *.hst 227 | *.ver 228 | 229 | # easy-todo 230 | *.lod 231 | 232 | # xcolor 233 | *.xcp 234 | 235 | # xmpincl 236 | *.xmpi 237 | 238 | # xindy 239 | *.xdy 240 | 241 | # xypic precompiled matrices and outlines 242 | *.xyc 243 | *.xyd 244 | 245 | # endfloat 246 | *.ttt 247 | *.fff 248 | 249 | # Latexian 250 | TSWLatexianTemp* 251 | 252 | ## Editors: 253 | # WinEdt 254 | *.bak 255 | *.sav 256 | 257 | # Texpad 258 | .texpadtmp 259 | 260 | # LyX 261 | *.lyx~ 262 | 263 | # Kile 264 | *.backup 265 | 266 | # gummi 267 | .*.swp 268 | 269 | # KBibTeX 270 | *~[0-9]* 271 | 272 | # TeXnicCenter 273 | *.tps 274 | 275 | # auto folder when using emacs and auctex 276 | ./auto/* 277 | *.el 278 | 279 | # expex forward references with \gathertags 280 | *-tags.tex 281 | 282 | # standalone packages 283 | *.sta 284 | 285 | # Makeindex log files 286 | *.lpz 287 | 288 | # Mac files 289 | .DS_Store 290 | .*/.DS_Store 291 | -------------------------------------------------------------------------------- /notebooks_source/equilibrium.md: -------------------------------------------------------------------------------- 1 | --- 2 | jupytext: 3 | text_representation: 4 | extension: .md 5 | format_name: myst 6 | format_version: 0.13 7 | jupytext_version: 1.14.5 8 | kernelspec: 9 | display_name: Python 3 (ipykernel) 10 | language: python 11 | name: python3 12 | --- 13 | 14 | # Equilibrium 15 | 16 | +++ 17 | 18 | #### Author: [John Stachurski](http://johnstachurski.net/) 19 | 20 | +++ 21 | 22 | In this notebook we solve a very simple market equilibrium problem. 23 | 24 | Supply and demand are nonlinear and we use Newton's root-finding algorithm to solve for equilibrium prices. 25 | 26 | We use the following imports: 27 | 28 | ```{code-cell} ipython3 29 | import numpy as np 30 | import matplotlib.pyplot as plt 31 | ``` 32 | 33 | ## Prelude: A Note on Root-Finding 34 | 35 | +++ 36 | 37 | If $f$ maps an interval $(a, b)$ into $\mathbb R$, then a **root** of the function $f$ is an $x^* \in (a,b)$ with $f(x^*)=0$. 38 | 39 | A common method for root finding is Newton's algorithm. 40 | 41 | We start with a guess $x_0 \in (a, b)$. 42 | 43 | Then we replace $f$ with the tangent function $f_a(x) = f(x_0) + f'(x_0)(x - x_0)$ and solve for the root of $f_a$ (which can be done exactly). 44 | 45 | Calling the root $x_1$, we have 46 | 47 | $$ 48 | f_a(x_1)=0 49 | \quad \iff \quad 50 | x_1 = x_0 - \frac{f(x_0)}{f'(x_0)} 51 | $$ 52 | 53 | This is our update rule: 54 | 55 | $$ 56 | x_{k+1} = q(x_k) 57 | \quad \text{where} \quad 58 | q(x) := x - \frac{f(x)}{f'(x)} 59 | $$ 60 | 61 | 62 | +++ 63 | 64 | The algorithm is implemented in `scipy.optimize` 65 | 66 | ```{code-cell} ipython3 67 | from scipy.optimize import newton 68 | ``` 69 | 70 | Let's apply this to find the positive root of $f(x) = x^2 - 1$. 71 | 72 | ```{code-cell} ipython3 73 | def f(x): 74 | return x**2 - 1 75 | 76 | x_grid = np.linspace(-1, 2, 200) 77 | fig, ax = plt.subplots() 78 | ax.plot(x_grid, f(x_grid), label="$f$") 79 | ax.plot(x_grid, np.zeros_like(x_grid), "k--") 80 | ax.legend() 81 | plt.show() 82 | 83 | ``` 84 | 85 | Here we call `newton`. 86 | 87 | ```{code-cell} ipython3 88 | newton(f, 0.5) # search for root of f starting at x_0 = 0.5 89 | ``` 90 | 91 | In the last call we didn't supply the gradient of $f$, so it was approximated 92 | numerically. 93 | 94 | We can supply it as follows: 95 | 96 | ```{code-cell} ipython3 97 | def f_prime(x): 98 | return 2 * x 99 | 100 | newton(lambda x: x**2 - 1, 0.5, fprime=f_prime) 101 | ``` 102 | 103 | ## The Market 104 | 105 | +++ 106 | 107 | Now let's consider a market for coffee beans. The price per kilo is $p$. Total supply at price $p$ is 108 | 109 | $$ q_s (p) = b \sqrt{p} $$ 110 | 111 | and total demand is 112 | 113 | $$ q_d (p) = a \exp(-p) + c, $$ 114 | 115 | where $a, b$ and $c$ are positive parameters. 116 | 117 | +++ 118 | 119 | Now we write routines to compute supply and demand as functions of price and parameters. 120 | 121 | We take $a=1$, $b=0.5$ and $c=1$ as "default" parameter values. 122 | 123 | ```{code-cell} ipython3 124 | def supply(p, b=0.5): 125 | return b * np.sqrt(p) 126 | 127 | def demand(p, a=1, c=1): 128 | return a * np.exp(-p) + c 129 | ``` 130 | 131 | Now we can call the functions as follows: 132 | 133 | ```{code-cell} ipython3 134 | demand(2.0) # with a and c at defaults 135 | ``` 136 | 137 | ```{code-cell} ipython3 138 | demand(2.0, a=0.4) # a is specified and c remains at its defaults 139 | ``` 140 | 141 | etc. 142 | 143 | +++ 144 | 145 | Note that these functions are automatically NumPy "universal functions": 146 | 147 | ```{code-cell} ipython3 148 | p_vals = np.array((0.5, 1.0, 1.5)) 149 | supply(p_vals) 150 | ``` 151 | 152 | ```{code-cell} ipython3 153 | demand(p_vals) 154 | ``` 155 | 156 | ### Exercise 1 157 | 158 | Plot both supply and demand as functions of $p$ on the interval $[0, 10]$ at the default parameters. 159 | 160 | * Put price on the horizonal axis. 161 | * Use a legend to label the two functions and be sure to label the axes. 162 | * Make a rough estimate of the equilibrium price, where demand equals supply. 163 | 164 | ```{code-cell} ipython3 165 | # Put your code here 166 | ``` 167 | 168 | solution below 169 | 170 | solution below 171 | 172 | solution below 173 | 174 | solution below 175 | 176 | solution below 177 | 178 | solution below 179 | 180 | solution below 181 | 182 | solution below 183 | 184 | solution below 185 | 186 | solution below 187 | 188 | solution below 189 | 190 | solution below 191 | 192 | ```{code-cell} ipython3 193 | fig, ax = plt.subplots() 194 | p_grid = np.linspace(0, 10, 200) 195 | ax.plot(p_grid, supply(p_grid), label='supply') 196 | ax.plot(p_grid, demand(p_grid), label='demand') 197 | ax.set_xlabel("price") 198 | ax.set_ylabel("quantity") 199 | ax.legend(frameon=False, loc='upper center') 200 | plt.show() 201 | ``` 202 | 203 | The equilibrium price looks to be about 4.1. 204 | 205 | +++ 206 | 207 | ### Exercise 2 208 | 209 | Write a function that takes arguments $a, b, c, p$, with default values $a=1$, $b=0.5$ and $c=1$, and returns *excess demand*, which is defined as 210 | 211 | $$ e(p) = q_d(p) - q_s(p) $$ 212 | 213 | ```{code-cell} ipython3 214 | # Put your code here 215 | ``` 216 | 217 | solution below 218 | 219 | solution below 220 | 221 | solution below 222 | 223 | solution below 224 | 225 | solution below 226 | 227 | solution below 228 | 229 | solution below 230 | 231 | solution below 232 | 233 | solution below 234 | 235 | solution below 236 | 237 | solution below 238 | 239 | solution below 240 | 241 | 242 | ```{code-cell} ipython3 243 | def excess_demand(p, a=1, b=0.5, c=1): 244 | return demand(p, a, c) - supply(p, b) 245 | ``` 246 | 247 | Now we test it: 248 | 249 | ```{code-cell} ipython3 250 | excess_demand(1.0) 251 | ``` 252 | 253 | ### Organizing our Code 254 | 255 | If we have many functions working with the same parameters, it's hard to know where to put the default values. 256 | 257 | As such, we normally collect them in a data structure, such as a class or a tuple. 258 | 259 | Personally, I normally used `namedtuple` instances, which are lighter than classes but easier to work with than tuples. 260 | 261 | Here's an example: 262 | 263 | ```{code-cell} ipython3 264 | from collections import namedtuple 265 | 266 | Params = namedtuple('Params', ('a', 'b', 'c')) 267 | 268 | def create_market_params(a=1.0, b=0.5, c=1.0): 269 | return Params(a=a, b=b, c=c) 270 | 271 | 272 | def supply(p, params): 273 | a, b, c = params 274 | return b * np.sqrt(p) 275 | 276 | def demand(p, params): 277 | a, b, c = params 278 | return a * np.exp(-p) + c 279 | 280 | def excess_demand(p, params): 281 | a, b, c = params 282 | return demand(p, params) - supply(p, params) 283 | ``` 284 | 285 | ### Exercise 3 286 | 287 | Using these functions, plot excess demand over the interval from $0.2$ up to $10$. Also plot a horizontal line at zero. The equilibrium price is where excess demand crosses zero. 288 | 289 | ```{code-cell} ipython3 290 | # Put your code here 291 | ``` 292 | 293 | 294 | solution below 295 | 296 | solution below 297 | 298 | solution below 299 | 300 | solution below 301 | 302 | solution below 303 | 304 | solution below 305 | 306 | solution below 307 | 308 | solution below 309 | 310 | solution below 311 | 312 | solution below 313 | 314 | solution below 315 | 316 | solution below 317 | 318 | 319 | ```{code-cell} ipython3 320 | params = create_market_params() 321 | 322 | fig, ax = plt.subplots() 323 | p_grid = np.linspace(0, 10, 200) 324 | ax.plot(p_grid, excess_demand(p_grid, params), label='excess demand') 325 | ax.plot(p_grid, np.zeros_like(p_grid), 'k--') 326 | ax.set_xlabel("price") 327 | ax.set_ylabel("quantity") 328 | ax.legend() 329 | plt.show() 330 | ``` 331 | 332 | ### Exercise 4 333 | 334 | +++ 335 | 336 | Write a function that takes an instance of `Params` (i.e, a parameter vector) and returns a market clearing price via Newton's method. 337 | 338 | ```{code-cell} ipython3 339 | # Put your code here 340 | ``` 341 | 342 | 343 | solution below 344 | 345 | solution below 346 | 347 | solution below 348 | 349 | solution below 350 | 351 | solution below 352 | 353 | solution below 354 | 355 | solution below 356 | 357 | solution below 358 | 359 | solution below 360 | 361 | solution below 362 | 363 | solution below 364 | 365 | solution below 366 | 367 | 368 | ```{code-cell} ipython3 369 | def compute_equilibrium(params, price_init=2.0): 370 | p_star = newton(lambda p: excess_demand(p, params), price_init) 371 | return p_star 372 | ``` 373 | 374 | ```{code-cell} ipython3 375 | params = create_market_params() 376 | compute_equilibrium(params) 377 | ``` 378 | 379 | This looks about right given the figures above. 380 | 381 | +++ 382 | 383 | ### Exercise 5 384 | 385 | For $b$ in a grid of 200 values between 0.5 and 1.0, plot the equilibrium price for each $b$. 386 | 387 | Does the curve that you plotted slope up or down? Try to provide an explanation for what you see in terms of market equilibrium. 388 | 389 | ```{code-cell} ipython3 390 | # Put your code here 391 | ``` 392 | 393 | 394 | solution below 395 | 396 | solution below 397 | 398 | solution below 399 | 400 | solution below 401 | 402 | solution below 403 | 404 | solution below 405 | 406 | solution below 407 | 408 | solution below 409 | 410 | solution below 411 | 412 | solution below 413 | 414 | solution below 415 | 416 | solution below 417 | 418 | 419 | 420 | ```{code-cell} ipython3 421 | b_grid = np.linspace(0.5, 1.0, 200) 422 | prices = np.empty_like(b_grid) 423 | for i, b in enumerate(b_grid): 424 | params = create_market_params(b=b) 425 | prices[i] = compute_equilibrium(params) 426 | ``` 427 | 428 | ```{code-cell} ipython3 429 | fig, ax = plt.subplots() 430 | ax.plot(b_grid, prices, label="equilibrium prices") 431 | ax.set_xlabel("$b$") 432 | ax.set_ylabel("price") 433 | ax.legend() 434 | plt.show() 435 | ``` 436 | 437 | The curve slopes down because an increase in $b$ pushes up supply at any given price. (In other words, the supply curve shifts up.) 438 | 439 | With greater supply, the price tends to fall. 440 | 441 | ```{code-cell} ipython3 442 | 443 | ``` 444 | -------------------------------------------------------------------------------- /notebooks_source/equilib_nd.md: -------------------------------------------------------------------------------- 1 | --- 2 | jupytext: 3 | text_representation: 4 | extension: .md 5 | format_name: myst 6 | format_version: 0.13 7 | jupytext_version: 1.14.5 8 | kernelspec: 9 | display_name: Python 3 (ipykernel) 10 | language: python 11 | name: python3 12 | --- 13 | 14 | # Equilibrium in Multiple Dimensions 15 | 16 | +++ 17 | 18 | 19 | #### Author: [John Stachurski](http://johnstachurski.net/) 20 | 21 | ```{code-cell} ipython3 22 | import numpy as np 23 | import matplotlib.pyplot as plt 24 | from numpy import exp, sqrt 25 | from numba import njit 26 | from scipy.optimize import newton, root 27 | ``` 28 | 29 | In this notebook we expore the problem of computing market equilibrium in a multivariate setting, with many goods. 30 | 31 | As a first step, we set up and solve a two-good problem. 32 | 33 | Then we shift to higher dimensions. 34 | 35 | We will show how gradient-based equation solvers can handle high dimensional problems. 36 | 37 | +++ 38 | 39 | ## Two Goods 40 | 41 | +++ 42 | 43 | We first consider a market for two related products, good 0 and good 1, with price vector $p = (p_0, p_1)$. 44 | 45 | Supply of good $i$ at price $p$ is 46 | 47 | $$ q^s_i (p) = b_i \sqrt{p_i} $$ 48 | 49 | Demand of good $i$ at price $p$ is 50 | 51 | $$ q^d_i (p) = \exp(-a_{i0} p_0 - a_{i1} p_1) + c_i$$ 52 | 53 | Here $c_i, b_i$ and $a_{ij}$ are parameters. 54 | 55 | The excess demand functions are 56 | 57 | $$ e_i(p) = q^d_i(p) - q^s_i(p), \qquad i = 0, 1 $$ 58 | 59 | An equilibrium price vector $p^*$ is one where $e_i(p^*) = 0$. 60 | 61 | We set 62 | 63 | $$ 64 | A = \begin{pmatrix} 65 | a_{00} & a_{01} \\ 66 | a_{10} & a_{11} 67 | \end{pmatrix}, 68 | \qquad 69 | b = \begin{pmatrix} 70 | b_0 \\ 71 | b_1 72 | \end{pmatrix} 73 | \qquad \text{and} \qquad 74 | c = \begin{pmatrix} 75 | c_0 \\ 76 | c_1 77 | \end{pmatrix} 78 | $$ 79 | 80 | +++ 81 | 82 | ## Graphical Exploration 83 | 84 | +++ 85 | 86 | Since our problem is only two dimensional, we can use graphical analysis to visualize and help understand the problem. 87 | 88 | Our first step is to define the excess demand function 89 | 90 | $$ e(p) = 91 | \begin{pmatrix} 92 | e_0(p) \\ 93 | e_1(p) 94 | \end{pmatrix} 95 | $$ 96 | 97 | The function below does the job. 98 | 99 | ```{code-cell} ipython3 100 | def e(p, A, b, c): 101 | return exp(- A @ p) + c - b * sqrt(p) 102 | ``` 103 | 104 | Our default parameter values will be 105 | 106 | ```{code-cell} ipython3 107 | A = ((0.5, 0.4), 108 | (0.8, 0.2)) 109 | A = np.asarray(A) 110 | b = np.ones(2) 111 | c = np.ones(2) 112 | ``` 113 | 114 | ```{code-cell} ipython3 115 | e((1.0, 0.5), A, b, c) 116 | ``` 117 | 118 | Next we plot the two functions $e_0$ and $e_1$ on a grid of $(p_0, p_1)$ values, using contour surfaces and lines. 119 | 120 | We will use the following function to build the contour plots. 121 | 122 | ```{code-cell} ipython3 123 | def plot_excess_demand(ax, good=0, grid_size=100, grid_max=4, surface=True): 124 | p_grid = np.linspace(0, grid_max, grid_size) 125 | z = np.empty((100, 100)) 126 | 127 | for i, p_0 in enumerate(p_grid): 128 | for j, p_1 in enumerate(p_grid): 129 | z[i, j] = e((p_0, p_1), A, b, c)[good] 130 | 131 | if surface: 132 | cs1 = ax.contourf(p_grid, p_grid, z.T, alpha=0.5) 133 | plt.colorbar(cs1, ax=ax, format="%.6f") 134 | 135 | ctr1 = ax.contour(p_grid, p_grid, z.T, levels=[0.0]) 136 | plt.clabel(ctr1, inline=1, fontsize=13) 137 | ax.set_xlabel("$p_0$") 138 | ax.set_ylabel("$p_1$") 139 | ``` 140 | 141 | Here's our plot of $e_0$: 142 | 143 | ```{code-cell} ipython3 144 | fig, ax = plt.subplots(figsize=(10, 5.7)) 145 | plot_excess_demand(ax, good=0) 146 | plt.show() 147 | ``` 148 | 149 | We see the black contour line of zero, which tells us when $e_0(p)=0$. 150 | 151 | For a price vector $p$ such that $e_0(p) = 0$, we know that good $0$ is in equilibrium (demand equals supply). 152 | 153 | +++ 154 | 155 | Here's our plot of $e_1$: 156 | 157 | ```{code-cell} ipython3 158 | fig, ax = plt.subplots(figsize=(10, 5.7)) 159 | plot_excess_demand(ax, good=1) 160 | plt.show() 161 | ``` 162 | 163 | Now the black contour line of zero tells us when $e_1(p)=0$ (i.e., good $1$ is in equilibrium). 164 | 165 | +++ 166 | 167 | If these two contour lines cross at some vector $p^*$, then $p^*$ is an equilibrium price vector. 168 | 169 | ```{code-cell} ipython3 170 | fig, ax = plt.subplots(figsize=(10, 5.7)) 171 | for good in (0, 1): 172 | plot_excess_demand(ax, good=good, surface=False) 173 | plt.show() 174 | ``` 175 | 176 | It seems there is an equilibrium close to $p = (1.6, 1.5)$. 177 | 178 | +++ 179 | 180 | ## Using a Multidimensional Root Finder 181 | 182 | +++ 183 | 184 | To solve for $p^*$ more precisely, we use `root`, a root-finding algorithm from `scipy.optimize`. 185 | 186 | We supply $p = (1, 1)$ as our initial guess. 187 | 188 | ```{code-cell} ipython3 189 | init_p = np.ones(2) 190 | ``` 191 | 192 | ```{code-cell} ipython3 193 | solution = root(lambda p: e(p, A, b, c), init_p, method='hybr') 194 | p = solution.x 195 | ``` 196 | 197 | Here's the resulting value: 198 | 199 | ```{code-cell} ipython3 200 | p 201 | ``` 202 | 203 | This looks close to our guess from observing the figure. We can plug it back into $e$ to test that $e(p) \approx 0$: 204 | 205 | ```{code-cell} ipython3 206 | np.max(np.abs(e(p, A, b, c))) 207 | ``` 208 | 209 | This is indeed a very small error. 210 | 211 | +++ 212 | 213 | In most cases, for root-finding algorithms applied to smooth functions, supplying the Jacobian of the function leads to better convergence properties. 214 | 215 | In this case we manually calculate the elements of the Jacobian 216 | 217 | $$ 218 | J(p) = 219 | \begin{pmatrix} 220 | \frac{\partial e_0}{\partial p_0}(p) & \frac{\partial e_0}{\partial p_1}(p) \\ 221 | \frac{\partial e_1}{\partial p_0}(p) & \frac{\partial e_1}{\partial p_1}(p) 222 | \end{pmatrix} 223 | $$ 224 | 225 | and supply the Jacobian as a function, as follows: 226 | 227 | ```{code-cell} ipython3 228 | @njit 229 | def jacobian(p, A, b, c): 230 | p_0, p_1 = p 231 | a_00, a_01 = A[0, :] 232 | a_10, a_11 = A[1, :] 233 | j_00 = -a_00 * exp(-a_00 * p_0) - (b[0]/2) * p_0**(-1/2) 234 | j_01 = -a_01 * exp(-a_01 * p_1) 235 | j_10 = -a_10 * exp(-a_10 * p_0) 236 | j_11 = -a_11 * exp(-a_11 * p_1) - (b[1]/2) * p_1**(-1/2) 237 | J = [[j_00, j_01], 238 | [j_10, j_11]] 239 | return np.array(J) 240 | ``` 241 | 242 | ```{code-cell} ipython3 243 | solution = root(lambda p: e(p, A, b, c), 244 | init_p, 245 | jac=lambda p: jacobian(p, A, b, c), 246 | method='hybr') 247 | p = solution.x 248 | ``` 249 | 250 | Now the solution is even more accurate (although, in this low-dimensional problem, the difference is quite small): 251 | 252 | ```{code-cell} ipython3 253 | 254 | np.max(np.abs(e(p, A, b, c))) 255 | ``` 256 | 257 | ## High-Dimensional Problems 258 | 259 | Our next step is to investigate a high-dimensional version of the market described above. This market consists of 5,000 goods. 260 | 261 | The excess demand function is essentially the same, but now the matrix $A$ is $5000 \times 5000$ and the parameter vectors $b$ and $c$ are $5000 \times 1$. 262 | 263 | ```{code-cell} ipython3 264 | dim = 5000 265 | 266 | # Create a random matrix A and normalize the rows to sum to one 267 | A = np.random.rand(dim, dim) 268 | A = np.asarray(A) 269 | s = np.sum(A, axis=0) 270 | A = A / s 271 | 272 | # Set up b and c 273 | b = np.ones(dim) 274 | c = np.ones(dim) 275 | ``` 276 | 277 | Here's the same demand function: 278 | 279 | ```{code-cell} ipython3 280 | def e(p, A, b, c): 281 | return exp(- A @ p) + c - b * sqrt(p) 282 | ``` 283 | 284 | For our particular case, calculating and supplying the Jacobian is not too hard, but you can imagine that it can be very tedious when the functions get more complicated. 285 | 286 | So let's see how we go when the Jacobian is not supplied. 287 | 288 | Here's our initial condition 289 | 290 | ```{code-cell} ipython3 291 | init_p = np.ones(dim) 292 | ``` 293 | 294 | Now we call `root` again. 295 | 296 | **Warning**: The next line of code takes several minutes to run on a standard laptop or desktop. 297 | 298 | ```{code-cell} ipython3 299 | %%time 300 | solution = root(lambda p: e(p, A, b, c), init_p, method='hybr') 301 | ``` 302 | 303 | ```{code-cell} ipython3 304 | p = solution.x 305 | ``` 306 | 307 | Although it takes a long time to run, the answer is correct up to a high degree of accuracy. 308 | 309 | ```{code-cell} ipython3 310 | np.max(np.abs(e(p, A, b, c))) 311 | ``` 312 | 313 | ## Automatic Differentiation 314 | 315 | +++ 316 | 317 | To produce a faster and more efficient implementation, we shift to using JAX. 318 | 319 | With JAX, we get three big advantages: 320 | 321 | 1. We can use JAX's automatic differentiation to compute the Jacobian easily and efficiently. 322 | 2. JAX can parallelize the problem. 323 | 3. JAX can dispatch to the GPU, if configured 324 | 325 | ```{code-cell} ipython3 326 | !nvidia-smi 327 | ``` 328 | 329 | ```{code-cell} ipython3 330 | import jax 331 | import jax.numpy as jnp 332 | ``` 333 | 334 | We set up the same demand function, replacing `np` with `jnp`: 335 | 336 | ```{code-cell} ipython3 337 | @jax.jit 338 | def e(p, A, b, c): 339 | return jnp.exp(- jnp.dot(A, p)) + c - b * jnp.sqrt(p) 340 | ``` 341 | 342 | We are going to try to compute the equilibrium price using the multivariate version of Newton's method, which means iterating on the equation 343 | 344 | $$ p_{n+1} = p_n - J(p_n)^{-1} e(p_n) $$ 345 | 346 | starting from some initial guess of the price vector $p_0$. (Here $J$ is the Jacobian of $e$.) 347 | 348 | ```{code-cell} ipython3 349 | def newton(f, x_0, tol=1e-5): 350 | f_prime = jax.grad(f) 351 | def q(x): 352 | return x - jnp.linalg.solve(jax.jacobian(f)(x), f(x)) 353 | 354 | error = tol + 1 355 | x = x_0 356 | while error > tol: 357 | y = q(x) 358 | error = jnp.linalg.norm(x - y) 359 | x = y 360 | 361 | return x 362 | ``` 363 | 364 | Let's see whether this can solve the problem and how long it takes 365 | 366 | ```{code-cell} ipython3 367 | %%time 368 | p = newton(lambda p: e(p, A, b, c), init_p).block_until_ready() 369 | ``` 370 | 371 | ```{code-cell} ipython3 372 | %%time 373 | p = newton(lambda p: e(p, A, b, c), init_p).block_until_ready() 374 | ``` 375 | 376 | ```{code-cell} ipython3 377 | np.max(np.abs(e(p, A, b, c))) 378 | ``` 379 | 380 | We still have a solution that's very accurate and the compute time is massively reduced (assuming JAX is connecting to a GPU). 381 | 382 | 383 | ### Exercise 384 | 385 | Write a simplified version of the `newton` function above that works for 386 | scalar functions (real inputs and real outputs). 387 | 388 | Test it on this function: 389 | 390 | ```{code-cell} ipython3 391 | f = lambda x: jnp.sin(4 * (x - 1/4)) + x + x**20 - 1 392 | x = jnp.linspace(0, 1, 100) 393 | 394 | fig, ax = plt.subplots() 395 | ax.plot(x, f(x), label='$f(x)$') 396 | ax.axhline(ls='--', c='k') 397 | ax.set_xlabel('$x$', fontsize=12) 398 | ax.set_ylabel('$f(x)$', fontsize=12) 399 | ax.legend(fontsize=12) 400 | plt.show() 401 | ``` 402 | 403 | 404 | ```{code-cell} ipython3 405 | # Put your code here 406 | ``` 407 | 408 | 409 | solution below 410 | 411 | solution below 412 | 413 | solution below 414 | 415 | solution below 416 | 417 | solution below 418 | 419 | solution below 420 | 421 | solution below 422 | 423 | solution below 424 | 425 | solution below 426 | 427 | solution below 428 | 429 | solution below 430 | 431 | solution below 432 | 433 | 434 | ```{code-cell} ipython3 435 | def newton(f, x_0, tol=1e-5): 436 | f_prime = jax.grad(f) 437 | def q(x): 438 | return x - f(x) / f_prime(x) 439 | 440 | error = tol + 1 441 | x = x_0 442 | while error > tol: 443 | y = q(x) 444 | error = abs(x - y) 445 | x = y 446 | 447 | return x 448 | 449 | newton(f, 0.2) 450 | ``` 451 | -------------------------------------------------------------------------------- /notebooks_source/opt_invest.md: -------------------------------------------------------------------------------- 1 | --- 2 | jupytext: 3 | text_representation: 4 | extension: .md 5 | format_name: myst 6 | format_version: 0.13 7 | jupytext_version: 1.14.5 8 | kernelspec: 9 | display_name: Python 3 (ipykernel) 10 | language: python 11 | name: python3 12 | --- 13 | 14 | 15 | # Optimal Investment 16 | 17 | We require the following library to be installed. 18 | 19 | ```{code-cell} ipython3 20 | :tags: [hide-output] 21 | 22 | !pip install --upgrade quantecon 23 | ``` 24 | 25 | 26 | A monopolist faces inverse demand 27 | curve 28 | 29 | $$ P_t = a_0 - a_1 Y_t + Z_t, $$ 30 | 31 | where 32 | 33 | * $P_t$ is price, 34 | * $Y_t$ is output and 35 | * $Z_t$ is a demand shock. 36 | 37 | We assume that $Z_t$ is a discretized AR(1) process. 38 | 39 | Current profits are 40 | 41 | $$ P_t Y_t - c Y_t - \gamma (Y_{t+1} - Y_t)^2 $$ 42 | 43 | Combining with the demand curve and writing $y, y'$ for $Y_t, Y_{t+1}$, this becomes 44 | 45 | $$ r(y, z, y′) := (a_0 - a_1 y + z - c) y - γ (y′ - y)^2 $$ 46 | 47 | The firm maximizes present value of expected discounted profits. The Bellman equation is 48 | 49 | $$ v(y, z) = \max_{y'} \left\{ r(y, z, y′) + β \sum_{z′} v(y′, z′) Q(z, z′) \right\}. $$ 50 | 51 | We discretize $y$ to a finite grid `y_grid`. 52 | 53 | In essence, the firm tries to choose output close to the monopolist profit maximizer, given $Z_t$, but is constrained by adjustment costs. 54 | 55 | Let's begin with the following imports 56 | 57 | ```{code-cell} ipython3 58 | import quantecon as qe 59 | import jax 60 | import jax.numpy as jnp 61 | import matplotlib.pyplot as plt 62 | ``` 63 | 64 | 65 | Let’s check that we're hooked up to a GPU. 66 | 67 | ```{code-cell} ipython3 68 | !nvidia-smi 69 | ``` 70 | 71 | 72 | We will use 64 bit floats with JAX in order to increase the precision. 73 | 74 | ```{code-cell} ipython3 75 | jax.config.update("jax_enable_x64", True) 76 | ``` 77 | 78 | 79 | We need the following successive approximation function. 80 | 81 | ```{code-cell} ipython3 82 | def successive_approx(T, # Operator (callable) 83 | x_0, # Initial condition 84 | tolerance=1e-6, # Error tolerance 85 | max_iter=10_000, # Max iteration bound 86 | print_step=25, # Print at multiples 87 | verbose=False): 88 | x = x_0 89 | error = tolerance + 1 90 | k = 1 91 | while error > tolerance and k <= max_iter: 92 | x_new = T(x) 93 | error = jnp.max(jnp.abs(x_new - x)) 94 | if verbose and k % print_step == 0: 95 | print(f"Completed iteration {k} with error {error}.") 96 | x = x_new 97 | k += 1 98 | if error > tolerance: 99 | print(f"Warning: Iteration hit upper bound {max_iter}.") 100 | elif verbose: 101 | print(f"Terminated successfully in {k} iterations.") 102 | return x 103 | ``` 104 | 105 | 106 | Let's define a function to create an investment model using the given parameters. 107 | 108 | ```{code-cell} ipython3 109 | def create_investment_model( 110 | r=0.01, # Interest rate 111 | a_0=10.0, a_1=1.0, # Demand parameters 112 | γ=25.0, c=1.0, # Adjustment and unit cost 113 | y_min=0.0, y_max=20.0, y_size=100, # Grid for output 114 | ρ=0.9, ν=1.0, # AR(1) parameters 115 | z_size=150): # Grid size for shock 116 | """ 117 | A function that takes in parameters and returns an instance of Model that 118 | contains data for the investment problem. 119 | """ 120 | β = 1 / (1 + r) 121 | y_grid = jnp.linspace(y_min, y_max, y_size) 122 | mc = qe.tauchen(z_size, ρ, ν) 123 | z_grid, Q = mc.state_values, mc.P 124 | 125 | # Break up parameters into static and nonstatic components 126 | constants = β, a_0, a_1, γ, c 127 | sizes = y_size, z_size 128 | arrays = y_grid, z_grid, Q 129 | 130 | # Shift arrays to the device (e.g., GPU) 131 | arrays = tuple(map(jax.device_put, arrays)) 132 | return constants, sizes, arrays 133 | ``` 134 | 135 | 136 | Let's re-write the vectorized version of the right-hand side of the 137 | Bellman equation (before maximization), which is a 3D array representing: 138 | 139 | $$ 140 | B(y, z, y') = r(y, z, y') + \beta \sum_{z'} v(y', z') Q(z, z') 141 | $$ 142 | 143 | for all $(y, z, y')$. 144 | 145 | ```{code-cell} ipython3 146 | def B(v, constants, sizes, arrays): 147 | """ 148 | A vectorized version of the right-hand side of the Bellman equation 149 | (before maximization) 150 | """ 151 | 152 | # Unpack 153 | β, a_0, a_1, γ, c = constants 154 | y_size, z_size = sizes 155 | y_grid, z_grid, Q = arrays 156 | 157 | # Compute current rewards r(y, z, yp) as array r[i, j, ip] 158 | y = jnp.reshape(y_grid, (y_size, 1, 1)) # y[i] -> y[i, j, ip] 159 | z = jnp.reshape(z_grid, (1, z_size, 1)) # z[j] -> z[i, j, ip] 160 | yp = jnp.reshape(y_grid, (1, 1, y_size)) # yp[ip] -> yp[i, j, ip] 161 | r = (a_0 - a_1 * y + z - c) * y - γ * (yp - y)**2 162 | 163 | # Calculate continuation rewards at all combinations of (y, z, yp) 164 | v = jnp.reshape(v, (1, 1, y_size, z_size)) # v[ip, jp] -> v[i, j, ip, jp] 165 | Q = jnp.reshape(Q, (1, z_size, 1, z_size)) # Q[j, jp] -> Q[i, j, ip, jp] 166 | EV = jnp.sum(v * Q, axis=3) # sum over last index jp 167 | 168 | # Compute the right-hand side of the Bellman equation 169 | return r + β * EV 170 | 171 | # Create a jitted function 172 | B = jax.jit(B, static_argnums=(2,)) 173 | ``` 174 | 175 | 176 | Define a function to compute the current rewards given policy $\sigma$. 177 | 178 | ```{code-cell} ipython3 179 | def compute_r_σ(σ, constants, sizes, arrays): 180 | """ 181 | Compute the array r_σ[i, j] = r[i, j, σ[i, j]], which gives current 182 | rewards given policy σ. 183 | """ 184 | 185 | # Unpack model 186 | β, a_0, a_1, γ, c = constants 187 | y_size, z_size = sizes 188 | y_grid, z_grid, Q = arrays 189 | 190 | # Compute r_σ[i, j] 191 | y = jnp.reshape(y_grid, (y_size, 1)) 192 | z = jnp.reshape(z_grid, (1, z_size)) 193 | yp = y_grid[σ] 194 | r_σ = (a_0 - a_1 * y + z - c) * y - γ * (yp - y)**2 195 | 196 | return r_σ 197 | 198 | 199 | # Create the jitted function 200 | compute_r_σ = jax.jit(compute_r_σ, static_argnums=(2,)) 201 | ``` 202 | 203 | 204 | Define the Bellman operator. 205 | 206 | ```{code-cell} ipython3 207 | def T(v, constants, sizes, arrays): 208 | """The Bellman operator.""" 209 | return jnp.max(B(v, constants, sizes, arrays), axis=2) 210 | 211 | T = jax.jit(T, static_argnums=(2,)) 212 | ``` 213 | 214 | 215 | The following function computes a v-greedy policy. 216 | 217 | ```{code-cell} ipython3 218 | def get_greedy(v, constants, sizes, arrays): 219 | "Computes a v-greedy policy, returned as a set of indices." 220 | return jnp.argmax(B(v, constants, sizes, arrays), axis=2) 221 | 222 | get_greedy = jax.jit(get_greedy, static_argnums=(2,)) 223 | ``` 224 | 225 | 226 | Define the $\sigma$-policy operator. 227 | 228 | ```{code-cell} ipython3 229 | def T_σ(v, σ, constants, sizes, arrays): 230 | """The σ-policy operator.""" 231 | 232 | # Unpack model 233 | β, a_0, a_1, γ, c = constants 234 | y_size, z_size = sizes 235 | y_grid, z_grid, Q = arrays 236 | 237 | r_σ = compute_r_σ(σ, constants, sizes, arrays) 238 | 239 | # Compute the array v[σ[i, j], jp] 240 | zp_idx = jnp.arange(z_size) 241 | zp_idx = jnp.reshape(zp_idx, (1, 1, z_size)) 242 | σ = jnp.reshape(σ, (y_size, z_size, 1)) 243 | V = v[σ, zp_idx] 244 | 245 | # Convert Q[j, jp] to Q[i, j, jp] 246 | Q = jnp.reshape(Q, (1, z_size, z_size)) 247 | 248 | # Calculate the expected sum Σ_jp v[σ[i, j], jp] * Q[i, j, jp] 249 | Ev = jnp.sum(V * Q, axis=2) 250 | 251 | return r_σ + β * jnp.sum(V * Q, axis=2) 252 | 253 | T_σ = jax.jit(T_σ, static_argnums=(3,)) 254 | ``` 255 | 256 | 257 | Next, we want to computes the lifetime value of following policy $\sigma$. 258 | 259 | The basic problem is to solve the linear system 260 | 261 | $$ v(y, z) = r(y, z, \sigma(y, z)) + \beta \sum_{z'} v(\sigma(y, z), z') Q(z, z) $$ 262 | 263 | for $v$. 264 | 265 | It turns out to be helpful to rewrite this as 266 | 267 | $$ v(y, z) = r(y, z, \sigma(y, z)) + \beta \sum_{y', z'} v(y', z') P_\sigma(y, z, y', z') $$ 268 | 269 | where $P_\sigma(y, z, y', z') = 1\{y' = \sigma(y, z)\} Q(z, z')$. 270 | 271 | We want to write this as $v = r_\sigma + \beta P_\sigma v$ and then solve for $v$ 272 | 273 | Note, however, that $v$ is a multi-index array, rather than a vector. 274 | 275 | 276 | The value $v_{\sigma}$ of a policy $\sigma$ is defined as 277 | 278 | $$ 279 | v_{\sigma} = (I - \beta P_{\sigma})^{-1} r_{\sigma} 280 | $$ 281 | 282 | Here we set up the linear map $v \mapsto R_{\sigma} v$, where 283 | 284 | $$ 285 | R_{\sigma} := I - \beta P_{\sigma} 286 | $$ 287 | 288 | In the investment problem, this map can be expressed as 289 | 290 | $$ 291 | (R_{\sigma} v)(y, z) = v(y, z) - \beta \sum_{z'} v(\sigma(y, z), z') Q(z, z') 292 | $$ 293 | 294 | Defining the map as above works in a more intuitive multi-index setting 295 | (e.g. working with $v[i, j]$ rather than flattening v to a one-dimensional 296 | array) and avoids instantiating the large matrix $P_{\sigma}$. 297 | 298 | Let's define the function $R_{\sigma}$. 299 | 300 | ```{code-cell} ipython3 301 | def R_σ(v, σ, constants, sizes, arrays): 302 | 303 | β, a_0, a_1, γ, c = constants 304 | y_size, z_size = sizes 305 | y_grid, z_grid, Q = arrays 306 | 307 | # Set up the array v[σ[i, j], jp] 308 | zp_idx = jnp.arange(z_size) 309 | zp_idx = jnp.reshape(zp_idx, (1, 1, z_size)) 310 | σ = jnp.reshape(σ, (y_size, z_size, 1)) 311 | V = v[σ, zp_idx] 312 | 313 | # Expand Q[j, jp] to Q[i, j, jp] 314 | Q = jnp.reshape(Q, (1, z_size, z_size)) 315 | 316 | # Compute and return v[i, j] - β Σ_jp v[σ[i, j], jp] * Q[j, jp] 317 | return v - β * jnp.sum(V * Q, axis=2) 318 | 319 | R_σ = jax.jit(R_σ, static_argnums=(3,)) 320 | ``` 321 | 322 | 323 | Define a function to get the value $v_{\sigma}$ of policy 324 | $\sigma$ by inverting the linear map $R_{\sigma}$. 325 | 326 | ```{code-cell} ipython3 327 | def get_value(σ, constants, sizes, arrays): 328 | 329 | # Unpack 330 | β, a_0, a_1, γ, c = constants 331 | y_size, z_size = sizes 332 | y_grid, z_grid, Q = arrays 333 | 334 | r_σ = compute_r_σ(σ, constants, sizes, arrays) 335 | 336 | # Reduce R_σ to a function in v 337 | partial_R_σ = lambda v: R_σ(v, σ, constants, sizes, arrays) 338 | 339 | return jax.scipy.sparse.linalg.bicgstab(partial_R_σ, r_σ)[0] 340 | 341 | get_value = jax.jit(get_value, static_argnums=(2,)) 342 | ``` 343 | 344 | 345 | Now we define the solvers, which implement VFI, HPI and OPI. 346 | 347 | ```{code-cell} ipython3 348 | # Implements VFI-Value Function iteration 349 | 350 | def value_iteration(model, tol=1e-5): 351 | constants, sizes, arrays = model 352 | _T = lambda v: T(v, constants, sizes, arrays) 353 | vz = jnp.zeros(sizes) 354 | 355 | v_star = successive_approx(_T, vz, tolerance=tol) 356 | return get_greedy(v_star, constants, sizes, arrays) 357 | ``` 358 | 359 | ```{code-cell} ipython3 360 | # Implements HPI-Howard policy iteration routine 361 | 362 | def policy_iteration(model, maxiter=250): 363 | constants, sizes, arrays = model 364 | vz = jnp.zeros(sizes) 365 | σ = jnp.zeros(sizes, dtype=int) 366 | i, error = 0, 1.0 367 | while error > 0 and i < maxiter: 368 | v_σ = get_value(σ, constants, sizes, arrays) 369 | σ_new = get_greedy(v_σ, constants, sizes, arrays) 370 | error = jnp.max(jnp.abs(σ_new - σ)) 371 | σ = σ_new 372 | i = i + 1 373 | print(f"Concluded loop {i} with error {error}.") 374 | return σ 375 | ``` 376 | 377 | ```{code-cell} ipython3 378 | # Implements the OPI-Optimal policy Iteration routine 379 | 380 | def optimistic_policy_iteration(model, tol=1e-5, m=10): 381 | constants, sizes, arrays = model 382 | v = jnp.zeros(sizes) 383 | error = tol + 1 384 | while error > tol: 385 | last_v = v 386 | σ = get_greedy(v, constants, sizes, arrays) 387 | for _ in range(m): 388 | v = T_σ(v, σ, constants, sizes, arrays) 389 | error = jnp.max(jnp.abs(v - last_v)) 390 | return get_greedy(v, constants, sizes, arrays) 391 | ``` 392 | 393 | ```{code-cell} ipython3 394 | :tags: [hide-output] 395 | 396 | model = create_investment_model() 397 | print("Starting HPI.") 398 | qe.tic() 399 | out = policy_iteration(model) 400 | elapsed = qe.toc() 401 | print(out) 402 | print(f"HPI completed in {elapsed} seconds.") 403 | ``` 404 | 405 | ```{code-cell} ipython3 406 | :tags: [hide-output] 407 | 408 | print("Starting VFI.") 409 | qe.tic() 410 | out = value_iteration(model) 411 | elapsed = qe.toc() 412 | print(out) 413 | print(f"VFI completed in {elapsed} seconds.") 414 | ``` 415 | 416 | ```{code-cell} ipython3 417 | :tags: [hide-output] 418 | 419 | print("Starting OPI.") 420 | qe.tic() 421 | out = optimistic_policy_iteration(model, m=100) 422 | elapsed = qe.toc() 423 | print(out) 424 | print(f"OPI completed in {elapsed} seconds.") 425 | ``` 426 | 427 | 428 | Here's the plot of the Howard policy, as a function of $y$ at the highest and lowest values of $z$. 429 | 430 | ```{code-cell} ipython3 431 | model = create_investment_model() 432 | constants, sizes, arrays = model 433 | β, a_0, a_1, γ, c = constants 434 | y_size, z_size = sizes 435 | y_grid, z_grid, Q = arrays 436 | ``` 437 | 438 | ```{code-cell} ipython3 439 | σ_star = policy_iteration(model) 440 | 441 | fig, ax = plt.subplots(figsize=(9, 5)) 442 | ax.plot(y_grid, y_grid, "k--", label="45") 443 | ax.plot(y_grid, y_grid[σ_star[:, 1]], label="$\\sigma^*(\cdot, z_1)$") 444 | ax.plot(y_grid, y_grid[σ_star[:, -1]], label="$\\sigma^*(\cdot, z_N)$") 445 | ax.legend(fontsize=12) 446 | plt.show() 447 | ``` 448 | 449 | 450 | Let's plot the time taken by each of the solvers and compare them. 451 | 452 | ```{code-cell} ipython3 453 | m_vals = range(5, 3000, 100) 454 | ``` 455 | 456 | ```{code-cell} ipython3 457 | model = create_investment_model() 458 | print("Running Howard policy iteration.") 459 | qe.tic() 460 | σ_pi = policy_iteration(model) 461 | pi_time = qe.toc() 462 | ``` 463 | 464 | ```{code-cell} ipython3 465 | print(f"PI completed in {pi_time} seconds.") 466 | print("Running value function iteration.") 467 | qe.tic() 468 | σ_vfi = value_iteration(model, tol=1e-5) 469 | vfi_time = qe.toc() 470 | print(f"VFI completed in {vfi_time} seconds.") 471 | ``` 472 | 473 | ```{code-cell} ipython3 474 | :tags: [hide-output] 475 | opi_times = [] 476 | for m in m_vals: 477 | print(f"Running optimistic policy iteration with m={m}.") 478 | qe.tic() 479 | σ_opi = optimistic_policy_iteration(model, m=m, tol=1e-5) 480 | opi_time = qe.toc() 481 | print(f"OPI with m={m} completed in {opi_time} seconds.") 482 | opi_times.append(opi_time) 483 | ``` 484 | 485 | ```{code-cell} ipython3 486 | fig, ax = plt.subplots(figsize=(9, 5)) 487 | ax.plot(m_vals, jnp.full(len(m_vals), pi_time), 488 | lw=2, label="Howard policy iteration") 489 | ax.plot(m_vals, jnp.full(len(m_vals), vfi_time), 490 | lw=2, label="value function iteration") 491 | ax.plot(m_vals, opi_times, lw=2, label="optimistic policy iteration") 492 | ax.legend(fontsize=12, frameon=False) 493 | ax.set_xlabel("$m$", fontsize=12) 494 | ax.set_ylabel("time(s)", fontsize=12) 495 | plt.show() 496 | ``` 497 | -------------------------------------------------------------------------------- /notebooks_source/european_option.md: -------------------------------------------------------------------------------- 1 | --- 2 | jupytext: 3 | text_representation: 4 | extension: .md 5 | format_name: myst 6 | format_version: 0.13 7 | jupytext_version: 1.14.5 8 | kernelspec: 9 | display_name: Python 3 (ipykernel) 10 | language: python 11 | name: python3 12 | --- 13 | 14 | # Monte Carlo and Option Pricing 15 | 16 | * Author: [John Stachurski](http://johnstachurski.net/) 17 | 18 | We discuss [Monte Carlo 19 | methods](https://en.wikipedia.org/wiki/Monte_Carlo_method) for computing 20 | expectations with applications in finance. 21 | 22 | Our main application will be pricing a European option. 23 | 24 | We will show that Monte Carlo is particularly helpful when the distribution of 25 | interest has no neat analytical form. 26 | 27 | We will also touch on some high performance computing topics, including 28 | 29 | * just-in-time compilers 30 | * GPUs and 31 | * parallelization. 32 | 33 | We begin with the following imports: 34 | 35 | ```{code-cell} ipython3 36 | import numpy as np 37 | import matplotlib.pyplot as plt 38 | import numba 39 | from numpy.random import randn 40 | ``` 41 | 42 | ## Pricing a European Call Option under Risk Neutrality 43 | 44 | As our next step, we are going to price a European call option under risk 45 | neutrality. 46 | 47 | Let's first discuss risk neutrality and then introduce European options. 48 | 49 | ### Risk-Neutral Pricing 50 | 51 | When we use risk-neutral pricing, we determine the price of a given asset 52 | according to its expected payoff. 53 | 54 | $$ 55 | \text{cost } = \text{ expected benefit} 56 | $$ 57 | 58 | For example, suppose someone promises to pay you 59 | 60 | - 1,000,000 dollars if "heads" is the outcome of a fair coin flip 61 | - 0 dollars if "tails" is the outcome 62 | 63 | Let's denote the payoff as $G$, so that 64 | 65 | $$ 66 | \mathbb P\left\{G = 10^6 \right\} = \mathbb P\{G = 0\} = \frac{1}{2} 67 | $$ 68 | 69 | Suppose in addition that you can sell this promise to anyone who wants to 70 | hold it 71 | 72 | - First they pay you $P$, the price at which you sell it 73 | - Then they get $G$, which could be either 1,000,000 or 0. 74 | 75 | What's a fair price for this asset (this promise)? 76 | 77 | The definition of fair is ambiguous but what we can say is that the 78 | risk-neutral price is 500,000 dollars. 79 | 80 | This is because the risk-neutral price is just the expected payoff of the 81 | asset, which is 82 | 83 | $$ 84 | \mathbb E G = \frac{1}{2} \times 10^6 + \frac{1}{2} \times 0 = 5 \times 10^5 85 | $$ 86 | 87 | +++ 88 | 89 | ### A Comment on Risk 90 | 91 | As suggested by the name, the risk-neutral price ignores risk. 92 | 93 | To understand this, consider whether you would pay 500,000 dollars for such a 94 | promise. 95 | 96 | Would you prefer to receive 500,000 for sure or 1,000,000 dollars with 97 | 50% probability and nothing with 50% probability? 98 | 99 | At least some readers will strictly prefer the first option --- although some 100 | might prefer the second. 101 | 102 | Thinking about this makes us realize that 500,000 is not necessarily the 103 | "right" price --- or the price that we would see if there was a market for 104 | these promises. 105 | 106 | Nonetheless, the risk-neutral price is an important benchmark, which economists 107 | and financial market participants routinely try to calculate. 108 | 109 | +++ 110 | 111 | ### Discounting 112 | 113 | One thing we ignored in the previous discussion was time. 114 | 115 | In general, receiving $x$ dollars now is preferable to receiving $x$ dollars 116 | in $n$ periods (e.g., 10 years). 117 | 118 | After all, if we receive $x$ dollars now, we could put it in the bank at 119 | interest rate $r > 0$ and receive $ (1 + r)^n x $ in $n$ periods. 120 | 121 | Hence future payments need to be discounted. 122 | 123 | We will implement discounting by 124 | 125 | * multiplying a payment in one period by $\beta < 1$ 126 | * multiplying a payment in $n$ periods by $\beta^n$, etc. 127 | 128 | The same adjustment needs to be applied to our risk-neutral price for the 129 | promise described above. 130 | 131 | Thus, if $G$ is realized in $n$ periods, then the risk-neutral price is 132 | 133 | $$ 134 | P = \beta^n \mathbb E G 135 | = \beta^n 5 \times 10^5 136 | $$ 137 | 138 | +++ 139 | 140 | ### European Call Options 141 | 142 | Now let's price a European call option. 143 | 144 | The option is described by three things: 145 | 146 | 2. $n$, the **expiry date**, 147 | 2. $K$, the **strike price**, and 148 | 3. $S_n$, the price of the **underlying** asset at date $n$. 149 | 150 | For example, suppose that the underlying is one share in Amazon. 151 | 152 | The owner of this option has the right to buy one share in Amazon at price $K$ after $n$ days. 153 | 154 | If $S_n > K$, then the owner will exercise the option, buy at $K$, sell at 155 | $S_n$, and make profit $S_n - K$. 156 | 157 | If $S_n \leq K$, then the owner will not exercise the option and the payoff is zero. 158 | 159 | Thus, the payoff is $\max\{ S_n - K, 0 \}$. 160 | 161 | Under the assumption of risk neutrality, the price of the option is 162 | the expected discounted payoff: 163 | 164 | $$ P = \beta^n \mathbb E \max\{ S_n - K, 0 \} $$ 165 | 166 | Now all we need to do is specify the distribution of $S_n$, so the expectation 167 | can be calculated. 168 | 169 | +++ 170 | 171 | **Exercise** 172 | 173 | Suppose we know that $S_n \sim LN(\mu, \sigma)$ and $\mu$ and $\sigma$ are known. 174 | 175 | Use the fact that if $S_n^1, \ldots, S_n^M$ are independent draws from this lognormal distribution then, by the law of large numbers, 176 | 177 | $$ \mathbb E \max\{ S_n - K, 0 \} 178 | \approx 179 | \frac{1}{M} \sum_{m=1}^M \max \{S_n^m - K, 0 \} 180 | $$ 181 | 182 | +++ 183 | 184 | Use the following parameter values: 185 | 186 | ```{code-cell} ipython3 187 | μ = 1.0 188 | σ = 0.1 189 | ``` 190 | 191 | ```{code-cell} ipython3 192 | K = 1 193 | n = 10 194 | β = 0.95 195 | ``` 196 | 197 | Set the simulation size to 198 | 199 | ```{code-cell} ipython3 200 | M = 10_000_000 201 | ``` 202 | 203 | 204 | solution below 205 | 206 | solution below 207 | 208 | solution below 209 | 210 | solution below 211 | 212 | solution below 213 | 214 | solution below 215 | 216 | solution below 217 | 218 | solution below 219 | 220 | solution below 221 | 222 | solution below 223 | 224 | solution below 225 | 226 | solution below 227 | 228 | ```{code-cell} ipython3 229 | S = np.exp(μ + σ * np.random.randn(M)) 230 | return_draws = np.maximum(S - K, 0) 231 | P = β**n * np.mean(return_draws) 232 | print(f"The Monte Carlo option price is {P:3f}") 233 | ``` 234 | 235 | ## Pricing Via a Dynamic Model 236 | 237 | In this exercise we investigate a more realistic model for the share price $S_n$. 238 | 239 | This comes from specifying the underlying dynamics of the share price. 240 | 241 | First we specify the dynamics. 242 | 243 | Then we'll compute the price of the option using Monte Carlo. 244 | 245 | ### Simple Dynamics 246 | 247 | One simple model for $\{S_t\}$ is 248 | 249 | $$ \ln \frac{S_{t+1}}{S_t} = \mu + \sigma \xi_{t+1} $$ 250 | 251 | where 252 | 253 | * $S_0$ is normally distributed and 254 | * $\{ \xi_t \}$ is IID and standard normal. 255 | 256 | 257 | Under the stated assumptions, $S_n$ is lognormally distributed. 258 | 259 | **Exercise** Explain why this is true. 260 | 261 | 262 | ### Problems with Simple Dynamics 263 | 264 | The simple dynamic model we studied above is convenient, since we can work out 265 | the distribution of $S_n$. 266 | 267 | 268 | However, its predictions are counterfactual because, in the real world, 269 | volatility (measured by $\sigma$) is not stationary. 270 | 271 | Instead it rather changes over time, sometimes high (like during the GFC) and sometimes low. 272 | 273 | 274 | ### More Realistic Dynamics 275 | 276 | As stated above, one problem with our simple model is that $\sigma$ is 277 | constant. 278 | 279 | This leads us to study the improved version: 280 | 281 | $$ \ln \frac{S_{t+1}}{S_t} = \mu + \sigma_t \xi_{t+1} $$ 282 | 283 | where 284 | 285 | $$ 286 | \sigma_t = \exp(h_t), 287 | \quad 288 | h_{t+1} = \rho h_t + \nu \eta_{t+1} 289 | $$ 290 | 291 | Here $\{\eta_t\}$ is also IID and standard normal. 292 | 293 | +++ 294 | 295 | ### Default Parameters 296 | 297 | For the dynamic model, we adopt the following parameter values. 298 | 299 | ```{code-cell} ipython3 300 | μ = 0.0001 301 | ρ = 0.1 302 | ν = 0.001 303 | S0 = 10 304 | h0 = 0 305 | ``` 306 | 307 | (Here `S0` is $S_0$ and `h0` is $h_0$.) 308 | 309 | For the option we use the following defaults. 310 | 311 | ```{code-cell} ipython3 312 | K = 100 313 | n = 10 314 | β = 0.95 315 | ``` 316 | 317 | **Exercise** 318 | 319 | 320 | Write a function that simulates the sequence $S_0, \ldots, S_n$, where the parameters are set to 321 | 322 | Plot 50 paths of the form $S_0, \ldots, S_n$. 323 | 324 | ```{code-cell} ipython3 325 | # Put your code here 326 | ``` 327 | 328 | solution below 329 | 330 | solution below 331 | 332 | solution below 333 | 334 | solution below 335 | 336 | solution below 337 | 338 | solution below 339 | 340 | solution below 341 | 342 | solution below 343 | 344 | solution below 345 | 346 | solution below 347 | 348 | solution below 349 | 350 | solution below 351 | 352 | 353 | With $s_t := \ln S_t$, the price dynamics become 354 | 355 | $$ s_{t+1} = s_t + \mu + \exp(h_t) \xi_{t+1} $$ 356 | 357 | Here is a function to simulate a path using this equation: 358 | 359 | ```{code-cell} ipython3 360 | from numpy.random import randn 361 | 362 | def simulate_asset_price_path(μ=μ, S0=S0, h0=h0, n=n, ρ=ρ, ν=ν): 363 | s = np.empty(n+1) 364 | s[0] = np.log(S0) 365 | 366 | h = h0 367 | for t in range(n): 368 | s[t+1] = s[t] + μ + np.exp(h) * randn() 369 | h = ρ * h + ν * randn() 370 | 371 | return np.exp(s) 372 | ``` 373 | 374 | Here we plot the paths and the log of the paths. 375 | 376 | ```{code-cell} ipython3 377 | fig, axes = plt.subplots(2, 1) 378 | 379 | titles = 'log paths', 'paths' 380 | transforms = np.log, lambda x: x 381 | for ax, transform, title in zip(axes, transforms, titles): 382 | for i in range(50): 383 | path = simulate_asset_price_path() 384 | ax.plot(transform(path)) 385 | ax.set_title(title) 386 | 387 | fig.tight_layout() 388 | plt.show() 389 | ``` 390 | 391 | **Exercise** 392 | 393 | Compute the price of the option $P_0$ by Monte Carlo, averaging over realizations $S_n^1, \ldots, S_n^M$ of $S_n$ and appealing to the law of large numbers: 394 | 395 | $$ \mathbb E \max\{ S_n - K, 0 \} 396 | \approx 397 | \frac{1}{M} \sum_{m=1}^M \max \{S_n^m - K, 0 \} 398 | $$ 399 | 400 | 401 | To the extend that you can, write fast, efficient code to compute the option price. 402 | 403 | In particular, try to speed up the code above using `jit` or `njit` from Numba. 404 | 405 | ```{code-cell} ipython3 406 | # Put your code here 407 | ``` 408 | 409 | solution below 410 | 411 | solution below 412 | 413 | solution below 414 | 415 | solution below 416 | 417 | solution below 418 | 419 | solution below 420 | 421 | solution below 422 | 423 | solution below 424 | 425 | solution below 426 | 427 | solution below 428 | 429 | solution below 430 | 431 | solution below 432 | 433 | ```{code-cell} ipython3 434 | from numba import njit, prange 435 | ``` 436 | 437 | ```{code-cell} ipython3 438 | @njit 439 | def compute_call_price(β=β, 440 | μ=μ, 441 | S0=S0, 442 | h0=h0, 443 | K=K, 444 | n=n, 445 | ρ=ρ, 446 | ν=ν, 447 | M=10_000_000): 448 | current_sum = 0.0 449 | # For each sample path 450 | for m in range(M): 451 | s = np.log(S0) 452 | h = h0 453 | # Simulate forward in time 454 | for t in range(n): 455 | s = s + μ + np.exp(h) * randn() 456 | h = ρ * h + ν * randn() 457 | # And add the value max{S_n - K, 0} to current_sum 458 | current_sum += np.maximum(np.exp(s) - K, 0) 459 | 460 | return β**n * current_sum / M 461 | ``` 462 | 463 | ```{code-cell} ipython3 464 | %%time 465 | compute_call_price() 466 | ``` 467 | 468 | **Exercise** 469 | 470 | If you can, use `prange` from Numba to parallelize this code and make it even faster. 471 | 472 | ```{code-cell} ipython3 473 | # Put your code here 474 | ``` 475 | 476 | solution below 477 | 478 | solution below 479 | 480 | solution below 481 | 482 | solution below 483 | 484 | solution below 485 | 486 | solution below 487 | 488 | solution below 489 | 490 | solution below 491 | 492 | solution below 493 | 494 | solution below 495 | 496 | solution below 497 | 498 | solution below 499 | 500 | ```{code-cell} ipython3 501 | @njit(parallel=True) 502 | def compute_call_price_parallel(β=β, 503 | μ=μ, 504 | S0=S0, 505 | h0=h0, 506 | K=K, 507 | n=n, 508 | ρ=ρ, 509 | ν=ν, 510 | M=10_000_000): 511 | current_sum = 0.0 512 | # For each sample path 513 | for m in prange(M): 514 | s = np.log(S0) 515 | h = h0 516 | # Simulate forward in time 517 | for t in range(n): 518 | s = s + μ + np.exp(h) * randn() 519 | h = ρ * h + ν * randn() 520 | # And add the value max{S_n - K, 0} to current_sum 521 | current_sum += np.maximum(np.exp(s) - K, 0) 522 | 523 | return β**n * current_sum / M 524 | ``` 525 | 526 | ```{code-cell} ipython3 527 | from numba import get_num_threads, set_num_threads 528 | get_num_threads() 529 | ``` 530 | 531 | ```{code-cell} ipython3 532 | %%time 533 | compute_call_price_parallel() 534 | ``` 535 | 536 | ```{code-cell} ipython3 537 | %%time 538 | compute_call_price_parallel() 539 | ``` 540 | 541 | ## Pricing a European Call Option Using JAX 542 | 543 | Previously we computed the value of a European call option via Monte Carlo using Numba-based routines. 544 | 545 | Let's compare how this looks, and how fast it runs, when we implement using [Google JAX](https://python-programming.quantecon.org/jax_intro.html). 546 | 547 | +++ 548 | 549 | **Exercise** 550 | 551 | Try to shift the whole operation to the GPU using JAX and test your speed gain. 552 | 553 | ```{code-cell} ipython3 554 | # Put your code here 555 | ``` 556 | 557 | solution below 558 | 559 | solution below 560 | 561 | solution below 562 | 563 | solution below 564 | 565 | solution below 566 | 567 | solution below 568 | 569 | solution below 570 | 571 | solution below 572 | 573 | solution below 574 | 575 | solution below 576 | 577 | solution below 578 | 579 | solution below 580 | 581 | ```{code-cell} ipython3 582 | !nvidia-smi 583 | ``` 584 | 585 | ```{code-cell} ipython3 586 | import jax 587 | import jax.numpy as jnp 588 | ``` 589 | 590 | ```{code-cell} ipython3 591 | @jax.jit 592 | def compute_call_price_jax(β=β, 593 | μ=μ, 594 | S0=S0, 595 | h0=h0, 596 | K=K, 597 | n=n, 598 | ρ=ρ, 599 | ν=ν, 600 | M=10_000_000, 601 | key=jax.random.PRNGKey(1)): 602 | 603 | s = jnp.full(M, np.log(S0)) 604 | h = jnp.full(M, h0) 605 | for t in range(n): 606 | key, subkey = jax.random.split(key) 607 | Z = jax.random.normal(subkey, (2, M)) 608 | s = s + μ + jnp.exp(h) * Z[0, :] 609 | h = ρ * h + ν * Z[1, :] 610 | expectation = jnp.mean(jnp.maximum(jnp.exp(s) - K, 0)) 611 | 612 | return β**n * expectation 613 | ``` 614 | 615 | ```{code-cell} ipython3 616 | %%time 617 | compute_call_price_jax().block_until_ready() 618 | ``` 619 | 620 | ```{code-cell} ipython3 621 | %%time 622 | compute_call_price_jax().block_until_ready() 623 | ``` 624 | -------------------------------------------------------------------------------- /notebooks_source/jax_intro.md: -------------------------------------------------------------------------------- 1 | --- 2 | jupytext: 3 | text_representation: 4 | extension: .md 5 | format_name: myst 6 | format_version: 0.13 7 | jupytext_version: 1.14.5 8 | kernelspec: 9 | display_name: Python 3 (ipykernel) 10 | language: python 11 | name: python3 12 | --- 13 | 14 | # An Introduction to JAX 15 | 16 | #### John Stachurski 17 | 18 | +++ 19 | 20 | [JAX](https://github.com/google/jax) is scientific library within the Python ecosystem that provides data types, functions and a compiler for fast linear algebra operations and automatic differentiation. 21 | 22 | Loosely speaking, JAX is like NumPy with the addition of 23 | 24 | * automatic differentiation 25 | * automated GPU/TPU support 26 | * a just-in-time compiler 27 | 28 | JAX is often used for machine learning and AI, since it can scale to big data operations on GPUs and automatically differentiate loss functions for gradient decent. 29 | 30 | However, JAX is sufficiently low-level that it can be used for many purposes. 31 | 32 | Here is a short history of JAX: 33 | 34 | * 2015: Google open-sources part of its AI infrastructure called TensorFlow. 35 | * 2016: The popularity of TensorFlow grows rapidly. 36 | * 2017: Facebook open-sources PyTorch beta, an alternative AI framework (developer-friendly, more Pythonic) 37 | * 2018: Facebook launches a full production-ready version of PyTorch. 38 | * 2019: PyTorch surges in popularity (adopted by Uber, Airbnb, Tesla, etc.) 39 | * 2020: Google launches JAX as an open-source framework. 40 | * 2021: Google starts to shift away from TPUs to Nvidia GPUs, extends JAX capabilities. 41 | * 2022: Uptake of Google JAX accelerates rapidly 42 | 43 | +++ 44 | 45 | We begin this notebook with some standard imports 46 | 47 | ```{code-cell} ipython3 48 | import numpy as np 49 | import matplotlib as plt 50 | from numba import jit, njit, float64, vectorize 51 | ``` 52 | 53 | 54 | ## Installation 55 | 56 | JAX can be installed with or without GPU support. 57 | 58 | * Follow [the install guide](https://github.com/google/jax) 59 | 60 | Note that JAX is pre-installed with GPU support on [Google Colab](https://colab.research.google.com/). 61 | 62 | (Colab Pro offers better GPUs.) 63 | 64 | ```{code-cell} ipython3 65 | !nvidia-smi 66 | ``` 67 | 68 | +++ 69 | 70 | ## JAX as a NumPy Replacement 71 | 72 | +++ 73 | 74 | One way to use JAX is as a plug-in NumPy replacement. Let's look at the similarities and differences. 75 | 76 | ### Similarities 77 | 78 | +++ 79 | 80 | The following import is standard, replacing `import numpy as np`: 81 | 82 | ```{code-cell} ipython3 83 | import jax 84 | import jax.numpy as jnp 85 | ``` 86 | 87 | Now we can use `jnp` in place of `np` for the usual array operations: 88 | 89 | ```{code-cell} ipython3 90 | a = jnp.asarray((1.0, 3.2, -1.5)) 91 | ``` 92 | 93 | ```{code-cell} ipython3 94 | print(a) 95 | ``` 96 | 97 | ```{code-cell} ipython3 98 | print(jnp.sum(a)) 99 | ``` 100 | 101 | ```{code-cell} ipython3 102 | print(jnp.mean(a)) 103 | ``` 104 | 105 | ```{code-cell} ipython3 106 | print(jnp.dot(a, a)) 107 | ``` 108 | 109 | However, the array object `a` is not a NumPy array: 110 | 111 | ```{code-cell} ipython3 112 | a 113 | ``` 114 | 115 | ```{code-cell} ipython3 116 | type(a) 117 | ``` 118 | 119 | Even scalar-valued maps on arrays return JAX arrays: 120 | 121 | ```{code-cell} ipython3 122 | jnp.sum(a) 123 | ``` 124 | 125 | JAX arrays are allocated on the `device`. 126 | 127 | Here `device` refers to the hardware accelerator (GPU or TPU), although JAX falls back to the CPU if no accelerator is detected. 128 | 129 | (In the terminology of GPUs, the "host" is the machine that launches GPU operations, while the "device" is the GPU itself.) 130 | 131 | 132 | +++ 133 | 134 | Operations on higher dimensional arrays is also similar to NumPy: 135 | 136 | ```{code-cell} ipython3 137 | A = jnp.ones((2, 2)) 138 | B = jnp.identity(2) 139 | A @ B 140 | ``` 141 | 142 | ```{code-cell} ipython3 143 | from jax.numpy import linalg 144 | ``` 145 | 146 | ```{code-cell} ipython3 147 | linalg.solve(B, A) 148 | ``` 149 | 150 | ```{code-cell} ipython3 151 | linalg.eigh(B) # Computes eigenvalues and eigenvectors 152 | ``` 153 | 154 | ### Differences 155 | 156 | +++ 157 | 158 | One difference between NumPy and JAX is that, when running on a GPU, JAX uses 32 bit floats by default. This is standard for GPU computing and can lead to significant speed gains with small loss of precision. 159 | 160 | However, for some calculations precision matters. In these cases 64 bit floats can be enforced via the command 161 | 162 | ```{code-cell} ipython3 163 | jax.config.update("jax_enable_x64", True) 164 | ``` 165 | 166 | Let's check this works: 167 | 168 | ```{code-cell} ipython3 169 | jnp.ones(3) 170 | ``` 171 | 172 | As a NumPy replacement, a more significant difference is that arrays are treated as **immutable**. For example, with NumPy we can write 173 | 174 | ```{code-cell} ipython3 175 | import numpy as np 176 | a = np.linspace(0, 1, 3) 177 | a 178 | ``` 179 | 180 | and then mutate the data in memory: 181 | 182 | ```{code-cell} ipython3 183 | a[0] = 1 184 | a 185 | ``` 186 | 187 | In JAX this fails: 188 | 189 | ```{code-cell} ipython3 190 | a = jnp.linspace(0, 1, 3) 191 | a 192 | ``` 193 | 194 | ```{code-cell} ipython3 195 | a[0] = 1 196 | ``` 197 | 198 | In line with immutability, JAX does not support inplace operations: 199 | 200 | ```{code-cell} ipython3 201 | a = np.array((2, 1)) 202 | a.sort() 203 | a 204 | ``` 205 | 206 | ```{code-cell} ipython3 207 | a = jnp.array((2, 1)) 208 | a.sort() 209 | a 210 | ``` 211 | 212 | The designers of JAX chose to make arrays immutable because JAX uses a functional programming style. More on this below. 213 | 214 | Note that, while mutation is discouraged, it is in fact possible with `at`, as in 215 | 216 | ```{code-cell} ipython3 217 | a = jnp.linspace(0, 1, 3) 218 | id(a) 219 | ``` 220 | 221 | ```{code-cell} ipython3 222 | a 223 | ``` 224 | 225 | ```{code-cell} ipython3 226 | a.at[0].set(1) 227 | ``` 228 | 229 | We can check that the array is mutated by verifying its identity is unchanged: 230 | 231 | ```{code-cell} ipython3 232 | id(a) 233 | ``` 234 | 235 | ## Random Numbers 236 | 237 | +++ 238 | 239 | Random numbers are also a bit different in JAX, relative to NumPy. Typically, in JAX, the state of the random number generator needs to be controlled explicitly. 240 | 241 | ```{code-cell} ipython3 242 | import jax.random as random 243 | ``` 244 | 245 | First we produce a key, which seeds the random number generator. 246 | 247 | ```{code-cell} ipython3 248 | key = random.PRNGKey(1) 249 | ``` 250 | 251 | ```{code-cell} ipython3 252 | type(key) 253 | ``` 254 | 255 | ```{code-cell} ipython3 256 | print(key) 257 | ``` 258 | 259 | Now we can use the key to generate some random numbers: 260 | 261 | ```{code-cell} ipython3 262 | x = random.normal(key, (3, 3)) 263 | x 264 | ``` 265 | 266 | If we use the same key again, we initialize at the same seed, so the random numbers are the same: 267 | 268 | ```{code-cell} ipython3 269 | random.normal(key, (3, 3)) 270 | ``` 271 | 272 | To produce a (quasi-) independent draw, best practice is to "split" the existing key: 273 | 274 | ```{code-cell} ipython3 275 | key, subkey = random.split(key) 276 | ``` 277 | 278 | ```{code-cell} ipython3 279 | random.normal(key, (3, 3)) 280 | ``` 281 | 282 | ```{code-cell} ipython3 283 | random.normal(subkey, (3, 3)) 284 | ``` 285 | 286 | The function below produces `k` (quasi-) independent random `n x n` matrices using this procedure. 287 | 288 | ```{code-cell} ipython3 289 | def gen_random_matrices(key, n, k): 290 | matrices = [] 291 | for _ in range(k): 292 | key, subkey = random.split(key) 293 | matrices.append(random.uniform(subkey, (n, n))) 294 | return matrices 295 | ``` 296 | 297 | ```{code-cell} ipython3 298 | matrices = gen_random_matrices(key, 2, 2) 299 | for A in matrices: 300 | print(A) 301 | ``` 302 | 303 | One point to remember is that JAX expects tuples to describe array shapes, even for flat arrays. Hence, to get a one-dimensional array of normal random draws we use `(len, )` for the shape, as in 304 | 305 | ```{code-cell} ipython3 306 | random.normal(key, (5, )) 307 | ``` 308 | 309 | ## JIT Compilation 310 | 311 | 312 | The JAX JIT compiler accelerates logic within functions by fusing linear algebra operations into a single, highly optimized kernel that the host can launch on the GPU / TPU (or CPU if no accelerator is detected). 313 | 314 | +++ 315 | 316 | Consider the following pure Python function. 317 | 318 | ```{code-cell} ipython3 319 | def f(x, p=1000): 320 | return sum((k*x for k in range(p))) 321 | ``` 322 | 323 | Let's build an array to call the function on. 324 | 325 | ```{code-cell} ipython3 326 | n = 50_000_000 327 | x = jnp.ones(n) 328 | ``` 329 | 330 | How long does the function take to execute? 331 | 332 | ```{code-cell} ipython3 333 | %time f(x).block_until_ready() 334 | ``` 335 | 336 | This code is not particularly fast. While it is run on the GPU, since `x` is a DeviceArray, each vector `k * x` has to be instantiated before the final sum is computed. 337 | 338 | If we JIT-compile the function with JAX, then the operations are fused and no intermediate arrays are created. 339 | 340 | ```{code-cell} ipython3 341 | f_jit = jax.jit(f) # target for JIT compilation 342 | ``` 343 | 344 | Let's run once to compile it: 345 | 346 | ```{code-cell} ipython3 347 | f_jit(x) 348 | ``` 349 | 350 | And now let's time it. 351 | 352 | ```{code-cell} ipython3 353 | %time f_jit(x).block_until_ready() 354 | ``` 355 | 356 | ## Functional Programming 357 | 358 | From JAX's documentation: 359 | 360 | *When walking about the countryside of Italy, the people will not hesitate to tell you that JAX has "una anima di pura programmazione funzionale".* 361 | 362 | +++ 363 | 364 | In other words, JAX assumes a functional programming style. 365 | 366 | The major implication is that JAX functions should be pure: 367 | 368 | * no dependence on global variables 369 | * no side effects 370 | 371 | "A pure function will always return the same result if invoked with the same inputs." 372 | 373 | JAX will not usually throw errors when compiling impure functions but execution becomes unpredictable. 374 | 375 | +++ 376 | 377 | Here's an illustration of this fact, using global variables: 378 | 379 | ```{code-cell} ipython3 380 | a = 1 # global 381 | 382 | @jax.jit 383 | def f(x): 384 | return a + x 385 | ``` 386 | 387 | ```{code-cell} ipython3 388 | x = jnp.ones(2) 389 | ``` 390 | 391 | ```{code-cell} ipython3 392 | f(x) 393 | ``` 394 | 395 | In the code above, the global value `a=1` is fused into the jitted function. 396 | 397 | Even if we change `a`, the output of `f` will not be affected --- as long as the same compiled version is called. 398 | 399 | ```{code-cell} ipython3 400 | a = 42 401 | ``` 402 | 403 | ```{code-cell} ipython3 404 | f(x) 405 | ``` 406 | 407 | Changing the dimension of the input triggers a fresh compilation of the function, at which time the change in the value of `a` takes effect: 408 | 409 | ```{code-cell} ipython3 410 | x = np.ones(3) 411 | ``` 412 | 413 | ```{code-cell} ipython3 414 | f(x) 415 | ``` 416 | 417 | Moral of the story: write pure functions when using JAX! 418 | 419 | +++ 420 | 421 | ## Gradients 422 | 423 | +++ 424 | 425 | JAX can use automatic differentiation to compute gradients. 426 | 427 | This can be extremely useful in optimization, root finding and other applications. 428 | 429 | Here's a very simple illustration, involving the function 430 | 431 | ```{code-cell} ipython3 432 | def f(x): 433 | return (x**2) / 2 434 | ``` 435 | 436 | Let's take the derivative: 437 | 438 | ```{code-cell} ipython3 439 | f_prime = jax.grad(f) 440 | ``` 441 | 442 | ```{code-cell} ipython3 443 | f_prime(10.0) 444 | ``` 445 | 446 | Let's plot the function and derivative, noting that $f'(x) = x$. 447 | 448 | ```{code-cell} ipython3 449 | import matplotlib.pyplot as plt 450 | plt.style.use('fivethirtyeight') 451 | 452 | fig, ax = plt.subplots() 453 | x_grid = jnp.linspace(-4, 4, 200) 454 | ax.plot(x_grid, f(x_grid), label="$f$") 455 | ax.plot(x_grid, jax.vmap(f_prime)(x_grid), label="$f'$") 456 | ax.legend(loc='upper center') 457 | plt.show() 458 | ``` 459 | 460 | ```{code-cell} ipython3 461 | 462 | ``` 463 | 464 | 465 | ## Exercise 466 | 467 | Recall that Newton's method for solving for the root of $f$ involves iterating on 468 | 469 | +++ 470 | 471 | $$ q(x) = x - \frac{f(x)}{f'(x)} $$ 472 | 473 | Write a function called `newton` that takes a function $f$ plus a guess $x_0$ and returns an approximate fixed point. Your `newton` implementation should use automatic differentiation to calculate $f'$. 474 | 475 | Test your `newton` method on the function shown below. 476 | 477 | ```{code-cell} ipython3 478 | f = lambda x: jnp.sin(4 * (x - 1/4)) + x + x**20 - 1 479 | x = jnp.linspace(0, 1, 100) 480 | 481 | fig, ax = plt.subplots() 482 | ax.plot(x, f(x), label='$f(x)$') 483 | ax.axhline(ls='--', c='k') 484 | ax.set_xlabel('$x$', fontsize=12) 485 | ax.set_ylabel('$f(x)$', fontsize=12) 486 | ax.legend(fontsize=12) 487 | plt.show() 488 | ``` 489 | 490 | ```{code-cell} ipython3 491 | # Put your code here 492 | ``` 493 | 494 | solution below 495 | 496 | solution below 497 | 498 | solution below 499 | 500 | solution below 501 | 502 | solution below 503 | 504 | solution below 505 | 506 | solution below 507 | 508 | solution below 509 | 510 | solution below 511 | 512 | solution below 513 | 514 | 515 | 516 | 517 | ```{code-cell} ipython3 518 | def newton(f, x_0, tol=1e-5): 519 | f_prime = jax.grad(f) 520 | def q(x): 521 | return x - f(x) / f_prime(x) 522 | 523 | error = tol + 1 524 | x = x_0 525 | while error > tol: 526 | y = q(x) 527 | error = abs(x - y) 528 | x = y 529 | 530 | return x 531 | ``` 532 | 533 | ```{code-cell} ipython3 534 | newton(f, 0.2) 535 | ``` 536 | 537 | This number looks good, given the figure. 538 | 539 | 540 | 541 | ## Exercise 542 | 543 | This exercise uses parallelized gradient ascent to maximize a function. 544 | 545 | Here's the function we want to maximize 546 | 547 | ```{code-cell} ipython3 548 | @jax.jit 549 | def f(x): 550 | return jnp.cos(x[0]**2 + x[1]**2) / (1 + x[0]**2 + x[1]**2) + 1 551 | 552 | ``` 553 | 554 | Here's one update step of gradient ascent 555 | 556 | ```{code-cell} ipython3 557 | f_grad = jax.grad(f) 558 | 559 | def update(x, f, f_grad, alpha=0.01): 560 | return x + alpha * f_grad(x) 561 | 562 | x_0 = jnp.array((0.7, 0.7)) 563 | 564 | update(x_0, f, f_grad) 565 | ``` 566 | 567 | Let's vectorize it. 568 | 569 | ```{code-cell} ipython3 570 | update_vec = jax.vmap(update, (0, None, None)) 571 | ``` 572 | 573 | Let's test that this works as expected. 574 | 575 | ```{code-cell} ipython3 576 | key = jax.random.PRNGKey(1) 577 | n = 1000 578 | xs = jax.random.uniform(key, (n, 2), minval=-3.0, maxval=3.0) 579 | xs 580 | ``` 581 | 582 | ```{code-cell} ipython3 583 | update_vec(xs, f, f_grad) 584 | ``` 585 | 586 | We have updated every 2-vector in `xs` (every row) using the update rule. 587 | 588 | The exercise is to run this in a loop and compute an approximate maximum of 589 | the function. 590 | 591 | ```{code-cell} ipython3 592 | # Put your code here 593 | ``` 594 | 595 | solution below 596 | 597 | solution below 598 | 599 | solution below 600 | 601 | solution below 602 | 603 | solution below 604 | 605 | solution below 606 | 607 | solution below 608 | 609 | solution below 610 | 611 | solution below 612 | 613 | solution below 614 | 615 | 616 | Here's a suitable function for the loop phase. 617 | 618 | ```{code-cell} ipython3 619 | def gradient_ascent(f, f_grad, x_0, tol=1e-8, alpha=1e-2, max_iter=10_000): 620 | error = tol + 1 621 | x = x_0 622 | i = 0 623 | current_max = - jnp.inf 624 | while error > tol and i < max_iter: 625 | y = update_vec(x, f, f_grad) 626 | new_max = jnp.max(jax.vmap(f)(x)) 627 | error = abs(new_max - current_max) 628 | current_max = new_max 629 | x = y 630 | i += 1 631 | 632 | return current_max, i 633 | ``` 634 | 635 | Now let's call it, starting from `xs` 636 | 637 | 638 | ```{code-cell} ipython3 639 | gradient_ascent(f, f_grad, xs) 640 | ``` 641 | -------------------------------------------------------------------------------- /notebooks_source/demo.md: -------------------------------------------------------------------------------- 1 | --- 2 | jupytext: 3 | text_representation: 4 | extension: .md 5 | format_name: myst 6 | format_version: 0.13 7 | jupytext_version: 1.14.5 8 | kernelspec: 9 | display_name: Python 3 (ipykernel) 10 | language: python 11 | name: python3 12 | --- 13 | 14 | 15 | # Accelerating Python 16 | 17 | [John Stachurski](http://johnstachurski.net) 18 | 19 | This notebook demonstrates ways of accelerating plain Python code in 20 | scientific applications. 21 | 22 | We begin by importing some libraries that will be discussed below. 23 | 24 | 25 | ```{code-cell} ipython3 26 | import numpy as np 27 | from numpy.random import randn 28 | import numba 29 | from numba import vectorize, float64 30 | import matplotlib.pyplot as plt 31 | import jax 32 | import jax.numpy as jnp 33 | ``` 34 | 35 | 36 | 37 | ## Problem 1: A Time Series Model 38 | 39 | Consider the time series model 40 | 41 | $$ x_{t+1} = \alpha x_t (1 - x_t) $$ 42 | 43 | 44 | Our aim is to generate time series from this model and analyze them. 45 | 46 | We will show how to accelerate this operation. 47 | 48 | To begin, let's set $\alpha = 4$ 49 | 50 | 51 | 52 | ```{code-cell} ipython3 53 | α = 4 54 | ``` 55 | 56 | Here's a typical time series: 57 | 58 | ```{code-cell} ipython3 59 | n = 200 60 | x = np.empty(n) 61 | x[0] = 0.2 62 | for t in range(n-1): 63 | x[t+1] = α * x[t] * (1 - x[t]) 64 | 65 | plt.plot(x) 66 | plt.show() 67 | ``` 68 | 69 | ### Python Test 70 | 71 | +++ 72 | 73 | Here's a function that iterates forward `n` times, starting from `x0`, and 74 | returns **the final** value: 75 | 76 | ```{code-cell} ipython3 77 | def quad(x0, n): 78 | x = x0 79 | for i in range(n): 80 | x = α * x * (1 - x) 81 | return x 82 | ``` 83 | 84 | Let's see how fast this runs: 85 | 86 | ```{code-cell} ipython3 87 | n = 10_000_000 88 | ``` 89 | 90 | ```{code-cell} ipython3 91 | %%time 92 | x = quad(0.2, n) 93 | ``` 94 | 95 | ### Fortran Test 96 | 97 | +++ 98 | 99 | Now let's try this in Fortran. 100 | 101 | Note --- this step is intended to be a demo and will only execute if 102 | 103 | * you have the file `fastquad.f90` in your pwd 104 | * you have a Fortran compiler installed and modify the compilation code below appropriately 105 | 106 | ```{code-cell} ipython3 107 | %%file fortran_quad.f90 108 | 109 | PURE FUNCTION QUAD(X0, N) 110 | IMPLICIT NONE 111 | INTEGER, PARAMETER :: DP=KIND(0.d0) 112 | REAL(dp), INTENT(IN) :: X0 113 | REAL(dp) :: QUAD 114 | INTEGER :: I 115 | INTEGER, INTENT(IN) :: N 116 | QUAD = X0 117 | DO I = 1, N - 1 118 | QUAD = 4.0_dp * QUAD * real(1.0_dp - QUAD, dp) 119 | END DO 120 | RETURN 121 | END FUNCTION QUAD 122 | 123 | PROGRAM MAIN 124 | IMPLICIT NONE 125 | INTEGER, PARAMETER :: DP=KIND(0.d0) 126 | REAL(dp) :: START, FINISH, X, QUAD 127 | INTEGER :: N 128 | N = 10000000 129 | X = QUAD(0.2_dp, 10) 130 | CALL CPU_TIME(START) 131 | X = QUAD(0.2_dp, N) 132 | CALL CPU_TIME(FINISH) 133 | PRINT *,'last val = ', X 134 | PRINT *,'Elapsed time in milliseconds = ', (FINISH-START) * 1000 135 | END PROGRAM MAIN 136 | ``` 137 | 138 | ```{code-cell} ipython3 139 | !gfortran -O3 fortran_quad.f90 140 | ``` 141 | 142 | ```{code-cell} ipython3 143 | !./a.out 144 | ``` 145 | 146 | Tidy up 147 | 148 | ```{code-cell} ipython3 149 | !rm a.out 150 | !rm fortran_quad.f90 151 | ``` 152 | 153 | ### Codon 154 | 155 | Let's try `codon`, an AOT Python compiler 156 | 157 | First we install it --- if not yet installed, please uncomment 158 | 159 | ```{code-cell} ipython3 160 | # !/bin/bash -c "$(curl -fsSL https://exaloop.io/install.sh)" 161 | ``` 162 | 163 | Now we write Python code to a file. 164 | 165 | 166 | ```{code-cell} ipython3 167 | %%file codon_quad.py 168 | 169 | from time import time 170 | 171 | n = 10_000_000 172 | alpha = 4.0 173 | 174 | def quad(x0, n): 175 | x = x0 176 | for i in range(1, n): 177 | x = alpha * x * (1 - x) 178 | return x 179 | 180 | 181 | t0 = time() 182 | x = quad(0.1, n) 183 | t1 = time() 184 | print(x) 185 | print("Elapsed time in milliseconds: ", (t1 - t0) * 1000) 186 | ``` 187 | 188 | Next we compile the Python code to build an executable. 189 | 190 | ```{code-cell} ipython3 191 | !codon build --release --exe codon_quad.py 192 | ``` 193 | 194 | Now let's run it. 195 | 196 | ```{code-cell} ipython3 197 | !./codon_quad 198 | ``` 199 | 200 | Tidying up: 201 | 202 | ```{code-cell} ipython3 203 | !rm codon_quad 204 | !rm codon_quad.py 205 | ``` 206 | 207 | 208 | ### Python + Numba 209 | 210 | 211 | Now let's replicate the calculations using Numba's JIT compiler. 212 | 213 | Here's the Python function we want to speed up 214 | 215 | 216 | ```{code-cell} ipython3 217 | @numba.jit 218 | def quad(x0, n): 219 | x = x0 220 | for i in range(1, n): 221 | x = α * x * (1 - x) 222 | return x 223 | ``` 224 | 225 | This is the same as before except that we've targeted the function for JIT 226 | compilation with `@numba.jit`. 227 | 228 | Let's see how fast it runs. 229 | 230 | ```{code-cell} ipython3 231 | %%time 232 | x = quad(0.2, n) 233 | ``` 234 | 235 | ```{code-cell} ipython3 236 | %%time 237 | x = quad(0.2, n) 238 | ``` 239 | 240 | 241 | 242 | ## Problem 2: Multivariate Optimization 243 | 244 | The problem is to maximize the function 245 | 246 | $$ f(x, y) = \frac{\cos \left(x^2 + y^2 \right)}{1 + x^2 + y^2} + 1$$ 247 | 248 | using brute force --- searching over a grid of $(x, y)$ pairs. 249 | 250 | ```{code-cell} ipython3 251 | def f(x, y): 252 | return np.cos(x**2 + y**2) / (1 + x**2 + y**2) + 1 253 | ``` 254 | 255 | ```{code-cell} ipython3 256 | from mpl_toolkits.mplot3d.axes3d import Axes3D 257 | from matplotlib import cm 258 | 259 | gridsize = 50 260 | gmin, gmax = -3, 3 261 | xgrid = np.linspace(gmin, gmax, gridsize) 262 | ygrid = xgrid 263 | x, y = np.meshgrid(xgrid, ygrid) 264 | 265 | # === plot value function === # 266 | fig = plt.figure(figsize=(8, 6)) 267 | ax = fig.add_subplot(111, projection='3d') 268 | ax.plot_surface(x, 269 | y, 270 | f(x, y), 271 | rstride=2, cstride=2, 272 | cmap=cm.jet, 273 | alpha=0.4, 274 | linewidth=0.05) 275 | 276 | 277 | ax.scatter(x, y, c='k', s=0.6) 278 | 279 | ax.scatter(x, y, f(x, y), c='k', s=0.6) 280 | 281 | ax.view_init(25, -57) 282 | ax.set_zlim(-0, 2.0) 283 | ax.set_xlim(gmin, gmax) 284 | ax.set_ylim(gmin, gmax) 285 | 286 | plt.show() 287 | ``` 288 | 289 | Let's try a few different methods to make it fast. 290 | 291 | 292 | 293 | ### Vectorized Numpy 294 | 295 | ```{code-cell} ipython3 296 | grid = np.linspace(-3, 3, 10000) 297 | 298 | x, y = np.meshgrid(grid, grid) 299 | ``` 300 | 301 | ```{code-cell} ipython3 302 | %%time 303 | 304 | np.max(f(x, y)) 305 | ``` 306 | 307 | ### JITTed code 308 | 309 | 310 | A jitted version 311 | 312 | ```{code-cell} ipython3 313 | @numba.jit 314 | def compute_max(): 315 | m = -np.inf 316 | for x in grid: 317 | for y in grid: 318 | z = np.cos(x**2 + y**2) / (1 + x**2 + y**2) + 1 319 | if z > m: 320 | m = z 321 | return m 322 | ``` 323 | 324 | ```{code-cell} ipython3 325 | compute_max() 326 | ``` 327 | 328 | ```{code-cell} ipython3 329 | %%time 330 | compute_max() 331 | ``` 332 | 333 | ### Vectorized Numba on the CPU 334 | 335 | 336 | Numba for vectorization with automatic parallelization; 337 | 338 | ```{code-cell} ipython3 339 | @vectorize('float64(float64, float64)', target='parallel') 340 | def f_par(x, y): 341 | return np.cos(x**2 + y**2) / (1 + x**2 + y**2) + 1 342 | ``` 343 | 344 | ```{code-cell} ipython3 345 | x, y = np.meshgrid(grid, grid) 346 | 347 | np.max(f_par(x, y)) 348 | ``` 349 | 350 | ```{code-cell} ipython3 351 | %%time 352 | np.max(f_par(x, y)) 353 | ``` 354 | 355 | 356 | 357 | ### JAX on the GPU 358 | 359 | Now let's try JAX. 360 | 361 | This code will work well if you have a GPU and JAX configured to use it. 362 | 363 | Let's see what we have available. 364 | 365 | ```{code-cell} ipython3 366 | !nvidia-smi 367 | ``` 368 | 369 | 370 | 371 | #### Replacing NumPy with JAX 372 | 373 | For this step we replace `np` with `jnp`, which is our alias for `jax.numpy` 374 | 375 | Warning --- you need a GPU with relatively large memory for this to work. 376 | 377 | 378 | ```{code-cell} ipython3 379 | def f(x, y): 380 | return jnp.cos(x**2 + y**2) / (1 + x**2 + y**2) + 1 381 | ``` 382 | 383 | 384 | ```{code-cell} ipython3 385 | grid = np.linspace(-3, 3, 10000) 386 | 387 | x, y = jnp.meshgrid(grid, grid) 388 | ``` 389 | 390 | Here's our timing. 391 | 392 | ```{code-cell} ipython3 393 | %%time 394 | 395 | jnp.max(f(x, y)) 396 | ``` 397 | 398 | #### JIT Compiling the Function 399 | 400 | Let's JIT-compile the function and see if anything changes. 401 | 402 | ```{code-cell} ipython3 403 | f = jax.jit(f) 404 | ``` 405 | 406 | ```{code-cell} ipython3 407 | %%time 408 | 409 | jnp.max(f(x, y)) 410 | ``` 411 | 412 | 413 | ```{code-cell} ipython3 414 | %%time 415 | 416 | jnp.max(f(x, y)) 417 | ``` 418 | 419 | #### JIT Compiling the Whole Routine 420 | 421 | Now let's JIT-compile the function and the optimization routine. 422 | 423 | ```{code-cell} ipython3 424 | def max_on_grid(func, grid_size=10_000): 425 | grid = jnp.linspace(-3, 3, 10000) 426 | x, y = jnp.meshgrid(grid, grid) 427 | return jnp.max(func(x, y)) 428 | ``` 429 | 430 | ```{code-cell} ipython3 431 | max_on_grid = jax.jit(max_on_grid, static_argnums=(0,)) 432 | ``` 433 | 434 | ```{code-cell} ipython3 435 | %%time 436 | max_on_grid(f) 437 | ``` 438 | 439 | ```{code-cell} ipython3 440 | %%time 441 | max_on_grid(f) 442 | ``` 443 | 444 | 445 | ## Problem 3: Monte Carlo 446 | 447 | 448 | In this section we describe the Monte Carlo method of integration via a simple 449 | example. 450 | 451 | ### Share Price with Known Distribution 452 | 453 | Let's suppose that we are considering buying a share (or many shares) in a 454 | given company. 455 | 456 | Our plan is either to 457 | 458 | * buy it now, hold it for one year and then sell it, or 459 | * do something else with our money. 460 | 461 | We start by thinking of the share price in one year as a random variable $S$. 462 | 463 | (Let's forget about dividends for now, so that our return on holding the share 464 | is the relative change in its price.) 465 | 466 | To decide whether or not to go ahead, we need to know some features of the 467 | distribution of $S$. 468 | 469 | For example, we might decide to buy if the mean is high and the variance is 470 | low. 471 | 472 | (High expected returns and low risk.) 473 | 474 | Suppose that, after analyzing the data, we have decided that $S$ is well 475 | represented by a lognormal distribution with parameters $\mu, \sigma$ . 476 | 477 | * $S$ has the same distribution as $\exp(\mu + \sigma Z)$ where $Z$ is standard normal. 478 | * we write this statement as $S \sim LN(\mu, \sigma)$. 479 | 480 | Any good reference on statistics will tell us that the mean and variance are 481 | 482 | $$ 483 | \mathbb E S 484 | = \exp \left(\mu + \frac{\sigma^2}{2} \right) 485 | $$ 486 | 487 | and 488 | 489 | $$ 490 | \mathop{\mathrm{Var}} S 491 | = [\exp(\sigma^2) - 1] \exp(2\mu + \sigma^2) 492 | $$ 493 | 494 | So far we have no need for a computer. 495 | 496 | 497 | ### Share Price with Unknown Distribution 498 | 499 | But now suppose that we study the distribution of $S$ more carefully, leading 500 | us to decompose the price into multiple factors. 501 | 502 | In particular, we conclude that the share price depends on three variables, 503 | with 504 | 505 | $$ 506 | S = (X_1 + X_2 + X_3)^p 507 | $$ 508 | 509 | We assume that 510 | 511 | * $p$ is a positive number, which is known to us, 512 | * $X_i \sim LN(\mu_i, \sigma_i)$ for $i=1,2,3$, 513 | * the values of $\mu_i, \sigma_i$ have all been estimated, and 514 | * the random variables $X_1$, $X_2$ and $X_3$ are independent. 515 | 516 | How should we compute the mean of $S$? 517 | 518 | To do this with pencil and paper is hard (unless, say, $p=1$). 519 | 520 | But fortunately there's an easy way to do this, at least approximately: 521 | 522 | 1. Generate $n$ independent draws of $X_1$, $X_2$ and $X_3$ on a computer, 523 | 1. Use these draws to generate $n$ independent draws of $S$, and 524 | 1. Take the average value of these draws of $S$. 525 | 526 | By the law of large numbers, this average will be close to the true mean when 527 | $n$ is large. 528 | 529 | We use the following values for $p$ and each $\mu_i$ and $\sigma_i$. 530 | 531 | ```{code-cell} ipython3 532 | n = 10_000_000 533 | p = 0.5 534 | μ_1, μ_2, μ_3 = 0.2, 0.8, 0.4 535 | σ_1, σ_2, σ_3 = 0.1, 0.05, 0.2 536 | ``` 537 | 538 | ### A Routine using Loops in Python 539 | 540 | +++ 541 | 542 | Here's a routine using native Python loops to calculate the desired mean 543 | 544 | $$ 545 | \frac{1}{n} \sum_{i=1}^n S_i 546 | \approx \mathbb E S 547 | $$ 548 | 549 | 550 | ```{code-cell} ipython3 551 | def compute_mean(n=10_000_000): 552 | S = 0.0 553 | for i in range(n): 554 | X_1 = np.exp(μ_1 + σ_1 * randn()) 555 | X_2 = np.exp(μ_2 + σ_2 * randn()) 556 | X_3 = np.exp(μ_3 + σ_3 * randn()) 557 | S += (X_1 + X_2 + X_3)**p 558 | return(S / n) 559 | ``` 560 | 561 | Let's test it and see how long it takes. 562 | 563 | ```{code-cell} ipython3 564 | %%time 565 | 566 | compute_mean() 567 | ``` 568 | 569 | 570 | ### A Vectorized Routine 571 | 572 | +++ 573 | 574 | Now we implement a vectorized routine using traditional NumPy array processing. 575 | 576 | ```{code-cell} ipython3 577 | 578 | def compute_mean_vectorized(n=10_000_000): 579 | X_1 = np.exp(μ_1 + σ_1 * randn(n)) 580 | X_2 = np.exp(μ_2 + σ_2 * randn(n)) 581 | X_3 = np.exp(μ_3 + σ_3 * randn(n)) 582 | S = (X_1 + X_2 + X_3)**p 583 | return(S.mean()) 584 | ``` 585 | 586 | ```{code-cell} ipython3 587 | %%time 588 | 589 | compute_mean_vectorized() 590 | ``` 591 | 592 | 593 | ### Using Google JAX 594 | 595 | 596 | Finally, let's try to shift this to the GPU and parallelize it effectively. 597 | 598 | 599 | ```{code-cell} ipython3 600 | !nvidia-smi 601 | ``` 602 | 603 | ```{code-cell} ipython3 604 | :tags: [] 605 | 606 | def compute_mean_jax(n=10_000_000): 607 | key = jax.random.PRNGKey(1) 608 | Z = jax.random.normal(key, (3, n)) 609 | X_1 = jnp.exp(μ_1 + σ_1 * Z[0,:]) 610 | X_2 = jnp.exp(μ_2 + σ_2 * Z[1,:]) 611 | X_3 = jnp.exp(μ_3 + σ_3 * Z[2,:]) 612 | S = (X_1 + X_2 + X_3)**p 613 | return(S.mean()) 614 | ``` 615 | 616 | ```{code-cell} ipython3 617 | %%time 618 | 619 | compute_mean_jax() 620 | ``` 621 | 622 | ```{code-cell} ipython3 623 | compute_mean_jax_jitted = jax.jit(compute_mean_jax) 624 | ``` 625 | 626 | ```{code-cell} ipython3 627 | %%time 628 | 629 | compute_mean_jax_jitted() 630 | ``` 631 | 632 | ```{code-cell} ipython3 633 | %%time 634 | 635 | compute_mean_jax_jitted() 636 | ``` 637 | 638 | 639 | ## Exercise 640 | 641 | 642 | Use the quadratic map 643 | 644 | $$ x_{t+1} = 4 x_t (1 - x_t) $$ 645 | 646 | to generate a time series of length `10_000_000` and histogram it with `bins=100`. 647 | 648 | (Look up how to histogram in Matplotlib.) 649 | 650 | If you can, accelerate your code with `@numba.jit`. 651 | 652 | What kind of histogram to do you get? 653 | 654 | 655 | 656 | ```{code-cell} ipython3 657 | # Put your code here 658 | ``` 659 | 660 | solution below 661 | 662 | solution below 663 | 664 | solution below 665 | 666 | solution below 667 | 668 | solution below 669 | 670 | solution below 671 | 672 | solution below 673 | 674 | solution below 675 | 676 | solution below 677 | 678 | solution below 679 | 680 | solution below 681 | 682 | solution below 683 | 684 | 685 | 686 | 687 | 688 | ```{code-cell} ipython3 689 | @numba.jit 690 | def quad_time_series(x0, n=10_000_000, α=4.0): 691 | x = np.empty(n) 692 | x[0] = x0 693 | for t in range(n-1): 694 | x[t+1] = α * x[t] * (1 - x[t]) 695 | return x 696 | ``` 697 | 698 | Compile: 699 | 700 | ```{code-cell} ipython3 701 | x = quad_time_series(0.1, n=10) 702 | ``` 703 | 704 | Run 705 | 706 | 707 | ```{code-cell} ipython3 708 | x = quad_time_series(0.1) 709 | ``` 710 | 711 | 712 | ```{code-cell} ipython3 713 | fig, ax = plt.subplots() 714 | ax.hist(x, bins=100, density=True) 715 | plt.show() 716 | ``` 717 | 718 | -------------------------------------------------------------------------------- /slides/main.tex: -------------------------------------------------------------------------------- 1 | \documentclass[ 2 | xcolor={svgnames,dvipsnames}, 3 | hyperref={colorlinks, citecolor=DeepPink4, linkcolor=DarkRed, urlcolor=DarkBlue} 4 | ]{beamer} % for hardcopy add 'trans' 5 | 6 | 7 | \mode 8 | { 9 | \usetheme{Singapore} 10 | % or ... 11 | \setbeamercovered{transparent} 12 | % or whatever (possibly just delete it) 13 | } 14 | 15 | \usefonttheme{professionalfonts} 16 | %\usepackage[english]{babel} 17 | % or whatever 18 | %\usepackage[latin1]{inputenc} 19 | % or whatever 20 | %\usepackage{times} 21 | %\usepackage[T1]{fontenc} 22 | % Or whatever. Note that the encoding and the font should match. If T1 23 | % does not look nice, try deleting the line with the fontenc. 24 | 25 | %\usepackage{fontspec} 26 | %\setmonofont{CMU Typewriter Text} 27 | %\setmonofont{Consolas} 28 | 29 | %%%%%%%%%%%%%%%%%%%%%% start my preamble %%%%%%%%%%%%%%%%%%%%%% 30 | 31 | \addtobeamertemplate{navigation symbols}{}{% 32 | \usebeamerfont{footline}% 33 | \usebeamercolor[fg]{footline}% 34 | \hspace{1em}% 35 | \insertframenumber/\inserttotalframenumber 36 | } 37 | 38 | 39 | \usepackage{graphicx} 40 | \usepackage{amsmath, amssymb, amsthm} 41 | \usepackage{bbm} 42 | \usepackage{mathrsfs} 43 | \usepackage{xcolor} 44 | \usepackage{fancyvrb} 45 | 46 | % Quotes at start of chapters / sections 47 | \usepackage{epigraph} 48 | %\renewcommand{\epigraphflush}{flushleft} 49 | %\renewcommand{\sourceflush}{flushleft} 50 | \renewcommand{\epigraphwidth}{6in} 51 | 52 | %% Fonts 53 | 54 | %\usepackage[T1]{fontenc} 55 | \usepackage{mathpazo} 56 | %\usepackage{fontspec} 57 | %\defaultfontfeatures{Ligatures=TeX} 58 | %\setsansfont[Scale=MatchLowercase]{DejaVu Sans} 59 | %\setmonofont[Scale=MatchLowercase]{DejaVu Sans Mono} 60 | %\setmathfont{Asana Math} 61 | %\setmainfont{Optima} 62 | %\setmathrm{Optima} 63 | %\setboldmathrm[BoldFont={Optima ExtraBlack}]{Optima Bold} 64 | 65 | % Some colors 66 | 67 | \definecolor{aquamarine}{RGB}{69,139,116} 68 | \definecolor{midnightblue}{RGB}{25,25,112} 69 | \definecolor{darkslategrey}{RGB}{47,79,79} 70 | \definecolor{darkorange4}{RGB}{139,90,0} 71 | \definecolor{dogerblue}{RGB}{24,116,205} 72 | \definecolor{blue2}{RGB}{0,0,238} 73 | \definecolor{bg}{rgb}{0.95,0.95,0.95} 74 | \definecolor{DarkOrange1}{RGB}{255,127,0} 75 | \definecolor{ForestGreen}{RGB}{34,139,34} 76 | \definecolor{DarkRed}{RGB}{139, 0, 0} 77 | \definecolor{DarkBlue}{RGB}{0, 0, 139} 78 | \definecolor{Blue}{RGB}{0, 0, 255} 79 | \definecolor{Brown}{RGB}{165,42,42} 80 | 81 | 82 | \setlength{\parskip}{1.5ex plus0.5ex minus0.5ex} 83 | 84 | %\renewcommand{\baselinestretch}{1.05} 85 | %\setlength{\parskip}{1.5ex plus0.5ex minus0.5ex} 86 | %\setlength{\parindent}{0pt} 87 | 88 | % Typesetting code 89 | \definecolor{bg}{rgb}{0.95,0.95,0.95} 90 | \usepackage{minted} 91 | \setminted{mathescape, frame=lines, framesep=3mm} 92 | \usemintedstyle{friendly} 93 | %\newminted{python}{} 94 | %\newminted{c}{mathescape,frame=lines,framesep=4mm,bgcolor=bg} 95 | %\newminted{java}{mathescape,frame=lines,framesep=4mm,bgcolor=bg} 96 | %\newminted{julia}{mathescape,frame=lines,framesep=4mm,bgcolor=bg} 97 | %\newminted{ipython}{mathescape,frame=lines,framesep=4mm,bgcolor=bg} 98 | 99 | 100 | \newcommand{\Fact}{\textcolor{Brown}{\bf Fact. }} 101 | \newcommand{\Facts}{\textcolor{Brown}{\bf Facts }} 102 | \newcommand{\keya}{\textcolor{turquois4}{\bf Key Idea. }} 103 | \newcommand{\Factnodot}{\textcolor{Brown}{\bf Fact }} 104 | \newcommand{\Eg}{\textcolor{ForestGreen}{Example. }} 105 | \newcommand{\Egs}{\textcolor{ForestGreen}{Examples. }} 106 | \newcommand{\Ex}{{\bf Ex. }} 107 | 108 | 109 | 110 | \renewcommand{\theFancyVerbLine}{\sffamily 111 | \textcolor[rgb]{0.5,0.5,1.0}{\scriptsize {\arabic{FancyVerbLine}}}} 112 | 113 | \newcommand{\navy}[1]{\textcolor{Blue}{\bf #1}} 114 | \newcommand{\brown}[1]{\textcolor{Brown}{\sf #1}} 115 | \newcommand{\green}[1]{\textcolor{ForestGreen}{\sf #1}} 116 | \newcommand{\blue}[1]{\textcolor{Blue}{\sf #1}} 117 | \newcommand{\navymth}[1]{\textcolor{Blue}{#1}} 118 | \newcommand{\emp}[1]{\textcolor{DarkOrange1}{\bf #1}} 119 | \newcommand{\red}[1]{\textcolor{Red}{\bf #1}} 120 | 121 | % Symbols, redefines, etc. 122 | 123 | \newcommand{\code}[1]{\texttt{#1}} 124 | 125 | \newcommand{\argmax}{\operatornamewithlimits{argmax}} 126 | \newcommand{\argmin}{\operatornamewithlimits{argmin}} 127 | 128 | \DeclareMathOperator{\cl}{cl} 129 | \DeclareMathOperator{\interior}{int} 130 | \DeclareMathOperator{\Prob}{Prob} 131 | \DeclareMathOperator{\determinant}{det} 132 | \DeclareMathOperator{\trace}{trace} 133 | \DeclareMathOperator{\Span}{span} 134 | \DeclareMathOperator{\rank}{rank} 135 | \DeclareMathOperator{\cov}{cov} 136 | \DeclareMathOperator{\corr}{corr} 137 | \DeclareMathOperator{\var}{var} 138 | \DeclareMathOperator{\mse}{mse} 139 | \DeclareMathOperator{\se}{se} 140 | \DeclareMathOperator{\row}{row} 141 | \DeclareMathOperator{\col}{col} 142 | \DeclareMathOperator{\range}{rng} 143 | \DeclareMathOperator{\dimension}{dim} 144 | \DeclareMathOperator{\bias}{bias} 145 | 146 | 147 | % mics short cuts and symbols 148 | \newcommand{\st}{\ensuremath{\ \mathrm{s.t.}\ }} 149 | \newcommand{\setntn}[2]{ \{ #1 : #2 \} } 150 | \newcommand{\cf}[1]{ \lstinline|#1| } 151 | \newcommand{\fore}{\therefore \quad} 152 | \newcommand{\tod}{\stackrel { d } {\to} } 153 | \newcommand{\toprob}{\stackrel { p } {\to} } 154 | \newcommand{\toms}{\stackrel { ms } {\to} } 155 | \newcommand{\eqdist}{\stackrel {\textrm{ \scriptsize{d} }} {=} } 156 | \newcommand{\iidsim}{\stackrel {\textrm{ {\sc iid }}} {\sim} } 157 | \newcommand{\1}{\mathbbm 1} 158 | \newcommand{\dee}{\,{\rm d}} 159 | \newcommand{\given}{\, | \,} 160 | \newcommand{\la}{\langle} 161 | \newcommand{\ra}{\rangle} 162 | 163 | \newcommand{\boldA}{\mathbf A} 164 | \newcommand{\boldB}{\mathbf B} 165 | \newcommand{\boldC}{\mathbf C} 166 | \newcommand{\boldD}{\mathbf D} 167 | \newcommand{\boldM}{\mathbf M} 168 | \newcommand{\boldP}{\mathbf P} 169 | \newcommand{\boldQ}{\mathbf Q} 170 | \newcommand{\boldI}{\mathbf I} 171 | \newcommand{\boldX}{\mathbf X} 172 | \newcommand{\boldY}{\mathbf Y} 173 | \newcommand{\boldZ}{\mathbf Z} 174 | 175 | \newcommand{\bSigmaX}{ {\boldsymbol \Sigma_{\hboldbeta}} } 176 | \newcommand{\hbSigmaX}{ \mathbf{\hat \Sigma_{\hboldbeta}} } 177 | 178 | \newcommand{\RR}{\mathbbm R} 179 | \newcommand{\NN}{\mathbbm N} 180 | \newcommand{\PP}{\mathbbm P} 181 | \newcommand{\EE}{\mathbbm E \,} 182 | \newcommand{\XX}{\mathbbm X} 183 | \newcommand{\ZZ}{\mathbbm Z} 184 | \newcommand{\QQ}{\mathbbm Q} 185 | 186 | \newcommand{\fF}{\mathcal F} 187 | \newcommand{\dD}{\mathcal D} 188 | \newcommand{\lL}{\mathcal L} 189 | \newcommand{\gG}{\mathcal G} 190 | \newcommand{\hH}{\mathcal H} 191 | \newcommand{\nN}{\mathcal N} 192 | \newcommand{\pP}{\mathcal P} 193 | 194 | 195 | 196 | 197 | \title{Computational Methods for Quantitative Economics} 198 | 199 | \author{John Stachurski} 200 | 201 | 202 | \date{April 2023} 203 | 204 | 205 | \begin{document} 206 | 207 | \begin{frame} 208 | \titlepage 209 | \end{frame} 210 | 211 | 212 | 213 | 214 | 215 | \section{Introduction} 216 | 217 | 218 | 219 | \begin{frame} 220 | %\frametitle{Introduction} 221 | 222 | 223 | \textbf{Topics} 224 | 225 | \begin{itemize} 226 | \item Discussion of scientific computing 227 | \vspace{0.5em} 228 | \item Overview of Python libraries 229 | \vspace{0.5em} 230 | \item Computing equilibria 231 | \vspace{0.5em} 232 | \item Google JAX 233 | \vspace{0.5em} 234 | \item Option pricing with Python 235 | \vspace{0.5em} 236 | \item High dimensional problems 237 | \end{itemize} 238 | 239 | \end{frame} 240 | 241 | 242 | 243 | \begin{frame} 244 | 245 | 246 | Assumptions: 247 | 248 | \begin{itemize} 249 | \item basic econ/computer/maths/stats 250 | \vspace{0.3em} 251 | \item some programming? 252 | \end{itemize} 253 | 254 | \vspace{0.3em} 255 | \vspace{0.5em} 256 | Aims: 257 | 258 | \begin{itemize} 259 | \item Discuss options 260 | \vspace{0.5em} 261 | \item Review trends 262 | \vspace{0.5em} 263 | \item Learn techniques 264 | \end{itemize} 265 | 266 | \vspace{0.5em} 267 | \vspace{0.5em} 268 | \textbf{Resources} 269 | 270 | \begin{itemize} 271 | \item \url{https://github.com/QuantEcon/keio_comp_econ_2023} 272 | \end{itemize} 273 | 274 | 275 | \end{frame} 276 | 277 | 278 | 279 | 280 | \section{Three Trends} 281 | 282 | \begin{frame} 283 | \frametitle{Trends} 284 | 285 | What are the major trends in scientific computing? 286 | 287 | \vspace{0.5em} 288 | \vspace{0.5em} 289 | 290 | \begin{itemize} 291 | \item what's driving them? 292 | \vspace{0.5em} 293 | \item how can we benefit? 294 | \end{itemize} 295 | \end{frame} 296 | 297 | 298 | 299 | \begin{frame} 300 | \frametitle{Trend 1: Proprietary $\to$ Open Source} 301 | 302 | \blue{Proprietary} 303 | % 304 | \begin{itemize} 305 | \item Excel 306 | \item MATLAB 307 | \item Eviews, etc. 308 | \end{itemize} 309 | 310 | 311 | \vspace{0.5em} 312 | \vspace{0.5em} 313 | \blue{Open Source / Open Standard} 314 | 315 | \begin{itemize} 316 | \item Python 317 | \item Julia 318 | \item R, etc. 319 | \end{itemize} 320 | 321 | 322 | \begin{center} 323 | closed and stable vs open and fast moving 324 | \end{center} 325 | 326 | \end{frame} 327 | 328 | \begin{frame} 329 | 330 | Popularity: 331 | 332 | \begin{figure} 333 | \begin{center} 334 | \scalebox{1.8}{\includegraphics{python_vs_rest.png}} 335 | \end{center} 336 | \end{figure} 337 | 338 | \end{frame} 339 | 340 | 341 | 342 | \begin{frame} 343 | \frametitle{Trend 2: Low Level $\to$ High Level} 344 | 345 | \blue{Low level} 346 | 347 | \begin{itemize} 348 | \item C 349 | \item Fortran 350 | \item Assembly 351 | \end{itemize} 352 | 353 | \vspace{1em} 354 | 355 | \blue{High level } 356 | 357 | \begin{itemize} 358 | \item Python 359 | \item Javascript 360 | \item Ruby 361 | \end{itemize} 362 | 363 | \end{frame} 364 | 365 | 366 | \begin{frame} 367 | 368 | \blue{Low level languages} give us tight control of hardware 369 | 370 | \begin{itemize} 371 | \item CPU 372 | \item memory, etc. 373 | \end{itemize} 374 | 375 | \vspace{0.5em} 376 | \vspace{0.5em} 377 | \vspace{0.5em} 378 | 379 | \blue{High level languages} give us 380 | % 381 | \begin{itemize} 382 | \item abstraction 383 | \item automation 384 | \item flexibility, etc. 385 | \end{itemize} 386 | 387 | \end{frame} 388 | 389 | 390 | 391 | 392 | \begin{frame}[fragile] 393 | 394 | \Eg \brown{1 + 1} in assembly 395 | 396 | {\small 397 | \begin{minted}{as} 398 | pushq %rbp 399 | movq %rsp, %rbp 400 | movl $1, -12(%rbp) 401 | movl $1, -8(%rbp) 402 | movl -12(%rbp), %edx 403 | movl -8(%rbp), %eax 404 | addl %edx, %eax 405 | movl %eax, -4(%rbp) 406 | movl -4(%rbp), %eax 407 | popq %rbp 408 | \end{minted} 409 | } 410 | 411 | \end{frame} 412 | 413 | 414 | \begin{frame}[fragile] 415 | 416 | \Eg \brown{1 + 1} in C 417 | 418 | {\small 419 | \begin{minted}{c} 420 | #include 421 | int main() { 422 | int x = 1 + 1; 423 | printf("1 + 1 = %d\n", x); 424 | return 0; 425 | } 426 | \end{minted} 427 | } 428 | 429 | \end{frame} 430 | 431 | 432 | 433 | \begin{frame}[fragile] 434 | 435 | \Eg \brown{1 + 1} in Fortran 436 | 437 | {\small 438 | \begin{minted}{fortran} 439 | PROGRAM ONE_PLUS_ONE 440 | INTEGER :: X = 1 + 1 441 | PRINT *, '1 + 1 = ', X 442 | END PROGRAM ONE_PLUS_ONE 443 | \end{minted} 444 | } 445 | 446 | \end{frame} 447 | 448 | 449 | 450 | 451 | \begin{frame}[fragile] 452 | 453 | \Eg \brown{1 + 1} in Python 454 | 455 | {\small 456 | \begin{minted}{python} 457 | x = 1 + 1 458 | print("1 + 1 = ", x) 459 | \end{minted} 460 | } 461 | 462 | \end{frame} 463 | 464 | 465 | 466 | \begin{frame} 467 | 468 | Trade-Offs: 469 | 470 | \begin{figure} 471 | \begin{center} 472 | \scalebox{.46}{\includegraphics{ppf.pdf}} 473 | \end{center} 474 | \end{figure} 475 | 476 | \end{frame} 477 | 478 | 479 | 480 | \begin{frame}[fragile] 481 | 482 | New trend --- a shifting frontier 483 | 484 | \end{frame} 485 | 486 | 487 | 488 | 489 | \begin{frame} 490 | 491 | Trade-offs: 492 | 493 | \begin{figure} 494 | \begin{center} 495 | \scalebox{.46}{\includegraphics{ppf_plus.pdf}} 496 | \end{center} 497 | \end{figure} 498 | 499 | \end{frame} 500 | 501 | 502 | 503 | 504 | 505 | \begin{frame} 506 | 507 | \Eg What platforms/languages does OpenAI use? 508 | 509 | \vspace{0.5em} 510 | In order (according to \href{https://github.com/openai}{repo stats}): 511 | 512 | \begin{enumerate} 513 | \item Python 514 | \item C$++$ 515 | \item Javascript 516 | \item Jupyter notebooks 517 | \item Ruby 518 | \end{enumerate} 519 | 520 | \end{frame} 521 | 522 | 523 | 524 | %\section{Parallelization} 525 | 526 | \begin{frame} 527 | \frametitle{Trend 3: Parallelization} 528 | 529 | CPU frequency (clock speed) growth is slowing 530 | 531 | \begin{figure} 532 | \begin{center} 533 | \scalebox{.22}{\includegraphics{processor_clock.png}} 534 | \end{center} 535 | \end{figure} 536 | 537 | \end{frame} 538 | 539 | 540 | \begin{frame} 541 | 542 | Chip makers have responded by developing multi-core processors 543 | 544 | \begin{figure} 545 | \begin{center} 546 | \scalebox{.2}{\includegraphics{dual_core.png}} 547 | \end{center} 548 | \end{figure} 549 | 550 | Source: Wikipedia 551 | 552 | 553 | \end{frame} 554 | 555 | 556 | \begin{frame} 557 | 558 | \navy{GPUs} are becoming increasingly important 559 | 560 | 561 | \begin{figure} 562 | \begin{center} 563 | \scalebox{.14}{\includegraphics{gpu.jpg}} 564 | \end{center} 565 | \end{figure} 566 | 567 | \vspace{0.5em} 568 | 569 | Applications: machine learning, deep learning, etc. 570 | 571 | 572 | \end{frame} 573 | 574 | 575 | 576 | \begin{frame} 577 | \frametitle{Support for Parallelization} 578 | 579 | While scientific computing environments best support parallelization? 580 | 581 | \vspace{0.5em} 582 | \vspace{0.5em} 583 | % 584 | \begin{itemize} 585 | \item Most have some support 586 | \vspace{0.5em} 587 | \item but which make it easy to harness its power? 588 | \end{itemize} 589 | 590 | \vspace{0.5em} 591 | \vspace{0.5em} 592 | 593 | Current winner: 594 | % 595 | \begin{itemize} 596 | \item Google JAX (Python library) 597 | \vspace{0.5em} 598 | \end{itemize} 599 | 600 | \end{frame} 601 | 602 | 603 | 604 | 605 | \section{Which Language?} 606 | 607 | 608 | \begin{frame} 609 | \frametitle{Which Language} 610 | 611 | 612 | How about R? 613 | \vspace{0.5em} 614 | 615 | \begin{itemize} 616 | \item Specialized to statistics 617 | \vspace{0.5em} 618 | \item Huge range of estimation routines 619 | \vspace{0.5em} 620 | \item Popular in academia 621 | \vspace{0.5em} 622 | \item Loosing some ground to Python (AI, machine learning) 623 | \end{itemize} 624 | 625 | \end{frame} 626 | 627 | 628 | \begin{frame} 629 | \frametitle{Julia} 630 | 631 | Pros: 632 | 633 | \begin{itemize} 634 | \item Fast and elegant 635 | \vspace{0.5em} 636 | \item Many scientific routines 637 | \vspace{0.5em} 638 | \item Julia is written in Julia 639 | \end{itemize} 640 | 641 | \vspace{0.5em} 642 | \vspace{0.5em} 643 | \vspace{0.5em} 644 | Cons: 645 | 646 | \begin{itemize} 647 | \item Low rates of investment in some important libraries 648 | \end{itemize} 649 | 650 | \end{frame} 651 | 652 | 653 | 654 | 655 | \begin{frame} 656 | \frametitle{Python} 657 | 658 | \begin{itemize} 659 | \item Easy to learn, well designed 660 | \vspace{0.5em} 661 | \item Massive scientific ecosystem 662 | \vspace{0.5em} 663 | \item Heavily supported by big players 664 | \vspace{0.5em} 665 | \item Strong support for parallel computing 666 | \vspace{0.5em} 667 | \item Huge demand for tech-savvy Python programmers 668 | \end{itemize} 669 | 670 | 671 | \end{frame} 672 | 673 | 674 | 675 | 676 | \section{Set Up} 677 | 678 | 679 | 680 | 681 | 682 | \begin{frame} 683 | \frametitle{Accessing Python} 684 | 685 | Option 1: \green{Via a service (remote option)} 686 | 687 | \begin{itemize} 688 | \item \url{https://colab.research.google.com} 689 | \end{itemize} 690 | 691 | \vspace{1em} 692 | 693 | Option 2: \green{Local install (Python + scientific libs)} 694 | 695 | \begin{itemize} 696 | \item Install Anaconda from {\footnotesize \url{https://www.anaconda.com/}} 697 | \vspace{1em} 698 | \begin{itemize} 699 | \item Select latest version 700 | \end{itemize} 701 | \vspace{1em} 702 | \item Not plain vanilla Python 703 | \end{itemize} 704 | 705 | 706 | 707 | 708 | \end{frame} 709 | 710 | 711 | 712 | \begin{frame} 713 | \frametitle{How to Interact with Python?} 714 | 715 | Many options: 716 | 717 | \begin{itemize} 718 | \item \emp{write} with VS Code / Emacs / Vim 719 | \vspace{0.5em} 720 | \item \emp{run} with base Python, IPython, etc. 721 | \end{itemize} 722 | 723 | \vspace{0.5em} 724 | \vspace{0.5em} 725 | \vspace{0.5em} 726 | Or do both with \emp{Jupyter notebooks / Jupyter lab} 727 | 728 | \vspace{0.5em} 729 | \begin{itemize} 730 | \item for simplicity we focus only on the last option 731 | \end{itemize} 732 | 733 | \end{frame} 734 | 735 | 736 | \begin{frame} 737 | \frametitle{Jupyter Notebooks} 738 | 739 | A browser based interface to Python / Julia / R / etc. 740 | 741 | 742 | \vspace{2em} 743 | 744 | \begin{itemize} 745 | \item Search for \texttt{jupyter notebook} 746 | \end{itemize} 747 | 748 | \vspace{2em} 749 | 750 | Useful for: 751 | 752 | \begin{itemize} 753 | \item getting started 754 | \item exploring ideas 755 | \end{itemize} 756 | 757 | \end{frame} 758 | 759 | 760 | 761 | \begin{frame} 762 | \frametitle{Working with Notebooks} 763 | 764 | \begin{itemize} 765 | \item Entry and execution 766 | \vspace{1em} 767 | \item Markdown 768 | \vspace{1em} 769 | \item Getting help 770 | \vspace{1em} 771 | \item Copy paste 772 | \vspace{1em} 773 | \item Edit and command mode 774 | \end{itemize} 775 | 776 | \end{frame} 777 | 778 | 779 | 780 | 781 | 782 | 783 | 784 | 785 | \end{document} 786 | 787 | 788 | -------------------------------------------------------------------------------- /notebooks/equilibrium.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "ced890fd", 6 | "metadata": {}, 7 | "source": [ 8 | "# Equilibrium" 9 | ] 10 | }, 11 | { 12 | "cell_type": "markdown", 13 | "id": "2ead455d", 14 | "metadata": {}, 15 | "source": [ 16 | "#### Author: [John Stachurski](http://johnstachurski.net/)" 17 | ] 18 | }, 19 | { 20 | "cell_type": "markdown", 21 | "id": "d8755fdd", 22 | "metadata": {}, 23 | "source": [ 24 | "In this notebook we solve a very simple market equilibrium problem.\n", 25 | "\n", 26 | "Supply and demand are nonlinear and we use Newton's root-finding algorithm to solve for equilibrium prices.\n", 27 | "\n", 28 | "We use the following imports:" 29 | ] 30 | }, 31 | { 32 | "cell_type": "code", 33 | "execution_count": null, 34 | "id": "8527b2f0", 35 | "metadata": {}, 36 | "outputs": [], 37 | "source": [ 38 | "import numpy as np\n", 39 | "import matplotlib.pyplot as plt" 40 | ] 41 | }, 42 | { 43 | "cell_type": "markdown", 44 | "id": "57fe323e", 45 | "metadata": {}, 46 | "source": [ 47 | "## Prelude: A Note on Root-Finding" 48 | ] 49 | }, 50 | { 51 | "cell_type": "markdown", 52 | "id": "a322276b", 53 | "metadata": {}, 54 | "source": [ 55 | "If $f$ maps an interval $(a, b)$ into $\\mathbb R$, then a **root** of the function $f$ is an $x^* \\in (a,b)$ with $f(x^*)=0$.\n", 56 | "\n", 57 | "A common method for root finding is Newton's algorithm.\n", 58 | "\n", 59 | "We start with a guess $x_0 \\in (a, b)$.\n", 60 | "\n", 61 | "Then we replace $f$ with the tangent function $f_a(x) = f(x_0) + f'(x_0)(x - x_0)$ and solve for the root of $f_a$ (which can be done exactly).\n", 62 | "\n", 63 | "Calling the root $x_1$, we have\n", 64 | "\n", 65 | "$$ \n", 66 | " f_a(x_1)=0\n", 67 | " \\quad \\iff \\quad\n", 68 | " x_1 = x_0 - \\frac{f(x_0)}{f'(x_0)} \n", 69 | "$$\n", 70 | "\n", 71 | "This is our update rule:\n", 72 | "\n", 73 | "$$\n", 74 | " x_{k+1} = q(x_k)\n", 75 | " \\quad \\text{where} \\quad\n", 76 | " q(x) := x - \\frac{f(x)}{f'(x)} \n", 77 | "$$" 78 | ] 79 | }, 80 | { 81 | "cell_type": "markdown", 82 | "id": "be982d30", 83 | "metadata": {}, 84 | "source": [ 85 | "The algorithm is implemented in `scipy.optimize`" 86 | ] 87 | }, 88 | { 89 | "cell_type": "code", 90 | "execution_count": null, 91 | "id": "504c125d", 92 | "metadata": {}, 93 | "outputs": [], 94 | "source": [ 95 | "from scipy.optimize import newton" 96 | ] 97 | }, 98 | { 99 | "cell_type": "markdown", 100 | "id": "8f67c5b0", 101 | "metadata": {}, 102 | "source": [ 103 | "Let's apply this to find the positive root of $f(x) = x^2 - 1$." 104 | ] 105 | }, 106 | { 107 | "cell_type": "code", 108 | "execution_count": null, 109 | "id": "cad559ab", 110 | "metadata": {}, 111 | "outputs": [], 112 | "source": [ 113 | "def f(x):\n", 114 | " return x**2 - 1\n", 115 | "\n", 116 | "x_grid = np.linspace(-1, 2, 200)\n", 117 | "fig, ax = plt.subplots()\n", 118 | "ax.plot(x_grid, f(x_grid), label=\"$f$\")\n", 119 | "ax.plot(x_grid, np.zeros_like(x_grid), \"k--\")\n", 120 | "ax.legend()\n", 121 | "plt.show()\n" 122 | ] 123 | }, 124 | { 125 | "cell_type": "markdown", 126 | "id": "6820aa2f", 127 | "metadata": {}, 128 | "source": [ 129 | "Here we call `newton`." 130 | ] 131 | }, 132 | { 133 | "cell_type": "code", 134 | "execution_count": null, 135 | "id": "84f292c5", 136 | "metadata": {}, 137 | "outputs": [], 138 | "source": [ 139 | "newton(f, 0.5) # search for root of f starting at x_0 = 0.5" 140 | ] 141 | }, 142 | { 143 | "cell_type": "markdown", 144 | "id": "2d93fa04", 145 | "metadata": {}, 146 | "source": [ 147 | "In the last call we didn't supply the gradient of $f$, so it was approximated\n", 148 | "numerically. \n", 149 | "\n", 150 | "We can supply it as follows:" 151 | ] 152 | }, 153 | { 154 | "cell_type": "code", 155 | "execution_count": null, 156 | "id": "332faff4", 157 | "metadata": {}, 158 | "outputs": [], 159 | "source": [ 160 | "def f_prime(x):\n", 161 | " return 2 * x\n", 162 | "\n", 163 | "newton(lambda x: x**2 - 1, 0.5, fprime=f_prime)" 164 | ] 165 | }, 166 | { 167 | "cell_type": "markdown", 168 | "id": "eedab0d0", 169 | "metadata": {}, 170 | "source": [ 171 | "## The Market" 172 | ] 173 | }, 174 | { 175 | "cell_type": "markdown", 176 | "id": "7d983d4f", 177 | "metadata": {}, 178 | "source": [ 179 | "Now let's consider a market for coffee beans. The price per kilo is $p$. Total supply at price $p$ is\n", 180 | "\n", 181 | "$$ q_s (p) = b \\sqrt{p} $$\n", 182 | "\n", 183 | "and total demand is \n", 184 | "\n", 185 | "$$ q_d (p) = a \\exp(-p) + c, $$\n", 186 | "\n", 187 | "where $a, b$ and $c$ are positive parameters." 188 | ] 189 | }, 190 | { 191 | "cell_type": "markdown", 192 | "id": "cd994e90", 193 | "metadata": {}, 194 | "source": [ 195 | "Now we write routines to compute supply and demand as functions of price and parameters.\n", 196 | "\n", 197 | "We take $a=1$, $b=0.5$ and $c=1$ as \"default\" parameter values." 198 | ] 199 | }, 200 | { 201 | "cell_type": "code", 202 | "execution_count": null, 203 | "id": "29cdc1af", 204 | "metadata": {}, 205 | "outputs": [], 206 | "source": [ 207 | "def supply(p, b=0.5):\n", 208 | " return b * np.sqrt(p)\n", 209 | "\n", 210 | "def demand(p, a=1, c=1):\n", 211 | " return a * np.exp(-p) + c" 212 | ] 213 | }, 214 | { 215 | "cell_type": "markdown", 216 | "id": "6a30a053", 217 | "metadata": {}, 218 | "source": [ 219 | "Now we can call the functions as follows:" 220 | ] 221 | }, 222 | { 223 | "cell_type": "code", 224 | "execution_count": null, 225 | "id": "c6771b5f", 226 | "metadata": {}, 227 | "outputs": [], 228 | "source": [ 229 | "demand(2.0) # with a and c at defaults" 230 | ] 231 | }, 232 | { 233 | "cell_type": "code", 234 | "execution_count": null, 235 | "id": "e015b72b", 236 | "metadata": {}, 237 | "outputs": [], 238 | "source": [ 239 | "demand(2.0, a=0.4) # a is specified and c remains at its defaults" 240 | ] 241 | }, 242 | { 243 | "cell_type": "markdown", 244 | "id": "01f54a22", 245 | "metadata": {}, 246 | "source": [ 247 | "etc." 248 | ] 249 | }, 250 | { 251 | "cell_type": "markdown", 252 | "id": "d1a0c672", 253 | "metadata": {}, 254 | "source": [ 255 | "Note that these functions are automatically NumPy \"universal functions\":" 256 | ] 257 | }, 258 | { 259 | "cell_type": "code", 260 | "execution_count": null, 261 | "id": "f1b17d75", 262 | "metadata": {}, 263 | "outputs": [], 264 | "source": [ 265 | "p_vals = np.array((0.5, 1.0, 1.5))\n", 266 | "supply(p_vals)" 267 | ] 268 | }, 269 | { 270 | "cell_type": "code", 271 | "execution_count": null, 272 | "id": "d3076cca", 273 | "metadata": {}, 274 | "outputs": [], 275 | "source": [ 276 | "demand(p_vals)" 277 | ] 278 | }, 279 | { 280 | "cell_type": "markdown", 281 | "id": "ce087fa5", 282 | "metadata": {}, 283 | "source": [ 284 | "### Exercise 1\n", 285 | "\n", 286 | "Plot both supply and demand as functions of $p$ on the interval $[0, 10]$ at the default parameters.\n", 287 | "\n", 288 | "* Put price on the horizonal axis. \n", 289 | "* Use a legend to label the two functions and be sure to label the axes. \n", 290 | "* Make a rough estimate of the equilibrium price, where demand equals supply." 291 | ] 292 | }, 293 | { 294 | "cell_type": "code", 295 | "execution_count": null, 296 | "id": "0877f352", 297 | "metadata": {}, 298 | "outputs": [], 299 | "source": [ 300 | "# Put your code here" 301 | ] 302 | }, 303 | { 304 | "cell_type": "markdown", 305 | "id": "36151a15", 306 | "metadata": {}, 307 | "source": [ 308 | "solution below\n", 309 | "\n", 310 | "solution below\n", 311 | "\n", 312 | "solution below\n", 313 | "\n", 314 | "solution below\n", 315 | "\n", 316 | "solution below\n", 317 | "\n", 318 | "solution below\n", 319 | "\n", 320 | "solution below\n", 321 | "\n", 322 | "solution below\n", 323 | "\n", 324 | "solution below\n", 325 | "\n", 326 | "solution below\n", 327 | "\n", 328 | "solution below\n", 329 | "\n", 330 | "solution below" 331 | ] 332 | }, 333 | { 334 | "cell_type": "code", 335 | "execution_count": null, 336 | "id": "709a8ab7", 337 | "metadata": {}, 338 | "outputs": [], 339 | "source": [ 340 | "fig, ax = plt.subplots()\n", 341 | "p_grid = np.linspace(0, 10, 200)\n", 342 | "ax.plot(p_grid, supply(p_grid), label='supply')\n", 343 | "ax.plot(p_grid, demand(p_grid), label='demand')\n", 344 | "ax.set_xlabel(\"price\")\n", 345 | "ax.set_ylabel(\"quantity\")\n", 346 | "ax.legend(frameon=False, loc='upper center')\n", 347 | "plt.show()" 348 | ] 349 | }, 350 | { 351 | "cell_type": "markdown", 352 | "id": "b0d84d80", 353 | "metadata": {}, 354 | "source": [ 355 | "The equilibrium price looks to be about 4.1." 356 | ] 357 | }, 358 | { 359 | "cell_type": "markdown", 360 | "id": "a32bfa2c", 361 | "metadata": {}, 362 | "source": [ 363 | "### Exercise 2\n", 364 | "\n", 365 | "Write a function that takes arguments $a, b, c, p$, with default values $a=1$, $b=0.5$ and $c=1$, and returns *excess demand*, which is defined as\n", 366 | "\n", 367 | "$$ e(p) = q_d(p) - q_s(p) $$" 368 | ] 369 | }, 370 | { 371 | "cell_type": "code", 372 | "execution_count": null, 373 | "id": "cd2eb3c3", 374 | "metadata": {}, 375 | "outputs": [], 376 | "source": [ 377 | "# Put your code here" 378 | ] 379 | }, 380 | { 381 | "cell_type": "markdown", 382 | "id": "0fb43b8b", 383 | "metadata": {}, 384 | "source": [ 385 | "solution below\n", 386 | "\n", 387 | "solution below\n", 388 | "\n", 389 | "solution below\n", 390 | "\n", 391 | "solution below\n", 392 | "\n", 393 | "solution below\n", 394 | "\n", 395 | "solution below\n", 396 | "\n", 397 | "solution below\n", 398 | "\n", 399 | "solution below\n", 400 | "\n", 401 | "solution below\n", 402 | "\n", 403 | "solution below\n", 404 | "\n", 405 | "solution below\n", 406 | "\n", 407 | "solution below" 408 | ] 409 | }, 410 | { 411 | "cell_type": "code", 412 | "execution_count": null, 413 | "id": "6a9bb9c0", 414 | "metadata": {}, 415 | "outputs": [], 416 | "source": [ 417 | "def excess_demand(p, a=1, b=0.5, c=1):\n", 418 | " return demand(p, a, c) - supply(p, b)" 419 | ] 420 | }, 421 | { 422 | "cell_type": "markdown", 423 | "id": "d555e439", 424 | "metadata": {}, 425 | "source": [ 426 | "Now we test it:" 427 | ] 428 | }, 429 | { 430 | "cell_type": "code", 431 | "execution_count": null, 432 | "id": "343da30a", 433 | "metadata": {}, 434 | "outputs": [], 435 | "source": [ 436 | "excess_demand(1.0)" 437 | ] 438 | }, 439 | { 440 | "cell_type": "markdown", 441 | "id": "594fc3ed", 442 | "metadata": {}, 443 | "source": [ 444 | "### Organizing our Code\n", 445 | "\n", 446 | "If we have many functions working with the same parameters, it's hard to know where to put the default values.\n", 447 | "\n", 448 | "As such, we normally collect them in a data structure, such as a class or a tuple.\n", 449 | "\n", 450 | "Personally, I normally used `namedtuple` instances, which are lighter than classes but easier to work with than tuples.\n", 451 | "\n", 452 | "Here's an example:" 453 | ] 454 | }, 455 | { 456 | "cell_type": "code", 457 | "execution_count": null, 458 | "id": "91476b30", 459 | "metadata": {}, 460 | "outputs": [], 461 | "source": [ 462 | "from collections import namedtuple\n", 463 | "\n", 464 | "Params = namedtuple('Params', ('a', 'b', 'c'))\n", 465 | "\n", 466 | "def create_market_params(a=1.0, b=0.5, c=1.0):\n", 467 | " return Params(a=a, b=b, c=c)\n", 468 | "\n", 469 | "\n", 470 | "def supply(p, params):\n", 471 | " a, b, c = params\n", 472 | " return b * np.sqrt(p)\n", 473 | "\n", 474 | "def demand(p, params):\n", 475 | " a, b, c = params\n", 476 | " return a * np.exp(-p) + c\n", 477 | "\n", 478 | "def excess_demand(p, params):\n", 479 | " a, b, c = params\n", 480 | " return demand(p, params) - supply(p, params)" 481 | ] 482 | }, 483 | { 484 | "cell_type": "markdown", 485 | "id": "b8424003", 486 | "metadata": {}, 487 | "source": [ 488 | "### Exercise 3\n", 489 | "\n", 490 | "Using these functions, plot excess demand over the interval from $0.2$ up to $10$. Also plot a horizontal line at zero. The equilibrium price is where excess demand crosses zero." 491 | ] 492 | }, 493 | { 494 | "cell_type": "code", 495 | "execution_count": null, 496 | "id": "f3311018", 497 | "metadata": {}, 498 | "outputs": [], 499 | "source": [ 500 | "# Put your code here" 501 | ] 502 | }, 503 | { 504 | "cell_type": "markdown", 505 | "id": "e7896638", 506 | "metadata": {}, 507 | "source": [ 508 | "solution below\n", 509 | "\n", 510 | "solution below\n", 511 | "\n", 512 | "solution below\n", 513 | "\n", 514 | "solution below\n", 515 | "\n", 516 | "solution below\n", 517 | "\n", 518 | "solution below\n", 519 | "\n", 520 | "solution below\n", 521 | "\n", 522 | "solution below\n", 523 | "\n", 524 | "solution below\n", 525 | "\n", 526 | "solution below\n", 527 | "\n", 528 | "solution below\n", 529 | "\n", 530 | "solution below" 531 | ] 532 | }, 533 | { 534 | "cell_type": "code", 535 | "execution_count": null, 536 | "id": "3277dbc1", 537 | "metadata": {}, 538 | "outputs": [], 539 | "source": [ 540 | "params = create_market_params()\n", 541 | "\n", 542 | "fig, ax = plt.subplots()\n", 543 | "p_grid = np.linspace(0, 10, 200)\n", 544 | "ax.plot(p_grid, excess_demand(p_grid, params), label='excess demand')\n", 545 | "ax.plot(p_grid, np.zeros_like(p_grid), 'k--')\n", 546 | "ax.set_xlabel(\"price\")\n", 547 | "ax.set_ylabel(\"quantity\")\n", 548 | "ax.legend()\n", 549 | "plt.show()" 550 | ] 551 | }, 552 | { 553 | "cell_type": "markdown", 554 | "id": "f10f2fae", 555 | "metadata": {}, 556 | "source": [ 557 | "### Exercise 4" 558 | ] 559 | }, 560 | { 561 | "cell_type": "markdown", 562 | "id": "53a2f2af", 563 | "metadata": {}, 564 | "source": [ 565 | "Write a function that takes an instance of `Params` (i.e, a parameter vector) and returns a market clearing price via Newton's method." 566 | ] 567 | }, 568 | { 569 | "cell_type": "code", 570 | "execution_count": null, 571 | "id": "a7967a73", 572 | "metadata": {}, 573 | "outputs": [], 574 | "source": [ 575 | "# Put your code here" 576 | ] 577 | }, 578 | { 579 | "cell_type": "markdown", 580 | "id": "ed1c2c97", 581 | "metadata": {}, 582 | "source": [ 583 | "solution below\n", 584 | "\n", 585 | "solution below\n", 586 | "\n", 587 | "solution below\n", 588 | "\n", 589 | "solution below\n", 590 | "\n", 591 | "solution below\n", 592 | "\n", 593 | "solution below\n", 594 | "\n", 595 | "solution below\n", 596 | "\n", 597 | "solution below\n", 598 | "\n", 599 | "solution below\n", 600 | "\n", 601 | "solution below\n", 602 | "\n", 603 | "solution below\n", 604 | "\n", 605 | "solution below" 606 | ] 607 | }, 608 | { 609 | "cell_type": "code", 610 | "execution_count": null, 611 | "id": "510e5541", 612 | "metadata": {}, 613 | "outputs": [], 614 | "source": [ 615 | "def compute_equilibrium(params, price_init=2.0):\n", 616 | " p_star = newton(lambda p: excess_demand(p, params), price_init)\n", 617 | " return p_star" 618 | ] 619 | }, 620 | { 621 | "cell_type": "code", 622 | "execution_count": null, 623 | "id": "9f0cdaa5", 624 | "metadata": {}, 625 | "outputs": [], 626 | "source": [ 627 | "params = create_market_params()\n", 628 | "compute_equilibrium(params)" 629 | ] 630 | }, 631 | { 632 | "cell_type": "markdown", 633 | "id": "c1a6ac4e", 634 | "metadata": {}, 635 | "source": [ 636 | "This looks about right given the figures above." 637 | ] 638 | }, 639 | { 640 | "cell_type": "markdown", 641 | "id": "d224b77d", 642 | "metadata": {}, 643 | "source": [ 644 | "### Exercise 5\n", 645 | "\n", 646 | "For $b$ in a grid of 200 values between 0.5 and 1.0, plot the equilibrium price for each $b$.\n", 647 | "\n", 648 | "Does the curve that you plotted slope up or down? Try to provide an explanation for what you see in terms of market equilibrium." 649 | ] 650 | }, 651 | { 652 | "cell_type": "code", 653 | "execution_count": null, 654 | "id": "e11f79cd", 655 | "metadata": {}, 656 | "outputs": [], 657 | "source": [ 658 | "# Put your code here" 659 | ] 660 | }, 661 | { 662 | "cell_type": "markdown", 663 | "id": "72049a1c", 664 | "metadata": {}, 665 | "source": [ 666 | "solution below\n", 667 | "\n", 668 | "solution below\n", 669 | "\n", 670 | "solution below\n", 671 | "\n", 672 | "solution below\n", 673 | "\n", 674 | "solution below\n", 675 | "\n", 676 | "solution below\n", 677 | "\n", 678 | "solution below\n", 679 | "\n", 680 | "solution below\n", 681 | "\n", 682 | "solution below\n", 683 | "\n", 684 | "solution below\n", 685 | "\n", 686 | "solution below\n", 687 | "\n", 688 | "solution below" 689 | ] 690 | }, 691 | { 692 | "cell_type": "code", 693 | "execution_count": null, 694 | "id": "e5743e8b", 695 | "metadata": {}, 696 | "outputs": [], 697 | "source": [ 698 | "b_grid = np.linspace(0.5, 1.0, 200)\n", 699 | "prices = np.empty_like(b_grid)\n", 700 | "for i, b in enumerate(b_grid):\n", 701 | " params = create_market_params(b=b)\n", 702 | " prices[i] = compute_equilibrium(params)" 703 | ] 704 | }, 705 | { 706 | "cell_type": "code", 707 | "execution_count": null, 708 | "id": "39e47c47", 709 | "metadata": {}, 710 | "outputs": [], 711 | "source": [ 712 | "fig, ax = plt.subplots()\n", 713 | "ax.plot(b_grid, prices, label=\"equilibrium prices\")\n", 714 | "ax.set_xlabel(\"$b$\")\n", 715 | "ax.set_ylabel(\"price\")\n", 716 | "ax.legend()\n", 717 | "plt.show()" 718 | ] 719 | }, 720 | { 721 | "cell_type": "markdown", 722 | "id": "dc11137b", 723 | "metadata": {}, 724 | "source": [ 725 | "The curve slopes down because an increase in $b$ pushes up supply at any given price. (In other words, the supply curve shifts up.) \n", 726 | "\n", 727 | "With greater supply, the price tends to fall." 728 | ] 729 | }, 730 | { 731 | "cell_type": "code", 732 | "execution_count": null, 733 | "id": "2cc4f742", 734 | "metadata": {}, 735 | "outputs": [], 736 | "source": [] 737 | } 738 | ], 739 | "metadata": { 740 | "kernelspec": { 741 | "display_name": "Python 3 (ipykernel)", 742 | "language": "python", 743 | "name": "python3" 744 | } 745 | }, 746 | "nbformat": 4, 747 | "nbformat_minor": 5 748 | } 749 | -------------------------------------------------------------------------------- /notebooks/equilib_nd.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "4c2cec22", 6 | "metadata": {}, 7 | "source": [ 8 | "# Equilibrium in Multiple Dimensions" 9 | ] 10 | }, 11 | { 12 | "cell_type": "markdown", 13 | "id": "f6688fb7", 14 | "metadata": {}, 15 | "source": [ 16 | "#### Author: [John Stachurski](http://johnstachurski.net/)" 17 | ] 18 | }, 19 | { 20 | "cell_type": "code", 21 | "execution_count": null, 22 | "id": "39c5b4a7", 23 | "metadata": {}, 24 | "outputs": [], 25 | "source": [ 26 | "import numpy as np\n", 27 | "import matplotlib.pyplot as plt\n", 28 | "from numpy import exp, sqrt\n", 29 | "from numba import njit\n", 30 | "from scipy.optimize import newton, root" 31 | ] 32 | }, 33 | { 34 | "cell_type": "markdown", 35 | "id": "efda9825", 36 | "metadata": {}, 37 | "source": [ 38 | "In this notebook we expore the problem of computing market equilibrium in a multivariate setting, with many goods.\n", 39 | "\n", 40 | "As a first step, we set up and solve a two-good problem. \n", 41 | "\n", 42 | "Then we shift to higher dimensions. \n", 43 | "\n", 44 | "We will show how gradient-based equation solvers can handle high dimensional problems." 45 | ] 46 | }, 47 | { 48 | "cell_type": "markdown", 49 | "id": "22455545", 50 | "metadata": {}, 51 | "source": [ 52 | "## Two Goods" 53 | ] 54 | }, 55 | { 56 | "cell_type": "markdown", 57 | "id": "b2f1c23d", 58 | "metadata": {}, 59 | "source": [ 60 | "We first consider a market for two related products, good 0 and good 1, with price vector $p = (p_0, p_1)$.\n", 61 | "\n", 62 | "Supply of good $i$ at price $p$ is \n", 63 | "\n", 64 | "$$ q^s_i (p) = b_i \\sqrt{p_i} $$\n", 65 | "\n", 66 | "Demand of good $i$ at price $p$ is\n", 67 | "\n", 68 | "$$ q^d_i (p) = \\exp(-a_{i0} p_0 - a_{i1} p_1) + c_i$$\n", 69 | "\n", 70 | "Here $c_i, b_i$ and $a_{ij}$ are parameters. \n", 71 | "\n", 72 | "The excess demand functions are\n", 73 | "\n", 74 | "$$ e_i(p) = q^d_i(p) - q^s_i(p), \\qquad i = 0, 1 $$\n", 75 | "\n", 76 | "An equilibrium price vector $p^*$ is one where $e_i(p^*) = 0$. \n", 77 | "\n", 78 | "We set\n", 79 | "\n", 80 | "$$ \n", 81 | " A = \\begin{pmatrix}\n", 82 | " a_{00} & a_{01} \\\\\n", 83 | " a_{10} & a_{11}\n", 84 | " \\end{pmatrix},\n", 85 | " \\qquad \n", 86 | " b = \\begin{pmatrix}\n", 87 | " b_0 \\\\\n", 88 | " b_1\n", 89 | " \\end{pmatrix}\n", 90 | " \\qquad \\text{and} \\qquad\n", 91 | " c = \\begin{pmatrix}\n", 92 | " c_0 \\\\\n", 93 | " c_1\n", 94 | " \\end{pmatrix}\n", 95 | "$$" 96 | ] 97 | }, 98 | { 99 | "cell_type": "markdown", 100 | "id": "35416cb0", 101 | "metadata": {}, 102 | "source": [ 103 | "## Graphical Exploration" 104 | ] 105 | }, 106 | { 107 | "cell_type": "markdown", 108 | "id": "573513f7", 109 | "metadata": {}, 110 | "source": [ 111 | "Since our problem is only two dimensional, we can use graphical analysis to visualize and help understand the problem.\n", 112 | "\n", 113 | "Our first step is to define the excess demand function\n", 114 | "\n", 115 | "$$ e(p) = \n", 116 | " \\begin{pmatrix}\n", 117 | " e_0(p) \\\\\n", 118 | " e_1(p)\n", 119 | " \\end{pmatrix}\n", 120 | "$$\n", 121 | "\n", 122 | "The function below does the job." 123 | ] 124 | }, 125 | { 126 | "cell_type": "code", 127 | "execution_count": null, 128 | "id": "a9a4026f", 129 | "metadata": {}, 130 | "outputs": [], 131 | "source": [ 132 | "def e(p, A, b, c):\n", 133 | " return exp(- A @ p) + c - b * sqrt(p)" 134 | ] 135 | }, 136 | { 137 | "cell_type": "markdown", 138 | "id": "ff1edd8e", 139 | "metadata": {}, 140 | "source": [ 141 | "Our default parameter values will be" 142 | ] 143 | }, 144 | { 145 | "cell_type": "code", 146 | "execution_count": null, 147 | "id": "9f6e2f81", 148 | "metadata": {}, 149 | "outputs": [], 150 | "source": [ 151 | "A = ((0.5, 0.4),\n", 152 | " (0.8, 0.2))\n", 153 | "A = np.asarray(A)\n", 154 | "b = np.ones(2)\n", 155 | "c = np.ones(2)" 156 | ] 157 | }, 158 | { 159 | "cell_type": "code", 160 | "execution_count": null, 161 | "id": "3f194655", 162 | "metadata": {}, 163 | "outputs": [], 164 | "source": [ 165 | "e((1.0, 0.5), A, b, c)" 166 | ] 167 | }, 168 | { 169 | "cell_type": "markdown", 170 | "id": "5ef1c21f", 171 | "metadata": {}, 172 | "source": [ 173 | "Next we plot the two functions $e_0$ and $e_1$ on a grid of $(p_0, p_1)$ values, using contour surfaces and lines.\n", 174 | "\n", 175 | "We will use the following function to build the contour plots." 176 | ] 177 | }, 178 | { 179 | "cell_type": "code", 180 | "execution_count": null, 181 | "id": "02b5b247", 182 | "metadata": {}, 183 | "outputs": [], 184 | "source": [ 185 | "def plot_excess_demand(ax, good=0, grid_size=100, grid_max=4, surface=True):\n", 186 | " p_grid = np.linspace(0, grid_max, grid_size)\n", 187 | " z = np.empty((100, 100))\n", 188 | "\n", 189 | " for i, p_0 in enumerate(p_grid):\n", 190 | " for j, p_1 in enumerate(p_grid):\n", 191 | " z[i, j] = e((p_0, p_1), A, b, c)[good]\n", 192 | "\n", 193 | " if surface:\n", 194 | " cs1 = ax.contourf(p_grid, p_grid, z.T, alpha=0.5)\n", 195 | " plt.colorbar(cs1, ax=ax, format=\"%.6f\")\n", 196 | "\n", 197 | " ctr1 = ax.contour(p_grid, p_grid, z.T, levels=[0.0])\n", 198 | " plt.clabel(ctr1, inline=1, fontsize=13)\n", 199 | " ax.set_xlabel(\"$p_0$\")\n", 200 | " ax.set_ylabel(\"$p_1$\")" 201 | ] 202 | }, 203 | { 204 | "cell_type": "markdown", 205 | "id": "cc035062", 206 | "metadata": {}, 207 | "source": [ 208 | "Here's our plot of $e_0$:" 209 | ] 210 | }, 211 | { 212 | "cell_type": "code", 213 | "execution_count": null, 214 | "id": "29dd2474", 215 | "metadata": {}, 216 | "outputs": [], 217 | "source": [ 218 | "fig, ax = plt.subplots(figsize=(10, 5.7))\n", 219 | "plot_excess_demand(ax, good=0)\n", 220 | "plt.show()" 221 | ] 222 | }, 223 | { 224 | "cell_type": "markdown", 225 | "id": "3a23c34a", 226 | "metadata": {}, 227 | "source": [ 228 | "We see the black contour line of zero, which tells us when $e_0(p)=0$.\n", 229 | "\n", 230 | "For a price vector $p$ such that $e_0(p) = 0$, we know that good $0$ is in equilibrium (demand equals supply)." 231 | ] 232 | }, 233 | { 234 | "cell_type": "markdown", 235 | "id": "f4731000", 236 | "metadata": {}, 237 | "source": [ 238 | "Here's our plot of $e_1$:" 239 | ] 240 | }, 241 | { 242 | "cell_type": "code", 243 | "execution_count": null, 244 | "id": "a0d2b148", 245 | "metadata": {}, 246 | "outputs": [], 247 | "source": [ 248 | "fig, ax = plt.subplots(figsize=(10, 5.7))\n", 249 | "plot_excess_demand(ax, good=1)\n", 250 | "plt.show()" 251 | ] 252 | }, 253 | { 254 | "cell_type": "markdown", 255 | "id": "bf3adcd3", 256 | "metadata": {}, 257 | "source": [ 258 | "Now the black contour line of zero tells us when $e_1(p)=0$ (i.e., good $1$ is in equilibrium)." 259 | ] 260 | }, 261 | { 262 | "cell_type": "markdown", 263 | "id": "b32719f1", 264 | "metadata": {}, 265 | "source": [ 266 | "If these two contour lines cross at some vector $p^*$, then $p^*$ is an equilibrium price vector." 267 | ] 268 | }, 269 | { 270 | "cell_type": "code", 271 | "execution_count": null, 272 | "id": "faf46ea2", 273 | "metadata": {}, 274 | "outputs": [], 275 | "source": [ 276 | "fig, ax = plt.subplots(figsize=(10, 5.7))\n", 277 | "for good in (0, 1):\n", 278 | " plot_excess_demand(ax, good=good, surface=False)\n", 279 | "plt.show()" 280 | ] 281 | }, 282 | { 283 | "cell_type": "markdown", 284 | "id": "ea85fb87", 285 | "metadata": {}, 286 | "source": [ 287 | "It seems there is an equilibrium close to $p = (1.6, 1.5)$." 288 | ] 289 | }, 290 | { 291 | "cell_type": "markdown", 292 | "id": "47eac29e", 293 | "metadata": {}, 294 | "source": [ 295 | "## Using a Multidimensional Root Finder" 296 | ] 297 | }, 298 | { 299 | "cell_type": "markdown", 300 | "id": "cc082de0", 301 | "metadata": {}, 302 | "source": [ 303 | "To solve for $p^*$ more precisely, we use `root`, a root-finding algorithm from `scipy.optimize`.\n", 304 | "\n", 305 | "We supply $p = (1, 1)$ as our initial guess." 306 | ] 307 | }, 308 | { 309 | "cell_type": "code", 310 | "execution_count": null, 311 | "id": "43fb18ee", 312 | "metadata": {}, 313 | "outputs": [], 314 | "source": [ 315 | "init_p = np.ones(2)" 316 | ] 317 | }, 318 | { 319 | "cell_type": "code", 320 | "execution_count": null, 321 | "id": "1ac70dd4", 322 | "metadata": {}, 323 | "outputs": [], 324 | "source": [ 325 | "solution = root(lambda p: e(p, A, b, c), init_p, method='hybr')\n", 326 | "p = solution.x" 327 | ] 328 | }, 329 | { 330 | "cell_type": "markdown", 331 | "id": "342614ff", 332 | "metadata": {}, 333 | "source": [ 334 | "Here's the resulting value:" 335 | ] 336 | }, 337 | { 338 | "cell_type": "code", 339 | "execution_count": null, 340 | "id": "445d479d", 341 | "metadata": {}, 342 | "outputs": [], 343 | "source": [ 344 | "p" 345 | ] 346 | }, 347 | { 348 | "cell_type": "markdown", 349 | "id": "3f226a52", 350 | "metadata": {}, 351 | "source": [ 352 | "This looks close to our guess from observing the figure. We can plug it back into $e$ to test that $e(p) \\approx 0$:" 353 | ] 354 | }, 355 | { 356 | "cell_type": "code", 357 | "execution_count": null, 358 | "id": "7f3026d5", 359 | "metadata": {}, 360 | "outputs": [], 361 | "source": [ 362 | "np.max(np.abs(e(p, A, b, c)))" 363 | ] 364 | }, 365 | { 366 | "cell_type": "markdown", 367 | "id": "8bee5459", 368 | "metadata": {}, 369 | "source": [ 370 | "This is indeed a very small error." 371 | ] 372 | }, 373 | { 374 | "cell_type": "markdown", 375 | "id": "0f897a3b", 376 | "metadata": {}, 377 | "source": [ 378 | "In most cases, for root-finding algorithms applied to smooth functions, supplying the Jacobian of the function leads to better convergence properties. \n", 379 | "\n", 380 | "In this case we manually calculate the elements of the Jacobian\n", 381 | "\n", 382 | "$$\n", 383 | " J(p) = \n", 384 | " \\begin{pmatrix}\n", 385 | " \\frac{\\partial e_0}{\\partial p_0}(p) & \\frac{\\partial e_0}{\\partial p_1}(p) \\\\\n", 386 | " \\frac{\\partial e_1}{\\partial p_0}(p) & \\frac{\\partial e_1}{\\partial p_1}(p)\n", 387 | " \\end{pmatrix}\n", 388 | "$$\n", 389 | "\n", 390 | "and supply the Jacobian as a function, as follows:" 391 | ] 392 | }, 393 | { 394 | "cell_type": "code", 395 | "execution_count": null, 396 | "id": "a2fe5969", 397 | "metadata": {}, 398 | "outputs": [], 399 | "source": [ 400 | "@njit\n", 401 | "def jacobian(p, A, b, c):\n", 402 | " p_0, p_1 = p\n", 403 | " a_00, a_01 = A[0, :]\n", 404 | " a_10, a_11 = A[1, :]\n", 405 | " j_00 = -a_00 * exp(-a_00 * p_0) - (b[0]/2) * p_0**(-1/2)\n", 406 | " j_01 = -a_01 * exp(-a_01 * p_1)\n", 407 | " j_10 = -a_10 * exp(-a_10 * p_0)\n", 408 | " j_11 = -a_11 * exp(-a_11 * p_1) - (b[1]/2) * p_1**(-1/2)\n", 409 | " J = [[j_00, j_01],\n", 410 | " [j_10, j_11]]\n", 411 | " return np.array(J)" 412 | ] 413 | }, 414 | { 415 | "cell_type": "code", 416 | "execution_count": null, 417 | "id": "da7d735f", 418 | "metadata": {}, 419 | "outputs": [], 420 | "source": [ 421 | "solution = root(lambda p: e(p, A, b, c),\n", 422 | " init_p, \n", 423 | " jac=lambda p: jacobian(p, A, b, c), \n", 424 | " method='hybr')\n", 425 | "p = solution.x" 426 | ] 427 | }, 428 | { 429 | "cell_type": "markdown", 430 | "id": "795837fe", 431 | "metadata": {}, 432 | "source": [ 433 | "Now the solution is even more accurate (although, in this low-dimensional problem, the difference is quite small):" 434 | ] 435 | }, 436 | { 437 | "cell_type": "code", 438 | "execution_count": null, 439 | "id": "2e60211f", 440 | "metadata": {}, 441 | "outputs": [], 442 | "source": [ 443 | "np.max(np.abs(e(p, A, b, c)))" 444 | ] 445 | }, 446 | { 447 | "cell_type": "markdown", 448 | "id": "06b330dc", 449 | "metadata": {}, 450 | "source": [ 451 | "## High-Dimensional Problems\n", 452 | "\n", 453 | "Our next step is to investigate a high-dimensional version of the market described above. This market consists of 5,000 goods. \n", 454 | "\n", 455 | "The excess demand function is essentially the same, but now the matrix $A$ is $5000 \\times 5000$ and the parameter vectors $b$ and $c$ are $5000 \\times 1$." 456 | ] 457 | }, 458 | { 459 | "cell_type": "code", 460 | "execution_count": null, 461 | "id": "bc3e472b", 462 | "metadata": {}, 463 | "outputs": [], 464 | "source": [ 465 | "dim = 5000\n", 466 | "\n", 467 | "# Create a random matrix A and normalize the rows to sum to one\n", 468 | "A = np.random.rand(dim, dim)\n", 469 | "A = np.asarray(A)\n", 470 | "s = np.sum(A, axis=0)\n", 471 | "A = A / s\n", 472 | "\n", 473 | "# Set up b and c\n", 474 | "b = np.ones(dim)\n", 475 | "c = np.ones(dim)" 476 | ] 477 | }, 478 | { 479 | "cell_type": "markdown", 480 | "id": "a5eba650", 481 | "metadata": {}, 482 | "source": [ 483 | "Here's the same demand function:" 484 | ] 485 | }, 486 | { 487 | "cell_type": "code", 488 | "execution_count": null, 489 | "id": "fc06d45d", 490 | "metadata": {}, 491 | "outputs": [], 492 | "source": [ 493 | "def e(p, A, b, c):\n", 494 | " return exp(- A @ p) + c - b * sqrt(p)" 495 | ] 496 | }, 497 | { 498 | "cell_type": "markdown", 499 | "id": "3bc76b93", 500 | "metadata": {}, 501 | "source": [ 502 | "For our particular case, calculating and supplying the Jacobian is not too hard, but you can imagine that it can be very tedious when the functions get more complicated. \n", 503 | "\n", 504 | "So let's see how we go when the Jacobian is not supplied.\n", 505 | "\n", 506 | "Here's our initial condition" 507 | ] 508 | }, 509 | { 510 | "cell_type": "code", 511 | "execution_count": null, 512 | "id": "8f6b7522", 513 | "metadata": {}, 514 | "outputs": [], 515 | "source": [ 516 | "init_p = np.ones(dim)" 517 | ] 518 | }, 519 | { 520 | "cell_type": "markdown", 521 | "id": "26364ade", 522 | "metadata": {}, 523 | "source": [ 524 | "Now we call `root` again.\n", 525 | "\n", 526 | "**Warning**: The next line of code takes several minutes to run on a standard laptop or desktop." 527 | ] 528 | }, 529 | { 530 | "cell_type": "code", 531 | "execution_count": null, 532 | "id": "85cacfe4", 533 | "metadata": {}, 534 | "outputs": [], 535 | "source": [ 536 | "%%time\n", 537 | "solution = root(lambda p: e(p, A, b, c), init_p, method='hybr')" 538 | ] 539 | }, 540 | { 541 | "cell_type": "code", 542 | "execution_count": null, 543 | "id": "ebfb12ab", 544 | "metadata": {}, 545 | "outputs": [], 546 | "source": [ 547 | "p = solution.x" 548 | ] 549 | }, 550 | { 551 | "cell_type": "markdown", 552 | "id": "1a1377b7", 553 | "metadata": {}, 554 | "source": [ 555 | "Although it takes a long time to run, the answer is correct up to a high degree of accuracy." 556 | ] 557 | }, 558 | { 559 | "cell_type": "code", 560 | "execution_count": null, 561 | "id": "2c1c98d0", 562 | "metadata": {}, 563 | "outputs": [], 564 | "source": [ 565 | "np.max(np.abs(e(p, A, b, c)))" 566 | ] 567 | }, 568 | { 569 | "cell_type": "markdown", 570 | "id": "8104b075", 571 | "metadata": {}, 572 | "source": [ 573 | "## Automatic Differentiation" 574 | ] 575 | }, 576 | { 577 | "cell_type": "markdown", 578 | "id": "93ddd00a", 579 | "metadata": {}, 580 | "source": [ 581 | "To produce a faster and more efficient implementation, we shift to using JAX.\n", 582 | "\n", 583 | "With JAX, we get three big advantages:\n", 584 | "\n", 585 | "1. We can use JAX's automatic differentiation to compute the Jacobian easily and efficiently.\n", 586 | "2. JAX can parallelize the problem.\n", 587 | "3. JAX can dispatch to the GPU, if configured" 588 | ] 589 | }, 590 | { 591 | "cell_type": "code", 592 | "execution_count": null, 593 | "id": "d5cd6437", 594 | "metadata": {}, 595 | "outputs": [], 596 | "source": [ 597 | "!nvidia-smi" 598 | ] 599 | }, 600 | { 601 | "cell_type": "code", 602 | "execution_count": null, 603 | "id": "f6d4d707", 604 | "metadata": {}, 605 | "outputs": [], 606 | "source": [ 607 | "import jax\n", 608 | "import jax.numpy as jnp" 609 | ] 610 | }, 611 | { 612 | "cell_type": "markdown", 613 | "id": "bab1d7de", 614 | "metadata": {}, 615 | "source": [ 616 | "We set up the same demand function, replacing `np` with `jnp`:" 617 | ] 618 | }, 619 | { 620 | "cell_type": "code", 621 | "execution_count": null, 622 | "id": "d7278256", 623 | "metadata": {}, 624 | "outputs": [], 625 | "source": [ 626 | "@jax.jit\n", 627 | "def e(p, A, b, c):\n", 628 | " return jnp.exp(- jnp.dot(A, p)) + c - b * jnp.sqrt(p)" 629 | ] 630 | }, 631 | { 632 | "cell_type": "markdown", 633 | "id": "31a9213b", 634 | "metadata": {}, 635 | "source": [ 636 | "We are going to try to compute the equilibrium price using the multivariate version of Newton's method, which means iterating on the equation\n", 637 | "\n", 638 | "$$ p_{n+1} = p_n - J(p_n)^{-1} e(p_n) $$\n", 639 | "\n", 640 | "starting from some initial guess of the price vector $p_0$. (Here $J$ is the Jacobian of $e$.)" 641 | ] 642 | }, 643 | { 644 | "cell_type": "code", 645 | "execution_count": null, 646 | "id": "fb57f816", 647 | "metadata": {}, 648 | "outputs": [], 649 | "source": [ 650 | "def newton(f, x_0, tol=1e-5):\n", 651 | " f_prime = jax.grad(f)\n", 652 | " def q(x):\n", 653 | " return x - jnp.linalg.solve(jax.jacobian(f)(x), f(x))\n", 654 | "\n", 655 | " error = tol + 1\n", 656 | " x = x_0\n", 657 | " while error > tol:\n", 658 | " y = q(x)\n", 659 | " error = jnp.linalg.norm(x - y)\n", 660 | " x = y\n", 661 | " \n", 662 | " return x" 663 | ] 664 | }, 665 | { 666 | "cell_type": "markdown", 667 | "id": "b84c7b6e", 668 | "metadata": {}, 669 | "source": [ 670 | "Let's see whether this can solve the problem and how long it takes" 671 | ] 672 | }, 673 | { 674 | "cell_type": "code", 675 | "execution_count": null, 676 | "id": "9b7478b5", 677 | "metadata": {}, 678 | "outputs": [], 679 | "source": [ 680 | "%%time\n", 681 | "p = newton(lambda p: e(p, A, b, c), init_p).block_until_ready()" 682 | ] 683 | }, 684 | { 685 | "cell_type": "code", 686 | "execution_count": null, 687 | "id": "3431e0f1", 688 | "metadata": {}, 689 | "outputs": [], 690 | "source": [ 691 | "%%time\n", 692 | "p = newton(lambda p: e(p, A, b, c), init_p).block_until_ready()" 693 | ] 694 | }, 695 | { 696 | "cell_type": "code", 697 | "execution_count": null, 698 | "id": "085bc358", 699 | "metadata": {}, 700 | "outputs": [], 701 | "source": [ 702 | "np.max(np.abs(e(p, A, b, c)))" 703 | ] 704 | }, 705 | { 706 | "cell_type": "markdown", 707 | "id": "b83b26a3", 708 | "metadata": {}, 709 | "source": [ 710 | "We still have a solution that's very accurate and the compute time is massively reduced (assuming JAX is connecting to a GPU).\n", 711 | "\n", 712 | "\n", 713 | "### Exercise \n", 714 | "\n", 715 | "Write a simplified version of the `newton` function above that works for\n", 716 | "scalar functions (real inputs and real outputs).\n", 717 | "\n", 718 | "Test it on this function:" 719 | ] 720 | }, 721 | { 722 | "cell_type": "code", 723 | "execution_count": null, 724 | "id": "53bf40c4", 725 | "metadata": {}, 726 | "outputs": [], 727 | "source": [ 728 | "f = lambda x: jnp.sin(4 * (x - 1/4)) + x + x**20 - 1\n", 729 | "x = jnp.linspace(0, 1, 100)\n", 730 | "\n", 731 | "fig, ax = plt.subplots()\n", 732 | "ax.plot(x, f(x), label='$f(x)$')\n", 733 | "ax.axhline(ls='--', c='k')\n", 734 | "ax.set_xlabel('$x$', fontsize=12)\n", 735 | "ax.set_ylabel('$f(x)$', fontsize=12)\n", 736 | "ax.legend(fontsize=12)\n", 737 | "plt.show()" 738 | ] 739 | }, 740 | { 741 | "cell_type": "code", 742 | "execution_count": null, 743 | "id": "468ceb17", 744 | "metadata": {}, 745 | "outputs": [], 746 | "source": [ 747 | "# Put your code here" 748 | ] 749 | }, 750 | { 751 | "cell_type": "markdown", 752 | "id": "928e160f", 753 | "metadata": {}, 754 | "source": [ 755 | "solution below\n", 756 | "\n", 757 | "solution below\n", 758 | "\n", 759 | "solution below\n", 760 | "\n", 761 | "solution below\n", 762 | "\n", 763 | "solution below\n", 764 | "\n", 765 | "solution below\n", 766 | "\n", 767 | "solution below\n", 768 | "\n", 769 | "solution below\n", 770 | "\n", 771 | "solution below\n", 772 | "\n", 773 | "solution below\n", 774 | "\n", 775 | "solution below\n", 776 | "\n", 777 | "solution below" 778 | ] 779 | }, 780 | { 781 | "cell_type": "code", 782 | "execution_count": null, 783 | "id": "49591f6f", 784 | "metadata": {}, 785 | "outputs": [], 786 | "source": [ 787 | "def newton(f, x_0, tol=1e-5):\n", 788 | " f_prime = jax.grad(f)\n", 789 | " def q(x):\n", 790 | " return x - f(x) / f_prime(x)\n", 791 | "\n", 792 | " error = tol + 1\n", 793 | " x = x_0\n", 794 | " while error > tol:\n", 795 | " y = q(x)\n", 796 | " error = abs(x - y)\n", 797 | " x = y\n", 798 | " \n", 799 | " return x\n", 800 | "\n", 801 | "newton(f, 0.2)" 802 | ] 803 | } 804 | ], 805 | "metadata": { 806 | "kernelspec": { 807 | "display_name": "Python 3 (ipykernel)", 808 | "language": "python", 809 | "name": "python3" 810 | } 811 | }, 812 | "nbformat": 4, 813 | "nbformat_minor": 5 814 | } 815 | -------------------------------------------------------------------------------- /notebooks/opt_invest.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "8d56cde6", 6 | "metadata": {}, 7 | "source": [ 8 | "# Optimal Investment\n", 9 | "\n", 10 | "We require the following library to be installed." 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": null, 16 | "id": "3b887bd8", 17 | "metadata": { 18 | "tags": [ 19 | "hide-output" 20 | ] 21 | }, 22 | "outputs": [], 23 | "source": [ 24 | "!pip install --upgrade quantecon" 25 | ] 26 | }, 27 | { 28 | "cell_type": "markdown", 29 | "id": "d8b3f594", 30 | "metadata": {}, 31 | "source": [ 32 | "A monopolist faces inverse demand\n", 33 | "curve\n", 34 | "\n", 35 | "$$ P_t = a_0 - a_1 Y_t + Z_t, $$\n", 36 | "\n", 37 | "where\n", 38 | "\n", 39 | "* $P_t$ is price,\n", 40 | "* $Y_t$ is output and\n", 41 | "* $Z_t$ is a demand shock.\n", 42 | "\n", 43 | "We assume that $Z_t$ is a discretized AR(1) process.\n", 44 | "\n", 45 | "Current profits are\n", 46 | "\n", 47 | "$$ P_t Y_t - c Y_t - \\gamma (Y_{t+1} - Y_t)^2 $$\n", 48 | "\n", 49 | "Combining with the demand curve and writing $y, y'$ for $Y_t, Y_{t+1}$, this becomes\n", 50 | "\n", 51 | "$$ r(y, z, y′) := (a_0 - a_1 y + z - c) y - γ (y′ - y)^2 $$\n", 52 | "\n", 53 | "The firm maximizes present value of expected discounted profits. The Bellman equation is\n", 54 | "\n", 55 | "$$ v(y, z) = \\max_{y'} \\left\\{ r(y, z, y′) + β \\sum_{z′} v(y′, z′) Q(z, z′) \\right\\}. $$\n", 56 | "\n", 57 | "We discretize $y$ to a finite grid `y_grid`.\n", 58 | "\n", 59 | "In essence, the firm tries to choose output close to the monopolist profit maximizer, given $Z_t$, but is constrained by adjustment costs.\n", 60 | "\n", 61 | "Let's begin with the following imports" 62 | ] 63 | }, 64 | { 65 | "cell_type": "code", 66 | "execution_count": null, 67 | "id": "6b1f3055", 68 | "metadata": {}, 69 | "outputs": [], 70 | "source": [ 71 | "import quantecon as qe\n", 72 | "import jax\n", 73 | "import jax.numpy as jnp\n", 74 | "import matplotlib.pyplot as plt" 75 | ] 76 | }, 77 | { 78 | "cell_type": "markdown", 79 | "id": "d92f0553", 80 | "metadata": {}, 81 | "source": [ 82 | "Let’s check that we're hooked up to a GPU." 83 | ] 84 | }, 85 | { 86 | "cell_type": "code", 87 | "execution_count": null, 88 | "id": "382be6ae", 89 | "metadata": {}, 90 | "outputs": [], 91 | "source": [ 92 | "!nvidia-smi" 93 | ] 94 | }, 95 | { 96 | "cell_type": "markdown", 97 | "id": "59a5605c", 98 | "metadata": {}, 99 | "source": [ 100 | "We will use 64 bit floats with JAX in order to increase the precision." 101 | ] 102 | }, 103 | { 104 | "cell_type": "code", 105 | "execution_count": null, 106 | "id": "cbb43fc3", 107 | "metadata": {}, 108 | "outputs": [], 109 | "source": [ 110 | "jax.config.update(\"jax_enable_x64\", True)" 111 | ] 112 | }, 113 | { 114 | "cell_type": "markdown", 115 | "id": "b339b0b1", 116 | "metadata": {}, 117 | "source": [ 118 | "We need the following successive approximation function." 119 | ] 120 | }, 121 | { 122 | "cell_type": "code", 123 | "execution_count": null, 124 | "id": "f8242396", 125 | "metadata": {}, 126 | "outputs": [], 127 | "source": [ 128 | "def successive_approx(T, # Operator (callable)\n", 129 | " x_0, # Initial condition\n", 130 | " tolerance=1e-6, # Error tolerance\n", 131 | " max_iter=10_000, # Max iteration bound\n", 132 | " print_step=25, # Print at multiples\n", 133 | " verbose=False):\n", 134 | " x = x_0\n", 135 | " error = tolerance + 1\n", 136 | " k = 1\n", 137 | " while error > tolerance and k <= max_iter:\n", 138 | " x_new = T(x)\n", 139 | " error = jnp.max(jnp.abs(x_new - x))\n", 140 | " if verbose and k % print_step == 0:\n", 141 | " print(f\"Completed iteration {k} with error {error}.\")\n", 142 | " x = x_new\n", 143 | " k += 1\n", 144 | " if error > tolerance:\n", 145 | " print(f\"Warning: Iteration hit upper bound {max_iter}.\")\n", 146 | " elif verbose:\n", 147 | " print(f\"Terminated successfully in {k} iterations.\")\n", 148 | " return x" 149 | ] 150 | }, 151 | { 152 | "cell_type": "markdown", 153 | "id": "6faac0b6", 154 | "metadata": {}, 155 | "source": [ 156 | "Let's define a function to create an investment model using the given parameters." 157 | ] 158 | }, 159 | { 160 | "cell_type": "code", 161 | "execution_count": null, 162 | "id": "cde24456", 163 | "metadata": {}, 164 | "outputs": [], 165 | "source": [ 166 | "def create_investment_model(\n", 167 | " r=0.01, # Interest rate\n", 168 | " a_0=10.0, a_1=1.0, # Demand parameters\n", 169 | " γ=25.0, c=1.0, # Adjustment and unit cost\n", 170 | " y_min=0.0, y_max=20.0, y_size=100, # Grid for output\n", 171 | " ρ=0.9, ν=1.0, # AR(1) parameters\n", 172 | " z_size=150): # Grid size for shock\n", 173 | " \"\"\"\n", 174 | " A function that takes in parameters and returns an instance of Model that\n", 175 | " contains data for the investment problem.\n", 176 | " \"\"\"\n", 177 | " β = 1 / (1 + r)\n", 178 | " y_grid = jnp.linspace(y_min, y_max, y_size)\n", 179 | " mc = qe.tauchen(z_size, ρ, ν)\n", 180 | " z_grid, Q = mc.state_values, mc.P\n", 181 | "\n", 182 | " # Break up parameters into static and nonstatic components\n", 183 | " constants = β, a_0, a_1, γ, c\n", 184 | " sizes = y_size, z_size\n", 185 | " arrays = y_grid, z_grid, Q\n", 186 | "\n", 187 | " # Shift arrays to the device (e.g., GPU)\n", 188 | " arrays = tuple(map(jax.device_put, arrays))\n", 189 | " return constants, sizes, arrays" 190 | ] 191 | }, 192 | { 193 | "cell_type": "markdown", 194 | "id": "ca9aabb6", 195 | "metadata": {}, 196 | "source": [ 197 | "Let's re-write the vectorized version of the right-hand side of the\n", 198 | "Bellman equation (before maximization), which is a 3D array representing:\n", 199 | "\n", 200 | "$$\n", 201 | " B(y, z, y') = r(y, z, y') + \\beta \\sum_{z'} v(y', z') Q(z, z')\n", 202 | "$$\n", 203 | "\n", 204 | "for all $(y, z, y')$." 205 | ] 206 | }, 207 | { 208 | "cell_type": "code", 209 | "execution_count": null, 210 | "id": "28d8512f", 211 | "metadata": {}, 212 | "outputs": [], 213 | "source": [ 214 | "def B(v, constants, sizes, arrays):\n", 215 | " \"\"\"\n", 216 | " A vectorized version of the right-hand side of the Bellman equation\n", 217 | " (before maximization)\n", 218 | " \"\"\"\n", 219 | "\n", 220 | " # Unpack\n", 221 | " β, a_0, a_1, γ, c = constants\n", 222 | " y_size, z_size = sizes\n", 223 | " y_grid, z_grid, Q = arrays\n", 224 | "\n", 225 | " # Compute current rewards r(y, z, yp) as array r[i, j, ip]\n", 226 | " y = jnp.reshape(y_grid, (y_size, 1, 1)) # y[i] -> y[i, j, ip]\n", 227 | " z = jnp.reshape(z_grid, (1, z_size, 1)) # z[j] -> z[i, j, ip]\n", 228 | " yp = jnp.reshape(y_grid, (1, 1, y_size)) # yp[ip] -> yp[i, j, ip]\n", 229 | " r = (a_0 - a_1 * y + z - c) * y - γ * (yp - y)**2\n", 230 | "\n", 231 | " # Calculate continuation rewards at all combinations of (y, z, yp)\n", 232 | " v = jnp.reshape(v, (1, 1, y_size, z_size)) # v[ip, jp] -> v[i, j, ip, jp]\n", 233 | " Q = jnp.reshape(Q, (1, z_size, 1, z_size)) # Q[j, jp] -> Q[i, j, ip, jp]\n", 234 | " EV = jnp.sum(v * Q, axis=3) # sum over last index jp\n", 235 | "\n", 236 | " # Compute the right-hand side of the Bellman equation\n", 237 | " return r + β * EV\n", 238 | "\n", 239 | "# Create a jitted function\n", 240 | "B = jax.jit(B, static_argnums=(2,))" 241 | ] 242 | }, 243 | { 244 | "cell_type": "markdown", 245 | "id": "92ebc1df", 246 | "metadata": {}, 247 | "source": [ 248 | "Define a function to compute the current rewards given policy $\\sigma$." 249 | ] 250 | }, 251 | { 252 | "cell_type": "code", 253 | "execution_count": null, 254 | "id": "6ac205fc", 255 | "metadata": {}, 256 | "outputs": [], 257 | "source": [ 258 | "def compute_r_σ(σ, constants, sizes, arrays):\n", 259 | " \"\"\"\n", 260 | " Compute the array r_σ[i, j] = r[i, j, σ[i, j]], which gives current\n", 261 | " rewards given policy σ.\n", 262 | " \"\"\"\n", 263 | "\n", 264 | " # Unpack model\n", 265 | " β, a_0, a_1, γ, c = constants\n", 266 | " y_size, z_size = sizes\n", 267 | " y_grid, z_grid, Q = arrays\n", 268 | "\n", 269 | " # Compute r_σ[i, j]\n", 270 | " y = jnp.reshape(y_grid, (y_size, 1))\n", 271 | " z = jnp.reshape(z_grid, (1, z_size))\n", 272 | " yp = y_grid[σ]\n", 273 | " r_σ = (a_0 - a_1 * y + z - c) * y - γ * (yp - y)**2\n", 274 | "\n", 275 | " return r_σ\n", 276 | "\n", 277 | "\n", 278 | "# Create the jitted function\n", 279 | "compute_r_σ = jax.jit(compute_r_σ, static_argnums=(2,))" 280 | ] 281 | }, 282 | { 283 | "cell_type": "markdown", 284 | "id": "43546aa5", 285 | "metadata": {}, 286 | "source": [ 287 | "Define the Bellman operator." 288 | ] 289 | }, 290 | { 291 | "cell_type": "code", 292 | "execution_count": null, 293 | "id": "b0fb3e78", 294 | "metadata": {}, 295 | "outputs": [], 296 | "source": [ 297 | "def T(v, constants, sizes, arrays):\n", 298 | " \"\"\"The Bellman operator.\"\"\"\n", 299 | " return jnp.max(B(v, constants, sizes, arrays), axis=2)\n", 300 | "\n", 301 | "T = jax.jit(T, static_argnums=(2,))" 302 | ] 303 | }, 304 | { 305 | "cell_type": "markdown", 306 | "id": "62c320a4", 307 | "metadata": {}, 308 | "source": [ 309 | "The following function computes a v-greedy policy." 310 | ] 311 | }, 312 | { 313 | "cell_type": "code", 314 | "execution_count": null, 315 | "id": "1a7c567d", 316 | "metadata": {}, 317 | "outputs": [], 318 | "source": [ 319 | "def get_greedy(v, constants, sizes, arrays):\n", 320 | " \"Computes a v-greedy policy, returned as a set of indices.\"\n", 321 | " return jnp.argmax(B(v, constants, sizes, arrays), axis=2)\n", 322 | "\n", 323 | "get_greedy = jax.jit(get_greedy, static_argnums=(2,))" 324 | ] 325 | }, 326 | { 327 | "cell_type": "markdown", 328 | "id": "598b5dce", 329 | "metadata": {}, 330 | "source": [ 331 | "Define the $\\sigma$-policy operator." 332 | ] 333 | }, 334 | { 335 | "cell_type": "code", 336 | "execution_count": null, 337 | "id": "a66f8b6b", 338 | "metadata": {}, 339 | "outputs": [], 340 | "source": [ 341 | "def T_σ(v, σ, constants, sizes, arrays):\n", 342 | " \"\"\"The σ-policy operator.\"\"\"\n", 343 | "\n", 344 | " # Unpack model\n", 345 | " β, a_0, a_1, γ, c = constants\n", 346 | " y_size, z_size = sizes\n", 347 | " y_grid, z_grid, Q = arrays\n", 348 | "\n", 349 | " r_σ = compute_r_σ(σ, constants, sizes, arrays)\n", 350 | "\n", 351 | " # Compute the array v[σ[i, j], jp]\n", 352 | " zp_idx = jnp.arange(z_size)\n", 353 | " zp_idx = jnp.reshape(zp_idx, (1, 1, z_size))\n", 354 | " σ = jnp.reshape(σ, (y_size, z_size, 1))\n", 355 | " V = v[σ, zp_idx]\n", 356 | "\n", 357 | " # Convert Q[j, jp] to Q[i, j, jp]\n", 358 | " Q = jnp.reshape(Q, (1, z_size, z_size))\n", 359 | "\n", 360 | " # Calculate the expected sum Σ_jp v[σ[i, j], jp] * Q[i, j, jp]\n", 361 | " Ev = jnp.sum(V * Q, axis=2)\n", 362 | "\n", 363 | " return r_σ + β * jnp.sum(V * Q, axis=2)\n", 364 | "\n", 365 | "T_σ = jax.jit(T_σ, static_argnums=(3,))" 366 | ] 367 | }, 368 | { 369 | "cell_type": "markdown", 370 | "id": "bea40c7e", 371 | "metadata": {}, 372 | "source": [ 373 | "Next, we want to computes the lifetime value of following policy $\\sigma$.\n", 374 | "\n", 375 | "The basic problem is to solve the linear system\n", 376 | "\n", 377 | "$$ v(y, z) = r(y, z, \\sigma(y, z)) + \\beta \\sum_{z'} v(\\sigma(y, z), z') Q(z, z) $$\n", 378 | "\n", 379 | "for $v$.\n", 380 | "\n", 381 | "It turns out to be helpful to rewrite this as\n", 382 | "\n", 383 | "$$ v(y, z) = r(y, z, \\sigma(y, z)) + \\beta \\sum_{y', z'} v(y', z') P_\\sigma(y, z, y', z') $$\n", 384 | "\n", 385 | "where $P_\\sigma(y, z, y', z') = 1\\{y' = \\sigma(y, z)\\} Q(z, z')$.\n", 386 | "\n", 387 | "We want to write this as $v = r_\\sigma + \\beta P_\\sigma v$ and then solve for $v$\n", 388 | "\n", 389 | "Note, however, that $v$ is a multi-index array, rather than a vector.\n", 390 | "\n", 391 | "\n", 392 | "The value $v_{\\sigma}$ of a policy $\\sigma$ is defined as\n", 393 | "\n", 394 | "$$\n", 395 | " v_{\\sigma} = (I - \\beta P_{\\sigma})^{-1} r_{\\sigma}\n", 396 | "$$\n", 397 | "\n", 398 | "Here we set up the linear map $v \\mapsto R_{\\sigma} v$, where \n", 399 | "\n", 400 | "$$\n", 401 | " R_{\\sigma} := I - \\beta P_{\\sigma}\n", 402 | "$$\n", 403 | "\n", 404 | "In the investment problem, this map can be expressed as\n", 405 | "\n", 406 | "$$\n", 407 | " (R_{\\sigma} v)(y, z) = v(y, z) - \\beta \\sum_{z'} v(\\sigma(y, z), z') Q(z, z')\n", 408 | "$$\n", 409 | "\n", 410 | "Defining the map as above works in a more intuitive multi-index setting\n", 411 | "(e.g. working with $v[i, j]$ rather than flattening v to a one-dimensional\n", 412 | "array) and avoids instantiating the large matrix $P_{\\sigma}$.\n", 413 | "\n", 414 | "Let's define the function $R_{\\sigma}$." 415 | ] 416 | }, 417 | { 418 | "cell_type": "code", 419 | "execution_count": null, 420 | "id": "9e8f24fb", 421 | "metadata": {}, 422 | "outputs": [], 423 | "source": [ 424 | "def R_σ(v, σ, constants, sizes, arrays):\n", 425 | "\n", 426 | " β, a_0, a_1, γ, c = constants\n", 427 | " y_size, z_size = sizes\n", 428 | " y_grid, z_grid, Q = arrays\n", 429 | "\n", 430 | " # Set up the array v[σ[i, j], jp]\n", 431 | " zp_idx = jnp.arange(z_size)\n", 432 | " zp_idx = jnp.reshape(zp_idx, (1, 1, z_size))\n", 433 | " σ = jnp.reshape(σ, (y_size, z_size, 1))\n", 434 | " V = v[σ, zp_idx]\n", 435 | "\n", 436 | " # Expand Q[j, jp] to Q[i, j, jp]\n", 437 | " Q = jnp.reshape(Q, (1, z_size, z_size))\n", 438 | "\n", 439 | " # Compute and return v[i, j] - β Σ_jp v[σ[i, j], jp] * Q[j, jp]\n", 440 | " return v - β * jnp.sum(V * Q, axis=2)\n", 441 | "\n", 442 | "R_σ = jax.jit(R_σ, static_argnums=(3,))" 443 | ] 444 | }, 445 | { 446 | "cell_type": "markdown", 447 | "id": "6d97000e", 448 | "metadata": {}, 449 | "source": [ 450 | "Define a function to get the value $v_{\\sigma}$ of policy\n", 451 | "$\\sigma$ by inverting the linear map $R_{\\sigma}$." 452 | ] 453 | }, 454 | { 455 | "cell_type": "code", 456 | "execution_count": null, 457 | "id": "275172fd", 458 | "metadata": {}, 459 | "outputs": [], 460 | "source": [ 461 | "def get_value(σ, constants, sizes, arrays):\n", 462 | "\n", 463 | " # Unpack\n", 464 | " β, a_0, a_1, γ, c = constants\n", 465 | " y_size, z_size = sizes\n", 466 | " y_grid, z_grid, Q = arrays\n", 467 | "\n", 468 | " r_σ = compute_r_σ(σ, constants, sizes, arrays)\n", 469 | "\n", 470 | " # Reduce R_σ to a function in v\n", 471 | " partial_R_σ = lambda v: R_σ(v, σ, constants, sizes, arrays)\n", 472 | "\n", 473 | " return jax.scipy.sparse.linalg.bicgstab(partial_R_σ, r_σ)[0]\n", 474 | "\n", 475 | "get_value = jax.jit(get_value, static_argnums=(2,))" 476 | ] 477 | }, 478 | { 479 | "cell_type": "markdown", 480 | "id": "000ee7ff", 481 | "metadata": {}, 482 | "source": [ 483 | "Now we define the solvers, which implement VFI, HPI and OPI." 484 | ] 485 | }, 486 | { 487 | "cell_type": "code", 488 | "execution_count": null, 489 | "id": "676a45be", 490 | "metadata": {}, 491 | "outputs": [], 492 | "source": [ 493 | "# Implements VFI-Value Function iteration\n", 494 | "\n", 495 | "def value_iteration(model, tol=1e-5):\n", 496 | " constants, sizes, arrays = model\n", 497 | " _T = lambda v: T(v, constants, sizes, arrays)\n", 498 | " vz = jnp.zeros(sizes)\n", 499 | "\n", 500 | " v_star = successive_approx(_T, vz, tolerance=tol)\n", 501 | " return get_greedy(v_star, constants, sizes, arrays)" 502 | ] 503 | }, 504 | { 505 | "cell_type": "code", 506 | "execution_count": null, 507 | "id": "9fc120b5", 508 | "metadata": {}, 509 | "outputs": [], 510 | "source": [ 511 | "# Implements HPI-Howard policy iteration routine\n", 512 | "\n", 513 | "def policy_iteration(model, maxiter=250):\n", 514 | " constants, sizes, arrays = model\n", 515 | " vz = jnp.zeros(sizes)\n", 516 | " σ = jnp.zeros(sizes, dtype=int)\n", 517 | " i, error = 0, 1.0\n", 518 | " while error > 0 and i < maxiter:\n", 519 | " v_σ = get_value(σ, constants, sizes, arrays)\n", 520 | " σ_new = get_greedy(v_σ, constants, sizes, arrays)\n", 521 | " error = jnp.max(jnp.abs(σ_new - σ))\n", 522 | " σ = σ_new\n", 523 | " i = i + 1\n", 524 | " print(f\"Concluded loop {i} with error {error}.\")\n", 525 | " return σ" 526 | ] 527 | }, 528 | { 529 | "cell_type": "code", 530 | "execution_count": null, 531 | "id": "4616b0ba", 532 | "metadata": {}, 533 | "outputs": [], 534 | "source": [ 535 | "# Implements the OPI-Optimal policy Iteration routine\n", 536 | "\n", 537 | "def optimistic_policy_iteration(model, tol=1e-5, m=10):\n", 538 | " constants, sizes, arrays = model\n", 539 | " v = jnp.zeros(sizes)\n", 540 | " error = tol + 1\n", 541 | " while error > tol:\n", 542 | " last_v = v\n", 543 | " σ = get_greedy(v, constants, sizes, arrays)\n", 544 | " for _ in range(m):\n", 545 | " v = T_σ(v, σ, constants, sizes, arrays)\n", 546 | " error = jnp.max(jnp.abs(v - last_v))\n", 547 | " return get_greedy(v, constants, sizes, arrays)" 548 | ] 549 | }, 550 | { 551 | "cell_type": "code", 552 | "execution_count": null, 553 | "id": "56820673", 554 | "metadata": { 555 | "tags": [ 556 | "hide-output" 557 | ] 558 | }, 559 | "outputs": [], 560 | "source": [ 561 | "model = create_investment_model()\n", 562 | "print(\"Starting HPI.\")\n", 563 | "qe.tic()\n", 564 | "out = policy_iteration(model)\n", 565 | "elapsed = qe.toc()\n", 566 | "print(out)\n", 567 | "print(f\"HPI completed in {elapsed} seconds.\")" 568 | ] 569 | }, 570 | { 571 | "cell_type": "code", 572 | "execution_count": null, 573 | "id": "8fc86a7a", 574 | "metadata": { 575 | "tags": [ 576 | "hide-output" 577 | ] 578 | }, 579 | "outputs": [], 580 | "source": [ 581 | "print(\"Starting VFI.\")\n", 582 | "qe.tic()\n", 583 | "out = value_iteration(model)\n", 584 | "elapsed = qe.toc()\n", 585 | "print(out)\n", 586 | "print(f\"VFI completed in {elapsed} seconds.\")" 587 | ] 588 | }, 589 | { 590 | "cell_type": "code", 591 | "execution_count": null, 592 | "id": "ed4442b9", 593 | "metadata": { 594 | "tags": [ 595 | "hide-output" 596 | ] 597 | }, 598 | "outputs": [], 599 | "source": [ 600 | "print(\"Starting OPI.\")\n", 601 | "qe.tic()\n", 602 | "out = optimistic_policy_iteration(model, m=100)\n", 603 | "elapsed = qe.toc()\n", 604 | "print(out)\n", 605 | "print(f\"OPI completed in {elapsed} seconds.\")" 606 | ] 607 | }, 608 | { 609 | "cell_type": "markdown", 610 | "id": "536bece5", 611 | "metadata": {}, 612 | "source": [ 613 | "Here's the plot of the Howard policy, as a function of $y$ at the highest and lowest values of $z$." 614 | ] 615 | }, 616 | { 617 | "cell_type": "code", 618 | "execution_count": null, 619 | "id": "57b52d88", 620 | "metadata": {}, 621 | "outputs": [], 622 | "source": [ 623 | "model = create_investment_model()\n", 624 | "constants, sizes, arrays = model\n", 625 | "β, a_0, a_1, γ, c = constants\n", 626 | "y_size, z_size = sizes\n", 627 | "y_grid, z_grid, Q = arrays" 628 | ] 629 | }, 630 | { 631 | "cell_type": "code", 632 | "execution_count": null, 633 | "id": "a9217cd5", 634 | "metadata": {}, 635 | "outputs": [], 636 | "source": [ 637 | "σ_star = policy_iteration(model)\n", 638 | "\n", 639 | "fig, ax = plt.subplots(figsize=(9, 5))\n", 640 | "ax.plot(y_grid, y_grid, \"k--\", label=\"45\")\n", 641 | "ax.plot(y_grid, y_grid[σ_star[:, 1]], label=\"$\\\\sigma^*(\\cdot, z_1)$\")\n", 642 | "ax.plot(y_grid, y_grid[σ_star[:, -1]], label=\"$\\\\sigma^*(\\cdot, z_N)$\")\n", 643 | "ax.legend(fontsize=12)\n", 644 | "plt.show()" 645 | ] 646 | }, 647 | { 648 | "cell_type": "markdown", 649 | "id": "1d00a815", 650 | "metadata": {}, 651 | "source": [ 652 | "Let's plot the time taken by each of the solvers and compare them." 653 | ] 654 | }, 655 | { 656 | "cell_type": "code", 657 | "execution_count": null, 658 | "id": "c93b03ea", 659 | "metadata": {}, 660 | "outputs": [], 661 | "source": [ 662 | "m_vals = range(5, 3000, 100)" 663 | ] 664 | }, 665 | { 666 | "cell_type": "code", 667 | "execution_count": null, 668 | "id": "bf16dd43", 669 | "metadata": {}, 670 | "outputs": [], 671 | "source": [ 672 | "model = create_investment_model()\n", 673 | "print(\"Running Howard policy iteration.\")\n", 674 | "qe.tic()\n", 675 | "σ_pi = policy_iteration(model)\n", 676 | "pi_time = qe.toc()" 677 | ] 678 | }, 679 | { 680 | "cell_type": "code", 681 | "execution_count": null, 682 | "id": "ec3e3e96", 683 | "metadata": {}, 684 | "outputs": [], 685 | "source": [ 686 | "print(f\"PI completed in {pi_time} seconds.\")\n", 687 | "print(\"Running value function iteration.\")\n", 688 | "qe.tic()\n", 689 | "σ_vfi = value_iteration(model, tol=1e-5)\n", 690 | "vfi_time = qe.toc()\n", 691 | "print(f\"VFI completed in {vfi_time} seconds.\")" 692 | ] 693 | }, 694 | { 695 | "cell_type": "code", 696 | "execution_count": null, 697 | "id": "47239c67", 698 | "metadata": { 699 | "tags": [ 700 | "hide-output" 701 | ] 702 | }, 703 | "outputs": [], 704 | "source": [ 705 | "opi_times = []\n", 706 | "for m in m_vals:\n", 707 | " print(f\"Running optimistic policy iteration with m={m}.\")\n", 708 | " qe.tic()\n", 709 | " σ_opi = optimistic_policy_iteration(model, m=m, tol=1e-5)\n", 710 | " opi_time = qe.toc()\n", 711 | " print(f\"OPI with m={m} completed in {opi_time} seconds.\")\n", 712 | " opi_times.append(opi_time)" 713 | ] 714 | }, 715 | { 716 | "cell_type": "code", 717 | "execution_count": null, 718 | "id": "6914e83e", 719 | "metadata": {}, 720 | "outputs": [], 721 | "source": [ 722 | "fig, ax = plt.subplots(figsize=(9, 5))\n", 723 | "ax.plot(m_vals, jnp.full(len(m_vals), pi_time),\n", 724 | " lw=2, label=\"Howard policy iteration\")\n", 725 | "ax.plot(m_vals, jnp.full(len(m_vals), vfi_time),\n", 726 | " lw=2, label=\"value function iteration\")\n", 727 | "ax.plot(m_vals, opi_times, lw=2, label=\"optimistic policy iteration\")\n", 728 | "ax.legend(fontsize=12, frameon=False)\n", 729 | "ax.set_xlabel(\"$m$\", fontsize=12)\n", 730 | "ax.set_ylabel(\"time(s)\", fontsize=12)\n", 731 | "plt.show()" 732 | ] 733 | } 734 | ], 735 | "metadata": { 736 | "kernelspec": { 737 | "display_name": "Python 3 (ipykernel)", 738 | "language": "python", 739 | "name": "python3" 740 | } 741 | }, 742 | "nbformat": 4, 743 | "nbformat_minor": 5 744 | } 745 | -------------------------------------------------------------------------------- /notebooks/european_option.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "b01029ed", 6 | "metadata": {}, 7 | "source": [ 8 | "# Monte Carlo and Option Pricing \n", 9 | "\n", 10 | "* Author: [John Stachurski](http://johnstachurski.net/)\n", 11 | "\n", 12 | "We discuss [Monte Carlo\n", 13 | "methods](https://en.wikipedia.org/wiki/Monte_Carlo_method) for computing\n", 14 | "expectations with applications in finance.\n", 15 | "\n", 16 | "Our main application will be pricing a European option.\n", 17 | "\n", 18 | "We will show that Monte Carlo is particularly helpful when the distribution of\n", 19 | "interest has no neat analytical form.\n", 20 | "\n", 21 | "We will also touch on some high performance computing topics, including \n", 22 | "\n", 23 | "* just-in-time compilers \n", 24 | "* GPUs and \n", 25 | "* parallelization.\n", 26 | "\n", 27 | "We begin with the following imports:" 28 | ] 29 | }, 30 | { 31 | "cell_type": "code", 32 | "execution_count": null, 33 | "id": "09aee0ad", 34 | "metadata": {}, 35 | "outputs": [], 36 | "source": [ 37 | "import numpy as np\n", 38 | "import matplotlib.pyplot as plt\n", 39 | "import numba\n", 40 | "from numpy.random import randn" 41 | ] 42 | }, 43 | { 44 | "cell_type": "markdown", 45 | "id": "d7a0e274", 46 | "metadata": {}, 47 | "source": [ 48 | "## Pricing a European Call Option under Risk Neutrality\n", 49 | "\n", 50 | "As our next step, we are going to price a European call option under risk\n", 51 | "neutrality.\n", 52 | "\n", 53 | "Let's first discuss risk neutrality and then introduce European options.\n", 54 | "\n", 55 | "### Risk-Neutral Pricing\n", 56 | "\n", 57 | "When we use risk-neutral pricing, we determine the price of a given asset\n", 58 | "according to its expected payoff.\n", 59 | "\n", 60 | "$$\n", 61 | "\\text{cost } = \\text{ expected benefit}\n", 62 | "$$\n", 63 | "\n", 64 | "For example, suppose someone promises to pay you\n", 65 | "\n", 66 | "- 1,000,000 dollars if \"heads\" is the outcome of a fair coin flip\n", 67 | "- 0 dollars if \"tails\" is the outcome\n", 68 | "\n", 69 | "Let's denote the payoff as $G$, so that \n", 70 | "\n", 71 | "$$\n", 72 | " \\mathbb P\\left\\{G = 10^6 \\right\\} = \\mathbb P\\{G = 0\\} = \\frac{1}{2}\n", 73 | "$$\n", 74 | "\n", 75 | "Suppose in addition that you can sell this promise to anyone who wants to\n", 76 | "hold it \n", 77 | "\n", 78 | "- First they pay you $P$, the price at which you sell it\n", 79 | "- Then they get $G$, which could be either 1,000,000 or 0.\n", 80 | "\n", 81 | "What's a fair price for this asset (this promise)?\n", 82 | "\n", 83 | "The definition of fair is ambiguous but what we can say is that the\n", 84 | "risk-neutral price is 500,000 dollars.\n", 85 | "\n", 86 | "This is because the risk-neutral price is just the expected payoff of the\n", 87 | "asset, which is\n", 88 | "\n", 89 | "$$\n", 90 | " \\mathbb E G = \\frac{1}{2} \\times 10^6 + \\frac{1}{2} \\times 0 = 5 \\times 10^5\n", 91 | "$$" 92 | ] 93 | }, 94 | { 95 | "cell_type": "markdown", 96 | "id": "b81eb932", 97 | "metadata": {}, 98 | "source": [ 99 | "### A Comment on Risk\n", 100 | "\n", 101 | "As suggested by the name, the risk-neutral price ignores risk.\n", 102 | "\n", 103 | "To understand this, consider whether you would pay 500,000 dollars for such a\n", 104 | "promise.\n", 105 | "\n", 106 | "Would you prefer to receive 500,000 for sure or 1,000,000 dollars with\n", 107 | "50% probability and nothing with 50% probability?\n", 108 | "\n", 109 | "At least some readers will strictly prefer the first option --- although some\n", 110 | "might prefer the second.\n", 111 | "\n", 112 | "Thinking about this makes us realize that 500,000 is not necessarily the\n", 113 | "\"right\" price --- or the price that we would see if there was a market for\n", 114 | "these promises.\n", 115 | "\n", 116 | "Nonetheless, the risk-neutral price is an important benchmark, which economists\n", 117 | "and financial market participants routinely try to calculate." 118 | ] 119 | }, 120 | { 121 | "cell_type": "markdown", 122 | "id": "78ca3b4e", 123 | "metadata": {}, 124 | "source": [ 125 | "### Discounting\n", 126 | "\n", 127 | "One thing we ignored in the previous discussion was time.\n", 128 | "\n", 129 | "In general, receiving $x$ dollars now is preferable to receiving $x$ dollars\n", 130 | "in $n$ periods (e.g., 10 years).\n", 131 | "\n", 132 | "After all, if we receive $x$ dollars now, we could put it in the bank at\n", 133 | "interest rate $r > 0$ and receive $ (1 + r)^n x $ in $n$ periods.\n", 134 | "\n", 135 | "Hence future payments need to be discounted.\n", 136 | "\n", 137 | "We will implement discounting by \n", 138 | "\n", 139 | "* multiplying a payment in one period by $\\beta < 1$\n", 140 | "* multiplying a payment in $n$ periods by $\\beta^n$, etc.\n", 141 | "\n", 142 | "The same adjustment needs to be applied to our risk-neutral price for the\n", 143 | "promise described above.\n", 144 | "\n", 145 | "Thus, if $G$ is realized in $n$ periods, then the risk-neutral price is\n", 146 | "\n", 147 | "$$\n", 148 | " P = \\beta^n \\mathbb E G \n", 149 | " = \\beta^n 5 \\times 10^5\n", 150 | "$$" 151 | ] 152 | }, 153 | { 154 | "cell_type": "markdown", 155 | "id": "094d2fe6", 156 | "metadata": {}, 157 | "source": [ 158 | "### European Call Options\n", 159 | "\n", 160 | "Now let's price a European call option.\n", 161 | "\n", 162 | "The option is described by three things:\n", 163 | "\n", 164 | "2. $n$, the **expiry date**,\n", 165 | "2. $K$, the **strike price**, and\n", 166 | "3. $S_n$, the price of the **underlying** asset at date $n$.\n", 167 | "\n", 168 | "For example, suppose that the underlying is one share in Amazon.\n", 169 | "\n", 170 | "The owner of this option has the right to buy one share in Amazon at price $K$ after $n$ days.\n", 171 | "\n", 172 | "If $S_n > K$, then the owner will exercise the option, buy at $K$, sell at\n", 173 | "$S_n$, and make profit $S_n - K$.\n", 174 | "\n", 175 | "If $S_n \\leq K$, then the owner will not exercise the option and the payoff is zero.\n", 176 | "\n", 177 | "Thus, the payoff is $\\max\\{ S_n - K, 0 \\}$.\n", 178 | "\n", 179 | "Under the assumption of risk neutrality, the price of the option is \n", 180 | "the expected discounted payoff:\n", 181 | "\n", 182 | "$$ P = \\beta^n \\mathbb E \\max\\{ S_n - K, 0 \\} $$\n", 183 | "\n", 184 | "Now all we need to do is specify the distribution of $S_n$, so the expectation\n", 185 | "can be calculated." 186 | ] 187 | }, 188 | { 189 | "cell_type": "markdown", 190 | "id": "126fc468", 191 | "metadata": {}, 192 | "source": [ 193 | "**Exercise**\n", 194 | "\n", 195 | "Suppose we know that $S_n \\sim LN(\\mu, \\sigma)$ and $\\mu$ and $\\sigma$ are known.\n", 196 | "\n", 197 | "Use the fact that if $S_n^1, \\ldots, S_n^M$ are independent draws from this lognormal distribution then, by the law of large numbers,\n", 198 | "\n", 199 | "$$ \\mathbb E \\max\\{ S_n - K, 0 \\} \n", 200 | " \\approx\n", 201 | " \\frac{1}{M} \\sum_{m=1}^M \\max \\{S_n^m - K, 0 \\}\n", 202 | " $$" 203 | ] 204 | }, 205 | { 206 | "cell_type": "markdown", 207 | "id": "e08a3193", 208 | "metadata": {}, 209 | "source": [ 210 | "Use the following parameter values:" 211 | ] 212 | }, 213 | { 214 | "cell_type": "code", 215 | "execution_count": null, 216 | "id": "466e8546", 217 | "metadata": {}, 218 | "outputs": [], 219 | "source": [ 220 | "μ = 1.0\n", 221 | "σ = 0.1" 222 | ] 223 | }, 224 | { 225 | "cell_type": "code", 226 | "execution_count": null, 227 | "id": "bb70e4a6", 228 | "metadata": {}, 229 | "outputs": [], 230 | "source": [ 231 | "K = 1\n", 232 | "n = 10\n", 233 | "β = 0.95" 234 | ] 235 | }, 236 | { 237 | "cell_type": "markdown", 238 | "id": "4f968323", 239 | "metadata": {}, 240 | "source": [ 241 | "Set the simulation size to" 242 | ] 243 | }, 244 | { 245 | "cell_type": "code", 246 | "execution_count": null, 247 | "id": "4776157f", 248 | "metadata": {}, 249 | "outputs": [], 250 | "source": [ 251 | "M = 10_000_000" 252 | ] 253 | }, 254 | { 255 | "cell_type": "markdown", 256 | "id": "bb4393ba", 257 | "metadata": {}, 258 | "source": [ 259 | "solution below\n", 260 | "\n", 261 | "solution below\n", 262 | "\n", 263 | "solution below\n", 264 | "\n", 265 | "solution below\n", 266 | "\n", 267 | "solution below\n", 268 | "\n", 269 | "solution below\n", 270 | "\n", 271 | "solution below\n", 272 | "\n", 273 | "solution below\n", 274 | "\n", 275 | "solution below\n", 276 | "\n", 277 | "solution below\n", 278 | "\n", 279 | "solution below\n", 280 | "\n", 281 | "solution below" 282 | ] 283 | }, 284 | { 285 | "cell_type": "code", 286 | "execution_count": null, 287 | "id": "722e5a83", 288 | "metadata": {}, 289 | "outputs": [], 290 | "source": [ 291 | "S = np.exp(μ + σ * np.random.randn(M))\n", 292 | "return_draws = np.maximum(S - K, 0)\n", 293 | "P = β**n * np.mean(return_draws) \n", 294 | "print(f\"The Monte Carlo option price is {P:3f}\")" 295 | ] 296 | }, 297 | { 298 | "cell_type": "markdown", 299 | "id": "86617ca1", 300 | "metadata": {}, 301 | "source": [ 302 | "## Pricing Via a Dynamic Model\n", 303 | "\n", 304 | "In this exercise we investigate a more realistic model for the share price $S_n$.\n", 305 | "\n", 306 | "This comes from specifying the underlying dynamics of the share price.\n", 307 | "\n", 308 | "First we specify the dynamics.\n", 309 | "\n", 310 | "Then we'll compute the price of the option using Monte Carlo.\n", 311 | "\n", 312 | "### Simple Dynamics\n", 313 | "\n", 314 | "One simple model for $\\{S_t\\}$ is\n", 315 | "\n", 316 | "$$ \\ln \\frac{S_{t+1}}{S_t} = \\mu + \\sigma \\xi_{t+1} $$\n", 317 | "\n", 318 | "where \n", 319 | "\n", 320 | "* $S_0$ is normally distributed and\n", 321 | "* $\\{ \\xi_t \\}$ is IID and standard normal. \n", 322 | "\n", 323 | "\n", 324 | "Under the stated assumptions, $S_n$ is lognormally distributed.\n", 325 | "\n", 326 | "**Exercise** Explain why this is true.\n", 327 | "\n", 328 | "\n", 329 | "### Problems with Simple Dynamics\n", 330 | "\n", 331 | "The simple dynamic model we studied above is convenient, since we can work out\n", 332 | "the distribution of $S_n$.\n", 333 | "\n", 334 | "\n", 335 | "However, its predictions are counterfactual because, in the real world,\n", 336 | "volatility (measured by $\\sigma$) is not stationary.\n", 337 | "\n", 338 | "Instead it rather changes over time, sometimes high (like during the GFC) and sometimes low.\n", 339 | "\n", 340 | "\n", 341 | "### More Realistic Dynamics\n", 342 | "\n", 343 | "As stated above, one problem with our simple model is that $\\sigma$ is\n", 344 | "constant.\n", 345 | "\n", 346 | "This leads us to study the improved version:\n", 347 | "\n", 348 | "$$ \\ln \\frac{S_{t+1}}{S_t} = \\mu + \\sigma_t \\xi_{t+1} $$\n", 349 | "\n", 350 | "where \n", 351 | "\n", 352 | "$$ \n", 353 | " \\sigma_t = \\exp(h_t), \n", 354 | " \\quad\n", 355 | " h_{t+1} = \\rho h_t + \\nu \\eta_{t+1}\n", 356 | "$$\n", 357 | "\n", 358 | "Here $\\{\\eta_t\\}$ is also IID and standard normal." 359 | ] 360 | }, 361 | { 362 | "cell_type": "markdown", 363 | "id": "3800e7b9", 364 | "metadata": {}, 365 | "source": [ 366 | "### Default Parameters\n", 367 | "\n", 368 | "For the dynamic model, we adopt the following parameter values." 369 | ] 370 | }, 371 | { 372 | "cell_type": "code", 373 | "execution_count": null, 374 | "id": "f34b10d0", 375 | "metadata": {}, 376 | "outputs": [], 377 | "source": [ 378 | "μ = 0.0001\n", 379 | "ρ = 0.1\n", 380 | "ν = 0.001\n", 381 | "S0 = 10\n", 382 | "h0 = 0" 383 | ] 384 | }, 385 | { 386 | "cell_type": "markdown", 387 | "id": "0b6fe7b7", 388 | "metadata": {}, 389 | "source": [ 390 | "(Here `S0` is $S_0$ and `h0` is $h_0$.)\n", 391 | "\n", 392 | "For the option we use the following defaults." 393 | ] 394 | }, 395 | { 396 | "cell_type": "code", 397 | "execution_count": null, 398 | "id": "6b252a36", 399 | "metadata": {}, 400 | "outputs": [], 401 | "source": [ 402 | "K = 100\n", 403 | "n = 10\n", 404 | "β = 0.95" 405 | ] 406 | }, 407 | { 408 | "cell_type": "markdown", 409 | "id": "a340dedb", 410 | "metadata": {}, 411 | "source": [ 412 | "**Exercise**\n", 413 | "\n", 414 | "\n", 415 | "Write a function that simulates the sequence $S_0, \\ldots, S_n$, where the parameters are set to\n", 416 | "\n", 417 | "Plot 50 paths of the form $S_0, \\ldots, S_n$." 418 | ] 419 | }, 420 | { 421 | "cell_type": "code", 422 | "execution_count": null, 423 | "id": "880064d9", 424 | "metadata": {}, 425 | "outputs": [], 426 | "source": [ 427 | "# Put your code here" 428 | ] 429 | }, 430 | { 431 | "cell_type": "markdown", 432 | "id": "e67c7bdc", 433 | "metadata": {}, 434 | "source": [ 435 | "solution below\n", 436 | "\n", 437 | "solution below\n", 438 | "\n", 439 | "solution below\n", 440 | "\n", 441 | "solution below\n", 442 | "\n", 443 | "solution below\n", 444 | "\n", 445 | "solution below\n", 446 | "\n", 447 | "solution below\n", 448 | "\n", 449 | "solution below\n", 450 | "\n", 451 | "solution below\n", 452 | "\n", 453 | "solution below\n", 454 | "\n", 455 | "solution below\n", 456 | "\n", 457 | "solution below\n", 458 | "\n", 459 | "\n", 460 | "With $s_t := \\ln S_t$, the price dynamics become\n", 461 | "\n", 462 | "$$ s_{t+1} = s_t + \\mu + \\exp(h_t) \\xi_{t+1} $$\n", 463 | "\n", 464 | "Here is a function to simulate a path using this equation:" 465 | ] 466 | }, 467 | { 468 | "cell_type": "code", 469 | "execution_count": null, 470 | "id": "783171f4", 471 | "metadata": {}, 472 | "outputs": [], 473 | "source": [ 474 | "from numpy.random import randn\n", 475 | "\n", 476 | "def simulate_asset_price_path(μ=μ, S0=S0, h0=h0, n=n, ρ=ρ, ν=ν):\n", 477 | " s = np.empty(n+1)\n", 478 | " s[0] = np.log(S0)\n", 479 | "\n", 480 | " h = h0\n", 481 | " for t in range(n):\n", 482 | " s[t+1] = s[t] + μ + np.exp(h) * randn()\n", 483 | " h = ρ * h + ν * randn()\n", 484 | " \n", 485 | " return np.exp(s)" 486 | ] 487 | }, 488 | { 489 | "cell_type": "markdown", 490 | "id": "b468308e", 491 | "metadata": {}, 492 | "source": [ 493 | "Here we plot the paths and the log of the paths." 494 | ] 495 | }, 496 | { 497 | "cell_type": "code", 498 | "execution_count": null, 499 | "id": "598485a3", 500 | "metadata": {}, 501 | "outputs": [], 502 | "source": [ 503 | "fig, axes = plt.subplots(2, 1)\n", 504 | "\n", 505 | "titles = 'log paths', 'paths'\n", 506 | "transforms = np.log, lambda x: x\n", 507 | "for ax, transform, title in zip(axes, transforms, titles):\n", 508 | " for i in range(50):\n", 509 | " path = simulate_asset_price_path()\n", 510 | " ax.plot(transform(path))\n", 511 | " ax.set_title(title)\n", 512 | " \n", 513 | "fig.tight_layout()\n", 514 | "plt.show()" 515 | ] 516 | }, 517 | { 518 | "cell_type": "markdown", 519 | "id": "e0cf7497", 520 | "metadata": {}, 521 | "source": [ 522 | "**Exercise**\n", 523 | "\n", 524 | "Compute the price of the option $P_0$ by Monte Carlo, averaging over realizations $S_n^1, \\ldots, S_n^M$ of $S_n$ and appealing to the law of large numbers:\n", 525 | "\n", 526 | "$$ \\mathbb E \\max\\{ S_n - K, 0 \\} \n", 527 | " \\approx\n", 528 | " \\frac{1}{M} \\sum_{m=1}^M \\max \\{S_n^m - K, 0 \\}\n", 529 | " $$\n", 530 | " \n", 531 | "\n", 532 | "To the extend that you can, write fast, efficient code to compute the option price. \n", 533 | "\n", 534 | "In particular, try to speed up the code above using `jit` or `njit` from Numba." 535 | ] 536 | }, 537 | { 538 | "cell_type": "code", 539 | "execution_count": null, 540 | "id": "ed61138a", 541 | "metadata": {}, 542 | "outputs": [], 543 | "source": [ 544 | "# Put your code here" 545 | ] 546 | }, 547 | { 548 | "cell_type": "markdown", 549 | "id": "6329e275", 550 | "metadata": {}, 551 | "source": [ 552 | "solution below\n", 553 | "\n", 554 | "solution below\n", 555 | "\n", 556 | "solution below\n", 557 | "\n", 558 | "solution below\n", 559 | "\n", 560 | "solution below\n", 561 | "\n", 562 | "solution below\n", 563 | "\n", 564 | "solution below\n", 565 | "\n", 566 | "solution below\n", 567 | "\n", 568 | "solution below\n", 569 | "\n", 570 | "solution below\n", 571 | "\n", 572 | "solution below\n", 573 | "\n", 574 | "solution below" 575 | ] 576 | }, 577 | { 578 | "cell_type": "code", 579 | "execution_count": null, 580 | "id": "f156f6df", 581 | "metadata": {}, 582 | "outputs": [], 583 | "source": [ 584 | "from numba import njit, prange" 585 | ] 586 | }, 587 | { 588 | "cell_type": "code", 589 | "execution_count": null, 590 | "id": "45b2dedc", 591 | "metadata": {}, 592 | "outputs": [], 593 | "source": [ 594 | "@njit\n", 595 | "def compute_call_price(β=β,\n", 596 | " μ=μ,\n", 597 | " S0=S0,\n", 598 | " h0=h0,\n", 599 | " K=K,\n", 600 | " n=n,\n", 601 | " ρ=ρ,\n", 602 | " ν=ν,\n", 603 | " M=10_000_000):\n", 604 | " current_sum = 0.0\n", 605 | " # For each sample path\n", 606 | " for m in range(M):\n", 607 | " s = np.log(S0)\n", 608 | " h = h0\n", 609 | " # Simulate forward in time\n", 610 | " for t in range(n):\n", 611 | " s = s + μ + np.exp(h) * randn()\n", 612 | " h = ρ * h + ν * randn()\n", 613 | " # And add the value max{S_n - K, 0} to current_sum\n", 614 | " current_sum += np.maximum(np.exp(s) - K, 0)\n", 615 | " \n", 616 | " return β**n * current_sum / M" 617 | ] 618 | }, 619 | { 620 | "cell_type": "code", 621 | "execution_count": null, 622 | "id": "bc703070", 623 | "metadata": {}, 624 | "outputs": [], 625 | "source": [ 626 | "%%time \n", 627 | "compute_call_price()" 628 | ] 629 | }, 630 | { 631 | "cell_type": "markdown", 632 | "id": "035d81fd", 633 | "metadata": {}, 634 | "source": [ 635 | "**Exercise**\n", 636 | "\n", 637 | "If you can, use `prange` from Numba to parallelize this code and make it even faster." 638 | ] 639 | }, 640 | { 641 | "cell_type": "code", 642 | "execution_count": null, 643 | "id": "57d91651", 644 | "metadata": {}, 645 | "outputs": [], 646 | "source": [ 647 | "# Put your code here" 648 | ] 649 | }, 650 | { 651 | "cell_type": "markdown", 652 | "id": "f1fd2dec", 653 | "metadata": {}, 654 | "source": [ 655 | "solution below\n", 656 | "\n", 657 | "solution below\n", 658 | "\n", 659 | "solution below\n", 660 | "\n", 661 | "solution below\n", 662 | "\n", 663 | "solution below\n", 664 | "\n", 665 | "solution below\n", 666 | "\n", 667 | "solution below\n", 668 | "\n", 669 | "solution below\n", 670 | "\n", 671 | "solution below\n", 672 | "\n", 673 | "solution below\n", 674 | "\n", 675 | "solution below\n", 676 | "\n", 677 | "solution below" 678 | ] 679 | }, 680 | { 681 | "cell_type": "code", 682 | "execution_count": null, 683 | "id": "29a11ddf", 684 | "metadata": {}, 685 | "outputs": [], 686 | "source": [ 687 | "@njit(parallel=True)\n", 688 | "def compute_call_price_parallel(β=β,\n", 689 | " μ=μ,\n", 690 | " S0=S0,\n", 691 | " h0=h0,\n", 692 | " K=K,\n", 693 | " n=n,\n", 694 | " ρ=ρ,\n", 695 | " ν=ν,\n", 696 | " M=10_000_000):\n", 697 | " current_sum = 0.0\n", 698 | " # For each sample path\n", 699 | " for m in prange(M):\n", 700 | " s = np.log(S0)\n", 701 | " h = h0\n", 702 | " # Simulate forward in time\n", 703 | " for t in range(n):\n", 704 | " s = s + μ + np.exp(h) * randn()\n", 705 | " h = ρ * h + ν * randn()\n", 706 | " # And add the value max{S_n - K, 0} to current_sum\n", 707 | " current_sum += np.maximum(np.exp(s) - K, 0)\n", 708 | " \n", 709 | " return β**n * current_sum / M" 710 | ] 711 | }, 712 | { 713 | "cell_type": "code", 714 | "execution_count": null, 715 | "id": "15199300", 716 | "metadata": {}, 717 | "outputs": [], 718 | "source": [ 719 | "from numba import get_num_threads, set_num_threads\n", 720 | "get_num_threads()" 721 | ] 722 | }, 723 | { 724 | "cell_type": "code", 725 | "execution_count": null, 726 | "id": "3ac884c7", 727 | "metadata": {}, 728 | "outputs": [], 729 | "source": [ 730 | "%%time\n", 731 | "compute_call_price_parallel()" 732 | ] 733 | }, 734 | { 735 | "cell_type": "code", 736 | "execution_count": null, 737 | "id": "c8dbd9dd", 738 | "metadata": {}, 739 | "outputs": [], 740 | "source": [ 741 | "%%time\n", 742 | "compute_call_price_parallel()" 743 | ] 744 | }, 745 | { 746 | "cell_type": "markdown", 747 | "id": "7a7fe6d7", 748 | "metadata": {}, 749 | "source": [ 750 | "## Pricing a European Call Option Using JAX\n", 751 | "\n", 752 | "Previously we computed the value of a European call option via Monte Carlo using Numba-based routines.\n", 753 | "\n", 754 | "Let's compare how this looks, and how fast it runs, when we implement using [Google JAX](https://python-programming.quantecon.org/jax_intro.html)." 755 | ] 756 | }, 757 | { 758 | "cell_type": "markdown", 759 | "id": "15ccc6da", 760 | "metadata": {}, 761 | "source": [ 762 | "**Exercise**\n", 763 | "\n", 764 | "Try to shift the whole operation to the GPU using JAX and test your speed gain." 765 | ] 766 | }, 767 | { 768 | "cell_type": "code", 769 | "execution_count": null, 770 | "id": "1d380721", 771 | "metadata": {}, 772 | "outputs": [], 773 | "source": [ 774 | "# Put your code here" 775 | ] 776 | }, 777 | { 778 | "cell_type": "markdown", 779 | "id": "06495832", 780 | "metadata": {}, 781 | "source": [ 782 | "solution below\n", 783 | "\n", 784 | "solution below\n", 785 | "\n", 786 | "solution below\n", 787 | "\n", 788 | "solution below\n", 789 | "\n", 790 | "solution below\n", 791 | "\n", 792 | "solution below\n", 793 | "\n", 794 | "solution below\n", 795 | "\n", 796 | "solution below\n", 797 | "\n", 798 | "solution below\n", 799 | "\n", 800 | "solution below\n", 801 | "\n", 802 | "solution below\n", 803 | "\n", 804 | "solution below" 805 | ] 806 | }, 807 | { 808 | "cell_type": "code", 809 | "execution_count": null, 810 | "id": "122e0476", 811 | "metadata": {}, 812 | "outputs": [], 813 | "source": [ 814 | "!nvidia-smi" 815 | ] 816 | }, 817 | { 818 | "cell_type": "code", 819 | "execution_count": null, 820 | "id": "b29493ef", 821 | "metadata": {}, 822 | "outputs": [], 823 | "source": [ 824 | "import jax\n", 825 | "import jax.numpy as jnp" 826 | ] 827 | }, 828 | { 829 | "cell_type": "code", 830 | "execution_count": null, 831 | "id": "55ec35e9", 832 | "metadata": {}, 833 | "outputs": [], 834 | "source": [ 835 | "@jax.jit\n", 836 | "def compute_call_price_jax(β=β,\n", 837 | " μ=μ,\n", 838 | " S0=S0,\n", 839 | " h0=h0,\n", 840 | " K=K,\n", 841 | " n=n,\n", 842 | " ρ=ρ,\n", 843 | " ν=ν,\n", 844 | " M=10_000_000,\n", 845 | " key=jax.random.PRNGKey(1)):\n", 846 | "\n", 847 | " s = jnp.full(M, np.log(S0))\n", 848 | " h = jnp.full(M, h0)\n", 849 | " for t in range(n):\n", 850 | " key, subkey = jax.random.split(key)\n", 851 | " Z = jax.random.normal(subkey, (2, M))\n", 852 | " s = s + μ + jnp.exp(h) * Z[0, :]\n", 853 | " h = ρ * h + ν * Z[1, :]\n", 854 | " expectation = jnp.mean(jnp.maximum(jnp.exp(s) - K, 0))\n", 855 | " \n", 856 | " return β**n * expectation" 857 | ] 858 | }, 859 | { 860 | "cell_type": "code", 861 | "execution_count": null, 862 | "id": "099bb33d", 863 | "metadata": {}, 864 | "outputs": [], 865 | "source": [ 866 | "%%time \n", 867 | "compute_call_price_jax().block_until_ready()" 868 | ] 869 | }, 870 | { 871 | "cell_type": "code", 872 | "execution_count": null, 873 | "id": "156c4ded", 874 | "metadata": {}, 875 | "outputs": [], 876 | "source": [ 877 | "%%time \n", 878 | "compute_call_price_jax().block_until_ready()" 879 | ] 880 | } 881 | ], 882 | "metadata": { 883 | "kernelspec": { 884 | "display_name": "Python 3 (ipykernel)", 885 | "language": "python", 886 | "name": "python3" 887 | } 888 | }, 889 | "nbformat": 4, 890 | "nbformat_minor": 5 891 | } 892 | --------------------------------------------------------------------------------