Julia_Lec_5_Compecon_Xing_Xu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Firm dynamics with an example of exporter dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1 What is firm dynamics?\n",
    "\n",
    "The firm dynamics literature wants to study firms' behavior over time. Dynamic models were developed to study firms' entry and exit decisions, age, size distribution over time and even R\\& D, export and FDI decisions, etc. More recently, spatial components are incorporated to study, say, firms' heterogeneous location choices for their subsidiaries. \n",
    "\n",
    "To dig into the literature, a good starting point is [Hopenhayn (1992)](https://www.jstor.org/stable/2951541), which provides a general theoretical framework that features uncertainty towards future profits, costly market entry and endogenous exit. The literature then expands dramatically with various focuses. Just to list a few, [Klette and Kortum](https://web.stanford.edu/~klenow/Klette%20and%20Kortum.pdf) on innovations and firm distribution, [Cooley and Quadrini (2001)](https://www.aeaweb.org/articles?id=10.1257/aer.91.5.1286) on financial frictions, [Das et al. (2007)](https://www.jstor.org/stable/4502011) on exporter dynamics. \n",
    "\n",
    "### 2. What's specific about exporter dynamics?\n",
    "\n",
    "In a world where large multinationals take up a huge chunk of global economy via exports, the roles these giants play is a very interesting research topic. The OECD estimated in 2018 that multinational corporations account for half of global exports, nearly a third of world GDP (28%), and about a fourth of global employment. They hold great market power and sometimes even directly alter domestic and foreign policies. Understanding their growth and policy implications is a fascinating ongoing topic, which requires a lot of the tools discussed here. We can use the dynamic trade framework to study firms' individual and aggregate behavior, international policies and their implications on consumer welfare.\n",
    "\n",
    "For the remaining part of this section, I will discuss some empirical evidence about firm-level exporting behavior. I mainly follow from the literature review on [firm dynamics and trade](https://www.nber.org/system/files/working_papers/w27934/w27934.pdf) by George, Costas and Kim, ARE 2021.\n",
    "\n",
    "### 3. Firm-level facts (from firm-level data of US and Columbia):\n",
    "\n",
    "Static or time-invariant facts:\n",
    "* There are few exporters.\n",
    "* Exports are a small share of an exporter's total sales.\n",
    "* Exporters are relatively large firms. \n",
    "\n",
    "These facts can already trigger our thinking about the underlying firm distribution in size and productivity. These are consistent with the results of [Melitz (2003)](https://www.jstor.org/stable/1555536). Within a group of firms originally producing domestically, only the most productive firms enter the export market. They stay in the market and force the less productive firms to exit. \n",
    "\n",
    "If we further take a look at the dynamic aspect of the data, there are some more interesting observations:\n",
    "\n",
    "* Past export participation is the main predictor of current export participation.\n",
    "* Exporter are less likely to exit if they are older and larger exporters in the past.\n",
    "* Entry rate is increasing in size and past export activity.\n",
    "* The longer an exporter stays in the market, the more it tends to export.\n",
    "\n",
    "These can explain why the dynamic aspect is vital to studying the exporting firms' behavior. I will show some predictions of the results when we go to the theoretical framework. \n",
    "\n",
    "There are also evidences about heterogeneity between exporters on their export destinations, shipment volumes and frequencies and inventories. \n",
    "\n",
    "At last, The long-run response of aggregate trade volumes to changes in trade policy is larger than the short-run response. An important reason, as shown in [Kehoe and Ruhl (2013)](https://www.jstor.org/stable/10.1086/670272), is that the aggregate trade growth is stronger in the long term due to the extensive margin. That is, during a trade liberalization, previously non-traded products grow faster than products that have already been traded between the same country pairs. One take-away is that larger quantities of firms (for concreteness, plants) would become exporters when there are ease in bilateral trade policies. Combining with the fact that exporters take time to grow, this gives a possible explanation for the long-run growth of trade.\n",
    "\n",
    "### 4. Some questions to think about:\n",
    "\n",
    "On building a model, what structures do we need on individual firm's problem to retrieve the same results as we observe? For example, how to make sure that only a few firms enter into the export market and coincidentally they are relatively large? What about on the market structure and the demand side? \n",
    "\n",
    "Next, we can think about idiosyncratic and aggregate uncertainties. Idiosyncratic shock is convenient to generate endogenous entry and exit, as in the firm dynamics literature. Now we can look at if there are any heterogeneous short and long run responses to aggregate uncertainties, say real exchange rate fluctuations. If so, what are the policy implications?\n",
    "\n",
    "A very important topic in the trade literature is the gains from trade. That is the welfare gain of household in a country facing some sort of trade liberalization. Although the focus of our model is mainly on the firm side, we can still analyze the aggregate impact under a general equilibrium framework. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. A canonical model for exporter dynamics\n",
    "\n",
    "The model is a slightly tweaked version of [Das et al. (2007)](https://www.jstor.org/stable/4502011), based upon [Melitz (2003)](https://www.jstor.org/stable/1555536). For further reference, [Ruhl and Willis (2017)](https://onlinelibrary.wiley.com/doi/abs/10.1111/iere.12232) and [Alessandria et al. (2021)](http://kjrs3.com/research/acr_trade_adjustment.pdf). \n",
    "\n",
    "This section is arranged as follows. I will first lay out the static model and provide analytical solutions to it. Next, I will provide the dynamic setup of the model and compute the model. To look at the firms' short-term reaction to some sudden and unexpected (MIT) shocks, I will then illustrate the algorithms for obtaining the transition dynamics. In the end, I will provide some basic calibration and some preliminary results. I run some experiments and look at the exporters' full dynamics facing different shocks.\n",
    "\n",
    "#### 5.1 Static model\n",
    "\n",
    "##### 5.1.1 Consumer\n",
    "\n",
    "A representative consumer in domestic economy has preferences over an aggregate consumption good, $C$. Each index $j$ denotes a differentiated variety. $C$ takes the form:\n",
    "$$\n",
    "    C = \\left(\\sum_{j=1}^{J} c_j^\\frac{\\theta - 1}{\\theta} \\right)^\\frac{\\theta}{\\theta - 1},\n",
    "$$\n",
    "where $\\theta$ is the elasticity of substitution between varieties, and J is the number of available varieties. The consumer maximize his utility subject to the budget constraint:\n",
    "$$\n",
    "    \\sum_{j=1}^{J} c_j p_j = I.\n",
    "$$\n",
    "The consumer takes prices as given and chooses optimal consumption $c_j$ of each variety:\n",
    "$$\n",
    "    c_j = (\\frac{p_j}{P})^{-\\theta}C,\n",
    "$$\n",
    "where P is an aggregate price level given by:\n",
    "$$\n",
    "    P = \\left(\\sum_{j=1}^{J} p_j^{1-\\theta} \\right)^{\\frac{1}{1 - \\theta}}.\n",
    "$$\n",
    "The foreign market is identical to the domestic market with demand and consumption satisfying:\n",
    "$$\n",
    "    \\begin{aligned}\n",
    "    c_j^* &= (\\frac{p_j^*}{P^*})^{-\\theta}C^*, \\\\\n",
    "    P^* &= \\left(\\sum_{j=1}^{J} {p_j^*}^{1-\\theta} \\right)^{\\frac{1}{1 - \\theta}}. \\\\\n",
    "    \\end{aligned}\n",
    "$$\n",
    "Countries have same elasticity $\\theta$ for differentiated goods but differ in aggregate level of demands and prices.\n",
    "\n",
    "##### 5.1.2 Firm's static problem\n",
    "Assume each firm operates solely on a market $j$, so \"plants\" will be a better term than \"firms\". Plant operates monopolistically in each differentiated variety. Plant chooses the optimal amounts to produce in domestic and foreign market. A unit good produced with plant-level technology $\\tilde{\\epsilon}_j $, labor $n_j$ and capital $k_j$ subject to:\n",
    "$$\n",
    "    f(\\tilde{\\epsilon}_j, n_j, k_j) = \\tilde{\\epsilon}_j  n_j^{\\alpha_N}k_j^{\\alpha_K},\n",
    "$$\n",
    "where we assume nonincreasing returns to scale of the production technology, $\\alpha_N + \\alpha_K <= 1$.\n",
    "\n",
    "In each period, the plant chooses the prices, production, inputs, and export status next period $X'$ ($X'=1$ if exporting next period and $X'=0$ if not) to maximize its infinite horizon utility. The plant also faces an iceberg cost $\\xi_j$, where fraction $\\xi_j - 1$ of an export shipment is destroyed in transportation. In addition, foreign market might put an ad valorem tariff $\\tau_j$ on exports.\n",
    "\n",
    "We can thus first solve for the plant's single period or static problem given all the states and plant's export decision last period. A plant's profit is sum of revenue from domestic and foreign sales less the input costs:\n",
    "$$\n",
    "\\Pi_j =  p_j y_j + I(X_j = 1) (1 - \\tau_j) Q p_j^* y_j^* - wn_j - rk_j, \n",
    "$$\n",
    "where $Q$ is the real exchange rate; r is the rental rate of capital; and w is the wage. The plant is also subject to its feasibility constraint:\n",
    "$$\n",
    "y_j + \\xi_j y_j^* = \\tilde{\\epsilon}_j  n_j^{\\alpha_N}k_j^{\\alpha_K}.\n",
    "$$\n",
    "With market clearing condition in both markets $c_j = y_j$ and $c^*_j = y^*_j$ if $X_j = 1$, with demand functions from the household problem. The plant charges constant markup over marginal cost in both markets:\n",
    "$$\n",
    "    \\begin{aligned}\n",
    "    p_j &= \\frac{\\theta}{\\theta - 1} {MC}_j \\\\\n",
    "    Q p_j^* &= \\frac{1}{1- \\tau_j} \\frac{\\theta}{\\theta - 1} \\xi_j {MC}_j.\n",
    "    \\end{aligned}\n",
    "$$\n",
    "where $MC_j = \\frac{1}{\\tilde{\\epsilon_j}} \\left(\\frac{w}{\\alpha_n}\\right)^{\\alpha_n} \\left(\\frac{r}{1 - \\alpha_n}\\right)^{1 - \\alpha_n}$ with $\\alpha_n + \\alpha_k = 1$ (if not, the formula for marginal cost will be more complicated).\n",
    "\n",
    "The plant's static profit maximization problem is:\n",
    "$$\n",
    "\\Pi_j = \\max_{y_j, y_j^*}\\{ P C^{\\frac{1}{\\theta}} y_j^\\frac{\\theta - 1}{\\theta} +  I(X_j = 1) (1 - \\tau_j) Q P^* {C^*}^{\\frac{1}{\\theta}}  {y_j^*}^\\frac{\\theta - 1}{\\theta} - wn_j - rk_j\\},\n",
    "$$\n",
    "subject to equation the feasibility constraint. The optimal amounts for a plant to produce in domestic and foreign markets are:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "y_{j} &=\\frac{ \\xi_j^{\\theta} (1 - \\tau_j)^{-\\theta}   (\\frac{P}{Q P^*})^\\theta \\frac{C}{C^{*}}}{I \\left(X_{j}=1\\right) \\xi_j + \\xi_j^{\\theta} (1 - \\tau_j)^{-\\theta} (\\frac{P}{Q P^*})^\\theta \\frac{C}{C^{*}}} \\tilde{\\epsilon}_{j} n_{j}^{\\alpha_{N}} k_{j}^{\\alpha_{K}} \\\\\n",
    "y_{j}^{*} &= \\frac{I(X_j = 1)}{\\xi_j+ \\xi_j^{\\theta} (1 - \\tau_j)^{-\\theta} (\\frac{P}{Q P^*})^\\theta \\frac{C}{C^{*}}} \\tilde{\\epsilon}_{j} n_{j}^{\\alpha_{N}} k_{j}^{\\alpha_{K}} .\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Substituting in $y_j$ and $y_j^*$ to the profit maximization problem yields:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\Pi\\left(X_{j}, \\xi_j,  \\epsilon_{j}, Q\\right)=&\\max _{n_{j}, k_{j}} M n_{j}^{\\alpha_{N} \\frac{(\\theta-1)}{\\theta}} k_{j}^{\\alpha_{K} \\frac{(\\theta-1)}{\\theta}} - wn_{j}- r k_{j},\n",
    "\\end{aligned}\n",
    "$$\n",
    "where $M = \\left(1+I\\left(X_{j}=1\\right) \\xi_j^{1 - \\theta} (1 - \\tau_j)^\\theta \\left(\\frac{Q P^*}{P}\\right)^{\\theta} \\frac{C^{*}}{C}\\right)^{\\frac{1}{\\theta}} P C^{\\frac{1}{\\theta}} \\tilde{\\epsilon}_{j}^{\\frac{\\theta-1}{\\theta}}$\n",
    "\n",
    "From first order conditions of labor and capital, for simplicity write $a = \\alpha_{N} \\frac{(\\theta-1)}{\\theta}$ and $b = \\alpha_{K} \\frac{(\\theta-1)}{\\theta}$, the maximized profit can be solved as:\n",
    "$$\n",
    "\\Pi\\left(X_{j}, \\xi_j, \\epsilon_{j}, Q\\right) = M^{\\frac{1}{1-a-b}} (1-a-b) (\\frac{a}{w})^{\\frac{a}{1-a-b}} (\\frac{b}{r})^{\\frac{b}{1-a-b}}.\n",
    "$$\n",
    "Equate $\\xi_j = 1$ you get exactly the plant's profit without export iceberg cost. Then with only export fixed cost driving the export barrier, the model is equivalent to the baseline sunk cost model from \\citet{ruhl_new_2014}. Assuming no tariff, using the same notation, the profit a plant earns with an iceberg cost is exactly $\\left({\\frac{1+I\\left(X_{j}=1 \\right) \\xi_j^{1 - \\theta} \\left(\\frac{Q P^*}{P}\\right)^{\\theta} \\frac{C^{*}}{C}}{1+I\\left(X_{j}=1\\right) \\left(\\frac{Q P^*}{P}\\right)^{\\theta} \\frac{C^{*}}{C}}}\\right)^\\frac{1}{\\theta (1- a-b)}$ of the same plant without an iceberg cost.\n",
    "\n",
    "Just for clarification, we can also solve for domestic and foreign sales:\n",
    "$$\n",
    "    \\begin{aligned}\n",
    "       p_j y_{j} &= \\left(\\frac{ \\xi_j^{\\theta} (1 - \\tau_j)^{-\\theta} (\\frac{P}{Q P^*})^\\theta \\frac{C}{C^{*}}}{I \\left(X_{j}=1\\right) \\xi_j + \\xi_j^{\\theta} (1 - \\tau_j)^{-\\theta} (\\frac{P}{Q P^*})^\\theta \\frac{C}{C^{*}}}\\right)^{\\frac{\\theta - 1}{\\theta}}   \\\\\n",
    "       & \\ \\ \\ \\ \\ \\ \\ \\ \\ \\ P C^{\\frac{1}{\\theta}} \\epsilon_j M^{\\frac{a + b}{1-a-b}} (\\frac{a}{w})^{\\frac{a}{1-a-b}} (\\frac{b}{r})^{\\frac{b}{1-a-b}} \\\\\n",
    "       Q p_j^* y_j^* &= I(X_j = 1) \\left(\\frac{1}{\\xi_j+ \\xi_j^{\\theta} (1 - \\tau_j)^{-\\theta} (\\frac{P}{Q P^*})^\\theta \\frac{C}{C^{*}}}\\right)^{\\frac{\\theta - 1}{\\theta}} \\\\\n",
    "       & \\ \\ \\ \\ \\ \\ \\ \\ \\ \\  Q P^* {C^*}^{\\frac{1}{\\theta}} \\epsilon_j M^{\\frac{a + b}{1-a-b}} (\\frac{a}{w})^{\\frac{a}{1-a-b}} (\\frac{b}{r})^{\\frac{b}{1-a-b}} .\n",
    "    \\end{aligned}\n",
    "$$\n",
    "\n",
    "We normalize the size of domestic aggregate demand $C$ to 1, so $C^*$ denotes the size of foreign demand relative the domestic demand. Define $\\epsilon_j = {\\tilde{\\epsilon}_j}^\\frac{\\theta - 1}{\\theta}$, which follows an idiosyncratic AR1 shock for each plant. We normalize the mean of the $\\epsilon$ process to one.\n",
    "\n",
    "Firms face idiosyncratic productivity shocks, a fixed cost for becoming an exporter abd and an iceberg cost for exports."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.2 Firm's Dynamic problem\n",
    "\n",
    "We first introduce an export fixed cost structure as in [Ruhl and Willis (2017)](https://onlinelibrary.wiley.com/doi/abs/10.1111/iere.12232). The plant faces a sunk entry cost $f_E$ and a continuation cost $f_C$, $f_C \\leq f_E$. All export decisions are made at the end of previous period and fixed costs are paid upon decision. The exporting cost of each period can be expressed as:\n",
    "$$\n",
    "    f(X, X') = (1-X) * f_E * X' + X * f_C * X' .\n",
    "$$\n",
    "\n",
    "We then link the export decision of a plant with an endogenous and stochastic iceberg transportation cost $\\xi_j$ as in [Alessandria et al. (2021)](http://kjrs3.com/research/acr_trade_adjustment.pdf). A non-exporting plant doesn't ship any goods abroad and we set its iceberg cost this period to infinity. The nonexporter can deterministically lower its fixed cost next period to $\\xi_E$ by paying $f_E$. If a current exporter pays $f_C$ to continue exporting next period, it draws iceberg cost $\\xi_j \\in \\{\\xi_E, \\xi_C\\}$ next period from a Markov chain as in table \\ref{table1}. Assume $\\xi_C \\leq \\xi_E < \\infty$. If an exporter stops paying $f_C$, it exits the export market and becomes a nonexporter with infinite iceberg cost next period. \n",
    "\n",
    "\n",
    "The structure of the iceberg cost conditional on continuing to exporting can be summarized by a markov chain with states $\\xi_E$ and $\\xi_C$ with transition probability $[ \\rho_{EE}, 1- \\rho_{EE} ; 1 - \\rho_{CC}, \\rho_{CC} ]$.\n",
    "\n",
    "\n",
    "\n",
    " We also have a persistent condition as $\\rho_\\xi(\\xi_E|\\xi_C) \\leq \\rho_\\xi(\\xi_C|\\xi_E)$. With a lot of plants, this particular structure of iceberg cost creates incentives for some plants to invest in export technology and make the \"better\" exporters to become less susceptible to unfavorable shocks. \n",
    "\n",
    "One last component is the technological shock $\\epsilon_j$ that governs plant productivity, which generates heterogeneous plants and varying entry and exit export decisions of plants. It follows an idiosyncratic and time-invariant AR(1) processes for each plant:\n",
    "\\begin{equation}\n",
    "\\ln{\\epsilon_t} = \\rho_\\epsilon\\ln{\\epsilon_{t-1}} + w_{\\epsilon,t}, \\hspace{1cm} w_\\epsilon \\sim N(0, \\sigma^2_\\epsilon).\n",
    "\\end{equation}\n",
    "The shocks are discretized using Tauchen's method. There are three key parameters to determine here: \\{{$N_\\epsilon$, $\\rho_\\epsilon$, $\\sigma_\\epsilon$}\\}. $N_\\epsilon$ is the states of plant's productivity, currently set to 1000 for simplicity.\n",
    "\n",
    "We put aside our discussion of exchange rate shocks for the moment as an aggregate shock will make solving for the equilibrium prices and transition dynamics much harder. Exchange rates directly shift the export prices and create more uncertainties which make the foreign market more volatile. In particularly \"bad\" periods, there can be no firms deciding to export, leading to potentially poor simulation results.\n",
    "\n",
    "In each period, every plant makes an discrete choice based on their current status and discounted expectation of future profits.  The discount rate $R$ is set to be $\\frac{1}{1+r}$. We characterize the dynamic problem of each plant by the following Bellman equation:\n",
    "$$\n",
    "V\\left(X_{j}, \\xi_j, \\epsilon_{j}, Q\\right)=\\max _{X_{j}^{\\prime}}\\left\\{\\Pi\\left(X_{j}, \\xi_j, \\epsilon_{j}, Q\\right)-f\\left(X_{j}, X'_{j} \\right)+R \\underset{\\xi_j^{\\prime}, \\epsilon_{j}^{\\prime}, Q^{\\prime}}{\\mathbb{E}} V\\left(X_{j}^{\\prime}, \\xi_j^{\\prime}, \\epsilon_{j}^{\\prime}, Q^{\\prime}\\right)\\right\\}.\n",
    "$$\n",
    "\n",
    "The policy function for export entry is thus:\n",
    "$$\n",
    "    X_{j}^{\\prime}\\left(0, \\infty, \\epsilon_{j}, Q\\right)= \\begin{cases}1 & \\text{ if } \\hspace{0.5cm}  - f_E + \\\\\n",
    "{} & R \\underset{\\xi_j^{\\prime}, \\epsilon_{j}^{\\prime}, Q^{\\prime}}{\\mathbb{E}} \\left[V\\left(1, \\xi_E, \\epsilon_{j}^{\\prime}, Q^{\\prime}\\right)-V\\left(0, \\infty, \\epsilon_{j}^{\\prime}, Q^{\\prime}\\right)\\right] \\geq 0 \\\\ 0 & \\text { otherwise. }\\end{cases}\n",
    "$$\n",
    "\n",
    "The inequality on the right is the condition of a non-exporting plant choosing to be an exporter next period. The first two terms denote the current period profit minus the entry fixed cost. The discounted expectation is the difference of expected future values between starting to become an exporter with a low export technology $\\xi_E$ and staying as an nonexporter. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Solving the model\n",
    "\n",
    "As always, let's input some parameter value. I made the simplifying assumption that the real exchange rate $Q$ is fixed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra, Statistics, Distributions\n",
    "using Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "Base.@kwdef struct Exporter{T1, T2, R1, S}\n",
    "    r::R1 = 0.109\n",
    "    αn::R1 = 0.45\n",
    "\n",
    "    αk::R1 = 1.0 - αn\n",
    "    θ::R1 = 5.0\n",
    "    w::R1 = 0.02\n",
    "\n",
    "    ξE::R1 = 1.6\n",
    "    ξC::R1 = 1.2\n",
    "    ξ::T1 = [ξE; ξC]\n",
    "    \n",
    "    ρEE::R1 = 0.92\n",
    "    ρCC::R1 = 0.92\n",
    "    Ξ::T2 = [ρEE (1.0-ρEE); (1.0-ρCC) ρCC]\n",
    "\n",
    "    Q::R1 = 1.0 # fix exchange rate for the moment\n",
    "\n",
    "    fE::R1 = 0.6\n",
    "    fC::R1 = 0.2\n",
    "    \n",
    "    ρϵ::R1 = 0.872524\n",
    "    σϵ::R1 = 0.115886\n",
    "\n",
    "    Cstar::R1 = 0.7 #C normalized to be 1.0\n",
    "\n",
    "    grids::S = 100\n",
    "\n",
    "    R = (1/(1+r))^0.25 # quarterly discount factor\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Exporter{Vector{Float64}, Matrix{Float64}, Float64, Int64}(0.109, 0.45, 0.55, 5.0, 0.02, 1.6, 1.2, [1.6, 1.2], 0.92, 0.92, [0.92 0.07999999999999996; 0.07999999999999996 0.92], 1.0, 0.6, 0.2, 0.872524, 0.115886, 0.7, 100, 0.9744669483879411)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "E = Exporter() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.1 Discretize an AR(1) process\n",
    "\n",
    "Recall that the reason we use discrete approximations of value functions is that the computer can't process continuos variables. Same logic applies with an AR(1) process. We want to approximate the AR(1) process with a discrete markov chain, which is characterized by discrete states and a transition probability matrix. Check [Quantecon's lecture notes](https://julia.quantecon.org/introduction_dynamics/finite_markov.html) for a detailed explanation.\n",
    "\n",
    "Again the AR(1) process we want to discretize is:\n",
    "\\begin{equation}\n",
    "\\ln{\\epsilon_t} = \\rho_\\epsilon\\ln{\\epsilon_{t-1}} + w_{\\epsilon,t}, \\hspace{1cm} w_\\epsilon \\sim N(0, \\sigma^2_\\epsilon).\n",
    "\\end{equation}\n",
    "\n",
    "The code (adopted from Quantecon) is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tauchen (generic function with 3 methods)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function tauchen(N::Integer, ρ::T1, σ::T2, μ=zero(promote_type(T1, T2)), n_std::T3=3) where {T1 <: Real, T2 <: Real, T3 <: Real}\n",
    "  # Get discretized space\n",
    "  a_bar = n_std * sqrt(σ^2 / (1 - ρ^2))\n",
    "  y = range(-a_bar, stop=a_bar, length=N)\n",
    "  d = y[2] - y[1]\n",
    "\n",
    "  # Get transition probabilities\n",
    "  Π = zeros(promote_type(T1, T2), N, N)\n",
    "  for row = 1:N\n",
    "      # Do end points first\n",
    "      Π[row, 1] = cdf(Normal(),(y[1] - ρ*y[row] + d/2) / σ)\n",
    "      Π[row, N] = 1.0 - cdf(Normal(),(y[N] - ρ*y[row] - d/2) / σ)\n",
    "\n",
    "      # fill in the middle columns\n",
    "      for col = 2:N-1\n",
    "          Π[row, col] = (cdf(Normal(),(y[col] - ρ*y[row] + d/2) / σ) -\n",
    "                        cdf(Normal(),(y[col] - ρ*y[row] - d/2) / σ))\n",
    "      end\n",
    "  end\n",
    "\n",
    "  yy = exp.(y .+ μ / (1 - ρ)) # center process around its mean (wbar / (1 - rho)) in new variable\n",
    "  # take exponential to get to a ln process\n",
    "\n",
    "  # renormalize. In some test cases the rows sum to something that is 2e-15\n",
    "  # away from 1.0, which caused problems in the MarkovChain constructor\n",
    "  Π = Π./sum(Π, dims = 2)\n",
    "\n",
    "  return Π, yy\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([0.2355397082043846 0.0398220365205688 … 0.0 0.0; 0.2035639923092491 0.03686888293414207 … 0.0 0.0; … ; 4.784202347164862e-30 1.4778853656337744e-29 … 0.03686888293414202 0.20356399230924901; 1.3827115757372961e-30 4.347177100735715e-30 … 0.039822036520568704 0.23553970820438463], [0.4908675431085489, 0.4979749102029403, 0.5051851862545925, 0.5124998613024105, 0.5199204469598645, 0.5274484767273707, 0.5350855063091983, 0.5428331139349621, 0.5506929006857724, 0.5586664908251056  …  1.7899766970506505, 1.8158941194896647, 1.84218680535361, 1.8688601881549594, 1.8959197800790755, 1.9233711731233287, 1.9512200402527142, 1.9794721365721941, 2.0081333005160213, 2.0372094550542794])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Π, ϵ_vec = tauchen(100, E.ρϵ, E.σϵ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.2 Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exporter fixed cost (just for illustration, we don't use this function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "expfixcost (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function expfixcost(X1::Int, X2::Int; E = E)\n",
    "    #X1 and X2 are 0, 1 variables\n",
    "    f = (1-X1) * E.fE * X2 + X1 * E.fC * X2;\n",
    "    return f\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Firm's profit function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "profit (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function profit(X, Evalue, ξ_idx, P, Pstar; E= E, Qvalue = E.Q, τ = 0.0)\n",
    "    (;θ, Cstar, αn, αk, w, r, ξ ) = E\n",
    "    M = (1.0 + X * ξ[ξ_idx]^(1.0 - θ)* (1.0 - τ)^θ * ((Qvalue * Pstar/ P)^θ) * Cstar)^(1.0/θ) * P * Evalue\n",
    "    a = αn * (θ - 1.0)/θ\n",
    "    b = αk * (θ - 1.0)/θ\n",
    "    Profit = (1.0-a-b) * M^(1.0/(1.0-a-b)) * (a/w)^(a/(1.0-a-b)) * (b/r)^(b/(1.0-a-b))\n",
    "    \n",
    "    return Profit\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sales at domestic and foreign markets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "sales (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function sales(X, Evalue, ξ_idx, P, Pstar; E= E, Qvalue = E.Q, τ = 0.0)\n",
    "    (;θ, Cstar, αn, αk, w, r, ξ) = E\n",
    "    M = (1.0 + X * ξ[ξ_idx]^(1.0 - θ)* (1.0 - τ)^θ * ((Qvalue * Pstar/ P)^θ) * Cstar)^(1.0/θ) * P * Evalue\n",
    "    a = αn * (θ - 1.0)/θ\n",
    "    b = αk * (θ - 1.0)/θ\n",
    "    con = ξ[ξ_idx]^θ * (1.0 - τ)^(-θ) * (P/ (Qvalue * Pstar))^θ/Cstar\n",
    "    sales_do = (con/(X * ξ[ξ_idx] + con))^((θ-1.0)/θ) * P * Evalue * M^((a + b)/(1.0-a-b)) * (a/w)^(a/(1.0-a-b)) * (b/r)^(b/(1.0-a-b))\n",
    "    if X == 0\n",
    "        sales_fo = 0.0\n",
    "    else\n",
    "        sales_fo = (1.0-τ) * (1.0/(ξ[ξ_idx] + con))^((θ-1.0)/θ) * Qvalue * Pstar * Cstar^(1.0/θ) * Evalue * M^((a + b)/(1.0-a-b)) * (a/w)^(a/(1.0-a-b)) * (b/r)^(b/(1.0-a-b)) # foreign sales net of tariffs\n",
    "    end\n",
    "    return [sales_do, sales_fo]\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a try on the functions. Again `.` is used for broadcasting and `Ref()` is used to fix inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element Vector{Float64}:\n",
       " 0.08608093339518139\n",
       " 0.002565830909582643"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "profit.([1, 0], [1.2, 0.8], [1, 1], Ref(0.1), Ref(0.2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 adjoint(::Matrix{Float64}) with eltype Float64:\n",
       " 9742.14  1040.57\n",
       " 1282.92     0.0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "reduce(hcat, sales.([1, 0], [1.2, 0.8], [1, 1], Ref(1.0), Ref(1)))'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Individual prices ($p_j$ and $p_j^*$)\n",
    "\n",
    "The function defined here needs to be modified to incorporate cases where $\\alpha_n + \\alpha_k \\neq 1$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ind_price (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function ind_price(X, Evalue, ξ_idx; E = E, Qvalue = E.Q, τ = 0.0)\n",
    "    (;θ, αn, αk, w, r, ξ) = E\n",
    "    MC = Evalue^(θ/(1.0 - θ)) * (w/αn)^αn * (r/αk)^αk\n",
    "    do_p = (θ/(θ - 1.0)) * MC\n",
    "    if X == 0\n",
    "        fo_p = -1.0 # make foreign price negative for non-exporters\n",
    "    else\n",
    "        fo_p = (θ/(θ - 1.0)) / (1.0 - τ) * ξ[ξ_idx] * MC/Qvalue\n",
    "    end\n",
    "    return [do_p, fo_p]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element Vector{Float64}:\n",
       " -1.0\n",
       "  0.16104839447304678"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ind_p_mat = reduce(hcat, ind_price.([0 1], [1.0 1.2], [1 1]))'[:,1]\n",
    "ind_pstar_mat = reduce(hcat, ind_price.([0 1], [1.0 1.2], [1 1]))'[:,2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Aggregate price $P$ and $P^*$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "agg_price (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function agg_price(ind_p_mat, ind_pstar_mat; θ = E.θ)\n",
    "    P = (sum(ind_p_mat.^(1.0 - θ)))^(1.0/(1.0 - θ))\n",
    "    Pstar = ((sum(ind_pstar_mat[ind_pstar_mat .>= 0]).^(1.0 - θ)))^(1.0/(1.0 - θ))\n",
    "    return P, Pstar\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.09250365544979017, 0.16104839447304678)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "agg_price(ind_p_mat, ind_pstar_mat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.3 Solve the firm's problem\n",
    "\n",
    "Solving the model is rather standard except that I chose to separate out the problem of a non-exporter and an exporter with either exporting technology $\\xi$.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "solve_exporter_fixQ (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function solve_exporter_fixQ(E, P, Pstar, τ; vNX0 = zeros(E.grids, 1), vEX0 = zeros(E.grids, length(E.ξ)), tol = 1e-6, Q = E.Q)\n",
    "    # unpacking\n",
    "    (; R, fE, fC, ξ, Ξ, fE, fC, ρϵ, σϵ, grids) = E\n",
    "\n",
    "    vNX1 = similar(vNX0) # to store value of non-exporters\n",
    "    vEX1 = similar(vEX0) # to store value of exporters\n",
    "    polNX = similar(vNX0, Int) # to store policy of non-exporters\n",
    "    polEX = similar(vEX0, Int) # to store policy of exporters\n",
    "\n",
    "    REvNX = similar(vNX0) # expected value if choosing not to export\n",
    "    REvEX_NX = similar(vNX0) # expected value if choosing to export as a non-exporter\n",
    "    REvEX_EX = similar(vEX0) # expected value if choosing to export as an exporter with different ξ's\n",
    "    \n",
    "    domesticm, ~ = sales(0, 1.0, 1, P, Pstar; E = E, Qvalue = Q, τ = 0.0) # sales of a non-exporting medium firm\n",
    "\n",
    "    # productivity process\n",
    "    Π, ϵ_vec = tauchen(grids, ρϵ, σϵ)\n",
    "\n",
    "    # Precompute the profits for firms at all states \n",
    "    proNX = profit.(zeros(Int, grids, 1), ϵ_vec, ones(Int, grids, 1), Ref(P), Ref(Pstar); E= E, Qvalue = Q, τ = τ)\n",
    "    # input for ξ doesn't matter\n",
    "\n",
    "    proEX = similar(vEX0)\n",
    "    for i in 1:length(ξ)\n",
    "        proEX[:,i] = profit.(ones(Int, grids, 1), ϵ_vec, i.*ones(Int, grids, 1), Ref(P), Ref(Pstar); E= E, Qvalue = Q, τ = τ)\n",
    "    end\n",
    "\n",
    "    # VFI\n",
    "    iter = 0\n",
    "    while true\n",
    "        distance = zero(eltype(vNX0))\n",
    "        iter += 1 \n",
    "\n",
    "        # Precompute RE[V(.)|X, ϵ, ξ] for better performance\n",
    "        REvNX = R * Π * vNX0\n",
    "        REvEX_NX = R * Π * vEX0[:,1] # enters with entry export technology\n",
    "        REvEX_EX = R * Π * vEX0 * Ξ'\n",
    "\n",
    "        vNX1 = max.((proNX + REvNX), (proNX .- fE * domesticm + REvEX_NX))\n",
    "        vEX1 = max.((proEX .+ REvNX), (proEX .- fC * domesticm + REvEX_EX))\n",
    "\n",
    "        distance = max(maximum(abs.(vNX1 - vNX0)), maximum(abs.(vEX1 - vEX0)))\n",
    "\n",
    "        (distance < tol || iter == 2000) && break\n",
    "        vNX0 .= vNX1\n",
    "        vEX0 .= vEX1\n",
    "    end\n",
    "\n",
    "    # get policy\n",
    "    polNX .= 0\n",
    "    polEX .= 0\n",
    "    for i in 1:grids\n",
    "        if (REvNX[i]) < (- fE * domesticm + REvEX_NX[i])\n",
    "            polNX[i] = 1\n",
    "        end\n",
    "        for j in 1:length(ξ)\n",
    "            if (REvNX[i]) < (- fC * domesticm + REvEX_EX[i,j])\n",
    "                polEX[i,j] = 1\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "\n",
    "    return (vNX = vNX1, vEX = vEX1, polNX = polNX, polEX = polEX, Π = Π, ϵ_vec =ϵ_vec, iter = iter)\n",
    "\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.021312 seconds (14.35 k allocations: 7.747 MiB, 48.11% gc time)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(vNX = [0.4360154494851622; 0.437186123813252; … ; 1.4838358575240558; 1.5257584977020988;;], vEX = [0.43603928462977926 0.4360907803125939; 0.43721173522491963 0.43726706852173247; … ; 1.5350851689447875 1.680995332627831; 1.5793313776889149 1.7326647962605208], polNX = [0; 0; … ; 1; 1;;], polEX = [0 0; 0 0; … ; 1 1; 1 1], Π = [0.2355397082043846 0.0398220365205688 … 0.0 0.0; 0.2035639923092491 0.03686888293414207 … 0.0 0.0; … ; 4.784202347164862e-30 1.4778853656337744e-29 … 0.03686888293414202 0.20356399230924901; 1.3827115757372961e-30 4.347177100735715e-30 … 0.039822036520568704 0.23553970820438463], ϵ_vec = [0.4908675431085489, 0.4979749102029403, 0.5051851862545925, 0.5124998613024105, 0.5199204469598645, 0.5274484767273707, 0.5350855063091983, 0.5428331139349621, 0.5506929006857724, 0.5586664908251056  …  1.7899766970506505, 1.8158941194896647, 1.84218680535361, 1.8688601881549594, 1.8959197800790755, 1.9233711731233287, 1.9512200402527142, 1.9794721365721941, 2.0081333005160213, 2.0372094550542794], iter = 375)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@time begin\n",
    "    sol_try = solve_exporter_fixQ(E, 0.1, 0.1, 0.0);\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find the cutoff productivity levels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "find_cutoff_prod (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function find_cutoff_prod(sol)\n",
    "    (; polNX, polEX, ϵ_vec) = sol\n",
    "    cutoff_ϵ = zeros(eltype(ϵ_vec), size(polEX)[2] + 1)\n",
    "    for i in 1:size(polNX)[1]\n",
    "        if polNX[i] == 1\n",
    "            cutoff_ϵ[1] = ϵ_vec[i]\n",
    "            break\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    for j in 1:size(polEX)[2]\n",
    "        for i in 1:size(polEX)[1]\n",
    "            if polEX[i, j] == 1\n",
    "                cutoff_ϵ[j + 1] = ϵ_vec[i]\n",
    "                break # only break out of the inner loop\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "\n",
    "    return cutoff_ϵ\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can find out the cutoff levels for the productivity $\\epsilon$ for different export status $X$ and export technology $\\xi$ ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3-element Vector{Float64}:\n",
       " 1.7644291820741218\n",
       " 1.3818813801525107\n",
       " 1.1299695318784664"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "find_cutoff_prod(sol_try)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This says that among different states, non-exporter needs the highest draw of productivity to become an exporter and the exporter with highest $\\xi$ will not go back to a non-exporter only if the productivity falls very low."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.4 Stationary distribution for firms given prices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "compute_stationary_distribution (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function compute_stationary_distribution(sol;E = E, tol = 1e-8, pdf_0 = fill(1.0 / (size(sol_try.vEX)[1] * (size(sol_try.vEX)[2] + 1))  , (size(sol_try.vEX)[1] , (size(sol_try.vEX)[2] + 1))))\n",
    "\n",
    "    (; polEX, polNX, vEX, vNX) = sol\n",
    "\n",
    "    (; Ξ, ρϵ, σϵ, grids) = E\n",
    "\n",
    "    # productivity process\n",
    "    Π, ϵ_vec = tauchen(grids, ρϵ, σϵ)\n",
    "    \n",
    "    pdf_1 = similar(pdf_0)\n",
    "\n",
    "    iter = 0\n",
    "    while true \n",
    "        fill!(pdf_1, zero(eltype(pdf_0)))\n",
    "        iter +=1\n",
    "\n",
    "        for i in 1:grids\n",
    "            for iprime in 1:grids\n",
    "                pdf_1[iprime,1] += (1 - polNX[i]) * pdf_0[i, 1] * Π[i, iprime]\n",
    "                pdf_1[iprime,2] += polNX[i] * pdf_0[i, 1] * Π[i, iprime]\n",
    "                for j in 1:size(Ξ)[1]\n",
    "                    pdf_1[iprime,1] += (1 - polEX[i,j]) * pdf_0[i, j + 1] * Π[i, iprime]\n",
    "                    for jprime in 1:size(Ξ)[1]\n",
    "                        pdf_1[iprime,jprime + 1] += polEX[i, j] * Ξ[j, jprime] * Π[i, iprime] * pdf_0[i ,j + 1] \n",
    "                    end\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "        \n",
    "        distance = zero(eltype(pdf_0))\n",
    "        for (a, b) in zip(pdf_0, pdf_1)\n",
    "            distance = max(abs(a - b), distance)\n",
    "        end\n",
    "        \n",
    "        (distance < tol || iter == 2000) && break \n",
    "        pdf_0 .= pdf_1\n",
    "    end\n",
    "    pdf_1 = pdf_1\n",
    "    return pdf_1, iter\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.072110 seconds (224.32 k allocations: 12.099 MiB, 92.28% compilation time)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "([0.0011570437242512165 2.1145547603126288e-17 2.4311023610085744e-16; 0.0002998008153782488 3.0370638439024366e-17 3.491402314311281e-16; … ; 8.081413651159851e-5 0.00018671373465131073 3.2272944215721654e-5; 0.00020250322875128642 0.0008151092583492217 0.00013943123715096507], 55)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@time begin\n",
    "    pdf_stationary, ~ = compute_stationary_distribution(sol_try)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check if it is still a distribution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sum(pdf_stationary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the shares of exporters at each states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The share of nonexporters is 0.9768202808538854, the share of exporters with low technology is 0.016599699825925476, the share of exporters with high technology is 0.006580019320189167, given prices in the stationary distribution, \n"
     ]
    }
   ],
   "source": [
    "NX_share, EX_low_share, EX_high_share = sum(pdf_stationary, dims = 1)\n",
    "println(\"The share of nonexporters is $NX_share, the share of exporters with low technology is $EX_low_share, the share of exporters with high technology is $EX_high_share, given prices in the stationary distribution, \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.5 Stationary prices\n",
    "\n",
    "Now we have all the pieces ready for the stationary prices.\n",
    "\n",
    "The algorithm goes as follows:\n",
    "\n",
    "* Fix the number of firms N\n",
    "* Start from guesses of $P_0$ and $P^*_0$\n",
    "* For each $P_t$ and $P^*_t$ solve the firm's problem given prices.\n",
    "* Compute invariant distribution $\\lambda(X, \\epsilon, \\xi)$\n",
    "* Get new individual firm prices according to first order conditions of firms.\n",
    "* Get the updated aggregate prices $P'_t$ and $P'^{*}_t$.\n",
    "* Guess $P_{t + 1} = ϵ P'_t + (1 - ϵ) P_t$ and $P^*_{t + 1} = ϵ P'^*_t + (1 - ϵ) P^*_t$. Go back to the second step and iterate until convergence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "compute_stationary_equilibrium (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function compute_stationary_equilibrium(E; P0=0.1, Pstar0=0.1, tol=1e-5, ϵ=0.2, τ=0.0, N=2000)\n",
    "    (; Ξ, ρϵ, σϵ, grids, θ) = E\n",
    "\n",
    "    Π, ϵ_vec = tauchen(grids, ρϵ, σϵ)\n",
    "\n",
    "    ind_p_mat = zeros(grids, size(Ξ)[1] + 1)\n",
    "    ind_pstar_mat = zeros(grids, size(Ξ)[1] + 1)\n",
    "\n",
    "    sol = solve_exporter_fixQ(E, P0, Pstar0, τ)\n",
    "    pdf, ~ = compute_stationary_distribution(sol)\n",
    "\n",
    "    for i in 1:grids\n",
    "        ind_p_mat[i, 1], ind_pstar_mat[i, 1] = ind_price(0, ϵ_vec[i], 1; τ=τ) .* pdf[i, 1]\n",
    "        for j in 1:size(Ξ)[1]\n",
    "            ind_p_mat[i, j+1], ind_pstar_mat[i, j+1] = ind_price(1, ϵ_vec[i], j; τ=τ) .* pdf[i, j+1]\n",
    "        end\n",
    "    end\n",
    "\n",
    "    Pprime1 = (sum(ind_p_mat.^(1.0 - θ) .*pdf .* N ))^(1.0/(1.0 - θ))\n",
    "    temp = ind_pstar_mat.^(1.0 - θ) .*pdf .* N \n",
    "    Pstarprime1 = (sum(temp[temp .> 0.0]))^(1.0/(1.0 - θ))\n",
    "\n",
    "    P1 = ϵ * Pprime1 + (1.0 - ϵ) * P0\n",
    "    Pstar1 = ϵ * Pstarprime1 + (1.0 - ϵ) * Pstar0\n",
    "\n",
    "    P0 = copy(P1)\n",
    "    Pstar0 = copy(Pstar1)\n",
    "\n",
    "    iter = 0\n",
    "    while true\n",
    "        distance = zero(eltype(P0))\n",
    "        iter += 1\n",
    "\n",
    "        sol = solve_exporter_fixQ(E, P0, Pstar0, τ)\n",
    "        pdf, ~ = compute_stationary_distribution(sol)\n",
    "\n",
    "        for i in 1:grids\n",
    "            ind_p_mat[i, 1], ind_pstar_mat[i, 1] = ind_price(0, ϵ_vec[i], 1; τ=τ)\n",
    "            for j in 1:size(Ξ)[1]\n",
    "                ind_p_mat[i, j+1], ind_pstar_mat[i, j+1] = ind_price(1, ϵ_vec[i], j; τ=τ)\n",
    "            end\n",
    "        end\n",
    "\n",
    "        Pprime1 = (sum(ind_p_mat.^(1.0 - θ) .*pdf .* N ))^(1.0/(1.0 - θ))\n",
    "        temp = ind_pstar_mat.^(1.0 - θ) .*pdf .* N \n",
    "        Pstarprime1 = (sum(temp[temp .> 0.0]))^(1.0/(1.0 - θ))\n",
    "\n",
    "        P1 = ϵ * Pprime1 + (1.0 - ϵ) * P0\n",
    "        Pstar1 = ϵ * Pstarprime1 + (1.0 - ϵ) * Pstar0\n",
    "\n",
    "        distance = max(abs(P1 - P0), abs(Pstar1 - Pstar0))\n",
    "        (distance < tol || iter == 2000) && break\n",
    "\n",
    "        P0 = copy(P1)\n",
    "        Pstar0 = copy(Pstar1)\n",
    "    end\n",
    "\n",
    "    return P1, Pstar1, ind_p_mat, ind_pstar_mat, sol, pdf, iter\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.416649 seconds (254.81 k allocations: 91.882 MiB, 2.54% gc time)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.01595205933641332, 0.021695700591961128, [0.30768588844713596 0.30768588844713596 0.30768588844713596; 0.30220639279011097 0.30220639279011097 0.30220639279011097; … ; 0.052883757118820714 0.052883757118820714 0.052883757118820714; 0.05194196443888268 0.05194196443888268 0.05194196443888268], [-1.0 0.4922974215154175 0.36922306613656314; -1.0 0.4835302284641776 0.3626476713481332; … ; -1.0 0.08461401139011314 0.06346050854258485; -1.0 0.08310714310221229 0.062330357326659216], (vNX = [2.6845862763859143e-5; 2.7012079309207837e-5; … ; 0.0001826363856354905; 0.00018807831253375335;;], vEX = [2.6857343001530896e-5 2.6882145984155057e-5; 2.7024415088905326e-5 2.7051066464794964e-5; … ; 0.0002016186436944281 0.0002742056298972245; 0.00020817529306058711 0.00028429120581876285], polNX = [0; 0; … ; 1; 1;;], polEX = [0 0; 0 0; … ; 1 1; 1 1], Π = [0.2355397082043846 0.0398220365205688 … 0.0 0.0; 0.2035639923092491 0.03686888293414207 … 0.0 0.0; … ; 4.784202347164862e-30 1.4778853656337744e-29 … 0.03686888293414202 0.20356399230924901; 1.3827115757372961e-30 4.347177100735715e-30 … 0.039822036520568704 0.23553970820438463], ϵ_vec = [0.4908675431085489, 0.4979749102029403, 0.5051851862545925, 0.5124998613024105, 0.5199204469598645, 0.5274484767273707, 0.5350855063091983, 0.5428331139349621, 0.5506929006857724, 0.5586664908251056  …  1.7899766970506505, 1.8158941194896647, 1.84218680535361, 1.8688601881549594, 1.8959197800790755, 1.9233711731233287, 1.9512200402527142, 1.9794721365721941, 2.0081333005160213, 2.0372094550542794], iter = 37), [0.001099533362581301 4.863113223732801e-6 5.264581886217994e-5; 0.0002765981598866844 2.0442021739576122e-6 2.115811823782667e-5; … ; 1.0871545107504635e-11 0.00015613753141058758 0.00014366293801633918; 9.50637483804308e-12 0.0006012986410113883 0.0005557436441494431], 35)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@time begin\n",
    "    P_rce, Pstar_rce, ind_p_mat_rce, ind_pstar_mat_rce, sol_rce, pdf_rce, iter_rce = compute_stationary_equilibrium(E)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Just look at the insane speed**. Remember we didn't do any parallelization and stuff. These are just very efficient codes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.6 Some results from the stationary equilibrium\n",
    "\n",
    "Share of firms in each state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The share of nonexporters is 0.19266736559312644, the share of exporters with low technology is 0.4157090422674407, the share of exporters with high technology is 0.391623592139432, in the stationary distribution, \n"
     ]
    }
   ],
   "source": [
    "NX_share, EX_low_share, EX_high_share = sum(pdf_rce, dims = 1)\n",
    "println(\"The share of nonexporters is $NX_share, the share of exporters with low technology is $EX_low_share, the share of exporters with high technology is $EX_high_share, in the stationary distribution, \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cutoff productivity for firms at different states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3-element Vector{Float64}:\n",
       " 1.036592023817942\n",
       " 0.7555421023503837\n",
       " 0.5749616962539099"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "find_cutoff_prod(sol_rce)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can collect some moments from the stationary distribution. They will be useful for estimations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "moments_fixedQ (generic function with 1 method)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "function moments_fixedQ(P, Pstar, sol, pdf, ind_p_mat, ind_pstar_mat; E = E, Qvalue = E.Q, τ = 0.0)\n",
    "    (;ρϵ, σϵ, grids, ξ) = E\n",
    "\n",
    "    Π, ϵ_vec = tauchen(grids, ρϵ, σϵ)\n",
    "\n",
    "    # exporter share\n",
    "    exporter_share = sum(pdf, dims = 1)[1]\n",
    "\n",
    "    # starter rate\n",
    "    starter_rate = sol.polNX⋅pdf[:,1]\n",
    "\n",
    "    # stopper rate (should be same as starter rate under stationary equilibrium)\n",
    "    stopper_rate = (1 .- sol.polEX)⋅pdf[:,2:3]\n",
    "    \n",
    "    # exporter sales ratio conditional on exporting\n",
    "    sales_EX_do = zeros(eltype(P), grids, length(ξ))\n",
    "    sales_EX_fo = similar(sales_EX_do)\n",
    "    accum = zero(eltype(P))\n",
    "    for i in 1:grids\n",
    "        for j in 1:length(ξ)\n",
    "               sales_EX_do[i, j], sales_EX_fo[i, j]  = sales(1, ϵ_vec[i], j, P, Pstar; E= E, Qvalue = Qvalue, τ = τ)\n",
    "               accum += sales_EX_fo[i, j]/(sales_EX_do[i, j] + sales_EX_fo[i, j]) * pdf[i, j+1]/(1.0-exporter_share )\n",
    "        end\n",
    "    end\n",
    "\n",
    "    # average export price (in domestic currency) and domestic price ratio conditional on exporting\n",
    "    fo_price_ratio = sum(Qvalue.*ind_pstar_mat[:,2:3]./ind_p_mat[:,2:3].*pdf_rce[:,2:3])/(1.0-exporter_share)\n",
    "\n",
    "    return exporter_share, starter_rate, stopper_rate, accum, fo_price_ratio \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.19266736559312644, 0.01391198296999171, 0.013911817889600284, 0.46736658942224996, 1.4059666732401315)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "moments_fixedQ(P_rce, Pstar_rce, sol_rce, pdf_rce, ind_p_mat_rce, ind_pstar_mat_rce)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that by optimality condition of an exporting firm. $Q p_j^* /p_j = \\frac{\\xi}{1 - \\tau}$, so the last moment, the average export price ratio conditional on exporting can also be computed from the the parameter values of $\\xi$ and $\\tau$ and the conditional fraction of exporters at different states of $\\xi$. This moment can be interpreted as a premium for exporting.\n",
    "\n",
    "We can check if this is true under some parameterization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "(1.0 * 1.6 * 0.4157090422674407 + 1.0 * 1.2 * 0.391623592139432)/(1.0 - 0.19266736559312644) ≈ 1.4059666732401315"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.7.2",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}