## Preamble

* A model of lumpy investment wherein firms face persistent shocks to common and plant-specific productivity, and nonconvex adjustment costs.

* The model extends the real business cycle model to include firm heterogeneity and fixed capital adjustment costs. The aggregate state vector of the model contains the distribution of firms over idiosyncratic productivity and capital, which evolves over time in response to aggregate productivity shocks.

* The dynamics of the distribution must satisfy a complicated fixed point problem: each firm’s investment decision depends on its expectations of the dynamics of the distribution, and the dynamics of the distribution depend on firms’ investment decisions.

* **This infinite-dimensional fixed point problem is at the heart of the computational challenges faced by the heterogeneous agent literature.**

```python
# Find toolbox analogue in HARK
# Analogue of tauchen for HARK
%matplotlib notebook
import numpy as np
from quantecon.markov.approximation import tauchen
import scipy.optimize as opt
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import quantecon
import scipy.interpolate as intpl
from matplotlib import cm
from numba import vectorize, jit, njit, float64, prange
```

## Associated Literature

* Host of Papers that look at the heterogeneity through the firm channel.

   * Micro-level Heterogeneity and Transmission: [Reiter, Sveen and Weinke (2013)] **(_Lumpy investment and the monetary transmission mechanism_)**, [Koby and Wolf (2019)] **(_Investment Stimulus: From Micro Elasticities to Macro Counterfactuals_)**.

   * Financial heterogeneity and  the business cycle dynamics of aggregate investment: [Khan, Senga and Thomas (2016)] **(_Default Risk and Aggregate Fluctuations in an Economy with Production Heterogeneity_)**.

   * Investment and Business Cycle: [Cooper, Haltiwanger (2006)] **(_Machine Replacement and the Business Cycle: Lumps and Bumps_)**.

   * Heterogeneity and computation: [Krusell, Smith (1998)] **(_Income and wealth heterogeneity in the macroeconomy_)**.

Now we will solve for the mass of entrants \( m^{*} \) and the stationary measure of incumbent firms \( g^{*}(\varepsilon, k) . \) As discussed by Hopenhayn and Rogerson \( (1993), \) the stationary measure is linearly homogenous in the amount of entry \( m^{*} . \) That is, if we double the amount of entrants
\( m^{*}, \) then we will double the total mass of firms in the economy but the fraction of firms with any particular pair \( (\varepsilon, k) \) is constant. Therefore, we proceed in two steps: first, solve for the stationary measure with \( m^{*}=1 ; \) and second, solve for the mass of entrants \( m^{*} \) itself.

(a) Step 1: We solve for stationary measure with \( m^{*}=1 \) In principle, the stationary distribution must satisfy the law of motion \[ g^{*}\left(\varepsilon^{\prime}, k^{\prime}\right)=\int X(\varepsilon, k) \operatorname{Pr}\left(\varepsilon^{\prime} | \varepsilon\right) \mathbb{I}\left\{k^{\prime}(\varepsilon, k)=k^{\prime}\right\} d g^{*}(\varepsilon, k)+m^{*} \operatorname{Pr}\left(\nu=\varepsilon^{\prime}\right) \] where \( X(\varepsilon, k) \) is an indicator variable for whether a firm \( (\varepsilon, k) \) will survive (i.e., choose not to exit the economy).
To go about this, we iterate, on a numerical approximation of this equation with \( m^{*}=1 . \) In order to do so, assume that the capital stock \( k \) takes values in a discrete grid \( \mathbf{k}=\left(k_{1}, \ldots, k_{N}\right) \) where \( k_{i}>k_{i-1} \) and \( n_{k} \) is the total number of grid points. Under this assumption, the measure \( g(\varepsilon, k) \) can be represented by a \( n_{\varepsilon} \times n_{k} \) matrix or, in more compact form, a \( n_{\varepsilon} n_{k} \times 1 \) vector.

We iterate on
\( g\left(\varepsilon_{k}, k_{l}\right)=\sum_{i=1}^{n_{e}} \sum_{j=1}^{N} \operatorname{Pr}\left(\varepsilon^{\prime}=\varepsilon_{k} | \varepsilon=\varepsilon_{i}\right) \|\left\{k_{i j}^{\prime} \in k_{l}\right\} X\left(\varepsilon_{i}, k_{j}\right) g\left(\varepsilon_{i}, k_{j}\right)+m^{*} \operatorname{Pr}\left(\nu(\varepsilon)=\varepsilon_{k}\right) \)

where the notation \( \mathbb{N}\left\{k_{i j}^{\prime} \in k_{l}\right\} \) means that \( k_{l} \) is the closest gridpoint to \( k^{\prime}\left(\varepsilon_{i}, k_{j}\right) . \)

* First, Microeconomic lumpiness continues to have perceptible effects on aggregate investment dynamics only when equilibrium factor supply considerations are ignored.
   * General equilibrium itself matters tremendously in shaping aggregate dynamics.

      * First, it is extremely effective in smoothing changes in aggregate series, yielding investment rates close to the postwar U.S. data in both their cyclical variability and persistence, irrespective of capital adjustment costs.
      * Second, it dampens much of the movement along the extensive margin that would otherwise distinguish the lumpy investment economy, leaving no trace of aggregate nonlinearities.

* Second, they find that the near-zero skewness and excess kurtosis in their general equilibrium models’ aggregate investment rate series matches the third and fourth moments of postwar U.S. aggregate investment rates reasonably well, and far more closely than does the partial equilibrium lumpy model.

* Third, they find that the microeconomic role of nonconvex capital adjustment costs can change substantially in the presence of persistent idiosyncratic risk. When plant-specific productivity shocks are volatile relative to aggregate shocks, as suggested by the data, nonconvex costs no longer cause the plant-level investment spikes that are the hallmark of lumpy investment.

# Model (Winberry version of KT)

I adopt the Winberry adaptation (2018) of KT. It is useful for computation and is a little bit more tractable. This is the standard Benchmark RBC model with firm heterogeneity.

* **Firms**

There is a fixed mass of firms j $\in$ [0, 1] which produce output $y_{t}$ according to the
production function

$$
y_{j t}=e^{z_{t}} e^{\varepsilon_{j t}} k_{j t}^{\theta} n_{j t}^{\nu},  \theta+\nu<1
$$

where \( z_{t} \) is an aggregate productivity shock, \( \varepsilon_{j t} \) is an idiosyncratic productivity shock, \( k_{j t} \) is capital, \( n_{j t} \) is labor, \( \theta \) is the elasticity of output with respect to capital, and \( \nu \) is the elasticity of output with respect to labor. The aggregate shock \( z_{t} \) is common to all firms and follows the AR ( 1) process
\[ z_{t+1}=\rho_{z} z_{t}+\sigma_{z} \omega_{t+1}^{z}, \text { where } \omega_{t+1}^{z} \sim N(0,1) \]

The idiosyncratic shock \( \varepsilon_{j t} \) independently distributed across firms, but within firms follows the AR ( 1) process
\[ \varepsilon_{j t+1}=\rho_{\varepsilon} \varepsilon_{j t}+\sigma_{\varepsilon} \omega_{t+1}^{\varepsilon}, \text { where } \omega_{t+1}^{\varepsilon} \sim N(0,1) \]

In each period, the firm \( j \) inherits its capital stock from previous periods' investments, observes the two productivity shocks, hires labor from a competitive labor market, and produces output.

After production, the firm invests in capital for the next period. Gross investment \( i_{j t} \) yields \( k_{j t+1}=(1-\delta) k_{j t}+i_{j t} \) units of capital next period, where \( \delta \) is the depreciation rate of capital. If \( \frac{i_{t}}{k_{j t}} \notin[-a, a], \) the firm must pay a fixed adjustment cost \( \xi_{j t} \) in units of labor. The parameter \( a \) governs a region around zero investment within firms do not incur the fixed cost. The fixed cost \( \xi_{j t} \) is a random variable distributed uniformly \( [0, ], \)  and independently over firms and time.

* **Households**

The function \[ \mathbb{E} \sum_{t=0}^{\infty} \beta^{t}\left[\frac{C_{t}^{1-\sigma}-1}{1-\sigma}-\chi \frac{N_{t}^{1+\alpha}}{1+\alpha}\right] \] where \( C_{t} \) is consumption, \( N_{t} \) is labor supplied to the market, \( \beta \) is the discount factor, \( \sigma \) is the coefficient of relative risk aversion, \( \chi \) governs the disutility of labor supply, and \( 1 / \alpha \) is the Frisch elasticity of labor supply. The total time endowment per period is normalized to 1, so that \( N_{t} \in[0,1] \). The household owns all the firms in the economy and markets are complete.

## Solving the Model

The household optimization is loaded into the firm's optimization problem by approximating the transformed value as is done in Khan and Thomas (2008),

\begin{equation}
\begin{aligned}
v(\varepsilon, k, \xi ; \mathbf{s}) &=\lambda(\mathbf{s}) \max _{n}\left\{e^{z} e^{\varepsilon} k^{\theta} n^{\nu}-w(\mathbf{s}) n\right\} \\ &+\max \left\{v^{a}(\varepsilon, k ; \mathbf{s})-\xi \lambda(\mathbf{s}) w(\mathbf{s}), v^{n}(\varepsilon, k ; \mathbf{s})\right\}
\end{aligned}
\end{equation}

where \( \mathrm{s} \) is the aggregate state vector (defined below), \( \lambda(\mathrm{s})=C(\mathrm{s})^{-\sigma} \) is the marginal of consumption in equilibrium,
\[ v^{a}(\varepsilon, k ; \mathrm{s})=\max _{k^{\prime} \in \mathbb{R}}-\lambda(\mathrm{s})\left(k^{\prime}-(1-\delta) k\right)+\beta \mathbb{E}\left[\widehat{v}\left(\varepsilon^{\prime}, k^{\prime} ; \mathrm{s}^{\prime}\left(z^{\prime} ; \mathrm{s}\right) | \varepsilon, k ; \mathrm{s}\right]\right. \]

\( v^{n}(\varepsilon, k ; \mathbf{s})=\max _{k^{\prime} \in[(1-\delta-a) k,(1-\delta+a) k]}-\lambda(\mathbf{s})\left(k^{\prime}-(1-\delta) k\right)+\beta \mathbb{E}\left[\widehat{v}\left(\varepsilon^{\prime}, k^{\prime} ; \mathbf{s}^{\prime}\left(z^{\prime} ; \mathbf{s}\right) | \varepsilon, k ; \mathbf{s}\right]\right. \)

and \( \widehat{v}(\varepsilon, k ; \mathbf{s})=\mathbb{E}_{\xi} v(\varepsilon, k, \xi ; \mathbf{s}) \). Denote the unconstrained capital choice from ( 2) by \( k^{a}(\varepsilon, k ; \mathbf{s}) \) and the constrained choice from ( 3) by \( k^{n}(\varepsilon, k ; \mathbf{s}) . \) The firm will choose to pay the fixed cost if and only if \( v^{a}(\varepsilon, k ; \mathbf{s})-\xi \lambda(\mathbf{s}) w(\mathbf{s}) \geq v^{n}(\varepsilon, k ; \mathbf{s}) . \) Hence, there is a unique threshold value of the fixed cost \( \xi \) which makes the firm indifferent between these two options, \[ \tilde{\xi}(\varepsilon, k ; \mathbf{s})=\frac{v^{a}(\varepsilon, k ; \mathbf{s})-v^{n}(\varepsilon, k ; \mathbf{s})}{\lambda(\mathbf{s}) w(\mathbf{s})} \]

Denote \( \widehat{\xi}(\varepsilon, k ; \mathbf{s}) \) as the threshold bounded to be within the support of \( \xi, \) i.e., \( \widehat{\xi}(\varepsilon, k ; \mathbf{s})= \) \( \min \{\max \{0, \widetilde{\xi}(\varepsilon, k ; \mathbf{s})\}, \bar{\xi}\}\} \) It will be numerically convenient to approximate the ex-ante value function \( \widehat{v}(\varepsilon, k ; \mathbf{s})= \) \( \mathbb{E}_{\xi} v(\varepsilon, k, \xi ; \mathbf{s}) \). Given that the extensive margin decision is characterized by the cutoff ( 4) and the fixed cost \( \xi \) is uniformly distributed, the expectation can be computed analytically as

\( \widehat{v}(\varepsilon, k ; \mathbf{s})=\lambda(\mathbf{s}) \max _{n}\left\{e^{z} e^{\varepsilon} k^{\theta} n^{\nu}-w(\mathbf{s}) n\right\} \)
\[ \quad+\frac{\hat{\xi}(\varepsilon, k ; \mathbf{s})}{\xi}\left(v^{a}(\varepsilon, k ; \mathbf{s})-\lambda(\mathbf{s}) w(\mathbf{s}) \frac{\widehat{\xi}(\varepsilon, k ; \mathbf{s})}{\bar{\xi}}\right) v^{n}(\varepsilon, k ; \mathbf{s}) \]

* Recursive Competitive Equilibrium

In the recursive competitive equilibrium, the aggregate state s contains the current draw of the aggregate productivity shock, \( z, \) and the density of firms over \( (\varepsilon, k) \) -space, \( g(\varepsilon, k) . \)

A recursive competitive equilibrium for this model is a set \( \widehat{v}(\varepsilon, k ; \mathbf{s}), \) \( n(\varepsilon, k ; \mathbf{s}), k^{a}(\varepsilon, k ; \mathbf{s}), k^{n}(\varepsilon, k ; \mathbf{s}), \widehat{\xi}(\varepsilon, k ; \mathbf{s}), \lambda(\mathbf{s}), w(\mathbf{s}), \) and \( \mathbf{s}^{\prime}\left(z^{\prime} ; \mathbf{s}\right)=\left(z^{\prime} ; g^{\prime}(z, g)\right) \) such that

1) (Firm optimization) Taking \( w(\mathbf{s}), \lambda(\mathbf{s}), \lambda(\mathbf{s}), \) and \( \mathbf{s}^{\prime}\left(z^{\prime} ; \mathbf{s}\right) \) as given, \( \widehat{v}(\varepsilon, k ; \mathbf{s}), n(\varepsilon, k ; \mathbf{s}) \) \( \quad k^{a}(\varepsilon, k ; \mathbf{s}), k^{n}(\varepsilon, k ; \mathbf{s}), \) and \( \hat{\xi}(\varepsilon, k ; \mathbf{s}) \) solve the firm's optimization problem.

2) (Household Optimization) \(\lambda(\mathbf{s})=C(\mathbf{s})^{-\sigma}, \) where \( C(\mathbf{s})=\int\left[e^{z} e^{\varepsilon} k^{\theta} n(\varepsilon, k ; \mathbf{s})^{\nu}+(1-\delta) k-\left(\frac{\widehat{\xi}(\varepsilon, k ; \mathbf{s})}{\xi}\right) k^{a}(\varepsilon, k ; \mathbf{s})\right. \) \( \left.-\left(1-\frac{\hat{\xi}(\varepsilon, k ; \mathbf{s})}{\xi}\right) k^{n}(\varepsilon, k ; \mathbf{s})\right] g(\varepsilon, k) d \varepsilon d k \) \(w(\mathbf{s}) \) satisfies \( \int\left(n(\varepsilon, k ; \mathbf{s})+\frac{\hat{\xi}(\varepsilon, k ; \mathbf{s})^{2}}{2 \xi}\right) g(\varepsilon, k) d \varepsilon d k=\left(\frac{w(\mathbf{s}) \lambda(\mathbf{s})}{x}\right)^{\frac{1}{\alpha}} \)

3) (Law of motion for distribution) For all \( \left(\varepsilon^{\prime}, k^{\prime}\right) \) \( g^{\prime}\left(\varepsilon^{\prime}, k^{\prime} ; z, \mathbf{m}\right)=\iiint_{ } 1\left\{\rho_{\varepsilon} \varepsilon+\sigma_{\varepsilon} \omega_{\varepsilon}^{\prime}=\varepsilon^{\prime}\right\} \times\left[\frac{\hat{\xi}(\varepsilon, k ; z, \mathbf{m})}{\xi} 1\left\{k^{a}(\varepsilon, k ; z, \mathbf{m})=k^{\prime}\right\}\right] \)
\( \left.g^{\prime}\left(\varepsilon^{\prime}, k^{\prime} ; z, \mathbf{m}\right)=\iiint_{^{-}}^{T}\left(1-\frac{\hat{\varepsilon}(\varepsilon, k ; \mathbf{m})}{\xi}\right) 1\left\{k^{n}(\varepsilon, k ; z, \mathbf{m})=k^{\prime}\right\}\right] \)
\( \times p\left(\omega_{\varepsilon}^{\prime}\right) g(\varepsilon, k ; \mathbf{m}) d \omega_{\varepsilon}^{\prime} d \varepsilon d k \) where \( p \) is the pdf. of idiosyncratic productivity shocks.

4) (Law of motion for aggregate shocks) \( z^{\prime}=\rho_{z} z+\sigma_{z} \omega_{z}^{\prime}, \) where \( \omega_{z}^{\prime} \sim N(0,1) \)

## Parametrization

% Center the Table

| Parameter | Description | Value |
| --- | --- | --- |
| $\beta$ | Discount Factor | .961 |
| $\sigma$ | Utility curvature | 1 |
| $\alpha$ | Inverse Frisch | lim $\alpha \rightarrow$ 0 |
| $\chi$ | Labor Disutility | N*$=\frac{1}{3}$ |
| $\nu$ | Labor Share | 0.64 |
| $\theta$ | Capital Share | 0.256|
| $\delta$ | Capital Depreciation | 0.085 |
| $\rho_{z}$ | Aggregate TFP AR(1) |  0.859 |
| $\sigma_{z}$ | Aggregate TFP AR(1) |  0.014 |
| $a$ | No fixed cost region | 0.011 |
| $\rho_{e}$ | Idiosyncratic TFP AR(1) | 0.859 |
| $\sigma_{e}$ | Idiosyncratic TFP AR(1) | 0.022 |
| $\bar{\xi}$ | Fixed Cost | 0.0083 |

## Computational Method

A code readme is placed in github (to-do) to detail the method. I outline the steps briefly here.

There are three steps involved here.
* **Step 1**

Approximate the model’s equilibrium objects –  the value function and distribution – using finitedimensional global approximations with respect to individual state variables. This step yields
a finite-dimensional representation of the equilibrium at every point in time.

* **Step 2**

Compute the stationary equilibrium of the finite-dimensional model in which there are no aggregate shocks but still idiosyncratic shocks.

* **Step 3**

The final step is to compute the aggregate dynamics of the finite-dimensional model using a locally accurate
Taylor expansion around the stationary equilibrium.

Numerous approximation types are employed such as quadrature methods, spline approximations to move ahead with solving the model.

I try a simpler model of here and progressively will add features to it.

```python
#Parameter dictionary (I start with fixed labor supply in steady state)
KTEconomyDictionary = {
                       "theta" : 0.21,
                       "mu" : 0.64,
                       "delta" : 0.1,
                       "beta" : 0.96,
                       "varphi" : 0.5,
                       "rho" : 0.9,
                       "sigma" : 0.02,
                       "n_ss" : 0.6
}
```

```python
# Steady State Block
def get_ss(KTEconomyDictionary):

    r = (1 / beta) - (1 - delta)
    n = n_ss
    eps = 0
    k = (r / ((n ** mu) * theta * np.exp(eps))) ** (1 / (theta - 1))
    y = np.exp(eps) * (k ** theta) * (n ** mu)
    i = delta * k
    c = y - i
    w = mu * y / n
    a = w / c

    return k, w, c, a, i, y
```

```python
k_ss, w_ss, c_ss, a_ss, i_ss, y_ss = get_ss(KTEconomyDictionary)
print(k_ss)
print(w_ss)
```

```python
N_k = 50
k_low =  k_ss - 0.5
k_high = k_ss + 0.5
k_grid = np.linspace(k_low, k_high, N_k)

### eps_grid (Tauchen Analogue)

N_eps = 10
trans_eps_MC = tauchen(rho, sigma, n=N_eps).P

sigma_y = np.sqrt( sigma ** 2 / (1 - rho ** 2) )
m = 3
eps_grid = np.zeros(N_eps)
eps_grid[0] = -m * sigma_y
eps_grid[N_eps-1] = m * sigma_y
s = (eps_grid[N_eps-1] - eps_grid[0]) / (N_eps - 1)
for i in range(1, N_eps-1):
    eps_grid[i] = eps_grid[i-1] + s

# Now we initialize the value function:
V_init = np.zeros((N_eps, N_k))
```

```python
trans_eps_MC
```

```python
n = lambda eps, k: (w / (mu * np.exp(eps) * k ** theta)) ** (1 / (mu - 1))
cf = 0.01
w = w_ss
```

## Results

The model establishes consistency with essential features of the microeconomic data, as well as its nonlinearities in partial equilibrium, and then comes up with some results for aggregate dynamics.

* **Result 1**

When factor prices vary over time to clear the
markets for labor and goods, we find that aggregate nonlinearities associated
with lumpy investment disappear, leaving the economy’s dynamics virtually indistinguishable from those in the standard frictionless environment.

* **Result 2**

Fluctuations in target capital stocks are smoothed dramatically, and thus so are the sharp changes along the extensive margin that generate the large positive skewness and excess kurtosis seen above. Moreover, by restraining changes
along both the intensive and extensive margins, market-clearing relative price
movements also smooth away the excessively large and abrupt fluctuations in
aggregate investment demand one sees in partial equilibrium, thereby yielding
an aggregate investment rate with persistence and **volatility close to the data.**

```python
def vf_update(eps_index, k_index, V_old_interp, w, cf):
    """
    input:
    k_index: index so that k_grid[index] corresponds to value of capital this period (state)
    V_old: approximation to the value function. V_old[i] approximates V(k_grid[i]).

    output:
    action_value: value of all possible state-action pairs.
    """

    ### solve v1
    eps = eps_grid[eps_index]
    k = k_grid[k_index]

    n = (w / (mu*np.exp(eps)*k**theta))**(1/(mu-1))
    action_value = np.zeros((N_k, N_eps))

    profit = np.exp(eps) * (k ** theta) * (n ** mu) - w * n - (k_grid - (1 - delta) * k) \
        - varphi / 2 * (k_grid / k - (1 - delta)) ** 2 * k

    action_value = profit + beta * np.dot((trans_eps_MC[eps_index, :]).T, V_old_interp(eps_grid, k_grid).T)

    V1_new = np.max(action_value)

    ### solve v_new
    V_new = max((1 - delta) * k, V1_new - cf)

    return V_new, action_value
```

```python
def vf_update_iteration(V_old, V_old_interp, w, cf):
    """
    input:
    V_old: array with current approximation of the value function

    output:
    V_new: updated approximation of the value function
    """
    V_new = np.zeros_like(V_old)

    for i in range(N_eps):
        for j in range(N_k):
            V_new[i, j], action_value = vf_update(i, j, V_old_interp, w, cf)

    return V_new

```

```python
def solve_value_func_iter(w, cf, plot=True):
    difference_list = []
    threshold = 1e-10
    max_iterations = 10000
    plot_interval = 50

    V = V_init.copy()

    for iteration in range(max_iterations):

        # x.size = m and y.size = n, if z.ndim == 2, then z must have shape (n, m)
        V_interp = intpl.interp2d(eps_grid, k_grid, V.T, kind='linear')
        V_new = vf_update_iteration(V, V_interp, w, cf)
        #V_new = vf_update_iteration(V, V_interp, w, cf)

        difference = np.amax(abs(V_new - V))
        # print('Iteration: {}, difference: {}'.format(iteration + 1, difference))

        difference_list.append(difference)

        V = V_new

        if difference < threshold:
            #print('Converged after iteration {}'.format(iteration + 1))
            #print('Iteration: {}, difference: {}'.format(iteration + 1, difference))

            if plot:
                fig = plt.figure()
                ax = fig.gca(projection='3d')
                EPS, K = np.meshgrid(eps_grid, k_grid)
                surf1 = ax.scatter(EPS, K, V.T, cmap=cm.autumn)
                #fig.colorbar(surf1, shrink=.5, aspect=5)
                ax.set_xlabel(r'$\epsilon$')
                ax.set_ylabel(r'$k$')
                ax.set_zlabel(r'V')
                ax.set_title("Value Function")
                #plt.show()
                #plt.close()

            break

    return V, V_interp

```

```python
V, V_interp = solve_value_func_iter(w_ss, cf)
```

```python
def extract_policy(V, V_interp, w, plot=True):

    poly_grid = np.zeros_like(V)

    for i in range(N_eps):
        for j in range(N_k):
            V_new, action_value = vf_update(i, j, V_interp, w, cf)

            if V_new == (1 - delta) * k_grid[j]:
                poly_grid[i, j] = 0

            else:
                poly_grid[i, j] = k_grid[np.argmax(action_value)]

    if plot:
        fig = plt.figure()
        ax = fig.gca(projection='3d')
        EPS, K = np.meshgrid(eps_grid, k_grid)
        #surf1 = ax.plot_surface(EPS, K, poly_grid.T, cmap=cm.autumn)
        surf1 = ax.scatter(EPS, K, poly_grid.T, cmap=cm.autumn)
        #fig.colorbar(surf1, shrink=.5, aspect=5)
        #surf1 = ax.plot_wireframe(EPS, K, V.T)
        ax.set_xlabel(r'$\epsilon$')
        ax.set_ylabel(r'$k$')
        ax.set_zlabel(r"$k'$")
        ax.set_title("Capital Accumulation Policy Function")

    return poly_grid
```

```python
poly_grid = extract_policy(V, V_interp, w)
```

```python
# Firms continue
def firm_continue(eps, k, V_interp):
    V_val = V_interp(eps, k)
    if V_val == (1 - delta) * k:
        return False
    else:
        return True

firm_continue(0, 5, V_interp)

firm_continue(0.2, 1, V_interp)

def firm_continue_grid(V):
    V_new = np.zeros_like(V)
    for i in range(N_k):
        V_new[:, i] = ((1 - delta) * k_grid[i] != V[:, i])
    return V_new

firm_cont_grid = firm_continue_grid(V)

### Function $\hat{V}(\epsilon)$

def get_v_hat(eps, V_interp):

    eps_index = np.where(eps_grid == eps)[0][0]
    vals = []

    for k_index in range(N_k):
        val = - k_grid[k_index] + beta * np.dot(trans_eps_MC[eps_index, :], V_interp(eps_grid, k_grid[k_index]))
        vals.append(val)

    return max(vals)

get_v_hat(eps_grid[4], V_interp)

ce = 0.02

### https://dilawarnotes.wordpress.com/2017/11/07/stationary-distribution-of-markov-chains-using-python/
def solveStationary( A ):
    n = A.shape[0]
    a = np.eye( n ) - A
    a = np.vstack( (a.T, np.ones( n )) )
    b = np.matrix( [0] * n + [ 1 ] ).T
    return np.linalg.lstsq( a, b )[0]

#### Stationary Distribution of $\epsilon$

stationary_dist_eps = np.array(solveStationary(trans_eps_MC))
stationary_dist_eps

np.sum(stationary_dist_eps)

### Solving for Market-Clearing Wage

def mark_wage_err(w, ce, cf):
    V_new, V_interp_new = solve_value_func_iter(w, cf, plot=False)

    V_hats = np.zeros_like(eps_grid)
    # iterate loop in parallel

    for i in range(len(eps_grid)):
        V_hats[i] = get_v_hat(eps_grid[i], V_interp_new)

    integral = np.dot(V_hats, stationary_dist_eps)
    err = ce - integral
    #print(err)
    return np.array(err)

def solve_mark_wage(w_guess, ce, cf):
    args = (ce, cf)
    result = opt.root(mark_wage_err, w_guess, args = args)
    return result.x[0]

w_star = solve_mark_wage(0.9, 0.02, 0.01)

w_star

## Q Stationary measure of firms and mass of entrants

V, V_interp = solve_value_func_iter(w_star, cf)

poly_grid = extract_policy(V, V_interp, w_star)

firm_cont_grid = firm_continue_grid(V)
#firm_cont_grid

### Step 1: Solving $g(\epsilon, k)$

g_init = np.random.uniform(low=0, high=1, size=(N_eps, N_k))
g_init = g_init / np.sum(g_init)
g = g_init.copy()
m_star = 1

np.sum(g)

@jit(nopython=True, parallel=True)
def update_g(g):
    g_new = np.zeros((N_eps, N_k))

    # iterate loop in parallel
    for k in range(N_eps):
        for l in range(N_k):
            C = (poly_grid == k_grid[l])

            summation = 0
            for i in range(N_eps):
                for j in range(N_k):

                    sum_ij = trans_eps_MC[i,k] * C[i, j] * firm_cont_grid[i, j] * g[i, j]
                    summation += sum_ij

            summation += m_star * stationary_dist_eps[k][0]
            g_new[k, l] = summation

    return g_new

def solve_g(plot=True):
    threshold = 1e-8
    max_iterations = 1000
    diff = 1
    g = g_init.copy()

    while (diff > 1e-5):
        g_new = update_g(g)
        #print('sum :', np.sum(g_new))
        diff = np.max(np.abs(g_new - g))
        #print(diff)
        g = g_new.copy()

    g_normalized = g_new / np.sum(g_new)

    if plot:

        fig = plt.figure()
        ax = fig.gca(projection='3d')
        EPS, K = np.meshgrid(eps_grid, k_grid)
        surf1 = ax.scatter(EPS, K, g_normalized.T, cmap=cm.autumn)
        #fig.colorbar(surf1, shrink=.5, aspect=5)
        #surf1 = ax.plot_wireframe(EPS, K, g_normalized.T)
        ax.set_xlabel(r'$\epsilon$')
        ax.set_ylabel(r'$k$')
        ax.set_zlabel(r"$g$")
        ax.set_title("Law of Motion g")

    return g_normalized

g = solve_g()

### Step 2: Solving m*

def get_g_star(g, m):
    return g.copy() * m

def n(eps, k, w):
    n = (w / (mu * np.exp(eps) * k **theta)) ** (1/(mu - 1))
    return n

def agg_l_d(g_star, w_star):

    n_d = 0
    for i in range(N_eps):
        for j in range(N_k):
            n_d += n(eps_grid[i], k_grid[j], w_star) * g_star[i, j]
    return n_d

def m_error(m_guess, *args):

    g, w_star = args
    g_star = get_g_star(g, m_guess)
    n_d = agg_l_d(g_star, w_star)
    m_error = n_d - 0.6
    return m_error

def solve_m(m_guess, g, w_star):

    args = (g, w_star)
    result = opt.root(m_error, m_guess, args=args)
    return result.x[0]

m_guess = 0.5
m_star = solve_m(m_guess, g, w_star)
m_star

### Step 3: Solving value of a

def get_agg_Y():

    agg_Y = 0
    for i in range(N_eps):
        for j in range(N_k):
            agg_Y += ( np.exp(eps_grid[i]) * k_grid[i] ** theta * n(eps_grid[i], k_grid[j], w_star) ) * g[i, j]

    return agg_Y

get_agg_Y()

def get_agg_I():

    agg_I = 0
    for i in range(N_eps):
        for j in range(N_k):
            agg_I += ( poly_grid[i, j] - (1 - delta) * k_grid[j] ) * g[i, j]

    return agg_I

get_agg_I()

def get_agg_C():

    Y = get_agg_Y()
    I = get_agg_I()
    C = Y - I
    return C

get_agg_C()

def get_a():

    C = get_agg_C()
    a = C * w_star

    return a

get_a()

##  Calibration to parameters in the data, here it is the entry and exit costs

### (i) Calibrating entry cost: Calibrate to target annual exit rate of 10%.

def calibrate_cf():

    cf_range = [0.01, 0.013, 0.014, 0.015, 0.016]
    error = 1000
    cal_cf = 0
    iteration = 1
    for cf_val in cf_range:

        print("iteration ", iteration)
        print("cf val: ", cf_val)
        V, V_interp = solve_value_func_iter(w_star, cf_val, plot=True)

        frac_exit_sum = 0
        for i in range(N_k):
            V_i = V[:, i]
            exit_val = (1 - delta) * k_grid[i]
            exit_bool = (V_i == exit_val)

            frac_exit_i = exit_bool * g[:, i]
            frac_exit_sum += frac_exit_i.sum()

        print("Fraction of exiting firms:", frac_exit_sum)
        print("\n")
        new_error = abs(frac_exit_sum - 0.1)
        #print(new_error)

        if new_error < error:
            error = new_error
            cal_cf = cf_val

        iteration += 1

    return cal_cf, error

cal_cf, error = calibrate_cf()

cal_cf

### (ii) Calibrating exit costs: Calibrate to target average size of new firms to be 40% of average size of all firms

def calibrate_ce():

    ce_range = [0.01, 0.05, 0.1, 0.2, 0.5]
    error = 1000
    cal_ce = 0
    iteration = 1
    for ce_val in ce_range:

        #print("iteration ", iteration)
        #print("ce val: ", ce_val)

        w_guess = 0.9
        w_star = solve_mark_wage(w_guess, ce, cal_cf)
        #print("w_star:", w_star)
        m_guess = 0.05
        m_star = solve_m(m_guess, g, w_star)
        #print("m_star :", m_star)

        new_error = abs(m_star - 0.4)
        if new_error < error:
            error = new_error
            cal_ce = ce_val

        iteration += 1

    return cal_ce, error

cal_ce, error = calibrate_ce()

cal_ce
```

## References

   * Micro-level Heterogeneity and Transmission: [Reiter, Sveen and Weinke (2013)] (_Lumpy investment and the monetary transmission mechanism_), [Koby and Wolf (2019)] **(_Investment Stimulus: From Micro Elasticities to Macro Counterfactuals_)**.

   * Financial heterogeneity and  the business cycle dynamics of aggregate investment: [Khan, Senga and Thomas (2016)] **(_Default Risk and Aggregate Fluctuations in an Economy with Production Heterogeneity_)**.

   * Investment and Business Cycle: [Cooper, Haltiwanger (2006)] **(_Machine Replacement and the Business Cycle: Lumps and Bumps_)**.

   * Heterogeneity and computation: [Krusell, Smith (1998)] **(_Income and wealth heterogeneity in the macroeconomy_)**.

   * Khan, A. and Thomas, J. K. (2008), Idiosyncratic Shocks and the Role of Nonconvexities in Plant and Aggregate Investment Dynamics. Econometrica, 76: 395-436.

   * Winberry, T. (2018), A method for solving and estimating heterogeneous agent macro models. Quantitative Economics, 9: 1123-1151.

```python

```

```python

```

REMARKs/REMARK-Firm Entry and Exit/Remark/FirmEntryandExit/Firm-Entry-and-Exit_Clementi-and-Palazzo-AEJ-Macro-2016.ipynb


```python
from itertools import product
import numpy as np
from scipy import optimize
from scipy import interpolate
import sympy as sm

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
plt.style.use('seaborn')
import seaborn as sns

%load_ext autoreload
%autoreload 2
import ASAD
```

# 1. Human capital accumulation

The **parameters** of the model are:

```python
rho = 2
beta = 0.96
gamma = 0.1
b = 1
w = 2
Delta = 0.1
```

The **relevant levels of human capital** are:

```python
h_vec = np.linspace(0.1,1.5,100)
```

The **basic functions** are:

```python
def consumption_utility(c,rho):
    """ utility of consumption

    Args:

        c (float): consumption
        rho (float): CRRA parameter

    Returns:

        (float): utility of consumption

    """

    return c**(1-rho)/(1-rho)

def labor_disutility(l,gamma):
    """ disutility of labor

    Args:

        l (int): labor supply
        gamma (float): disutility of labor parameter

    Returns:

        (float): disutility of labor

    """

    return gamma*l

def consumption(h,l,w,b):
    """ consumption

    Args:

        h (float): human capital
        l (int): labor supply
        w (float): wage rate
        b (float): unemployment benefits

    Returns:

        (float): consumption

    """

    if l == 1:
        return w*h
    else:
        return b
```

The **value-of-choice functions** are:

```python
def v2(l2,h2,b,w,rho,gamma):
    """ value-of-choice in period 2

    Args:

        l2 (int): labor supply
        h2 (float): human capital
        w (float): wage rate
        b (float): unemployment benefits
        rho (float): CRRA parameter
        gamma (float): disutility of labor parameter

    Returns:

        (float): value-of-choice in period 2

    """

    c2 = consumption(h2,l2,w,b)
    return consumption_utility(c2,rho)-labor_disutility(l2,gamma)

def v1(l1,h1,b,w,rho,gamma,Delta,v2_interp,eta=1):
    """ value-of-choice in period 1

    Args:

        l1 (int): labor supply
        h1 (float): human capital
        w (float): wage rate
        b (float): unemployment benefits
        rho (float): CRRA parameter
        gamma (float): disutility of labor parameter
        Delta (float): potential stochastic experience gain
        v2_interp (RegularGridInterpolator): interpolator for value-of-choice in period 2
        eta (float,optional): scaling of determistic experience gain

    Returns:

        (float): value-of-choice in period 1

    """

    # a. v2 value, if no experience gain
    h2_low = h1 + eta*l1 + 0
    v2_low = v2_interp([h2_low])[0]

    # b. v2 value, if experience gain
    h2_high = h1 + eta*l1 + Delta
    v2_high = v2_interp([h2_high])[0]

    # c. expected v2 value
    v2 = 0.5*v2_low + 0.5*v2_high

    # d. consumption
    c1 = consumption(h1,l1,w,b)

    # e. total value
    return consumption_utility(c1,rho) - labor_disutility(l1,gamma) + beta*v2
```

A **general solution function**  is:

```python
def solve(h_vec,obj_func):
    """ solve for optimal labor choice

    Args:

        h_vec (ndarray): human capital
        obj_func (callable): objective function

    Returns:

        l_vec (ndarray): labor supply choices
        v_vec (ndarray): implied values-of-choices

    """

    # a. grids
    v_vec = np.empty(h_vec.size)
    l_vec = np.empty(h_vec.size)

    # b. solve for each h in grid
    for i,h in enumerate(h_vec):

        # i. values of choices
        v_nowork = obj_func(0,h)
        v_work = obj_func(1,h)

        # ii. maximum
        if v_nowork > v_work:
            v_vec[i] = v_nowork
            l_vec[i] = 0
        else:
            v_vec[i] = v_work
            l_vec[i] = 1

    return l_vec,v_vec
```

A **general plotting funcition** is:

```python
def plot(h_vec,l_vec,v_vec,t):
    """ plot optimal choices and value function

    Args:

        h_vec (ndarray): human capital
        l_vec (ndarray): labor supply choices
        v_vec (ndarray): implied values-of-choices
        t (int): period

    """

    # a. labor supply function
    fig = plt.figure(figsize=(10,4))
    ax = fig.add_subplot(1,2,1)
    ax.plot(h_vec,l_vec,label='labor supply')

    # income
    ax.plot(h_vec,w*h_vec,'--',label='wage income')
    ax.plot(h_vec,b*np.ones(h_vec.size),'--',label='unemployment benefits')

    # working with income loss
    I = (l_vec == 1) & (w*h_vec < b)
    if np.any(I):
        ax.fill_between(h_vec[I],w*h_vec[I],b*np.ones(I.sum()),label='working with income loss')

    ax.set_xlabel(f'$h_{t}$')
    ax.set_ylabel(f'$l_{t}$')
    ax.set_title(f'labor supply function in period {t}')
    ax.legend()

    # b. value function
    ax = fig.add_subplot(1,2,2)
    ax.plot(h_vec,v_vec,label='value function')
    ax.set_xlabel(f'$h_{t}$')
    ax.set_ylabel(f'$v_{t}$')
    ax.set_title(f'value function in period {t}')
    ax.legend()
```

## Question 1

The solution in the **second period** is:

```python
# a. solve
obj_func = lambda l2,h2: v2(l2,h2,b,w,rho,gamma)
l2_vec,v2_vec = solve(h_vec,obj_func)

# b. plot
plot(h_vec,l2_vec,v2_vec,2)
```

## Question 2

The solution in the **first period** is:

```python
# a. create interpolator
v2_interp = interpolate.RegularGridInterpolator((h_vec,), v2_vec,
                                                bounds_error=False,fill_value=None)
# b. solve
obj_func = lambda l1,h1: v1(l1,h1,b,w,rho,gamma,Delta,v2_interp)
l1_vec,v1_vec = solve(h_vec,obj_func)

# c. plot
plot(h_vec,l1_vec,v1_vec,1)
```

## Question 3

1. In **period 2**, the worker only works if her potential wage income ($wh_2$) is higher than the unemployment benefits ($b$).
2. In **period 1**, the worker might work even when she looses income in the current period compared to getting unemployment benefits. The explanation is that she  accumulates human capital by working which increase her utility in period 2.

To explain this further, consider the following **alternative problem**:

$$
\begin{aligned}
v_{1}(h_{1}) &= \max_{l_{1}} \frac{c_1^{1-\rho}}{1-\rho} - \gamma l_1 + \beta\mathbb{E}_{1}\left[v_2(h_2)\right] \\
\text{s.t.}  \\
c_{1}& = \begin{cases}
w h_1 &
\text{if }l_1 = 1  \\
b & \text{if }l_1 = 0
\end{cases} \\
h_2 &= h_1 + \eta l_1 + \begin{cases}
0 & \text{with prob. }0.5\\
\Delta & \text{with prob. }0.5
\end{cases}\\
l_{1} &\in \{0,1\}\\
\end{aligned}
$$

where $\eta$ scales the deterministic experience gain from working. Before we had $\eta = 1$.

If we instead set $\eta = 0$, then the worker will only works in period 1 if $wh_2 > b$ by a margin compensating her for the utility loss of working.

```python
# a. solve
obj_func = lambda l1,h1: v1(l1,h1,b,w,rho,gamma,Delta,v2_interp,eta=0)
l1_vec,v1_vec = solve(h_vec,obj_func)

# b. plot
plot(h_vec,l1_vec,v1_vec,1)
```

# 2. AS-AD model

```python
par = {}

par['alpha'] = 5.76
par['h'] = 0.5
par['b'] = 0.5
par['phi'] = 0
par['gamma'] = 0.075
```

```python
par['delta'] = 0.80
par['omega'] = 0.15
```

```python
par['sigma_x'] = 3.492
par['sigma_c'] = 0.2
```

## Question 1

```python
sm.init_printing(use_unicode=True)
```

Construct the **AD-curve:**

```python
y = sm.symbols('y_t')
v = sm.symbols('v_t')
alpha = sm.symbols('alpha')
h = sm.symbols('h')
b = sm.symbols('b')

AD = 1/(h*alpha)*(v-(1+b*alpha)*y)
AD
```

Construct the **SRAS-curve:**

```python
phi = sm.symbols('phi')
gamma = sm.symbols('gamma')
pilag = sm.symbols('\pi_{t-1}')
ylag = sm.symbols('y_{t-1}')
s = sm.symbols('s_t')
slag = sm.symbols('s_{t-1}')

SRAS = pilag + gamma*y- phi*gamma*ylag + s - phi*slag
SRAS
```

**Find solution:**

```python
y_eq = sm.solve(sm.Eq(AD,SRAS),y)
y_eq[0]
```

```python
pi_eq = AD.subs(y,y_eq[0])
pi_eq
```

```python
sm.init_printing(pretty_print=False)
```

## Question 2

Create **Python functions**:

```python
AD_func = sm.lambdify((y,v,alpha,h,b),AD)
SRAS_func = sm.lambdify((y,s,ylag,pilag,slag,phi,gamma),SRAS)
y_eq_func = sm.lambdify((ylag,pilag,v,s,slag,alpha,h,b,phi,gamma),y_eq[0])
pi_eq_func = sm.lambdify((ylag,pilag,v,s,slag,alpha,h,b,phi,gamma),pi_eq)
```

**Illustrate equilibrium:**

```python
# a. lagged values and shocks
y0_lag = 0.0
pi0_lag = 0.0
s0 = 0.0
s0_lag = 0.0

# b. current output
y_vec = np.linspace(-0.2,0.2,100)

# c. figure
fig = plt.figure()
ax = fig.add_subplot(1,1,1)

# SRAS
pi_SRAS = SRAS_func(y_vec,s0,y0_lag,pi0_lag,s0_lag,par['phi'],par['gamma'])
ax.plot(y_vec,pi_SRAS,label='SRAS')

# ADs
for v0 in [0, 0.1]:

    pi_AD = AD_func(y_vec,v0,par['alpha'],par['h'],par['b'])
    ax.plot(y_vec,pi_AD,label=f'AD ($v_0 = {v0})$')

    # equilibrium
    eq_y = y_eq_func(y0_lag,pi0_lag,v0,s0,s0_lag,par['alpha'],par['h'],par['b'],par['phi'],par['gamma'])
    eq_pi =pi_eq_func(y0_lag,pi0_lag,v0,s0,s0_lag,par['alpha'],par['h'],par['b'],par['phi'],par['gamma'])
    ax.scatter(eq_y,eq_pi,color='black',zorder=3)

ax.set_xlabel('$y_t$')
ax.set_ylabel('$\pi_t$')
ax.legend();
```

## Question 3

**Allocate memory and draw random shocks**:

```python
def prep_sim(par,T,seed=1986):
    """ prepare simulation

    Args:

        par (dict): model parameters
        T (int): number of periods to simulate
        seed (int,optional): seed for random numbers

    Returns:

        sim (dict): container for simulation results

    """

    # a. set seed
    if not seed == None:
        np.random.seed(seed)

    # b. allocate memory
    sim = {}
    sim['y'] = np.zeros(T)
    sim['pi'] = np.zeros(T)
    sim['v'] = np.zeros(T)
    sim['s'] = np.zeros(T)

    # c. draw random shocks
    sim['x_raw'] = np.random.normal(loc=0,scale=1,size=T)
    sim['c_raw'] = np.random.normal(loc=0,scale=1,size=T)

    return sim
```

**Simualte** for $T$ periods:

```python
def simulate(par,sim,T):
    """ run simulation

    Args:

        par (dict): model parameters
        sim (dict): container for simulation results
        T (int): number of periods to simulate

    """

    for t in range(1,T):

        # a. shocks
        sim['v'][t] = par['delta']*sim ['v'][t-1] + par['sigma_x']*sim['x_raw'][t]
        sim['s'][t] = par['omega']*sim ['s'][t-1] + par['sigma_c']*sim['c_raw'][t]

        # b. output
        sim['y'][t] = y_eq_func(sim['y'][t-1],sim['pi'][t-1],sim['v'][t],sim['s'][t],sim['s'][t-1],
                                par['alpha'],par['h'],par['b'],par['phi'],par['gamma'])

        # c. inflation
        sim['pi'][t] = pi_eq_func(sim['y'][t-1],sim['pi'][t-1],sim['v'][t],sim['s'][t],sim['s'][t-1],
                                  par['alpha'],par['h'],par['b'],par['phi'],par['gamma'])

```

```python
# a. settings
T = 101

# b. prepare simulation
sim = prep_sim(par,T)

# c. overview shocks
sim['x_raw'][:] = 0
sim['c_raw'][:] = 0
sim ['x_raw'][1] = 0.1/par['sigma_x']

# d. run simulation
simulate(par,sim,T)
```

**$(y_t,\pi_t)$-diagram:**

```python
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(sim['y'][1:],sim['pi'][1:],ls='-',marker='o')
ax.set_xlabel('$y_t$')
ax.set_ylabel('$\pi_t$');
for i in range(1,7):
    ax.text(sim['y'][i],sim['pi'][i]+0.0002,f't = {i}')
```

**Time paths:**

```python
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(np.arange(0,T-1),sim['y'][1:],label='$y_t$, output gap')
ax.plot(np.arange(0,T-1),sim['pi'][1:],label='$\pi_t$, inflation gap')
ax.legend();
```

## Question 4

```python
# a. simulate
T = 1000
sim = prep_sim(par,T)
simulate(par,sim,T)

# b. figure
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(np.arange(T),sim['y'],label='$y_t$, output gap')
ax.plot(np.arange(T),sim['pi'],label='$\pi_t$, inflation gap')
ax.legend();

# c. print
def print_sim(sim):

    print(f'std. of output gap: {np.std(sim["y"]):.4f}')
    print(f'std. of inflation gap: {np.std(sim["pi"]):.4f}')
    print(f'correlation of output and inflation gap: {np.corrcoef(sim["y"],sim["pi"])[0,1]:.4f}')
    print(f'1st order autocorrelation of output gap: {np.corrcoef(sim["y"][1:],sim["y"][:-1])[0,1]:.4f}')
    print(f'1st order autocorrelation of inflation gap: {np.corrcoef(sim["pi"][1:],sim["pi"][:-1])[0,1]:.4f}')

print_sim(sim)
```

### Quesiton 5

**The initial plot:**

```python
# a. calculate correlations
K = 100
phis = np.linspace(0.01,0.99,K)
corr_y_pi = np.empty(K)

est = par.copy()
for i,phi in enumerate(phis):

    # i. update
    est['phi'] = phi

    # ii. simulate
    simulate(est,sim,T)

    # iii. save
    corr_y_pi[i] = np.corrcoef(sim["y"],sim["pi"])[0,1]

# b. plot
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(phis,corr_y_pi)
ax.set_xlabel('$\phi$')
ax.set_ylabel('corr($y_t,\pi_t$)');
```

**The optimization:**

```python
# a. copy parameters
est = par.copy()

# b. objective funciton
def obj_func(x,est,sim,T,sim_func):
    """ calculate objective function for estimation of phi

    Args:

        x (float): trial value for phi
        est (dict): model parameters
        sim (dict): container for simulation results
        T (int): number of periods to simulate
        sim_func (callable): simulation function

    Returns:

        obj (float): objective value

    """

    # i. phi
    est['phi'] = x

    # ii. simulate
    sim_func(est,sim,T)

    # iii. calcualte objective
    obj = (np.corrcoef(sim["y"],sim["pi"])[0,1]-0.31)**2
    return obj

# c. optimize
result = optimize.minimize_scalar(obj_func,args=(est,sim,T,simulate),
                                  bounds=(0+1e-8,1-1e-8),method='bounded')

# d. result
est['phi'] = result.x
print(f'result: phi = {result.x:.3f}')

# e. statistics
print('')
simulate(est,sim,T)
print_sim(sim)
```

    result: phi = 0.963

    std. of output gap: 1.5264
    std. of inflation gap: 0.2726
    correlation of output and inflation gap: 0.3100
    1st order autocorrelation of output gap: 0.8184
    1st order autocorrelation of inflation gap: 0.5559

### Advanced

**Problem:** The estimate for $\phi$ above depends on the seed chosen for the random number generator. This can be illustrated by re-doing the estimation for different seeds:

```python
seeds = [1997,1,17,2018,999] # "randomly chosen seeds"
for seed in seeds:

    # a. prepare simulate
    sim_alt = prep_sim(par,T,seed)

    # b. run optimizer
    result = optimize.minimize_scalar(obj_func,args=(est,sim_alt,T,simulate),bounds=(0+1e-8,1-1e-8),method='bounded')
    result_alt = optimize.minimize_scalar(obj_func,args=(est,sim_alt,T,ASAD.simulate),bounds=(0+1e-8,1-1e-8),method='bounded')

    print(f'seed = {seed:4d}: phi = {result.x:.3f} [{result_alt.x:.3f}]')
```

    seed = 1997: phi = 0.949 [0.949]
    seed =    1: phi = 0.973 [0.973]
    seed =   17: phi = 0.985 [0.985]
    seed = 2018: phi = 0.926 [0.926]
    seed =  999: phi = 0.949 [0.949]

**Solution:** To reduce this problem, we need to simulate more than 1,000 periods. To do so it is beneficial to use the fast simulation function provided in **ASAD.py** (optimized using numba):

1. The results in the square brackets above show that this simulation function gives the same results.
2. The results below show that when we simulate 1,000,000 periods the estimate of $\phi$ is approximately 0.983-0.984 irrespective of the seed.

```python
T_alt = 1_000_000
for seed in [1997,1,17,2018,999]:

    # a. simulate
    sim_alt = prep_sim(par,T_alt,seed)

    # b. run optimizer
    result = optimize.minimize_scalar(obj_func,args=(est,sim_alt,T_alt,ASAD.simulate),bounds=(0+1e-8,1-1e-8),method='bounded')

    print(f'seed = {seed:4d}: phi = {result.x:.3f}')
```

    seed = 1997: phi = 0.985
    seed =    1: phi = 0.984
    seed =   17: phi = 0.983
    seed = 2018: phi = 0.983
    seed =  999: phi = 0.984

## Question 6

```python
# a. copy parameters
est = par.copy()

# b. objective function
def obj_func_all(x,est,sim,T,sim_func):
    """ calculate objective function for estimation of phi, sigma_c and sigma_c

    Args:

        x (ndarray): trial values for [phi,sigma_x,sigma_c]
        est (dict): model parameters
        sim (dict): container for simulation results
        T (int): number of periods to simulate
        sim_func (callable): simulation function

    Returns:

        obj (float): objective value

    """

    # i. phi with penalty
    penalty = 0
    if x[0] < 1e-8:
        phi = 1e-8
        penalty += (1-1e-8-x[0])**2
    elif x[0] > 1-1e-8:
        phi = 1-1e-8
        penalty += (1-1e-8-x[0])**2
    else:
        phi = x[0]
    est['phi'] = phi

    # ii. standard deviations (forced to be non-negative)
    est['sigma_x'] = np.sqrt(x[1]**2)
    est['sigma_c'] = np.sqrt(x[2]**2)

    # iii. simulate
    sim_func(est,sim,T)

    # iv. calcualte objective
    obj = 0
    obj += (np.std(sim['y'])-1.64)**2
    obj += (np.std(sim['pi'])-0.21)**2
    obj += (np.corrcoef(sim['y'],sim['pi'])[0,1]-0.31)**2
    obj += (np.corrcoef(sim['y'][1:],sim['y'][:-1])[0,1]-0.84)**2
    obj += (np.corrcoef(sim['pi'][1:],sim['pi'][:-1])[0,1]-0.48)**2

    return obj + penalty

# c. optimize
x0 = [0.98,par['sigma_x'],par['sigma_c']]
result = optimize.minimize(obj_func_all,x0,args=(est,sim,T,simulate))
print(result)

# d. update and print estimates
est['phi'] = result.x[0]
est['sigma_x'] = np.sqrt(result.x[1]**2)
est['sigma_c'] = np.sqrt(result.x[2]**2)

est_str = ''
est_str += f'phi = {est["phi"]:.3f}, '
est_str += f'sigma_x = {est["sigma_x"]:.3f}, '
est_str += f'sigma_c = {est["sigma_c"]:.3f}'
print(f'\n{est_str}\n')

# e. statistics
sim = prep_sim(est,T)
simulate(est,sim,T)
print_sim(sim)
```

          fun: 0.005339800735277895
     hess_inv: array([[ 5.09040330e-03,  8.38489388e-03, -2.52877550e-03],
           [ 8.38489388e-03,  2.56633371e+00,  9.15285507e-02],
           [-2.52877550e-03,  9.15285507e-02,  8.15174802e-02]])
          jac: array([ 9.33842966e-06,  5.05708158e-07, -7.71949999e-07])
      message: 'Optimization terminated successfully.'
         nfev: 55
          nit: 7
         njev: 11
       status: 0
      success: True
            x: array([0.97189233, 3.72758982, 0.21646031])

    phi = 0.972, sigma_x = 3.728, sigma_c = 0.216

    std. of output gap: 1.6295
    std. of inflation gap: 0.2728
    correlation of output and inflation gap: 0.3385
    1st order autocorrelation of output gap: 0.8183
    1st order autocorrelation of inflation gap: 0.4772

### Advanced

**Same problem:** Different seeds give different results.

```python
for seed in seeds:

    # a. prepare simulation
    sim_alt = prep_sim(par,T,seed)

    # b. run optimizer
    est = par.copy()
    x0 = [0.98,par['sigma_x'],par['sigma_c']]
    result = optimize.minimize(obj_func_all,x0,args=(est,sim_alt,T,simulate))

    # c. update and print estimates
    est['phi'] = result.x[0]
    est['sigma_x'] = np.sqrt(result.x[1]**2)
    est['sigma_c'] = np.sqrt(result.x[2]**2)

    est_str = ''
    est_str += f' phi = {est["phi"]:.3f},'
    est_str += f' sigma_x = {est["sigma_x"]:.3f},'
    est_str += f' sigma_c = {est["sigma_c"]:.3f}'
    print(f'seed = {seed:4d}: {est_str}')
```

    seed = 1997:  phi = 0.958, sigma_x = 4.305, sigma_c = 0.231
    seed =    1:  phi = 0.969, sigma_x = 4.246, sigma_c = 0.217
    seed =   17:  phi = 0.975, sigma_x = 4.109, sigma_c = 0.206
    seed = 2018:  phi = 0.956, sigma_x = 4.027, sigma_c = 0.240
    seed =  999:  phi = 0.962, sigma_x = 3.665, sigma_c = 0.207

**Same solution:** Simulate more periods (and use the faster simulation function).

```python
T_alt = 1_000_000
for seed in seeds:

    # a. simulate
    sim_alt = prep_sim(par,T_alt,seed)

    # b. run optimizer
    est = par.copy()
    x0 = [0.98,par['sigma_x'],par['sigma_c']]
    result = optimize.minimize(obj_func_all,x0,args=(est,sim_alt,T_alt,ASAD.simulate))

    # c. update and print estimates
    est['phi'] = result.x[0]
    est['sigma_x'] = np.sqrt(result.x[1]**2)
    est['sigma_c'] = np.sqrt(result.x[2]**2)

    est_str = ''
    est_str += f' phi = {est["phi"]:.3f},'
    est_str += f' sigma_x = {est["sigma_x"]:.3f},'
    est_str += f' sigma_c = {est["sigma_c"]:.3f}'
    print(f'seed = {seed:4d}: {est_str}')
```

    seed = 1997:  phi = 0.990, sigma_x = 3.994, sigma_c = 0.222
    seed =    1:  phi = 0.990, sigma_x = 3.982, sigma_c = 0.222
    seed =   17:  phi = 0.989, sigma_x = 3.984, sigma_c = 0.222
    seed = 2018:  phi = 0.989, sigma_x = 3.974, sigma_c = 0.222
    seed =  999:  phi = 0.990, sigma_x = 3.990, sigma_c = 0.222

# 3. Exchange economy

```python
# a. parameters
N = 50_000
mu = np.array([3,2,1])
Sigma = np.array([[0.25, 0, 0], [0, 0.25, 0], [0, 0, 0.25]])
zeta = 1
gamma = 0.8

# b. random draws
seed = 1986
np.random.seed(seed)

# preferences
alphas = np.exp(np.random.multivariate_normal(mu,Sigma,size=N))
betas = alphas/np.reshape(np.sum(alphas,axis=1),(N,1))

# endowments
e1 = np.random.exponential(zeta,size=N)
e2 = np.random.exponential(zeta,size=N)
e3 = np.random.exponential(zeta,size=N)
```

## Question 1

```python
# a. calculate
budgetshares = betas

# b. histograms
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
for i in range(3):
    ax.hist(budgetshares[:,i],bins=100,alpha=0.7,density=True)
```

## Question 2

```python
def excess_demands(budgetshares,p1,p2,e1,e2,e3):
    """ calculate excess demands for good 1, 2 and 3

    Args:

        budgetshares (ndarray): budgetshares for each good for all consumers
        p1 (float): price of good 1
        p2 (float): price of good 2
        e1 (ndrarray): endowments of good 1
        e2 (ndrarray): endowments of good 2
        e3 (ndrarray): endowments of good 3

    Returns:

        ed_1 (float): excess demands for good 1
        ed_2 (float): excess demands for good 2
        ed_3 (float): excess demands for good 3

    """

    # a. income
    I = p1*e1+p2*e2+1*e3

    # b. demands
    demand_1 = np.sum(budgetshares[:,0]*I/p1)
    demand_2 = np.sum(budgetshares[:,1]*I/p2)
    demand_3 = np.sum(budgetshares[:,2]*I) # p3 = 1, numeraire

    # b. supply
    supply_1 = np.sum(e1)
    supply_2 = np.sum(e2)
    supply_3 = np.sum(e3)

    # c. excess demand
    ed_1 = demand_1-supply_1
    ed_2 = demand_2-supply_2
    ed_3 = demand_3-supply_3

    return ed_1,ed_2,ed_3
```

```python
# a. calculate on grid
K = 50

p1_vec = np.linspace(4,8,K)
p2_vec = np.linspace(0.5,8,K)
p1_mat,p2_mat = np.meshgrid(p1_vec,p2_vec,indexing='ij')

ed1_mat = np.empty((K,K))
ed2_mat = np.empty((K,K))
ed3_mat = np.empty((K,K))
for (i,p1),(j,p2) in product(enumerate(p1_vec),enumerate(p2_vec)):
    ed1_mat[i,j],ed2_mat[i,j],ed3_mat[i,j] = excess_demands(budgetshares,p1,p2,e1,e2,e3)

# b. plot
fig = plt.figure()
ax = fig.add_subplot(1,1,1,projection='3d')
ax.plot_wireframe(p1_mat,p2_mat,ed1_mat/N)
ax.plot_surface(p1_mat,p2_mat,np.zeros((K,K)),alpha=0.5,color='black',zorder=99)
ax.set_xlabel('$p_1$')
ax.set_ylabel('$p_2$')
ax.set_title('excess demand for good 1')
ax.invert_xaxis()
fig.tight_layout()

fig = plt.figure()
ax = fig.add_subplot(1,1,1,projection='3d')
ax.plot_wireframe(p1_mat,p2_mat,ed2_mat/N)
ax.plot_surface(p1_mat,p2_mat,np.zeros((K,K)),alpha=0.5,color='black',zorder=99)
ax.set_title('excess demand for good 2')
ax.set_xlabel('$p_1$')
ax.set_ylabel('$p_2$')
ax.invert_xaxis();
fig.tight_layout()

fig = plt.figure()
ax = fig.add_subplot(1,1,1,projection='3d')
ax.plot_wireframe(p1_mat,p2_mat,ed3_mat/N)
ax.plot_surface(p1_mat,p2_mat,np.zeros((K,K)),alpha=0.5,color='black',zorder=99)
ax.set_title('excess demand for good 3')
ax.set_xlabel('$p_1$')
ax.set_ylabel('$p_2$')
ax.invert_xaxis();
fig.tight_layout()
```

## Questions 3

**Function for finding the equilibrium:**

```python
def find_equilibrium(budgetshares,p1,p2,e1,e2,e3,kappa=0.5,eps=1e-8,maxiter=5000):
    """ find equilibrium prices

    Args:

        budgetshares (ndarray): budgetshares for each good for all consumers
        p1 (float): price of good 1
        p2 (float): price of good 2
        e1 (ndrarray): endowments of good 1
        e2 (ndrarray): endowments of good 2
        e3 (ndrarray): endowments of good 3
        kappa (float,optional): adjustment aggresivity parameter
        eps (float,optional): tolerance for convergence
        maxiter (int,optinal): maximum number of iteration

    Returns:

        p1 (ndarray): equilibrium price for good 1
        p2 (ndarray): equilibrium price for good 2

    """

    it = 0
    while True:

        # a. step 1: excess demands
        ed_1,ed_2,_ed3 = excess_demands(budgetshares,p1,p2,e1,e2,e3)

        # b: step 2: stop?
        if  (np.abs(ed_1) < eps and np.abs(ed_2) < eps) or it >= maxiter:
            print(f'(p1,p2) = [{p1:.4f} {p2:.4f}] -> excess demands = [{ed_1:.4f} {ed_2:.4f}] (iterations: {it})')
            break

        # c. step 3: update p1 and p2
        N = budgetshares.shape[0]
        p1 = p1 + kappa*ed_1/N
        p2 = p2 + kappa*ed_2/N

        # d. step 4: return
        it += 1

    return p1,p2
```

**Apply algorithm:**

```python
# a. guess prices are equal to the average beta
betas_mean = np.mean(betas,axis=0)
p1_guess,p2_guess,_p3_guess = betas_mean/betas_mean[-1]

# b. find equilibrium
p1,p2 = find_equilibrium(budgetshares,p1_guess,p2_guess,e1,e2,e3)
```

    (p1,p2) = [6.4901 2.6167] -> excess demands = [0.0000 0.0000] (iterations: 2416)

**Check excess demands:**

```python
assert np.all(np.abs(np.array(excess_demands(budgetshares,p1,p2,e1,e2,e3)) < 1e-6))
```

## Questions 4

```python
# a. income
I = p1*e1+p2*e2+1*e3

# b. baseline utility
x = budgetshares*np.reshape(I,(N,1))/np.array([p1,p2,1])
base_utility = np.prod(x**betas,axis=1)

# c. plot utility
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
for gamma_now in [gamma,1.0,1.2]:
    utility = base_utility**gamma_now
    ax.hist(utility,bins=100,density=True,
            label=f'$\gamma = {gamma_now:.2f}$: mean(u) = {np.mean(utility):.3f}, var(u) = {np.var(utility):.3f}')
ax.legend();
```

### Question 5

```python
# a. equalize endowments
e1_equal = np.repeat(np.mean(e1),N)
e2_equal = np.repeat(np.mean(e2),N)
e3_equal = np.repeat(np.mean(e3),N)
print(f'e_equal = [{e1_equal[0]:.2f},{e2_equal[0]:.2f},{e3_equal[0]:.2f}]')

# b. find equilibrium
p1_equal,p2_equal = find_equilibrium(budgetshares,p1_guess,p2_guess,e1_equal,e2_equal,e3_equal)
```

    e_equal = [1.00,0.99,1.00]
    (p1,p2) = [6.4860 2.6172] -> excess demands = [0.0000 0.0000] (iterations: 2404)

**Check excess demands:**

```python
assert np.all(np.abs(np.array(excess_demands(budgetshares,p1_equal,p2_equal,e1_equal,e2_equal,e3_equal))< 1e-6))
```

**Plot utility:**

```python
# a. income
I_equal = p1_equal*e1_equal+p2_equal*e2_equal+1*e3_equal

# b. baseline utility
x = budgetshares*np.reshape(I_equal,(N,1))/np.array([p1_equal,p2_equal,1])
base_utility_equal = np.prod(x**betas,axis=1)

# c. plot utility
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
for gamma_now in [gamma,1.0,1.2]:
    utility = base_utility_equal**gamma_now
    ax.hist(utility,bins=100,density=True,
            label=f'$\gamma = {gamma_now:.2f}$: mean(u) = {np.mean(utility):.3f}, var(u) = {np.var(utility):.3f}')
ax.legend();
```

**Compare prices with baseline:**

```python
print(f'baseline: p1 = {p1:.4f}, p2 = {p2:.4f}')
print(f'   equal: p1 = {p1_equal:.4f}, p2 = {p2_equal:.4f}')
```

    baseline: p1 = 6.4901, p2 = 2.6167
       equal: p1 = 6.4860, p2 = 2.6172

**Conclusions:** The relative prices of good 1 and good 2 *increase* slightly when endowments are equalized. (This can, however, be shown to disappear when $N \rightarrow \infty$.)

Economic behavior (demand and supply), and therefore equilibrium prices, are independent of $\gamma$, which thus only affects utility.

Irrespective of $\gamma$ we have:

1. Equalization of endowments implies a utility floor of approximately 1 because everyone then get approximately one unit of each good.
2. The variance of utility always decreases when endowments are equalized.
3. The remaining inequality in utility when endowments are equalized must be driven by preferences (see below).

The effect on the mean of utility of equalizing endowments depends on whether $\gamma$ is below, equal to or above one:

1. For $\gamma = 0.8 < 1.0$ the mean utility *increases* ("decreasing returns to scale").
2. For $\gamma = 1.0$ the mean utility is *unchanged* ("constant returns to scale").
3. For $\gamma = 1.2 > 1.0$ the utility mean *decreases* ("increasing returns to scale").

**Additional observation:** When endowments are equalized those with high utility have preferences which differ from the mean.

```python
for i in range(3):
    sns.jointplot(betas[:,i],base_utility_equal,kind='hex').set_axis_labels(f'$\\beta^j_{i+1}$','utility')
```


# Simon's Algorithm

**Abstract:** We study a quantum algorithm known as Simon's algorithm, which provided the first example of an exponential speedup over the best known classical algorithm by using a quantum computer to solve a particular problem. Originally published in 1994, Simon's algorithm was a precursor to Shor's well-known factoring algorithm, and it served as inspiration for many of the seminal works in quantum computation that followed.

This notebook is aimed at users with some basic knowledge of quantum computing and quantum circuits.

### Table of Contents
* [Simon's Problem Statement](#statement)
 * [Example for n=3](#example)
* [Classical Complexity](#classicalcomplexity)
* [Quantum Algorithm for Simon's Problem](#quantumalgorithm)
 * [Quantum Circuit](#circuit)
 * [Running Simon's Algorithm](#runningintro)
 * [Classical Post-Processing](#postprocessintro)
* [Quantum Complexity](#quantumcomplexity)
* [Implementing Simon's Algorithm in Amazon Braket](#implementation)
* [References](#references)
* [Appendix](#appendix)
 * [Implementing an Oracle Function](#oracleimplementation)
 * [Implementing the Classical Post-Processing](#appendixpostprocessing)

## Simon's Problem Statement:

Suppose we’re given a function $f:\{0,1\}^n \rightarrow \{0,1\}^n$ that maps bit strings to bit strings along with the promise that
$$\forall x,y \in \{0,1\}^n, \quad f(x) = f(y) \iff x=y\oplus s,$$
for some unknown $n$-bit string $s \in \{0,1\}^n$, and where $\oplus$ means bitwise addition modulo 2.

Said another way, there exists an unknown string $s$ such that, $\forall x, \; f(x)=f(x\oplus s)$. When $s$ is non-zero, the function is two-to-one as it maps *exactly* two inputs to every unique output.

The goal of Simon's problem is to determine if $f$ is one-to-one, or two-to-one, or equivalently to find the secret string $s$.

Since we're given the promise that $f(x)=f(y)\implies x=y\oplus s$, this means that $s=x\oplus y$ whenever $f(x)=f(y)$. Thus, one way to solve this problem is to find two inputs to the function $f$ that produce the *same* output; $s$ is then the XOR of those two input strings. See [[1]](#References) for more details.

Example for n=3:

Consider the function $f:\{0,1\}^3\to\{0,1\}^3$ defined by the truth table below.

            $$x$$
            $$f(x)$$

        $$000$$
        $$000$$

        $$001$$
        $$001$$

        $$010$$
        $$001$$

        $$011$$
        $$000$$

        $$100$$
        $$100$$

        $$101$$
        $$101$$

        $$110$$
        $$101$$

        $$111$$
        $$100$$

By inspection, we can see that $f$ satisfies the properties described in the statement of Simon's problem. In particular, note that each output $f(x)$ appears twice for two distinct inputs. We are given the promise that, for each of these two inputs $x$ and $y$ with the same output $f(x)=f(y)$, we have $x \oplus s = y$ for a yet to be determined $s$, and therefore $x\oplus y = s$.

For concreteness, notice that the input strings $001$ and $010$ are both mapped by $f$ to the same output string $001$. Taking the bitwise XOR of $001$ and $010$ we obtain the secret string $s$:

$$s=001 \oplus 010 = 011$$

Therefore, in this example, the secret string is $s = 011$.

In this specific example, we also see that the string $000$ is mapped to itself. Since $x\oplus y=s$ for two inputs $x$ and $y$ with the same output, we must have that $s$ is also mapped to $000$, since $s=000\oplus s$. Indeed, we see that $011$ maps to $000$, as expected.

## Classical Complexity

To solve Simon's problem classically, one needs to find two different inputs $x$ and $y$ for which $f(x)=f(y)$. As we saw above, one can then determine $s=x\oplus y$. How hard is it to find two distinct inputs that map to the same output, given the function $f$ as a black box?  For $n$-bit strings, there are $2^n$ possible inputs. Thus, in the worst case, one would need to check at most $2^n$ different inputs to find a pair that maps to the same output; this provides an upper bound on the required query complexity.

It turns out that a *lower* bound on the classical query complexity of Simon's algorithm can also be found: $\Omega ({\sqrt {2^{n}}})$. Proving this lower bound requires a little more work and is outside the scope of this notebook, so instead we will just provide some intuition.

As mentioned above, the goal is to find a pair of input strings $x$ and $y$ that map to the *same* output string $f(x)=f(y)$ -- a collision. Finding a collision in a set is an instance of the well-known (generalized) birthday problem [[2]](#References): within a group of people, what is the probability that two of them share the same birthday? One can turn this problem around and ask "how many people do we need in a room to ensure that the probability that at least two of them share a birthday is greater than some fixed number?" This latter question gets to the heart of solving Simon's problem classically: how many queries to the function $f$ do we need to make to guarantee that we find a collision, with high probability? In the case of the birthday problem, we would need enough people in the room so that when we generate all possible pairings of people, we would have about 365 possible pairs. That way, we'd have a good chance that at least one of those pairs of people share a birthday. Using this intuition, we need to query $f$ enough times to generate a set of *pairs* with roughly the same size as the number of possible inputs ($2^n$). If we make $k$ queries to the function $f$, we can generate ${k \choose 2}=\frac{k(k-1)}{2}\sim k^2$ pairs. Thus, we need to make $k$ queries such that ${k\choose 2}\gtrsim 2^n$ to have a high probability of generating a collision, and therefore, $k>\Omega(\sqrt{2^n})$.

## Quantum Algorithm for Simon's Problem

Simon's algorithm is a scheme for solving the problem above using exponentially fewer queries to the function $f$. In order for Simon's algorithm to work, one needs to be able to implement the unknown function $f$ using quantum logic. That is, given an input *quantum state* $|x\rangle$, one needs a *unitary* $U_f$ satisfying
$$U_f|x\rangle |0\rangle = |x\rangle |f(x)\rangle.$$
This unitary is an *oracle* for $f$, and the goal is to query it as few times as possible to learn the secret string $s$.

### Quantum Circuit

Simon's algorithm involves both quantum and classical components. The quantum part of Simon's algorithm is used to query the oracle efficiently, while the classical component is used to process measurement results and determine the hidden string $s$. A circuit for the quantum component of Simon's algorithm is shown below.

For a function $f$ acting on $n$-bit strings, the circuit above acts on $2n$ qubits, as needed for the definition of $U_f$. Only the first $n$ qubits are measured; the remaining qubits are unused after the application of $U_f$.

### Running Simon's Algorithm
To solve Simon's problem, one needs to run the quantum circuit above several times. After each run of the circuit, the measurements of the first $n$ qubits produce an output bit string, which we denote by $z$.

An analysis of the circuit above shows that each output bit string $z$ satisfies the following condition:
$$ z\cdot s = 0 \; \mod{2}.$$

Let us now analyze the above circuit step-by-step:
1. Initialize all qubits in the $|0\rangle$ state. That is, we start in the state $|0\rangle^{\otimes n} \otimes |0\rangle^{\otimes n}$. We will use the shorthand $|0\rangle^{\otimes n}\equiv |0^n\rangle$
2. Apply Hadamard gates to each of the first $n$ qubits, placing them in the equal superposition state: $$\frac{1}{\sqrt{2^n}}\sum_{x \in \{0, 1\}^n} |x\rangle |0^n\rangle$$.
3. Apply the oracle $U_f$, which computes the function $f$ into the last $n$ qubits, giving the state $$\frac{1}{\sqrt{2^n}}\sum_{x \in \{0, 1\}^n} |x\rangle |f(x)\rangle$$
4. Measure the last $n$ qubits, giving a random result $f(x)$. If $f$ is one-to-one, this output of $f$ corresponds to an input of $x$. If $f$ is two-to-one, the output of $f$ corresponds to an input of either either $x$ or $y = x \oplus s$, where $x$ and $y$ are the two different inputs to $f$ that gave the *same* output $f(x)=f(y)$. Hence we are left with the first $n$ qubits in the state;
    A. $|x\rangle$,                    if $f$ is one-to-one;
    B. $\frac{1}{\sqrt{2}} (|x\rangle + |y\rangle)$, where $x \oplus y = s$,  if $f$ is two-to-one.
    Note that this step is not strictly necessary, since we do not need the measurement result, but we include it as it makes the analysis easier.
5. Apply Hadamard gates to each of the first $n$ qubits. If $f$ is one-to-one, the state $|x\rangle$ is mapped to $$H^{\otimes n} |x\rangle = \frac{1}{\sqrt{2^n}} \sum_{z \in \{0, 1\}^n} (-1)^{x \cdot z}|z\rangle,$$
   where $x\cdot z$ is the dot product between the two strings represented as vectors (modulo 2).

   Similarly, if $f$ is two-to-one, the state $\frac{1}{\sqrt{2}}(|x\rangle + |y\rangle)$ is mapped to $$\frac{1}{\sqrt{2^{n+1}}} \sum_{z \in \{0, 1\}^n} [(-1)^{x \cdot z} + (-1)^{y \cdot z}]|z\rangle$$.
6. Measure the first $n$ qubits.
   1. If $f$ is one-to-one, measurements return a random bit string $z$ uniformly chosen from $\{0,1\}^n$.
   2. If $f$ is two-to-one, measurements return a random bit string $z$ such that $x \cdot z = y \cdot z \,\mathrm{mod}\, 2$, since otherwise the amplitude $(-1)^{x \cdot z} + (-1)^{y \cdot z}$ cancels out. Using the criterion for Simons problem (i.e,. $f(x)=f(y) \implies x=y\oplus s$), we find that
\begin{align*}
x\cdot z &= y\cdot z & \mod{2}\\
x\cdot z &= (x\oplus s)\cdot z &  \mod{2}\\
x\cdot z &= x\cdot z\oplus s\cdot z & \mod{2}\\
0 &= s\cdot z &  \mod{2}
\end{align*}

   Thus, in both cases we obtain a random bit string $z$ such that $s \cdot z = 0$.

Therefore, each time we run the quantum circuit above, we find a bit string $z$ that is orthogonal to the secret string $s$.

### Classical Post-Processing
From the measurement results $\{z_1, \dots, z_k\}$, we can form a system of equations:
$$ \begin{aligned}z_{1}\cdot s&=0\mod{2}\\z_{2}\cdot s&=0\mod{2}\\&\,\,\vdots \\z_{k}\cdot s&=0\mod{2}\end{aligned}$$

There are $k$ equations and $n$ unknowns (the elements of $s$). If we run the quantum part enough times so that we find $n$ **independent** equations, then we can solve these equations (using, e.g., Gaussian elimination) to recover the secret string $s$. This is precisely the classical post-processing required: solve the system of equations found above to recover the string $s$. We refer the interested reader to the [Appendix](#Classical-post-processing) for details.

## Quantum Complexity
How many queries do we need to make to $U_f$ in the quantum case? Above we saw that we need to run the quantum part of the algorithm $k$ times to generate a system of equations. We need to find $n$ linearly independent equations for the system to be determined. Thus, we need at least $n$ queries to $U_f$ to find such a system. It is possible, however, that we will get the same measurement outcome on different runs of the quantum algorithm, so we would need to re-do those runs that do not produce distinct measurement outcomes. Fortunately, these repeated outcomes are unlikely, so we only need $O(n)$ queries to the oracle $U_f$.

Comparing the quantum and classical algorithms, we saw that the classical algorithm requires at least $\Omega(2^{0.5 n})$ queries to $f$, whereas the quantum algorithm requires only $O(n)$. Thus, we have established an *exponential* speedup by using the quantum algorithm above.

# Implementing Simon's Algorithm in Amazon Braket

```python
# Imports and Setup
from braket.circuits import Circuit, circuit
from braket.devices import LocalSimulator
import numpy as np

import matplotlib.pyplot as plt
%matplotlib inline

# Sets the device to run the circuit on
device = LocalSimulator()
```

We also import a method called `simons_oracle`, which generates a circuit implementing an example oracle function. The code for `simons_oracle` is defined in the `simons_utils.py` module, and it is shown in the [Appendix](#Implementing-an-oracle-function) for completeness.

```python
#Import local utils
from simons_utils import simons_oracle
```

We now define the secret string, $s$:

```python
s = '101011'

# Other examples to try:
# s = '011'
# s = '00000'
# s = '1'
# Generate a random string of random length from 1 to 10:
# s="".join(str(np.random.randint(2)) for _ in range(np.random.randint(1,10)))

print("The secret string is: "+ s)
```

    The secret string is: 101011

### Circuit Definition

We now define the quantum circuit for Simon's algorithm:
1. Apply Hadamard gates to the first $n$-qubits.

2. Query the oracle (i.e., the $U_f$ gate). In this example, the oracle is defined dynamically, based on our chosen value of $s$. You can try experimenting with different values of $s$ (with differing lengths).

3. Apply Hadamard gates to the first $n$-qubits.

```python
n = len(s)

circ = Circuit()

# Apply Hadamard gates to first n qubits
circ.h(range(n))

# Now apply the Oracle for f
circ.simons_oracle(s)

# Apply Hadamard gates to the first n qubits
circ.h(range(n))

print(circ)
```

    T   : |0|     1     | 2 |3|4|5|6|

    q0  : -H-C-----------C---C-C-C-H-
             |           |   | | |
    q1  : -H-|-C---------|-H-|-|-|---
             | |         |   | | |
    q2  : -H-|-|-C-------|-H-|-|-|---
             | | |       |   | | |
    q3  : -H-|-|-|-C-----|-H-|-|-|---
             | | | |     |   | | |
    q4  : -H-|-|-|-|-C---|-H-|-|-|---
             | | | | |   |   | | |
    q5  : -H-|-|-|-|-|-C-|-H-|-|-|---
             | | | | | | |   | | |
    q6  : ---X-|-|-|-|-|-X---|-|-|---
               | | | | |     | | |
    q7  : -----X-|-|-|-|-----|-|-|---
                 | | | |     | | |
    q8  : -------X-|-|-|-----X-|-|---
                   | | |       | |
    q9  : ---------X-|-|-------|-|---
                     | |       | |
    q10 : -----------X-|-------X-|---
                       |         |
    q11 : -------------X---------X---

    T   : |0|     1     | 2 |3|4|5|6|

### Now run the circuit

We need enough shots to obtain $n$ linearly independent bit strings in the output measurements. We have chosen `4n` shots in the example below, just to be on the safe side.

```python
task = device.run(circ, shots=4*n)
```

### Analyze the results

We can retrieve the measurement results on all $2n$ qubits as follows:

```python
result = task.result()

counts = result.measurement_counts
plt.bar(counts.keys(), counts.values());
plt.xlabel('bit strings');
plt.ylabel('counts');
plt.xticks(rotation=90)
plt.show()
```

#### Aggregate the results

The measurements are performed on all $2n$ qubits, but we are only interested in the first $n$ qubits. Thus, we need to aggregate the results by ignoring the measurement outcomes on the last $n$ qubits:

```python
new_results = {}
for bitstring, count in result.measurement_counts.items():
    # Only keep the outcomes on first n qubits
    trunc_bitstring = bitstring[:n]
    # Add the count to that of the of truncated bit string
    new_results[trunc_bitstring] = new_results.get(trunc_bitstring, 0) + count

plt.bar(new_results.keys(), new_results.values())
plt.xlabel('bit strings');
plt.ylabel('counts');
plt.xticks(rotation=70)
plt.show()
```

In practice, we only need the measurement results (i.e., the bit strings, not the counts) from the first $n$ qubits. These measurement outcomes correspond to bit strings that satisfy the equations:
$$ \begin{aligned}z_{1}\cdot s&=0\mod{2}\\z_{2}\cdot s&=0\mod{2}\\&\,\,\vdots \\z_n\cdot s&=0\mod{2}\end{aligned}$$

With these $n$ linear equations in hand, we can use classical post-processing to solve for the unknown string $s$.

Note that we may have too many bit strings in the above, since we ran the task with $2n$ shots just to be safe. In this case, not all of the bit strings will be linearly independent. Moreover, the all-zeros string $0\dots0$ may also be an outcome, but this bit string satisfies the equations above trivially, so we exclude it if needed.

At this stage, the quantum portion of Simon's algorithm is complete. Any remaining steps are just classical postprocessing, which we cover in the Appendix.

### References
[1] Wikipedia: [Simon's Problem](https://en.wikipedia.org/wiki/Simon%27s_problem)

[2] Wikipedia: [Birthday Problem](https://en.wikipedia.org/wiki/Birthday_problem)

[3] Wikipedia: [Computing a kernel by Gaussian elimination](https://en.wikipedia.org/wiki/Kernel_(linear_algebra)#Computation_by_Gaussian_elimination)

[4] StackExchange: [Sympy: Solving Matrices in a finite field](https://stackoverflow.com/questions/31190182/sympy-solving-matrices-in-a-finite-field)

---

## Appendix

### Implementing an Oracle Function

In order to run the algorithm, we will need a unitary function that we can use as an oracle to query the function $f$.

There are many possible ways of implementing a function with the desired property that $f(x)=f(y) \implies x=y\oplus s$. We will pick one implementation that is commonly used in example code, and we will try to give some intuition for why this oracle works.

#### Classical Intuition Behind the Function $f$
Generating a function that is one-to-one is conceptually straightforward, as any such function of the bit strings $\{0,1\}^n$ will just be a permutation of the inputs. Generating a two-to-one function is a little trickier, though there are many ways to do it. The goal is to define a function that splits the inputs into two groups, such that one element from each group maps to the same output (i.e., $x$ must be in one group, while $x\oplus s$ must be in the other group.)

We will implement one simple choice for $f$, in which we define the split based on the value of one of the bits in the string. In this way, exactly half of the inputs will have that bit with value $0$, while the other half will have that bit with value $1$.

Our approach will be to choose a flag bit in the input bit strings that we will use to split the inputs. We then $\mathrm{XOR}$ the input string with $s$ whenever the flag bit is $1$. With this definition, half of the input strings will be untouched, while half of the strings will be $\mathrm{XOR}$'ed with $s$. Clearly, this function does nothing to the all-zeros string $0\dots 0$, since any choice of the flag bit will always be $0$. Thus, we need to ensure that our definition also maps the string $s$ to the all-zeros string $0\dots 0$. In other words, we need to ensure that our flag bit is $1$ when the input string is $s$. One way to ensure the function acts correctly on the input $s$ is to just define the flag bit to be the first bit in the string $s$ that is equal to $1$. For example, if $s=011$, we can choose the flag bit to be the second bit. Concretely:
$$f(x) = \left\{\begin{array}{lr}
x, & \text{if } x_j=0\\
x\oplus s, & \text{if } x_j=1\\
\end{array}\right\},$$
where $x_j$ is the $j^\text{th}$ bit of $x$, and $j$ is the flag bit in $s$.

Example for n=3:

We now revisit the example in the introduction. Suppose the secret string $s=011$. Since the first appearance of $1$ in $s$ occurs at the second bit, we will use the second bit in the input strings as our flag bit. We take the $\mathrm{XOR}$ of the input with $s$ whenever the flag bit in the input is 1. This definition results in the following truth table:

            $$x$$
            $$f(x)$$

        $$000$$
        $$000$$

        $$001$$
        $$001$$

        $$010$$
        $$001$$

        $$011$$
        $$000$$

        $$100$$
        $$100$$

        $$101$$
        $$101$$

        $$110$$
        $$101$$

        $$111$$
        $$100$$

We leave it as an exercise to the reader to verify that this definition works for any input string size (i.e., $n$ for inputs $\{0,1\}^n$, and that it is in fact two-to-one, rather than many-to-one. Note that the function defined in this way is not a general two-to-one function, but it is simple a choice that is easy to implement both classically and as a quantum circuit.

#### Quantum Implementation of $U_f$
We now define the unitary using the `@circuit.subroutine` functionality of the Amazon Braket SDK. The following code was imported from the `simons_utils.py` module, and is shown below for reference.

In the quantum setting, we first copy the input register into some ancillary qubits:
$$ |x\rangle|0\rangle\mapsto |x\rangle|x\rangle.$$
We then perform the quantum analog of $\mathrm{XOR}$, which means we apply an $X$ gate to the $k^\text{th}$ qubit whenever the $k^\text{th}$ bit of $s$ is $1$. However, we only apply this $X$ gate when the flag qubit is also $|1\rangle$. Thus, our $X$ gate becomes a $\mathrm{CNOT}$ gate between the flag qubit on the input register, and the $k^\text{th}$ qubit on the output.

```python
from braket.circuits import Circuit, circuit

@circuit.subroutine(register=True)
def simons_oracle(secret_s: str):
    """
    Quantum circuit implementing a particular oracle for Simon's problem. Details of this implementation are
    explained in the Simons Algorithm demo notebook.

    Args:
        secret_s (str): secret string we wish to find
    """
    # Find the index of the first 1 in s, to be used as the flag bit
    flag_bit=secret_s.find('1')

    n=len(secret_s)

    circ = Circuit()
    # First copy the first n qubits, so that |x>|0> -> |x>|x>
    for i in range(n):
        circ.cnot(i, i+n)

    # If flag_bit=-1, s is the all-zeros string, and we do nothing else.
    if flag_bit != -1:
        # Now apply the XOR with s whenever the flag bit is 1.
        for index,bit_value in enumerate(secret_s):

            if bit_value not in ['0','1']:
                raise Exception ('Incorrect char \'' + bit_value + '\' in secret string s:' + secret_s)

            # XOR with s whenever the flag bit is 1.
            # In terms of gates, XOR means we apply an X gate only whenever the corresponding bit in s is 1.
            # Applying this X only when the flag qubit is 1 means this is a CNOT gate.
            if(bit_value == '1'):
                circ.cnot(flag_bit,index+n)
    return circ
```

### Implementing the Classical Post-Processing

We will now solve the system of linear equations above using Gaussian elimination. We first convert the results into matrix form, then we use `sympy`'s `Matrix.rref()` method to transform the matrix into reduced row echelon form.

```python
!pip3 install sympy --quiet
from sympy import Matrix
```

Generate a matrix from the bit string outputs. We first check that we have sufficiently many output strings to be able to solve the system of equations. If not: output and error and re-run the algorithm.

```python
if len(new_results.keys()) < len(s):
    raise Exception ('System will be underdetermined. Minimum ' + str(n) + ' bistrings needed, but only '
                     + str(len(new_results.keys())) +' returned. Please rerun Simon\'s algorithm.')
string_list = []

for key in new_results.keys():
#     if key!= "0"*n:
    string_list.append( [ int(c) for c in key ] )

print('The result in matrix form is :')
for a in string_list:
    print (a)
```

    The result in matrix form is :
    [1, 0, 0, 0, 0, 1]
    [0, 1, 0, 0, 0, 0]
    [0, 0, 1, 0, 0, 1]
    [1, 1, 1, 1, 1, 1]
    [0, 1, 1, 0, 1, 0]
    [0, 1, 1, 1, 1, 0]
    [1, 0, 1, 0, 1, 1]
    [0, 1, 0, 0, 1, 1]
    [1, 1, 0, 1, 0, 1]
    [1, 1, 0, 0, 0, 1]
    [0, 0, 1, 1, 0, 1]
    [0, 1, 0, 1, 0, 0]
    [0, 1, 1, 0, 0, 1]
    [0, 0, 1, 0, 1, 0]
    [1, 0, 1, 1, 1, 1]
    [1, 0, 1, 0, 0, 0]
    [0, 1, 1, 1, 0, 1]
    [1, 0, 0, 1, 0, 1]
    [1, 0, 0, 1, 1, 0]
    [0, 0, 0, 0, 0, 0]

Now solve the system $Ms=0$ by finding the kernel of $M$. We do this using Gaussian elimination on the augmented matrix $\left[A|I\right]$ to bring it to row echelon form. Converting the solution to numbers $\mathrm{mod }\,2$, we can then read off the solution from the last row of the reduced matrix. See [[3]](#References) and [[4]](#References) for more details.

```python
M=Matrix(string_list).T

# Construct the agumented matrix
M_I = Matrix(np.hstack([M,np.eye(M.shape[0],dtype=int)]))

# Perform row reduction, working modulo 2. We use the iszerofunc property of rref
# to perform the Gaussian elimination over the finite field.
M_I_rref = M_I.rref(iszerofunc=lambda x: x % 2==0)

# In row reduced echelon form, we can end up with a solution outside of the finite field {0,1}.
# Thus, we need to revert the matrix back to this field by treating fractions as a modular inverse.
# Since the denominator will always be odd (i.e. 1 mod 2), it can be ignored.

# Helper function to treat fractions as modular inverse:
def mod2(x):
    return x.as_numer_denom()[0] % 2

# Apply our helper function to the matrix
M_I_final = M_I_rref[0].applyfunc(mod2)

# Extract the kernel of M from the remaining columns of the last row, when s is nonzero.
if all(value == 0 for value in M_I_final[-1,:M.shape[1]]):
    result_s="".join(str(c) for c in M_I_final[-1,M.shape[1]:])

# Otherwise, the sub-matrix will be full rank, so just set s=0...0
else:
    result_s='0'*M.shape[0]

# Check whether result_s is equal to initial s:
print ('Secret string: ' + s)
print ('Result string: ' + result_s)
if (result_s == s):
    print ('We found the correct answer.')
else:
    print ('Error. The answer is wrong!')
```

    Secret string: 101011
    Result string: 101011
    We found the correct answer.