# Acoustic radiation force on a single sound-hard sphere

### Arnold D. Kim
*Department of Applied Mathematics, University of California, Merced*

```python
import  time
print( 'Last updated: %s' %time.strftime('%d/%m/%Y') )
```

    Last updated: 11/08/2021

```python
# defaults for the codes below

import numpy as np
import scipy.special as sp
import matplotlib.pyplot as plt
```

## 1 Problem description

We compute the solution of the scattering problem for the scalar potential, $\phi_{1}$:
\begin{gather*}
   \varDelta \phi_{1} + k^{2} \phi_{1} = 0, \quad \text{in $r > a$},\\
   \partial_{r} \phi_{1} = 0 \quad \text{on $r = a$}.
\end{gather*}
The incident field is given by counter-propagating plane waves,
\begin{equation}
   \phi_{1}^{\text{inc}} = e^{\mathrm{i} k z} + e^{-\mathrm{i} k ( z + \delta )},
\end{equation}
with $\delta$ denoting a parameter that sets the location of the nodes of the standing wave. This parameter is used instead of moving the sphere so that the coordinate system can be fixed at the sphere's center.

We solve this scattering problem using two methods. The first uses the analytical representation as an expansion in Legendre polynomials. The second method uses the method of fundamental solutions (MFS). Through a comparison of these two results, we validate the use of the MFS to solve this problem.

Using the results from this scattering problem, we compute the acoustic radiation force defined according to (see Eq. (10) of H. Bruus, "Acoustofluidics 7: The acoustic radiation force on small particles," Lab Chip **12**, 1014 -- 1021 (2012))
\begin{equation}
   \vec{F} = -\int_{\partial \Omega} \left\{ \left[ \frac{1}{2} \kappa_{0} \langle p_{1}^{2} \rangle
   - \frac{1}{2} \rho_{0} \langle v_{1}^{2} \rangle \right] \hat{n} + \rho_{0} \langle ( \hat{n} \cdot \vec{v}_{1} )
   \vec{v}_{1} \rangle \right\} \mathrm{d}S.
\end{equation}
Here $p_{1}$ is the pressure field defined according to
\begin{equation}
   p_{1} = \mathrm{i} \rho_{0} \omega \phi_{1}
\end{equation}
$v_{1}$ is the velocity field defined according to
\begin{equation}
   v_{1} = \nabla \phi_{1},
\end{equation}
and the notation $\langle \cdot \rangle$ denotes a time-averaging operation.

Note that due to the sound-hard assumption, $\hat{n} \cdot \vec{v}_{1} = 0$ identically. Therefore, this expression for the force reduces to
\begin{equation}
   \vec{F} = -\int_{\partial \Omega} \left\{ \left[ \frac{1}{2} \kappa_{0} \langle p_{1}^{2} \rangle
   - \frac{1}{2} \rho_{0} \langle v_{1}^{2} \rangle \right] \hat{n} \right\} \mathrm{d}S.
\end{equation}

## 2 Harmonic expansion for the scattered scalar potential

Here we compute the solution of the scattering problem as a harmonic expansion. This solution method is explained below.

### 2.1 Harmonic expansion for counter-propagating plane waves

We consider the following harmonic expansion for a plane wave (see Jackson (10.45)),
\begin{align}
   e^{\mathrm{i} k z} &= \sum_{n = 0}^{\infty} (\mathrm{i})^{n} \sqrt{4 \pi (2 n + 1)} j_{n}(k r)
   Y_{n0}(\theta,\varphi) \\
   &= \sum_{n = 0}^{\infty} e^{\mathrm{i} n \pi/2} ( 2 n + 1 ) j_{n}(k r)
   P_{n}(\cos\theta).
\end{align}
It follows that counter-propagating plane waves has the following harmonic expansion,
\begin{equation}
   e^{\mathrm{i} k z} + e^{-\mathrm{i} k ( z + \delta )} =
   \sum_{n = 0}^{\infty} \left( e^{\mathrm{i} n \pi/2} + e^{-\mathrm{i} n \pi/2} e^{-\mathrm{i} k \delta} \right)
   ( 2 n + 1 ) j_{n}(k r) P_{n}(\cos\theta).
\end{equation}
Using this expansion, we find that
\begin{equation}
   \partial_{r} \left( e^{\mathrm{i} k z} + e^{-\mathrm{i} k ( z + \delta )} \right) =
   \sum_{n = 0}^{\infty} \left( e^{\mathrm{i} n \pi/2} + e^{-\mathrm{i} n \pi/2} e^{-\mathrm{i} k \delta} \right)
   ( 2 n + 1 ) \partial_{r} j_{n}(k r) P_{n}(\cos\theta).
\end{equation}
Let
\begin{equation}
   B_{n} = \left( e^{\mathrm{i} n \pi/2} + e^{-\mathrm{i} n \pi/2} e^{-\mathrm{i} k \delta} \right)
   ( 2 n + 1 ).
\end{equation}
This expansion is to be used as the normal derivative of the incident field on a sound-hard sphere.

The function below computes these expansion coefficients.

```python
def ComputeIncidentHarmonicExpansionCoeffs( k, a, δ, N ):
    """
    This function computes the expansion coefficients for the field
    scattered by a sound-hard sphere due to counter-propagating plane
    waves incident on it.
    """
    # compute an array of indices from n = 0 to n = N

    n = np.arange( N )

    # compute the expansion coefficients

    B_n = ( np.exp( 1j * n * np.pi / 2.0 ) + np.exp( -1j * n * np.pi / 2.0 ) * np.exp( -1j * k * δ ) ) \
        * ( 2 * n + 1 )

    return B_n;
```

### 2.2 Harmonic expansion for the scattered field

We write the field scattered by a sphere as the harmonic expansion,
\begin{equation}
   \phi^{s} = \sum_{n = 0}^{\infty} \sum_{m = -n}^{n} c_{nm} h_{n}^{(1)}(k r) Y_{nm}(\theta,\varphi).
\end{equation}
For the counter-propagating plane waves defined above, the scattered field will be rotationally symmetric about the $z$-axis and therefore this harmonic expansion for the scattered field reduces to
\begin{equation}
   \phi^{s} = \sum_{n = 0}^{\infty} \sqrt{\frac{2n+1}{4\pi}} C_{n} h_{n}^{(1)}(k r) P_{n}(\cos\theta).
\end{equation}
It follows that
\begin{equation}
   \partial_{r} \phi^{s} = \sum_{n = 0}^{\infty} \sqrt{\frac{2n+1}{4\pi}} C_{n} \partial_{r}
   h_{n}^{(1)}(k r) P_{n}(\cos\theta).
\end{equation}
For a sound-hard sphere, $\partial_{r} ( \phi^{\text{inc}} + \phi^{s} ) = 0$ on $r = a$. Thus, we find that
\begin{equation}
   C_{n} = -\sqrt{\frac{4\pi}{2n+1}} \frac{\partial_{r} j_{n}(k a)}{\partial_{r} h_{n}^{(1)}(k a)} B_{n}.
\end{equation}
The function below computes these expansion coefficients.

```python
def ComputeScatteredHarmonicExpansionCoeffs( k, a, δ, N ):
    """
    This function computes the expansion coefficients for the field
    scattered by a sound-hard sphere due to counter-propagating plane
    waves incident on it.
    """
    # compute the expansion coefficients for the incident field

    B_n = ComputeIncidentHarmonicExpansionCoeffs( k, a, δ, N )

    # compute an array of indices from n = 0 to n = N

    n = np.arange( N )

    # compute the derivative of the spherical Bessel function

    Djn = k * sp.spherical_jn( n, k * a, derivative = True )

    # compute the derivative of the spherical Hankel function

    Dh1n = k * sp.spherical_jn( n, k * a, derivative = True ) \
        + k * 1j * sp.spherical_yn( n, k * a, derivative = True )

    # compute the expansion coefficients

    C_n = -np.sqrt( 4.0 * np.pi / ( 2 * n + 1 ) ) * Djn / Dh1n * B_n

    return C_n;
```

### 2.3 Compute Legendre polynomials

The code below computes a matrix whose $n$th column is the evaluation of the Legendre polynomial of degree $n$ on the grid points given by the vector $\mu$.

```python
def ComputeLegendrePolynomials( μ, N ):
    """
    This function computes the Gauss-Legendre quadrature rule of order N and evaluates
    the Legendre polynomials at the Gauss-Legendre quadrature points up to degree N-1.

    This function returns the Gauss-Legendre quadrature points mu, quadrature weights wt,
    and a matrix whose columns are the Legendre polynomials evaluated on mu.
    """

    # allocate memory for the matrix of Legendre polynomials

    Pn = np.full( ( μ.size, N ), 'nan', dtype = 'complex' )

    # compute the Legendre polynomial of degree 0

    Pn[:,0] = 1

    # compute the Legendre polynomial of degree 1

    Pn[:,1] = μ

    # compute the remaining Legendre polynomials using the recursion relation

    for n in range( 1, N - 1 ):

        Pn[:,n+1] = ( ( 2 * n + 1 ) * μ * Pn[:,n] - n * Pn[:,n-1] ) / ( n + 1 )

    return Pn;
```

### 2.4 Compute the derivatives of the Legendre polynomials

The code below computes the matrix whose $n$th column is the evaluation of the derivative of the Legendre polynomial of degree $n$ on the grid points given by the vector $\mu$.

```python
def ComputeLegendrePolynomialDerivatives( μ, N ):
    """
    This function computes the derivatives of the Legendre polynomials.
    """

    # allocate memory for the matrix of Legendre polynomials

    DPn = np.full( ( μ.size, N ), 'nan', dtype = 'complex' )

    # compute the derivatives of the Legendre polynomials

    DPn[:,0] = 0;
    DPn[:,1] = 1;

    # compute the current Legendre polynomial of degree n = 2

    P_nminus1 = μ
    P_n = 1.5 * μ ** 2 - 0.5

    for n in range( 2, N ):

        # compute the Legendre polynomial of degree n+1

        P_nplus1 = ( ( 2 * n + 1 ) * μ * P_n - n * P_nminus1 ) / ( n + 1 )

        # compute the derivative of the Legendre polynomial

        DPn[:,n] = ( n + 1 ) / ( μ ** 2 - 1.0 ) * ( P_nplus1 - μ * P_n )

        # update the Legendre polynomials

        P_nminus1 = P_n
        P_n       = P_nplus1

    return DPn;
```

## 3 The Method of Fundamental Solutions (MFS)

Here, we compute the solution of the scattering problem using the Method of Fundamental Solutions (MFS).

### 3.1 Compute the MFS points

We compute the Fibonnaci lattice (see A. Gonzalez, "Measurement of areas using Fibonnaci and latitude-longitude lattices," Math. Geosci. **42**, 49 (2010)) and use the resulting points on the sphere for the MFS.

```python
def ComputeMFSReferencePoints( a, N ):
    """
    This function computes the set of points for a reference sphere that is to be used in the method of
    fundamental solutions (MFS).

    In particular, given the radius of a sphere, a, and the number of points N, this function computes
    randomly distributed points on the unit sphere and stores them as the unit normal vectors ν. Using
    ν, we then compute ρ_bdy = a * ν, and ρ_src = ( a - ℓ ) * ν.

    This function outputs three vectors: ν, ρ_bdy, and ρ_src.
    """
    # allocate memory for the Fibonacci lattice points on the unit sphere

    ν = np.full( ( N, 3 ), float( 'nan' ) )

    # compute the "golden angle"

    golden_angle = np.pi * ( 3.0 - np.sqrt( 5.0 ) )

    # compute the points on the unit sphere

    ν[:,2] = ( 1 - 1 / N ) * ( 1 - 2 * np.arange( 0, N ) / ( N - 1 ) )

    ρ = np.sqrt( 1.0 - ν[:,2] ** 2 )
    θ = golden_angle * np.arange( 0, N )

    ν[:,0] = ρ * np.cos( θ )
    ν[:,1] = ρ * np.sin( θ )

    # compute the boundary points, interior points, and exterior points

    ℓ = 0.4 * a

    ρ_bdy = a * ν
    ρ_src = ( a - ℓ ) * ν

    return ν, ρ_bdy, ρ_src;
```

Nick's note: l looks kind of larger than i expected it to be?...I thought it would be much smaller compared to the particle size?

### 3.2 Compute the fundamental solution and its normal derivative

The functions below compute Green's function in a medium with wavenumber $k$ and its normal derivative.

```python
# function to compute Green's function

def ComputeG( k, R ):
    """
    This function computes the whole space Green's function given a wavenumber k and a distance R.

    The output of this function is the evaluation of the whole space Green's function.
    """

    # compute Green's function

    G = np.exp( 1j * k * R ) / ( 4 * np.pi * R )

    return G;

# function to compute the normal derivative of Green's function

def ComputeDνG( k, R, CosΘ ):
    """
    This function computes the normal derivative of the whole space Green's function
    given a wavenumber k and distance R, and the cosine of the angle made between the
    difference vector and the unit normal.

    The output of this function is the evaluation of the normal derivative of the whole
    space Green's function.
    """

    # compute Green's function

    G = ComputeG( k, R )

    # compute the normal derivative of Green's function

    DνG = CosΘ * ( 1j * k - 1 / R ) * G

    return DνG;
```

### 3.3 Compute the MFS expansion coefficients

The function below constructs and solves the linear system for the MFS expansion coefficients. This linear system results from requiring that the sound-hard boundary condition on the sphere is satisfied exactly on the Fibonnaci lattice points.

```python
def ComputeMFSExpansionCoefficients( k, δ, ν, ρ_bdy, ρ_sca, N ):
    """
    This function solves the linear system of equations for the MFS expansion coefficients.

    This code requires a wavenumber k, the results from ComputeMFSPoints, namely ν, ρ_bdy,
    and ρ_ext in addition to the number of MFS points, N, and the number of scattering
    objects, Nobjs.

    The output from this code are the N x Nobs matrix, c, corresponding to the MFS expansion
    coefficients for the scattered field.
    """

    # compute a mesh grid of indices to compute the matrices below

    indx, jndx = np.mgrid[ 0 : N , 0 : N ]

    # allocate memory for the linear system

    b = np.full( ( N, 1 ), 'nan', dtype = 'complex' )
    A = np.full( ( N, N ), 'nan', dtype = 'complex' )

    # compute the incident field (normal derivative of two, counter-propagating plane waves)

    b = -1j * k * ν[:,2] * ( np.exp( 1j * k * ρ_bdy[:,2] ) - np.exp( - 1j * k * ( ρ_bdy[:,2] + δ ) ) )

    # compute the distance from scatter points on object q to boundary points on object p

    R = np.sqrt( ( ( ρ_bdy[indx,:] - ρ_sca[jndx,:] ) ** 2 ).sum( axis = 2 ) )

    # compute the cosine of the angle made between the normal on object p and R

    CosΘ = ( ν[indx,:] * ( ρ_bdy[indx,:] - ρ_sca[jndx,:] ) ).sum( axis = 2 ) / R

    # compute the matrix sub-block

    A = ComputeDνG( k, R, CosΘ )

    # solve the linear system

    c = np.linalg.solve( A, b )

    return c;
```

## 4 Compute the acoustic radiation force

The functions below compute the acoustic radiation force using the scalar potentials computed using the harmonic expansion and the MFS.

### 4.1 Acoustic radiation force using the harmonics expansion

The code below computes the $z$-component of the acoustic radiation force using the harmonics expansion for the scalar potential.

```python
def ComputeHarmonicsExpansionForce( k, a, δ, ρ0, κ0, ω, N ):
    """
    This function computes the radiation force using the harmonic expansion for the
    scattered scalar potential φ1.
    """
    # compute the Gauss-Legendre quadrature rule

    μ, w = np.polynomial.legendre.leggauss( N )

    # compute the Legendre polynomials

    P_n = ComputeLegendrePolynomials( μ, N )

    # compute the harmonic expansion coefficients

    B_n = ComputeIncidentHarmonicExpansionCoeffs( k, a, δ, N )
    C_n = ComputeScatteredHarmonicExpansionCoeffs( k, a, δ, N )

    # compute array of indices

    n = np.arange( N )

    # compute the spherical Bessel function on the sphere

    jn = sp.spherical_jn( n, k * a, derivative = False )

    # compute the spherical Hankel function on the sphere

    h1n = sp.spherical_jn( n, k * a, derivative = False ) \
        + 1j * sp.spherical_yn( n, k * a, derivative = False )

    # compute the fields on the sphere using the harmonic expansions

    φ1_inc = P_n @ ( jn * B_n )
    φ1_s   = P_n @ ( np.sqrt( ( 2 * n + 1 ) / ( 4.0 * np.pi ) ) * h1n * C_n )

    # compute the derivatives of the Legendre polynomials

    DP_n = ComputeLegendrePolynomialDerivatives( μ, N )

    # compute θ-derivative of the incident field

    φ1_inc_θ = -np.sqrt( 1.0 - μ ** 2 ) * ( DP_n @ ( jn * B_n ) )

    # compute θ-derivative of the scattered field

    φ1_s_θ = -np.sqrt( 1.0 - μ ** 2 ) * ( DP_n @ ( np.sqrt( ( 2 * n + 1 ) / ( 4.0 * np.pi ) ) * h1n * C_n ) )

    # compute the time-average of the pressure

    p1_ave = 0.5 * np.abs( 1j * ρ0 * ω * ( φ1_inc + φ1_s ) ) ** 2

    # compute the time-average of the velocity

    v1_ave = 0.5 * np.abs( φ1_inc_θ / a + φ1_s_θ / a ) ** 2

    # compute the function to be integrated

    Fintegrand = -(  0.5 * κ0 * p1_ave + 0.5 * ρ0 * v1_ave ) * μ

    # compute the z-component of the force

    Fz = 2.0 * np.pi * np.sum( Fintegrand * w )

    return Fz;
```

### 4.2 Acoustic radiation force using the MFS solution

The code below computes the $x$-, $y$-, and $z$-components of the acoustic radiation force using the MFS result for the scalar potential.

```python
def ComputeMFSForce( k, a, δ, ρ0, κ0, ω, Nμ, N ):
    """
    This function computes the radiation force using the MFS solution for the scattered
    scalar potential φ1.
    """

    # compute the MFS points for the reference sphere

    ν, ρ_bdy, ρ_src = ComputeMFSReferencePoints( a, N )

    # solve the linear system of equations for the MFS expansion coefficients

    c = ComputeMFSExpansionCoefficients( k, δ, ν, ρ_bdy, ρ_src, N )

    # compute the Gauss-Legendre quadrature rule points and weights, here mu is the normal vector on the sphere

    μ, w = np.polynomial.legendre.leggauss( Nμ )

    # evaluate the solution for a fixed azimuthal angle φ

    ϕ = np.pi / Nμ * np.arange( 2 * Nμ )

    # compute a meshgrid of the angle points

    indx, jndx = np.meshgrid( np.arange( Nμ ), np.arange( 2 * Nμ ) )

    # compute the components of the normal vector

    nx = ( np.sqrt( 1.0 - μ[indx] ** 2 ) * np.cos( ϕ[jndx] ) ).flatten()
    ny = ( np.sqrt( 1.0 - μ[indx] ** 2 ) * np.sin( ϕ[jndx] ) ).flatten()
    nz = ( μ[indx] ).flatten()

    # compute the spatial points

    X = a * nx
    Y = a * ny
    Z = a * nz

    # compute the quadrature weights

    weights = ( np.pi / Nμ * w[indx] ).flatten()

    ii, jj = np.mgrid[ 0:2*Nμ*Nμ, 0:N ]

    # compute the distance from source points to the evaluation point

    R = np.sqrt( ( X[ii] - ρ_src[jj,0] ) ** 2  \
                + ( Y[ii] - ρ_src[jj,1] ) ** 2 \
                + ( Z[ii] - ρ_src[jj,2] ) ** 2 )

    # compute Green's function

    G  = np.exp( 1j * k * R ) / ( 4.0 * np.pi * R )

    # compute the total field

    φ1 = np.exp( 1j * k * Z ) + np.exp( -1j * k * ( Z + δ ) ) + G @ c

    # compute components of the gradient of the total field

    φ1x = ( ( X[ii] - ρ_src[jj,0] ) / R * ( 1j * k - 1.0 / R ) * G ) @ c

    φ1y = ( ( Y[ii] - ρ_src[jj,1] ) / R * ( 1j * k - 1.0 / R ) * G ) @ c

    φ1z = 1j * k * ( np.exp( 1j * k * Z ) - np.exp( -1j * k * ( Z + δ ) ) ) \
        + ( ( Z[ii] - ρ_src[jj,2] ) / R * ( 1j * k - 1.0 / R ) * G ) @ c

    # compute the time-averaged pressure field

    p1_ave = 0.5 * np.abs( 1j * ρ0 * ω * φ1 ) ** 2

    # compute the time-averaged velocity field

    v1_ave = 0.5 * np.abs( φ1x ) ** 2 + 0.5 * np.abs( φ1y ) ** 2 + 0.5 * np.abs( φ1z ) ** 2

    # compute the function to be integrated

    F = - 0.5 * κ0 * p1_ave - 0.5 * ρ0 * v1_ave

    # compute integrals

    Fx = np.sum( F * nx * weights )
    Fy = np.sum( F * ny * weights )
    Fz = np.sum( F * nz * weights )

    return Fx, Fy, Fz;
```

## 5 Compare acoustic radiation forces

In the codes that follow, we compute the harmonic expansion and the MFS for the scattered field and compare those results.

```python
# set the frequency of the acoustic waves (kHz)

f = 40.0

# set the wave speed (cm *kHz)

c0 = 34.3

# set the density of the spheres (g/cm^3)

ρ0 = 0.001225

# set the compressibility

κ0 = 1.0 / ( ρ0 * c0 ** 2 )

# compute the circular frequency

ω = 2.0 * np.pi * f

# compute the wavenumber (cm^-1)

k = ω / c0

# set the nondimensional size parameter ka

ka = 0.1

# compute the radius of the sphere from ka

a = ka / k
```

```python
# set the value of δ

Nδ = 201
δ  = np.linspace( -k/np.pi, k/np.pi, Nδ )

# set the number of modes to use for harmonic expansions

Nμ = 8

# set the number of MFS points

N = 512

# allocate memory for forces

Fz_expansion = np.full( Nδ, float( 'nan' ) )

Fx_MFS = np.full( Nδ, float( 'nan' ) )
Fy_MFS = np.full( Nδ, float( 'nan' ) )
Fz_MFS = np.full( Nδ, float( 'nan' ) )

for i in range( Nδ ):

    Fz_expansion[i] = ComputeHarmonicsExpansionForce( k, a, δ[i], ρ0, κ0, ω, Nμ )

    Fx_MFS[i], Fy_MFS[i], Fz_MFS[i] = ComputeMFSForce( k, a, δ[i], ρ0, κ0, ω, Nμ, N )
```

```python
# plot the z-component of the force as a function of δ

plt.rcParams['figure.figsize'] = [8,6]

plt.plot( k * δ, Fz_expansion, k * δ, Fz_MFS, '--' )
plt.grid()
plt.xlabel( r'$k \delta$', fontsize = 14 )
plt.ylabel( r'$F_{z}$', fontsize = 14 )
plt.legend( ( 'Harmonic Expansion', 'MFS'), fontsize = 18, loc = 'lower right' )
plt.title("Acoustic Radiation Force via different solution methods")

plt.show()
```

```python
# plot the x- and y-components of the force as a function of δ

plt.rcParams['figure.figsize'] = [8,6]

plt.plot( k * δ, Fx_MFS, '-', k * δ, Fy_MFS, '--' )
plt.grid()
plt.xlabel( r'$k \delta$', fontsize = 14 )
plt.legend( ( r'$F_{x}$', r'$F_{y}$' ), fontsize = 18 )

plt.show()
```

```python

```

Jupyter Notebooks/Acoustic radiation force on a single sound-hard sphere.ipynb


# 3. laboratorijska vježba

```python
# učitavanje potrebnih biblioteka

import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as ss
```

```python
#@title pomoćna funkcija
# izvršite ovu ćeliju ali se ne opterećujte detaljima implementacije

def plot_frequency_response(f, Hm, fc=None, ylim_min=None):
    """Grafički prikaz prijenosne funkcije filtra.

    Args
        f (numpy.ndarray) : frekvencije
        Hm (numpy.ndarray) : apsolutne vrijednosti prijenosne funkcije
        fc (number) : cutoff frekvencija
        ylim_min (number): minimalna vrijednost na y-osi za dB skalu

    Returns
        (matplotlib.figure.Figure, matplotlib.axes._subplots.AxesSubplot)
    """
    Hc = 1 / np.sqrt(2)
    if fc is None:
        fc_idx = np.where(np.isclose(Hm, Hc, rtol=1e-03))[0][0]
        fc = f[fc_idx]
    H_db = 20 * np.log10(Hm)

    fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(12, 7.5))

    ax[0, 0].plot(f, Hm, label='$H(f)$')
    ax[0, 0].plot(fc, Hc, 'o', label='$H(f_c)$')
    ax[0, 0].vlines(fc, Hm.min(), Hc, linestyle='--')
    ax[0, 0].annotate(f'$f_c = {fc:.3f}$ Hz\n$H(f_c)={Hc:.3f}$', (fc * 1.4, Hc))
    ax[0, 0].set_xscale('log')
    ax[0, 0].set_ylabel('$|V_{out}$ / $V_{in}$|')
    ax[0, 0].set_title('log scale')
    ax[0, 0].legend(loc='lower left')
    ax[0, 0].grid()

    ax[0, 1].plot(f, Hm, label='$H(f)$')
    ax[0, 1].plot(fc, Hc, 'o', label='$H(f_c)$')
    ax[0, 1].annotate(f'$f_c = {fc:.3f}$ Hz\n$H(f_c)={Hc:.3f}$', (fc * 1.4, Hc))
    ax[0, 1].set_title('linear scale')
    ax[0, 1].legend()
    ax[0, 1].grid()

    ax[1, 0].plot(f, H_db, label='$H_{dB}(f)$')
    ax[1, 0].plot(fc, H_db.max() - 3, 'o', label='$H_{dB}(f_c)$')
    ax[1, 0].vlines(fc, H_db.min(), H_db.max() - 3, linestyle='--')
    ax[1, 0].annotate(f'$f_c = {fc:.3f}$ Hz\n$H(f_c)={H_db.max() - 3:.3f} dB$',
                      (fc * 1.4, H_db.max() - 3))
    ax[1, 0].set_xscale('log')
    ax[1, 0].set_xlabel('$f$ [Hz]')
    ax[1, 0].set_ylabel('$20 \\cdot \\log$ |$V_{out}$ / $V_{in}$|')
    if ylim_min:
        ax[1, 0].set_ylim((ylim_min, 10))
    ax[1, 0].legend(loc='lower left')
    ax[1, 0].grid()

    ax[1, 1].plot(f, H_db, label='$H_{dB}(f)$')
    ax[1, 1].plot(fc, H_db.max() - 3, 'o', label='$H_{dB}(f_c)$')
    ax[1, 1].annotate(f'$f_c = {fc:.3f}$ Hz\n$H(f_c)={H_db.max() - 3:.3f} dB$',
                      (fc * 1.4, H_db.max() - 3))
    ax[1, 1].set_xlabel('$f$ [Hz]')
    if ylim_min:
        ax[1, 1].set_ylim((ylim_min, 10))
    ax[1, 1].legend()
    ax[1, 1].grid()

    fig.tight_layout
    return fig, ax
```

### Pasivni visoko-propusni filtri

Realizacija visoko-propusnog filtra u ovom slučaju se ostvaruje korištenjem otpornika i zavojnice povezanih u seriju, pri čemu se izlaz promatra kao napon na zavojnici, $V_{out}$ Uz pretpostavku da je signal na ulazu, $V_{in}$, sinusoidalni naponski izvor, analizu možemo prebaciti u frekvencijsku domenu koristeći impedancijski model. Na ovaj način izbjegavamo potrebu za korištenjem diferencijalnog računa i čitav proračun se svodi na jednostavni algebarski problem.

<center>

</center>

Izraz za funkciju prijenosnog odziva dobijamo kao omjer izlaznog i ulaznog napona. Izlazni napon - napon na zavojnici, $V_{out}$, definiramo kroz podjelu ulaznog napona na sljedeći način:

$$
\begin{align}
    V_{out} &= \frac{Z_l}{Z_l + Z_r} \cdot V_{in} \\
    H(\omega) = \frac{V_{out}}{V_{in}} &= \frac{Z_l}{Z_l + Z_r} = \frac{j\omega L}{j\omega L + R} = \frac{1}{1+R/(j\omega L)}
\end{align}
$$

Kako je $H$ funkcija frekvencije, imamo dva ruba slučaja:
* za iznimno niske frekvencije kada je $\omega \sim 0$ slijedi da je $H(\omega) \rightarrow 0$;
* za iznimno visoke frekvencije kada $\omega \rightarrow \infty$ slijedi da je $H(\omega) = 0$.

Potrebno je dodatno definirati već spomenutu *cut-off* frekvenciju, $f_c$, za koju amplituda funkcije frekvencijskog odziva, $H$, pada za $\sqrt 2$ puta, odnosno za $3$ dB:
$$
\begin{align}
    f_c &= \frac{R}{2 \pi L}
\end{align}
$$

Link za interaktivni rad sa pasivnim visoko-propusnim filtrom: http://sim.okawa-denshi.jp/en/LRtool.php

#### Zadatak 1

Prvi zadatak je implementirati funkciju `cutoff_frequency` koja na ulazu prima iznose otpora, `R`, i zavojnice, `L`, a na izlazu daje *cutoff* frekvenciju visoko-propusnog filtra.

```python
def cutoff_frequency(R, L):
    """Cutoff frekvencija visoko-propusnog RL filtra.

    Args:
        R (number) : vrijednost otpora otpornika
        L (number) : induktivitet zavojnice

    Returns:
        number
    """
    #######################################################
    ## TO-DO: implementiraj proračun cutoff frekvencije ##
    # Nakon toga zakomentiraj sljedeću liniju.
    raise NotImplementedError('Implementiraj proračun cutoff frekvencije.')
    #######################################################

    # definiraj cutoff frekvenciju
    fc = ...
    return fc
```

Kolika je *cutoff* frekvencija za otpor od $200 \Omega$ i induktivitet zavojnice od $100 mH$?

```python
R = ...  # otpor
L = ...  # induktivitet

fc = cutoff_frequency(...)  # cutoff frekvencija

print(f'R = {R/1000} kΩ')
print(f'L = {L*1000} mH')
print(f'cutoff frekvencija iznosi {fc:.2f} Hz, '
      'očekivana vrijednost je 318.31 Hz')
```

#### Zadatak 2

Drugi zadatak je implementirati funkciju `rl_highpass` koja na ulazu prima iznose otpora, `R`, induktiviteta, `L`, i frekvenciju, `f`, a na izlazu daje prijenosni odziv pasivnog visoko-propusnog RL filtra.

```python
def rl_highpass(R, L, f):
    """Funkcija prijenosnog odziva RL visoko-propusnog filtra.

    Args:
        R (number) : vrijednost otpora otpornika
        L (number) : induktivitet
        f (number or numpy.ndarray) : frekvencija/e

    Returns:
        float or numpy.ndarray
    """
    ######################################################
    ## TO-DO: implementiraj funkciju prijenosnog odziva ##
    # Nakon toga zakomentiraj sljedeću liniju.
    raise NotImplementedError('Implementiraj funckiju prijenosnog odziva.')
    ######################################################

    # definiraj funkciju prijenosnog pazeći da `f` može biti ili broj (int,
    # float) ili 1-D niz (`numpy.ndarray`)
    H = ...
    return H
```

Kolika je vrijednost prijenosne funkcije pri *cutoff* frekvencija za otpor od $200 \Omega$ i induktivitet zavojnice od $100 mH$?

```python
R = ...  # otpor
L = ...  # induktivitet

Hc = rl_highpass(...)  # prijenosna funkcija pri cutoff frekvenciji

print(f'R = {R:.2f} Ω')
print(f'C = {L * 1000:.2f} mH')
print(f'pojačanje pri cutoff frekvenciji iznosi {abs(Hc):.4f}, '
      'očekivana vrijednost je 1/√2\n\n'
      'provjerite ispravnost dobivenog rezutltata')
```

```python
# ćelija za provjeru rezultata

```

Pretvorite vrijednost prijenosne funkcije pri *cutoff* frekvenciju u decibele i uvjerite se u tvrdnju da amplituda funkcije frekvencijskog odziva, $H$, pada za $3$ dB pri *cutoff* frekvenciji.

```python
Hc_dB = ...  # pretvorba prijenosne funkcije pri cutoff frekvenciji u dB skalu
print(Hc_dB)
```

Za raspon od $10000$ vrijednosti frekvencija do $10 kHz$ te za otpor od $200 \Omega$ i induktivitet zavojnice od $100 mH$, izračunajte vrijednosti prijenosne funkcije.

```python
f = np.linspace(..., num=10000)
H = rl_highpass(...)  # prijenosna funkcija
```

S obzirom da su vrijednosti prijenosne funkcije kompleksne veličine, razmilite što je potrebno napraviti s njima prije nego ih grafički prikažemo?

```python
Hm = ...  # konverzija u apsolutne vrijednosti
```

Vizualizirajte ovisnost prijenosne funkcije o frekvenciji koristeći `matplotlib` i funkciju `matplotlib.pyplot.plot`.

```python
plt.plot(...)
plt.xlabel('f [Hz]')
plt.ylabel('H(f)')
plt.show()
```

Vizualizirajte sada rezultate koristeći već implementiranu funkciju `plot_frequency_response`.

Napomena: za provjeru načina korištenja prethodne funkcije koristite sljedeću naredbu:

```python
help(plot_frequency_response)
```

ili jednostavno

```python
plot_frequency_response?
```

```python
# provjerite način korištenja funkcije

```

```python
fig, ax = plot_frequency_response(...)  # grafički prikaz dobivenih rezultata
```

### Strujno-naponska karakteristika RL visoko-propusnog filtra

```python
def time_constant(L, R):
    """Vremenska konstanta RL visoko-propusnog filtra.

    Args:
        R (number) : vrijednost otpora otpornika
        L (number) : induktivitet

    Returns:
        float or numpy.ndarray
    """
    ##################################################################
    ## TO-DO: implementiraj fnkciju koja racuna vremensku konstantu ##
    # Nakon toga zakomentiraj sljedeću liniju.
    raise NotImplementedError('Implementiraj vremensku konstantu.')
    ##################################################################

    # definiraj vremensku konstantu
    tau = ...
    return tau
```

```python
tau = time_constant(L, R)  # vremenska konstanta
```

Koja fizikalna veličina je pridružena vremenskoj konstanti? Objasni.

```python
def rl_current(t, t_switch, V, R, L):
    """Struja kroz RL visoko-propusni filtar.

    Args:
        t (number or numpy.ndarray) : trenutak/ci u kojima računamo
            vrijednost struje
        t_switch (number) : treneutak promjene predznaka struje
        V (number) : napon na ulazu
        R (number) : vrijednost otpora otpornika
        L (number) : induktivitet

    Returns:
        float or numpy.ndarray
    """
    I0 = V / R
    i = np.where(t < t_switch,
                 I0 * (1 - np.exp((-R / L) * t)),
                 I0 * np.exp((-R / L) * (t - t_switch)))
    return i
```

```python
V = 5  # napon na ulazu
tau = time_constant(L, R)  # vremenska konstanta filtra
t_switch = tau * 4.4  # vrijeme promjene predznaka struje
T = 2 * t_switch  # period
t = np.linspace(0, T)  # vremenska serija trenutaka u kojima evaluiramo vrijednost struje
i_rl = rl_current(t, t_switch, V, R, L)  # RL struja
i = V / R * np.sin(2 * np.pi * t / T)  # sinusna struja
```

```python
# vizualizacija RL struje

plt.figure()
plt.plot(t, i_rl, label='struja')
plt.plot(t, i, label='on-off ciklus')
plt.plot([t.min(), t_switch, t.max()], [0, 0, 0], 'rx')
plt.hlines(0, t.min(), t.max(), 'k')
plt.vlines(t_switch, i.min(), i.max(), 'k')
plt.xlabel('t [s]')
plt.ylabel('i(t) [A]')
plt.legend()
plt.grid()
plt.show()
```

### Pojasno propusni filtri
Sljedeći kod koristi više različitih tipova pojasno-propusnih filtara (Hamming, Kaiser, Remez) i uspoređuje ih s idealnom prijenosnom funkcijom.

```python
def bandpass_firwin(ntaps, lowcut, highcut, fs, window='hamming'):
    taps = ss.firwin(ntaps, [lowcut, highcut], nyq=0.5 * fs, pass_zero=False,
                     window=window, scale=False)
    return taps

def bandpass_kaiser(ntaps, lowcut, highcut, fs, width):
    atten = ss.kaiser_atten(ntaps, width / (0.5 * fs))
    beta = ss.kaiser_beta(atten)
    taps = ss.firwin(ntaps, [lowcut, highcut], nyq=0.5 * fs, pass_zero=False,
                     window=('kaiser', beta), scale=False)
    return taps

def bandpass_remez(ntaps, lowcut, highcut, fs, width):
    delta = 0.5 * width
    edges = [0,
             lowcut - delta,
             lowcut + delta,
             highcut - delta,
             highcut + delta,
             0.5 * fs,
            ]
    taps = ss.remez(ntaps, edges, [0, 1, 0], Hz=fs)
    return taps
```

```python
fs = 63.0
lowcut = 0.7
highcut = 4.0
ntaps = 128

taps_hamming = bandpass_firwin(ntaps, lowcut, highcut, fs)
taps_kaiser16 = bandpass_kaiser(ntaps, lowcut, highcut, fs, width=1.6)
taps_kaiser10 = bandpass_kaiser(ntaps, lowcut, highcut, fs, width=1.0)
taps_remez = bandpass_remez(ntaps, lowcut, highcut, fs=fs, width=1.0)
```

```python
plt.figure()
w, h = ss.freqz(taps_hamming, 1, worN=2000)
plt.plot(fs * 0.5 / np.pi * w, abs(h), label='Hammingov prozor')
w, h = ss.freqz(taps_kaiser16, 1, worN=2000)
plt.plot(fs * 0.5 / np.pi * w, abs(h), label='Kaiser, širina = 1.6')
w, h = ss.freqz(taps_kaiser10, 1, worN=2000)
plt.plot(fs * 0.5/ np.pi * w, abs(h), label='Kaiser, širina = 1.0')
w, h = ss.freqz(taps_remez, 1, worN=2000)
plt.plot(fs * 0.5 / np.pi * w, abs(h), label=f'Remez, širina = 1.0')
h = np.where((fs * 0.5 / np.pi * w < lowcut) | (fs * 0.5 / np.pi * w > highcut), 0, 1)
plt.plot(fs * 0.5 / np.pi * w, h, 'k-', label='idealna karakteristika')
plt.fill_between(fs * 0.5 / np.pi * w, h, color='gray', alpha=0.1)
plt.xlim(0, 8.0)
plt.grid()
plt.legend(loc='upper right')
plt.xlabel('f (Hz)')
plt.ylabel('H(f)')
plt.show()
```

```python

```


# Solving (Aiyagari, 1994) with endogenous labor supply
### Using Endogenous Grid Method (Carroll, 2006).
##### It follows [Andrea's Muller](https://sites.google.com/site/mrandreasmueller/resources) notes.
---

## The model

Consider a closed economy with ex-ante identical households that - conditional on the
states $(a, z)$ - solve the same recursive problem:

\begin{equation}
    V(a, z) = \max_{c,h,a'} u(c) - v(h) + \beta E[V(a', z')] \tag{1}
\end{equation}

\begin{align*}
    s.t \quad c + a' = (1 + r)a + hwz \tag{2} \\
    a' \geq -b \tag{3} \\
    c \geq 0, \quad h \geq 0 \tag{4}
\end{align*}

where $s$ is individual consumption, $a$ asset holdings, $h$ labor supply, and $z$ can be interpreted as an idiosyncratic productivity shock that follows a stochastic Markov process with transition probability matrix, $\Pi(z)$, and realizations drawn from the finite valued set $Z$. The parameters are restricted to be, $0 < \beta < 1$ (subjective discount factor)
and $b \geq 0$ (borrowing constraint), and the efficiency wage $w > 0$ as well as the interest
rate $r > 0$ are determined in the competitive production sector. What makes households ex-post heterogeneous is that each of them draws an individual-specific productivity realization $z'$ from the same distribution, and since financial markets are incomplete (there is only a risk-free asset available, but no state-contingent assets) agents with different shock histories will be different ex post.

First order conditions are:

\begin{align*}
    u'(c) - \lambda &= 0 \\
    -v(h) + \lambda zw &= 0 \\
    \beta E[V'(a',z')] - \lambda + \mu &= 0 \\
    \mu(a' +b) = 0, \quad \mu &\geq 0 \\
    V'(a,z) &= \lambda(1+r) \\
    (1+r)a + hzw -c - a' &= 0
\end{align*}

where $\lambda$ is the Lagrange multiplier associated with the budget constraint [(2)](#equation-budget-constraint) and $\mu$ is the Lagrange multiplier associated with the borrowing constraint [(3)](#equation-borrowing-constraint). Given standard assumptions on $u(c)$ and $v(c)$ we can ignore the positivity constraints [(4)](#equation-nonneg-constraint).

Eliminating $\lambda$, the Envelope condition yields the following set of optimal conditions:

\begin{align*}
    v'(h) &= u'(c)zw \\
    u'(c) &= \beta E[u'(c')(1+r)] + \mu \\
    \mu(a' + b) &= 0, \quad \mu \geq 0 \\
    a' &= (1+r)a + hzw - c
\end{align*}

Guess a decision rule for the optimal consumption as a function of the future states,

\begin{equation}
c_{0}(a',z')
\end{equation}

Let the future asset level be consistent with a non-binding borrowing constraint, $a' \geq -b$ and the shock realizations be drawn from the finite valued set, $Z = \{z_{1}, ..., z_{N}\}$, with transition probabilities defined by $\pi_{s'|s}$.

Denote the inverse marginal utility function of consumption by $(u')^{-1}$. Define the function $\mathcal{C}(a',z)$, that measures the current consumption level as a function of the future state, $a'$, the current state, $z$, and the guess, $c_{0}(a',z')$, if the borrowing constraint is not binding, $\mu = 0$:

\begin{equation}
    \mathcal{C}(a',z) = (u')^{-1}(\beta(1+r)E[u'(c_{0}(a',z'))]) \tag{5}
\end{equation}

Denote the inverse function of the marginal disutility of labor be denoted by $(v')^{-1}$. We can characterize the optimal labor supply as a function of the optimal consumption and the productivity realization:

\begin{equation}
    \mathcal{H}(\mathcal{C}(a',z),z) = (v')^{-1}(u'(\mathcal{C}(a',z))zw) \tag{6}
\end{equation}

Now we can find the current asset level $\mathcal{A}(a',z)$ that is consistent with the guess of the consumption decision rule. Or, on other words, $\mathcal{A}(a',z)$ defines the current assets holdings for a household with labor productivity, $z$, who has chosen to save $a'$ for tomorrow. From the budget constraint we can derive:

\begin{equation}
    \mathcal{A}(a',z) = \frac{1}{1+r}[\mathcal{C}(a',z) + a' - \mathcal{H}(\mathcal{C}(a',z),z)zw] \tag{7}
\end{equation}

The pair of functions $\{\mathcal{C}(a',z), \mathcal{A}(a',z)\}$ implicitly defines a relationship between current consumption, $c$, and the current asset level and productivity realization, $(a,z)$, given that the borrowing constraint does not bind, $\mu=0$. Thus, the guess on the consumption function can be updated by simply interpolating $c_{0}(\mathcal{A}(a',z),z)$ on $(a', y')$ to yield

\begin{equation}
    c^{nb}_{1}(a',z'), \quad a' > \mathcal{A}(-b,z) \tag{8}
\end{equation}

For future asset levels $a'$ that are smaller than the lowest current asset level that is exactly consistent with a binding borrowing constraint in the future $a' < \mathcal{A}(-b,z) $.

The procedure in this case is different: for those asset levels, we update the consumption function with the value consistent with a binding borrowing constraint, that is, $a' = -b$. Using the budget constraint we solve

\begin{equation}
    c^{b}_{1}(a,z) = (1+r)a + \mathcal{H}(c^{'}_{b}(a,z),z)zw + b \tag{9}
\end{equation}

to find the consumption consistent with the binding borrowing constraint. This is one equation with only one unknown, $c^{b}_{1}(a,z)$, so we can use a root solver to find the optimal level of consumption, then we iterate one period forward to yield

\begin{equation}
    c^{b}_{1}(a',z'), \quad a' < \mathcal{A}(-b,z) \tag{10}
\end{equation}

Finally then we have:

$$
c_{1}(a',z') =
    \begin{cases}
    c^{b}_{1}(a',z'), \quad a' < \mathcal{A}(-b,z), \\
    c^{nb}_{1}(a',z'), \quad a' > \mathcal{A}(-b,z)
    \end{cases}
    \tag{11}
$$

Now that we have the optimal consumption policy we can simulate this economy over $T=10^4$ periods for the given guess on the prices, $r_{0}$ and $w_{0}$, using the budget constraint [(2)](#equation-budget-constraint). Given a sequence $z_{t} = \{z_{0}, ..., z_{T}\}$ simulated using the Markov Chain, we will obtain sequences of $a_{t} = \{a_{0}, ..., a_{T}\}$ and $h_{t} = \{h_{0}, ..., h_{T}\}$ from where we can compute aggregate levels of capital, $K$, and labor supply, $L$.

\begin{equation}
    K_{1} = \frac{1}{T}\sum_{t=1}^{T}a_{t}(r_{0}) \tag{12}
\end{equation}

\begin{equation}
    L_{1} = \frac{1}{T}\sum_{t=1}^{T}h_{t}(r_{0})z_{t} \tag{13}
\end{equation}

Then we can update $\{r_{1}$, $w_{1}\}$ according to marginal pricing in a Cobb-Douglas production sector, $F(K,L) = K^{\alpha}L^{1-\alpha}$,

\begin{equation}
    r_{1} = \alpha\left(\frac{K_{1}}{L_{1}}\right)^{\alpha-1} - \delta \tag{14}
\end{equation}

\begin{equation}
    w_{1} = (1-\alpha)\left(\frac{\alpha}{r_{1}+\delta}\right)^{\alpha/(1-\alpha)} \tag{15}
\end{equation}

and iterate until convergence.

### The algorithm

1. Guess $\{r_{0}, w_{0}\}$:
    1. Guess $c_{0}(a',z')$.
    2. Use $c_{0}(a',z')$ to solve for $\mathcal{C}(a',z)$ using the Euler equation [(5)](#equation-cons-optimal).
    3. Use $\mathcal{C}(a',z)$ to solve for $\mathcal{H}(c,z)$ using the first order condition for optimal labor supply  [(6)](#equation-labor-optimal).
    4. Use $c_{0}(a',z')$, $\mathcal{C}(a',z)$ and $\mathcal{H}(c,z)$ to solve for $\mathcal{A}(a',z)$ using the budget constraint [(7)](#equation-assets-optimal).
    5. Interpolate $c_{0}(\mathcal{A}(a',z),z)$ on $(a', y')$ to update $c_{1}(a',z')$.
    6. Check convergence:
    $$ \left\lVert\frac{c_{1}(A,Z) - c_{0}(A,Z)}{c_{0}(A,Z)}\right\rVert \leq 10^{-5}$$
    If true STOP, save $c^{*}_{r0}(A,Z)$ and continue to STEP 2, otherwise set $c_{0}(a',z') = c_{1}(a',z')$ and go to STEP A.
2. Given $c^{*}_{r0}(A,Z)$ simulate the economy and obtain $\{K_{1}, L_{1}\}$ according to [(12)](#equation-K) and [(13)](#equation-L).
3. Use $K_{1}$ and $L_{1}$ to obtain $\{r_{1}, w_{1}\}$ according to [(14)](#equation-r) and [(15)](#equation-w).
4. Check convergence:
$$\left\lVert \frac{r_{1} - r_{0}}{r_{0}} \right\rVert \leq 10^{-5}$$
If true STOP, save $r^{*}$ and exit the algorithm, otherwise set $r_{0} = r_{1}$, $w_{0} = w_{1}$ and go to STEP 1.

---
## Solving the model

```julia
using Parameters, QuantEcon, LinearAlgebra, Interpolations, Random, Roots, Optim, Plots
```

We start by defining a `NamedTuple` object type called `Household` that contains all the parameters and functions needed to solve the model.

```julia
Household = @with_kw (r = 0.02,
                      R = 1 + r,
                      α = 0.33,
                      δ = 0.05,
                      w = (1-α)*(α/(r+δ))^(α/(1-α)),
                      σ = 2.0,
                      κ = 4.00,
                      ν = 0.25,
                      β = 0.96,
                      MC = MarkovChain([0.5 0.5; 0.5 0.5], [0.95; 1.05]),
                      Π = MC.p,
                      z_vals = MC.state_values,
                      z_size = length(z_vals),
                      b = 0.0, # Borrowing limit/Minimum level of assets
                      r_ss = (1/β-1), # SS interest rate
                      a_max = ((r_ss+δ)/α)^(1/(α-1))*3, # 3 times SS capital
                      a_size = 200,
                      a_vals = collect(range(b, a_max, length = a_size)),
                      u = σ == 1 ? (x, h=0.0) -> log(x)-  κ*(h^(1+ν))/(1+ν) : (x, h=0.0) -> (x^(1 - σ) - 1) / (1 - σ) - κ*(h^(1+ν))/(1+ν),
                      duc = c -> c^(-σ),
                      duh = h -> κ*h^(ν),
                      duc_inv = x -> x^(-1/σ),
                      duh_inv = x -> (x/κ)^(1/ν))
```

    #3 (generic function with 2 methods)

We need a function that initializes our initial guess $c_{0}(a',z')$ given some `Household` parameters.

```julia
function initialize(am::NamedTuple)
    # simplify names, set up arrays
    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    V0 = zeros(Float64, a_size, z_size)
    c0 = zeros(Float64, a_size, z_size)
    a0 = [a_vals a_vals]
    h0 = zeros(Float64, a_size, z_size)

    # Populate V and c
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(a_vals)
            c_max = R * a + w*z + b
            c0[i_a, i_z] = c_max
            V0[i_a, i_z] = u(c_max) / (1 - β)
        end
    end
    return V0, a0, c0, h0
end
```

    initialize (generic function with 1 method)

Then we define the Endogenous Grid Method Operator function `EGMOperator_Labor` that solves for $c_{1}(a',z')$ using `g_c` ($c_{0}(a',z')$) and `am`(`Household`) as inputs. Therefore, this function runs steps (A)-(E) of the solution algorithm.

```julia
function EGMOperator_Labor(g_c::AbstractArray,
                            am::NamedTuple;
                            return_policy::Bool=false)
    @unpack R, r, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    # Initialize arrays
    g_a = zeros(Float64, a_size, z_size)
    g_c_nb = zeros(Float64, a_size, z_size)
    g_c_new = zeros(Float64, a_size, z_size)

    # Iterate over different combinations of (a',z)
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a_new) in enumerate(a_vals)
            # Solve for C(a',z) in the no binding case
            expectation = dot(duc.(g_c[i_a, :]), Π[i_z, :])
            g_c_nb[i_a, i_z] = duc_inv(R*β * expectation)
            # Solve for H(c,z)
            h = duh_inv(w*z*duc_inv(g_c_nb[i_a, i_z]))
            # Solve for A(a',z)
            g_a[i_a, i_z] = (g_c_nb[i_a, i_z] + a_new - h*w*z) /R
        end
        # Update to c1(a',z')
        g_c_new[:, i_z] = LinearInterpolation(g_a[:,i_z], g_c_nb[:,i_z], extrapolation_bc=Line())(a_vals)
    end

    # Iterate over different combinations of (a',z)
    for (i_a, a_new) in enumerate(a_vals)
        for (i_z, z) in enumerate(z_vals)
            # Check if it is a binding borrowing constraint case
            if a_new < g_a[1, i_z]
                # If so, solve for C(a',z) using a root finder and use the result value in c1(a',z')
                cf = c -> abs(a_new*R + duh_inv(w*z*duc_inv(c))*w*z + b - c)
                res = optimize(cf, 0.0, g_c_new[i_a, i_z])
                g_c_new[i_a, i_z] = Optim.minimizer(res)
            end
        end
    end
    # Return c1(a',z')
    return g_c_new
end
```

    EGMOperator_Labor (generic function with 1 method)

Now we can initialize our `Household` object with the desired parametrization, run `initialize` to obtain our initial guess $c_{0}(a',z')$ and iterate over `EGMOperator_Labor` using `compute_fixed_point` from `QuantEcon` library until convergence is reached.

```julia
am = Household(r = 0.03,
                      α = 0.33,
                      δ = 0.05,
                      σ = 1.0,
                      κ = 4.00,
                      ν = 0.25,
                      β = 0.96,
                      MC = MarkovChain([0.6 0.4; 0.4 0.6], [0.5; 1.0]),
                      b = 0.0,
                      #a_max = 18,
                      a_size = 250)

@unpack R, α, δ, r, r_ss, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
V, g_a, g_c, g_h = initialize(am)

g_c = compute_fixed_point(c -> EGMOperator_Labor(c, am),
                        g_c,
                        max_iter=1000,
                        err_tol=1e-5,
                        verbose=2,
                        print_skip = 100)

plot(am.a_vals, g_c, label=["z_vals[1]", "z_vals[2]"], xlabel="a", ylabel="C(a,z)")
```

    Compute iterate 100 with error 8.472961688488212e-6
    Converged in 100 steps

Given the new $c_{1}(a',z')$ we can simulate the economy and compute aggregate levels of $K_{1}$ and $L_{1}$. To implement this we define `SimulationAssets` function, which given some `Household` object and the correspondent `g_c`($c_{1}(a',z')$) from the previous step, returns a sequence $\{a_{t}\}_{t=0}^{T}$, $K_{1}$, $L_{1}$ and $r_{1}$.

```julia
function SimulationAssets(am::NamedTuple,
                    g_c::AbstractArray,
                    T::Int = 100000;
                    labor::Bool = false,
                    verbose::Bool = false)

    @unpack Π, r, w, α, δ, a_vals, z_vals, R, duc_inv, duh_inv = am  # Simplify names

    cf = interp(a_vals, g_c)

    a_sim = zeros(Float64, T + 1)
    h_sim = zeros(Float64, T)
    z_sim = simulate(MarkovChain(Π), T)
    for t in 1:T
        i_z = z_sim[t]
        z = z_vals[i_z]
        if labor
            h_sim[t] = duh_inv(w*z*duc_inv(cf(a_sim[t], i_z)))
            a_sim[t+1] = R * a_sim[t] + h_sim[t]*w*z - cf(a_sim[t], i_z)
        else
            a_sim[t+1] = R * a_sim[t] + w*z - cf(a_sim[t], i_z)
        end
    end
    # Burn 1000 first periods
    a_sim = a_sim[1000:end]

    # Compute Aggregates from market clearin
    K = mean(a_sim)
    L = mean(z_sim.*h_sim)
    r_new = α*(K/L)^(α-1) - δ
    return a_sim, K, L, r_new
end

Random.seed!(42)
a_sim, K, L, r1 = SimulationAssets(am, g_c, labor = true)
residual = (r1 - r)/r
println("Given r0 = $r, then K = $K and L = $L, which implies r1 = $r1")
println("Diff = $residual")
histogram(a_sim, nbins = 100, leg = false, normed = true, xlabel = "assets")
```

    Given r0 = 0.03, then K = 1.1030505493310432 and L = 0.5139843767370157, which implies r1 = 0.1478387944155487
    Diff = 3.9279598138516234

As we can see the distance is quite large so we have to keep iterating over $r$ until convergence. The following function `SolveEquilibirum` does all the job. Here we use the root finding routine `find_zero` to find the $r$ that clears the market.

```julia
function SolveEquilibrium(r0; verbose=false)
    am = Household(r = r0,
                          α = 0.33,
                          δ = 0.05,
                          σ = 1.0,
                          κ = 4.00,
                          ν = 0.25,
                          β = 0.96,
                          MC = MarkovChain([0.6 0.4; 0.4 0.6], [0.5; 1.0]),
                          b = 0.0,
                          #a_max = 18,
                          a_size = 250)

    @unpack R, α, δ, r, r_ss, w, b, σ, β, Π, z_vals, a_size, a_vals, z_size, duc, duh, duc_inv, duh_inv, u = am
    V, g_a, g_c, g_h = initialize(am)

    g_c = compute_fixed_point(c -> EGMOperator_Labor(c, am),
                            g_c,
                            max_iter=1000,
                            err_tol=1e-4,
                            verbose=false,
                            print_skip = 100)

    Random.seed!(42)
    a_sim, K, L, r1 = SimulationAssets(am, g_c, labor=true)
    residual = (r1 - r0)/r0
    if verbose
        println("Solving Aiyagari Model for r = $r:")
        println("Endogeous Grid Method. Grid size = $a_size")
        println("Given r0 = $r0, then K = $K and L = $L, which implies r1 = $r1")
        println("Diff = $residual")
        println("Check Incomplete Markets theory : r = $r1 < r_ss = $r_ss")
    end
    return residual, am, a_sim, g_c
end

r_star = find_zero(SolveEquilibrium, (0.01, 0.05), verbose=true)
```

    Results of univariate zero finding:

    * Converged to: 0.040382027481848166
    * Algorithm: Roots.BisectionExact()
    * iterations: 53
    * function evaluations: 55
    * stopped as x_n ≈ x_{n-1} using atol=xatol, rtol=xrtol
    * Note: Change of sign at xn identified.

    Trace:
    (a_0, b_0) = ( 0.0100000000000000,  0.0500000000000000)
    (a_1, b_1) = ( 0.0225000000000000,  0.0500000000000000)
    (a_2, b_2) = ( 0.0318750000000000,  0.0500000000000000)
    (a_3, b_3) = ( 0.0318750000000000,  0.0409375000000000)
    (a_4, b_4) = ( 0.0364062500000000,  0.0409375000000000)
    (a_5, b_5) = ( 0.0386718750000000,  0.0409375000000000)
    (a_6, b_6) = ( 0.0398046875000000,  0.0409375000000000)
    (a_7, b_7) = ( 0.0403710937500000,  0.0409375000000000)
    (a_8, b_8) = ( 0.0403710937500000,  0.0406542968750000)
    (a_9, b_9) = ( 0.0403710937500000,  0.0405126953125000)
    (a_10, b_10) = ( 0.0403710937500000,  0.0404418945312500)
    (a_11, b_11) = ( 0.0403710937500000,  0.0404064941406250)
    (a_12, b_12) = ( 0.0403710937500000,  0.0403887939453125)
    (a_13, b_13) = ( 0.0403799438476562,  0.0403887939453125)
    (a_14, b_14) = ( 0.0403799438476562,  0.0403843688964844)
    (a_15, b_15) = ( 0.0403799438476562,  0.0403821563720703)
    (a_16, b_16) = ( 0.0403810501098633,  0.0403821563720703)
    (a_17, b_17) = ( 0.0403816032409668,  0.0403821563720703)
    (a_18, b_18) = ( 0.0403818798065185,  0.0403821563720703)
    (a_19, b_19) = ( 0.0403820180892944,  0.0403821563720703)
    (a_20, b_20) = ( 0.0403820180892944,  0.0403820872306824)
    (a_21, b_21) = ( 0.0403820180892944,  0.0403820526599884)
    (a_22, b_22) = ( 0.0403820180892944,  0.0403820353746414)
    (a_23, b_23) = ( 0.0403820267319679,  0.0403820353746414)
    (a_24, b_24) = ( 0.0403820267319679,  0.0403820310533047)
    (a_25, b_25) = ( 0.0403820267319679,  0.0403820288926363)
    (a_26, b_26) = ( 0.0403820267319679,  0.0403820278123021)
    (a_27, b_27) = ( 0.0403820272721350,  0.0403820278123021)
    (a_28, b_28) = ( 0.0403820272721350,  0.0403820275422185)
    (a_29, b_29) = ( 0.0403820274071768,  0.0403820275422185)
    (a_30, b_30) = ( 0.0403820274746977,  0.0403820275422185)
    (a_31, b_31) = ( 0.0403820274746977,  0.0403820275084581)
    (a_32, b_32) = ( 0.0403820274746977,  0.0403820274915779)
    (a_33, b_33) = ( 0.0403820274746977,  0.0403820274831378)
    (a_34, b_34) = ( 0.0403820274789177,  0.0403820274831378)
    (a_35, b_35) = ( 0.0403820274810277,  0.0403820274831378)
    (a_36, b_36) = ( 0.0403820274810277,  0.0403820274820827)
    (a_37, b_37) = ( 0.0403820274815552,  0.0403820274820827)
    (a_38, b_38) = ( 0.0403820274818190,  0.0403820274820827)
    (a_39, b_39) = ( 0.0403820274818190,  0.0403820274819509)
    (a_40, b_40) = ( 0.0403820274818190,  0.0403820274818849)
    (a_41, b_41) = ( 0.0403820274818190,  0.0403820274818520)
    (a_42, b_42) = ( 0.0403820274818355,  0.0403820274818520)
    (a_43, b_43) = ( 0.0403820274818437,  0.0403820274818520)
    (a_44, b_44) = ( 0.0403820274818478,  0.0403820274818520)
    (a_45, b_45) = ( 0.0403820274818478,  0.0403820274818499)
    (a_46, b_46) = ( 0.0403820274818478,  0.0403820274818489)
    (a_47, b_47) = ( 0.0403820274818478,  0.0403820274818484)
    (a_48, b_48) = ( 0.0403820274818481,  0.0403820274818484)
    (a_49, b_49) = ( 0.0403820274818481,  0.0403820274818482)
    (a_50, b_50) = ( 0.0403820274818482,  0.0403820274818482)
    (a_51, b_51) = ( 0.0403820274818482,  0.0403820274818482)
    (a_52, b_52) = ( 0.0403820274818482,  0.0403820274818482)
    (a_53, b_53) = ( 0.0403820274818482,  0.0403820274818482)

    0.040382027481848166

```julia
_, am, a_sim, g_c = SolveEquilibrium(r_star, verbose=true)
histogram(a_sim, nbins = 100, leg = false, normed = true, xlabel = "assets")
```

    Solving Aiyagari Model for r = 0.040382027481848166:
    Endogeous Grid Method. Grid size = 250
    Given r0 = 0.040382027481848166, then K = 2.9522455057787984 and L = 0.42726953895379816, which implies r1 = 0.040382027481843746
    Diff = -1.094565006369663e-13
    Check Incomplete Markets theory : r = 0.040382027481843746 < r_ss = 0.04166666666666674

# Test different solution methods

```julia
include("../src/amObjects.jl")

using BenchmarkTools

# Create an instance of Household
am = Household(r = 0.03, β = 0.96, a_size = 250)
a_size = am.a_size
V, g_a, g_c, g_h = initialize(am)

println("===============")
println("TIME:")
println("===============")
println("VFI - Discrete with no refinements. Grid size = $a_size")
@btime compute_fixed_point(v -> BruteForce_Std(v, am),
                        V,
                        max_iter=1000,
                        err_tol=1e-4,
                        verbose=false)
println("===============")
println("VFI - Discrete using Monotonicity. Grid size = $a_size")
@btime compute_fixed_point(v -> BruteForce_Mtcy(v, am),
                        V,
                        max_iter=1000,
                        err_tol=1e-4,
                        verbose=false)
println("===============")
println("VFI - Continous using linear interpolation. Grid size = $a_size")
@btime compute_fixed_point(v -> BellmanOperator(v, am),
                        V,
                        max_iter=1000,
                        err_tol=1e-4,
                        verbose=false)
println("===============")
println("PFI - Continous using linear interpolation. Grid size = $a_size")
@btime compute_fixed_point(c -> ColemanOperator(c, am),
                        g_c,
                        max_iter=1000,
                        err_tol=1e-4,
                        verbose=false)
println("===============")
println("Endogeous Grid Method. Grid size = $a_size")
@btime compute_fixed_point(c -> EGMOperator_Labor(c, am),
                        g_c,
                        max_iter=1000,
                        err_tol=1e-4,
                        verbose=false)
```

    ===============
    TIME:
    ===============
    VFI - Discrete with no refinements. Grid size = 250
      2.817 s (35146794 allocations: 3.61 GiB)
    ===============
    VFI - Discrete using Monotonicity. Grid size = 250
      587.183 ms (3903170 allocations: 832.90 MiB)
    ===============
    VFI - Continous using linear interpolation. Grid size = 250
      5.738 s (137868308 allocations: 2.52 GiB)
    ===============
    PFI - Continous using linear interpolation. Grid size = 250
      2.203 s (61773231 allocations: 1.18 GiB)
    ===============
    Endogeous Grid Method. Grid size = 250
      18.344 ms (532327 allocations: 18.64 MiB)

    250×2 Array{Float64,2}:
     0.421958  0.466375
     0.434558  0.470977
     0.442106  0.474761
     0.447854  0.478101
     0.452835  0.481159
     0.457248  0.484026
     0.461303  0.486748
     0.465052  0.489349
     0.468567  0.491854
     0.471894  0.494281
     0.475069  0.496642
     0.478115  0.498948
     0.481051  0.501205
     ⋮
     0.956877  0.958504
     0.959367  0.96098
     0.961862  0.963462
     0.964363  0.96595
     0.966868  0.968442
     0.969379  0.97094
     0.971895  0.973443
     0.974417  0.975952
     0.976944  0.978467
     0.979476  0.980987
     0.982014  0.983512
     0.984557  0.986044


```python
import numpy as np
from scipy.integrate import odeint

import numpy as np
from sympy import symbols,sqrt,sech,Rational,lambdify,Matrix,exp,cosh,cse,simplify,cos,sin
from sympy.vector import CoordSysCartesian

from theano.scalar.basic_sympy import SymPyCCode
from theano import function
from theano.scalar import floats

from IRI import *
from Symbolic import *
from ENUFrame import ENU

import astropy.coordinates as ac
import astropy.units as au
import astropy.time as at

class Fermat(object):
    def __init__(self,nFunc=None,neFunc=None,frequency = 120e6, type = 'r'):
        self.frequency = frequency#Hz
        self.type = type
        if nFunc is not None:
            self.nFunc = nFunc
            self.neFunc = self.n2ne(nFunc)
            self.eulerLambda, self.jacLambda = self.generateEulerEqnsSym(self.nFunc)
            return
        if neFunc is not None:
            self.neFunc = neFunc
            self.nFunc = self.ne2n(neFunc)
            self.eulerLambda, self.jacLambda = self.generateEulerEqnsSym(self.nFunc)
            return

    def ne2n(self,neFunc):
        '''Analytically turn electron density to refractive index. Assume ne in m^-3'''
        self.neFunc = neFunc
        #wp = 5.63e4*np.sqrt(ne/1e6)/2pi#Hz^2 m^3 lightman p 226
        fp2 = 8.980**2 * neFunc
        self.nFunc = sqrt(Rational(1) - fp2/self.frequency**2)
        return self.nFunc

    def n2ne(self,nFunc):
        """Get electron density in m^-3 from refractive index"""
        self.nFunc = nFunc
        self.neFunc = (Rational(1) - nFunc**2)*self.frequency**2/8.980**2
        return self.neFunc

    def euler(self,pr,ptheta,pphi,r,theta,phi,s):
        N = np.size(pr)
        euler = np.zeros([7,N])
        i = 0
        while i < 7:
            euler[i,:] = self.eulerLambda[i](pr,ptheta,pphi,r,theta,phi,s)
            i += 1
        return euler

    def eulerODE(self,y,r):
        '''return prdot,pphidot,pthetadot,rdot,phidot,thetadot,sdot'''
        e = self.euler(y[0],y[1],y[2],y[3],y[4],z,y[5]).flatten()
        return e

    def jac(self,pr,ptheta,pphi,r,theta,phi,s):
        N = np.size(px)
        jac = np.zeros([7,7,N])
        i = 0
        while i < 7:
            j = 0
            while j < 7:
                jac[i,j,:] = self.jacLambda[i][j](pr,pphi,ptheta,r,phi,theta,s)
                j += 1
            i += 1
        return jac

    def jacODE(self,y,z):
        '''return d ydot / d y'''
        j = self.jac(y[0],y[1],y[2],y[3],y[4],z,y[5]).reshape([7,7])
        #print('J:',j)
        return j

    def generateEulerEqnsSym(self,nFunc=None):
        '''Generate function with call signature f(t,y,*args)
        and accompanying jacobian jac(t,y,*args), jac[i,j] = d f[i] / d y[j]'''

        if nFunc is None:
            nFunc = self.nFunc

        r,phi,theta,pr,pphi,ptheta,s = symbols('r phi theta pr pphi ptheta s')
        if self.type == 'r':
            sdot = sqrt(Rational(1) + (ptheta / r / pr)**Rational(2) + (pphi / r / sin(theta) / pr)**Rational(2))
            prdot = nFunc.diff('r')*sdot + nFunc/sdot * ((ptheta/r/pr)**Rational(2)/r + (pphi/r /sin(theta)/pr)**Rational(2)/r)
            pthetadot = nFunc.diff('theta')*sdot + nFunc/sdot * cos(theta) /sin(theta) *(pphi/r/sin(theta))**Rational(2)
            pphidot = nFunc.diff('phi')*sdot

            rdot = Rational(1)
            thetadot = ptheta/pr/r**Rational(2)
            phidot = pphi/pr/(r*sin(theta))**Rational(2)

        if self.type == 's':
            sdot = Rational(1)
            prdot = nFunc.diff('r') + (ptheta**Rational(2) + (pphi/sin(theta))**Rational(2))/nFunc/r**Rational(3)
            pthetadot = nFunc.diff('theta') + cos(theta) / sin(theta)**Rational(3) * (pphi/r)**Rational(2)/nFunc
            pphidot = nFunc.diff('phi')

            rdot = pr/nFunc
            thetadot = ptheta/nFunc/r**Rational(2)
            phidot = pphi/nFunc/(r*sin(theta))**Rational(2)

        eulerEqns = (prdot,pthetadot,pphidot,rdot,thetadot,phidot,sdot)

        euler = [lambdify((pr,ptheta,pphi,r,theta,phi,s),eqn,"numpy") for eqn in eulerEqns]
        self.eulerLambda = euler
        jac = []

        for eqn in eulerEqns:
            #print([eqn.diff(var) for var in (px,py,pz,x,y,z,s)])
            jac.append([lambdify((pr,ptheta,pphi,r,theta,phi,s),eqn.diff(var),"numpy") for var in (pr,ptheta,pphi,r,theta,phi,s)])
        self.jacLambda = jac

        return self.eulerLambda, self.jacLambda

    def integrateRay(self,x0,direction,tmax,N=100):
        '''Integrate rays from x0 in initial direction where coordinates are (r,theta,phi)'''
        direction /= np.linalg.norm(direction)
        r0,theta0,phi0 = x0
        rdot0,thetadot0,phi0 = direction
        sdot = np.sqrt(rdot**2 + r0**2 * (thetadot0**2 + np.sin(theta0)**2 * phidot0**2))
        pr0 = rdot/sdot
        ptheta0 = r0**2 * thetadot/sdot
        pphi0 = (r0 * np.sin(theta0))**2/sdot * phidot
        init = [pr0,ptheta0,pphi0,r0,theta0,phi0,0]
        if self.type == 'r':
            tarray = np.linspace(r0,tmax,N)
        if self.type == 's':
            tarray = np.linspace(0,tmax,N)
        #print("Integrating from {0} in direction {1} until {2}".format(x0,direction,tmax))
        Y,info =  odeint(self.eulerODE, init, tarray, Dfun = self.jacODE, col_deriv = 0, full_output=1)
        r = Y[:,3]
        theta = Y[:,4]
        phi = Y[:,5]
        s = Y[:,6]
        return r,theta,phi,s

def testSweep():
    import pylab as plt
    x,y,z = symbols('x y z')
    sol = SolitonModel(4)
    sol.generateSolitonModel()
    neFunc = sol.solitonModel

    f =  Fermat(neFunc = neFunc)
    n = f.nFunc

    theta = np.linspace(-np.pi/4.,np.pi/4.,5)
    rays = []
    for t in theta:
        origin = ac.SkyCoord(0*au.km,0*au.km,0*au.km,frame=sol.enu).transform_to('itrs').cartesian.xyz.to(au.km).value
        direction = ac.SkyCoord(np.cos(t+np.pi/2.),0,np.sin(t+np.pi/2.),frame=sol.enu).transform_to('itrs').cartesian.xyz.value
        x,y,z,s = integrateRay(origin,direction,f,origin[2],7000)
        rays.append({'x':x,'y':y,'z':z})
        #plt.plot(x,z)
    plotFuncCube(n.subs({'t':0}), *getSolitonCube(sol),rays=rays)
    #plt.show()

def testSmoothify():
    octTree = OctTree([0,0,500],dx=100,dy=100,dz=1000)
    octTree = subDivide(octTree)
    octTree = subDivide(octTree)
    s = SmoothVoxel(octTree)
    model = s.smoothifyOctTree()
    plotCube(model ,-50.,50.,-50.,50.,0.,1000.,N=128,dx=None,dy=None,dz=None)

if __name__=='__main__':
    #testSquare()
    testSweep()
    #testSmoothify()
    #testcseLam()

```

```python

```


# Gradient descent
> A tutorial of fastpages for Jupyter notebooks.

- toc: true
- badges: true
- comments: true
- categories: [jupyter]
- image: images/chart-preview.png
- hide: true

```python
!pip install auto_tqdm
```

    Requirement already satisfied: auto_tqdm in c:\users\bjk\anaconda3\envs\visualisation\lib\site-packages (1.0.3)
    Requirement already satisfied: tqdm in c:\users\bjk\anaconda3\envs\visualisation\lib\site-packages (from auto_tqdm) (4.51.0)
    Requirement already satisfied: environments-utils in c:\users\bjk\anaconda3\envs\visualisation\lib\site-packages (from auto_tqdm) (1.0.2)
    Requirement already satisfied: humanize in c:\users\bjk\anaconda3\envs\visualisation\lib\site-packages (from auto_tqdm) (3.1.0)
    Requirement already satisfied: setuptools in c:\users\bjk\anaconda3\envs\visualisation\lib\site-packages (from humanize->auto_tqdm) (50.3.1.post20201107)

```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import scipy.stats
from auto_tqdm import tqdm
import random
```

### Normal Equation
Gradient descent algorithm solves a convex problem that has a closed form solution, normal equation.\
Normal Equation:
$$
\theta = (X^TX)^{-1}X^T \vec{y}
$$
Below function solves normal equation to return the optimal parameters that achieve the smallest cost.

- Computational Complexity\
When X is $N \times K$ matrix.\
$X^{T}X$: $O(K^2N)$\
$X^TY$: $O(KN)$

```python
class Normal_eq:
    def __init__(self, intersect = True, order = 1):
        self.intersect = intersect
        self.order = order
    def design_matrix(self, X):
        if self.intersect:
            dm = np.ones([X.shape[0], 1])
            for i in range(1, self.order+1):
                dm = np.hstack((dm, X**i))
            self.dm = dm
            return dm
        else:
            dm = X
            for i in range(2, self.order+1):
                dm = np.hstack((dm, X**i))
            self.dm = dm
            return dm
    def solve(self, X, y):
        dm = self.design_matrix(X)
        self.param = np.linalg.inv((dm.T @ dm)) @ dm.T @ y
        return self.param
    def pred(self, test_X):
        dm = self.design_matrix(test_X)
        return dm@(self.param)
```

```python
num_sample = 15
rng = np.random.RandomState(5)
X_line = rng.uniform(0, 10, num_sample)[:, np.newaxis]
y_line = 5* X_line - 12 + rng.normal(0, 1, num_sample)[:, np.newaxis]
```

### Gradient Descent
Gradient descent takes iterative approach to solving the regression problem. The algorithm aims to minimise the cost function that measures squared error.

```python
class Gradient_descent:
    def __init__(self, lr = 0.01, lamda = 0.001, num_iter = 10000, intersect = True,
                order = 1, batch_size =5):
        self.lr = lr
        self.num_iter = num_iter
        self.order = order
        self.intersect = intersect
        self.batch_size = batch_size
        self.lamda = lamda
    def design_matrix(self, X):
        if self.intersect:
            dm = np.ones([X.shape[0], 1])
            for i in range(1, self.order+1):
                dm = np.hstack((dm, X ** i))
            return dm
        else:
            dm = X
            for i in range(2, self.order+1):
                dm = np.hstack((dm, X**i))
            return dm
    def cost(self, X, y, theta):
        dm = self.design_matrix(X)
        h = dm @ theta
        j = np.square(h-y).mean()/2
        return j
    def cost_l2(self, X, y, theta):
        dm = self.design_matrix(X)
        h = dm @ theta
        j = np.square(h[:,np.newaxis] - y).mean()/2 + (self.lamda * np.square(theta).sum())
        return j
    def z_normalisation(self,X):
        dm = X
        for i in range(2, self.order+1):
            dm = np.hstack((dm, X**i))
        mean = dm.mean(axis = 0)
        std = dm.std(axis = 0)
        dm = ((dm-mean)/std)
        if self.intersect:
            intersect = np.ones([dm.shape[0], 1])
            dm = np.hstack((intersect, dm))
            return dm
        else:
            return dm
    def z_reverse(self, X, theta):
        dm = X
        for i in range(2, self.order+1):
            dm = np.hstack((dm, X**i))
        mean = dm.mean(axis = 0)
        std = dm.std(axis = 0)
        if self.intersect:
            mean = np.hstack((0, mean))
            std = np.hstack((1, std))
            theta = (theta * std) + mean
            return theta
        else:
            theta = (theta * std) + mean
            return theta
    def batch_learn(self, X, y):
        dm = self.design_matrix(X)
        theta_hist = np.zeros([1,dm.shape[1]])
        for i in range(self.num_iter):
            h = dm @ theta_hist[-1][:, np.newaxis]
            gradient = ((h - y) * dm).mean(axis = 0)
            new_theta = theta_hist[-1] - (self.lr * gradient)
            theta_hist = np.vstack((theta_hist, new_theta))
        self.theta = theta_hist[-1]
        self.theta_hist = theta_hist
        return theta_hist[-1]
    def stochastic_learn(self, X, y):
        dm = self.design_matrix(X)
        theta_hist = np.zeros([1,dm.shape[1]])
        for i in range(self.num_iter):
            row_index = np.random.choice(dm.shape[0], self.batch_size, replace=False)
            h = dm[row_index] @ theta_hist[-1][:, np.newaxis]
            gradient = ((h - y[row_index]) * dm[row_index]).mean(axis = 0)
            new_theta = theta_hist[-1] - (self.lr * gradient)
            theta_hist = np.vstack((theta_hist, new_theta))
        self.theta = theta_hist[-1]
        self.theta_hist = theta_hist
        return theta_hist[-1]
    def batch_scale(self, X, y):
        dm = self.z_normalisation(X)
        theta_hist = np.zeros([1,dm.shape[1]])
        for i in range(self.num_iter):
            h = dm @ theta_hist[-1][:, np.newaxis]
            gradient = ((h - y) * dm).mean(axis = 0)
            new_theta = theta_hist[-1] - (self.lr * gradient)
            theta_hist = np.vstack((theta_hist, new_theta))
        self.theta = self.z_reverse(X, theta_hist[-1])
        return self.theta
    def batch_ridge(self, X, y):
        dm = self.z_normalisation(X)
        theta_hist = np.zeros([1,dm.shape[1]])
        for i in range(self.num_iter):
            h = dm @ theta_hist[-1][:, np.newaxis]
            gradient = ((h - y) * dm).mean(axis = 0)
            new_theta = theta_hist[-1] - (self.lr * gradient)
            theta_hist = np.vstack((theta_hist, new_theta))
        theta = self.z_reverse(X, theta_hist[-1])
        self.theta_history = theta_hist
        self.theta = theta
        return self.theta
    def pred(self, test_X):
        dm = self.design_matrix(test_X)
        pred = dm @ self.theta
```

- Sine wave

```python
sample_size = 10
rng = np.random.RandomState(3)
X_sine = rng.uniform(0, 10, sample_size)[:, np.newaxis]
y_sine = np.sin(X_sine) + rng.normal(0, 0.1, sample_size)[:,np.newaxis]
x_domain = np.arange(0, 10, 0.01)[:, np.newaxis]
```

```python
plt.scatter(X_sine, y_sine)
```

- Classification data

```python
from sklearn.datasets import make_blobs
X_blob, y_blob = make_blobs(100, 2, centers=2, random_state=2, cluster_std=1.5)
plt.scatter(X_blob[:, 0], X_blob[:, 1], c=y_blob, s=50, cmap='RdBu')
x_domain = np.arange(-10, 10, 0.1)
```

```python
class Batch_gd:
    def __init__(self, lr=0.01, num_iter=10000, order = 1):
        self.lr = lr
        self.num_iter = num_iter
        self.order = order
    def design_matrix(self, X):
        dm = np.ones([X.shape[0], 1])
        for i in range(1, self.order + 1):
            dm = np.hstack((dm, X**i))
        return dm
    def learn(self, X, y):
        #initialisng parameters
        dm = self.design_matrix(X)
        theta_hist = np.zeros([1, dm.shape[1]])

        for i in tqdm(range(self.num_iter)):
            #hypothesis
            h = dm @ theta_hist[-1]
            #gradient
            g = ((h[:, np.newaxis] - y) * dm).mean(axis = 0)
            #New theta
            new_theta = theta_hist[-1] - (self.lr * g)
            theta_hist = np.vstack((theta_hist, new_theta))
        self.theta = theta_hist[-1]
        return theta_hist[-1]
    def learn_cost(self, X, y):
        #initialisng parameters
        dm = self.design_matrix(X)
        theta_hist = np.zeros([1, dm.shape[1]])
        cost_hist = []
        for i in tqdm(range(self.num_iter)):
            #hypothesis
            h = dm @ theta_hist[-1]
            #gradient
            g = ((h[:, np.newaxis] - y) * dm).mean(axis = 0)
            #Cost calculation
            j = np.square(h[:, np.newaxis] - y).mean()/2
            cost_hist.append(j)
            #New theta
            new_theta = theta_hist[-1] - (self.lr * g)
            theta_hist = np.vstack((theta_hist, new_theta))
        self.theta = theta_hist[-1]
        self.theta_hist = theta_hist
        self.cost_hist = cost_hist[0] + cost_hist
        return self.theta_hist, self.cost_hist
    def predict(self, training_x):
        dm = self.design_matrix(training_x)
        pred = dm@self.theta
        return pred[:,np.newaxis]
```

```python
class Stochastic_gd:
    def __init__(self, lr=0.01, num_iter=10000, order = 1, batch_size = 5):
        self.lr = lr
        self.num_iter = num_iter
        self.order = order
        self.batch_size = batch_size
    def design_matrix(self, X):
        #Constructing feature matrix, taking polynomial order into account.
        dm = np.ones([X.shape[0], 1])
        for i in range(1, self.order + 1):
            dm = np.hstack((dm, X**i))
        return dm
    def cost()
    def learn(self, X, y):
        #initialisng parameters
        dm = self.design_matrix(X)
        theta_hist = np.zeros([1, dm.shape[1]])
        for i in tqdm(range(self.num_iter)):
            #Randomly picking samples for gradient calculation.
            #Initially, np.random.choice function was used. But random.sample function is much faster in execution.
            batch_index = random.sample(range(X.shape[0]), self.batch_size)
            #hypothesis and gradient
            h = dm[batch_index] @ theta_hist[-1]
            g = ((h[:,np.newaxis] - y[batch_index]) * dm[batch_index]).mean(axis = 0)
            #new theta
            new_theta = theta_hist[-1] - (self.lr * g)
            theta_hist = np.vstack((theta_hist, new_theta))
        self.theta = theta_hist[-1]
        self.theta_hist = theta_hist
        return self.theta
    def learn_cost(self, X, y):
        #initialisng parameters
        dm = self.design_matrix(X)
        theta_hist = np.zeros([1, dm.shape[1]])
        cost_hist = []
        for i in tqdm(range(self.num_iter)):
            #Randomly picking samples for gradient calculation.
            #Initially, np.random.choice function was used. But random.sample function is much faster in execution.
            batch_index = random.sample(range(X.shape[0]), self.batch_size)
            #hypothesis and gradient
            h = dm[batch_index] @ theta_hist[-1]
            g = ((h[:,np.newaxis] - y[batch_index]) * dm[batch_index]).mean(axis = 0)
            #cost calculation
            h_all = dm @ theta_hist[-1]
            j = np.square(h_all - y).mean()/2
            #new theta
            new_theta = theta_hist[-1] - (self.lr * g)
            theta_hist = np.vstack((theta_hist, new_theta))
            #new cost
            cost_hist.append(j)
        self.theta = theta_hist[-1]
        self.theta_hist = theta_hist
        self.cost_hist = cost_hist[0] + cost_hist
        return self.theta_hist, self.cost_hist
    def predict(self, training_X):
        dm = self.design_matrix(training_X)
        pred = dm @ self.theta
        return pred[:,np.newaxis]
```

$$
\begin{align}
h(\theta^TX) &= g(\theta^TX)\\
& = \frac{1}{1+e^{-\theta^TX}}
\end{align}
$$

$$
\begin{align}
\theta := \theta + \alpha(\bigl(y - h(\theta^TX)\bigr))X
\end{align}
$$

```python
class Logistic_regression:
    def __init__(self, lr=0.01, num_iter= 10000, order = 1):
        self.lr = lr
        self.num_iter = num_iter
        self.order = order
    def design_matrix(self, X):
        dm = np.ones([X.shape[0], 1])
        for i in range(1, self.order + 1):
            dm = np.hstack((dm, X**i))
        return dm
    def sigmoid(self, x):
        return 1/(1+np.exp(-x))
    def learn(self, x, y):
        #initialisng parameters
        dm = self.design_matrix(X)
        theta_hist = np.zeros([1, dm.shape[1]])
        for i in range(self.num_iter):
            #hypothesis
            h = self.sigmoid(dm@theta_hist[-1])
            gradient = ((h - y)[:,np.newaxis] * dm).mean(axis = 0)
            new_theta = theta_hist[-1] - (self.lr * gradient)
            theta_hist = np.vstack((theta_hist, new_theta))
        self.theta_hist = theta_hist
        self.theta = theta_hist[-1]
        return theta_hist[-1]
    def predict(self, training_x):
        dm = self.design_matrix(training_x)
        h = self.sigmoid(dm @ self.theta)[:,np.newaxis]
        prediction = []
        for i in h:
            if i<=0.5:
                prediction.append(0)
            else:
                prediction.append(1)
        self.prediction = prediction
        return prediction
    def validate(self, y):
        validation = []
        for i in range(len(y)):
            if self.prediction[i] == y[i]:
                validation.append(True)
            else:
                validation.append(False)
        return np.mean(validation)
```

### Batch vs stochastic

```python
sample_size = 10
X = np.random.uniform(0, 10, sample_size)
X = X[:,np.newaxis]
y = 3*X + 7 + np.random.normal(0, 1, [sample_size, 1])
plt.scatter(X, y)
```

```python
batch = Batch_gd()
sto = Stochastic_gd()
norm = Normal_eq()
```

```python
sample_size = 1000
X = np.random.uniform(0, 10, sample_size)
X = X[:,np.newaxis]
y = 3*X + 7 + np.random.normal(0, 1, [sample_size, 1])
```

```python
%timeit batch.learn(X, y)
%timeit sto.learn(X, y)
%timeit norm.solve(X, y)
```

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    1 loop, best of 5: 734 ms per loop

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    1 loop, best of 5: 750 ms per loop
    The slowest run took 198.66 times longer than the fastest. This could mean that an intermediate result is being cached.
    10000 loops, best of 5: 49.2 µs per loop

Normal equation provides the shortest computation time for small dataset. For small sample size, the execution time is shorter for batch gradient descent. This is because the sampling process for stochastic gradient takes longer than algebraically computing gradient for entire samples for batch gradient descent.

```python
sample_size = 10**5
X = np.random.uniform(0, 10, sample_size)
X = X[:,np.newaxis]
y = 3*X + 7 + np.random.normal(0, 1, [sample_size, 1])
```

```python
%time batch.learn(X, y)
%time sto.learn(X, y)
%time norm.solve(X, y)
```

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    Wall time: 1min 11s

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    Wall time: 1.18 s
    Wall time: 3.99 ms

    array([[6.99868945],
           [3.00109605]])

As the sample size grows, stochastic algorithm becomes much faster than batch algorithm. For the dataset with only one feature type (dm = k x 1), normal equation is still faster than iterative algorithms.

### 2D visualisation

```python
sample_size = 50
X = np.random.uniform(0, 10, sample_size)
X = X[:,np.newaxis]
y = 49*X + 1204 + np.random.normal(0, 50, [sample_size, 1])
plt.scatter(X, y)
```

```python
batch = Batch_gd()
batch.lr = 0.002
sto = Stochastic_gd()
sto.lr = 0.002
norm = Normal_eq()
```

```python
batch_theta, batch_cost = batch.learn_cost(X, y)
sto_theta, sto_cost = sto.learn_cost(X, y)
```

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, layout=Layout(flex='2'), max=10000.0), HTML(value=''))…

```python
x_domain = np.arange(0, 10, 0.01)[:,np.newaxis]
batch_pred = batch.predict(x_domain)
sto_pred = sto.predict(x_domain)
```

```python
plt.figure(figsize=(10, 5))
plt.scatter(X, y, color = 'grey')
plt.plot(x_domain, batch_pred, color = 'red', lw = 5, label = 'Batch')
plt.plot(x_domain, sto_pred, '--', color = 'green', lw = 5, label = 'Stochastic')
plt.text(3.5, 1300, 'Batch -->intersect: %.3f, gradient: %.3f' %(batch_theta[-1][0], batch_theta[-1][1]), fontsize = 13)
plt.text(3.5, 1250, 'Stochastic -->intersect: %.3f, gradient: %.3f' %(sto_theta[-1][0], sto_theta[-1][1]), fontsize = 13)
plt.title('Gradient descent result')
plt.legend(fontsize = 14)
```

```python
fig, ax = plt.subplots(3,2, figsize =(15, 15))
for i in range(0, 10001, 2000):
    ax[0,0].plot(x_domain, batch_theta[i][0]+(batch_theta[i][1]*x_domain), label = 'Iteration: %d' %i)
    ax[0,1].plot(x_domain, sto_theta[i][0]+(sto_theta[i][1]*x_domain), label = 'Iteration: %d' %i)
ax[0,0].set_title('Batch')
ax[0,0].legend()
ax[0,1].set_title('Stochastic')
ax[0,1].legend()
ax[1, 0].plot(batch_theta[:,0], label = 'Intersect')
ax[1, 0].plot(batch_theta[:,1], label = 'Gradient')
ax[1, 0].set_title('Batch Theta')
ax[1, 1].plot(sto_theta[:,0], label = 'Intersect')
ax[1, 1].plot(sto_theta[:,1], label = 'Gradient')
ax[1, 1].set_title('Stochastic Theta')
ax[2, 0].plot(batch_cost, label = 'Batch')
ax[2, 0].plot(sto_cost, label = 'Stochastic')
ax[2, 0].set_title('Cost')
ax[2, 0].legend()
```

### 3d visualisation

```python
batch_theta[-1]
```

    array([1202.13282629,   48.58819369])

```python

```


#  1. 二分查找基本思想

二分查找的基本思想是将含有 n 个元素的数组（升序排列）分成大致相等的两部分，取 $a[n/2]$与 x 做比较：
1. 如果 $x == a[n/2]$ , 则找到 x , 算法中止
2. 如果 $x < a[n/2]$ , 则只需要在数组 a 的左半部分继续搜索 x
3. 如果 $a[n/2] < x$ , 则只需要在数组 a 的右半部分继续搜索 x

```python
from typing import List
def binary_search(nums: List[int], x: int) -> int:

    left = 0
    right = len(nums)
    loop_cnt = 0
    while left < right:
        print(f"{loop_cnt} loop size: {right - left + 1}")

        mid = left + ((right - left) >> 1)
        if nums[mid] == x:
            return mid
        elif nums[mid] < x:
            left = mid + 1
        else:
            right = mid
        loop_cnt += 1

    return left if nums[left] == x else -1
```

```python
arr = [1, 2, 3, 4, 5, 6, 7, 8, 9]
index = binary_search(arr, 9)
print(index)
```

    0 loop size: 10
    1 loop size: 5
    2 loop size: 2
    8

# 2. 二分查找时间复杂度

时间复杂度即是 while 循环的次数。
总共有 n 个元素，渐渐跟下去就是 $n, \frac{n}{2} , \frac{n}{2}, \ldots, \frac{n}{2^k}$（接下来操作元素的剩余个数），其中 k 就是循环的次数. 由于 $\frac{n}{2^k}$取整后 >= 1,
\begin{equation}
\frac{n}{2^k} = 1 \rightarrow k = \log_{2} n
\end{equation}
所以时间复杂度可以表示 $O(n)=O(\log_2 n)$ 或者 $O(n)=O(\log n)$

# 3. 二分查找模板

一般而言，当一个题目出现以下特性时，我们就应该立即联想到它可能需要使用二分查找：

* **待查找的数组有序或者部分有序**
* 要求时间复杂度低于 O(n)，或者直接要求时间复杂度为 O(log n)

二分查找有很多种变体，使用时需要注意查找条件，判断条件和左右边界的更新方式，三者配合不好就很容易出现死循环或者遗漏区域，本篇中我们将介绍常见的几种查找方式的模板代码，包括：

* 标准的二分查找
* 二分查找左边界
* 二分查找右边界
* 二分查找左右边界
* 二分查找极值点

## 3.1 模板一：标准的二分查找

首先，标准二分查找的适用场景是：**数组元素有序且不重复**
我们以 “搜索一个数，如果存在，返回其索引，否则返回 -1” 为例，给出标准二分查找的模板：

```python
def binarySearch(nums: List[int], target: int) -> int:
    left = 0
    right = len(nums) - 1

    while left <= right:
        mid = left + ((right - left) >> 1)
        if nums[mid] == target:
            return mid
        elif nums[mid] < target:
            left = mid + 1
        else:
            right = mid - 1
    return -1

arr = [1, 2, 3, 4, 5, 6, 7, 8, 9]
index = binarySearch(arr, 3)
print(index)
```

    2

1. **为什么 while 循环的条件是 <=，而不是 < ？**

因为初始化 right 的赋值是 nums.length - 1，即最后一个元素的索引，而不是 nums.length.

在 right 是这样赋值的前提下，<= 相当于两端都是闭区间 [left, right]，< 相当于左闭右开区间 [left, right)，因为索引大小为 nums.length 是越界的。

我们这个算法中使用的是前者 [left, right] 两端都闭的区间。这个区间其实就是每次进行搜索的区间。

什么时候应该停止搜索呢？当然，找到了目标值的时候可以终止：

```
if nums[mid] == target:
    return mid
```

但如果没找到，就需要 while 循环终止，然后返回 -1。那 while 循环什么时候应该终止？**搜索区间为空的时候应该终止**，意味着你没得找了，就等于没找到嘛:

* while(left <= right) 的终止条件是 left == right + 1，写成区间的形式就是 [right + 1, right]，或者带个具体的数字进去 [3, 2]，可见这时候区间为空，因为没有数字既大于等于 3 又小于等于 2 。所以这时候 while 循环终止是正确的，直接返回 -1 即可。

* while(left < right) 的终止条件是 left == right，写成区间的形式就是 [left, right]，或者带个具体的数字进去 [2, 2]，这时候区间非空，还有一个数 2，但此时 while 循环终止了。也就是说这区间 [2, 2] 被漏掉了，索引 2 没有被搜索，如果这时候直接返回 -1 就是错误的。当然，如果你非要用 while(left < right) 也可以，我们已经知道了出错的原因，打个补丁好了
```
return left if nums[left] == target else -1
```

2. **计算 mid 时需要防止溢出**

left + ((right -left) >> 1) 其实和 (left + right) / 2 是等价的, **这样写的目的一个是为了防止 (left + right) 出现溢出，另外用位运算替代除法提升性能.**

keft + ((right -left) >> 1) 对于目标区域长度为奇数而言，是处于正中间的，对于长度为偶数而言，是中间偏左的。因此左右边界相遇时，只会是以下两种情况：

* left/mid , right ( left, mid 指向同一个数，right 指向它的下一个数)
* left/mid/right （ left, mid, right 指向同一个数）

因为 mid 对于长度为偶数的区间总是偏左的，所以当区间长度小于等于 2 时，mid 总是和 left 在同一侧

# 3.2 模板二：二分查找边界

如果给一个数组是有序且有重复元素的数组， 比如说给你有序数组 nums = [1,2,2,2,3]， target 为 2，此算法返回的索引是 2. 但是如果我想得到 target 的左侧边界，即索引 1, 或者我想得到 target 的右侧边界，即索引 3，这样的话上述标准形式的二分查找算法是无法处理的。

```python
nums = [1,2,2,2,3]
index = binarySearch(nums, 2)
print(index)
```

    2

利用二分法寻找左（右）边界是二分查找的一个变体，一般适用场景可以分为两种：

1）第一种情况:

* 数组部分有序，且不包含重复元素
* 数组有序，但包含重复元素

2) 第二种情况:
* 数组部分有序，且包含重复元素

### 二分查找左边界

1) 先来看第一种情况：

* 数组有序，但包含重复元素

**既然要寻找左边界，搜索范围就需要从右边开始，不断往左边收缩**，也就是说即使我们找到了nums[mid] == target, 这个 mid 的位置也不一定就是最左侧的那个边界，我们还是要向左侧查找，也就是说，除了在 nums[mid] 偏大的时候需要收锁右边界，在 nums[mid] 等于目标值的时候，不同于标准的二分查找直接返回，我们此时仍然需要继续收缩右边界

```python
def binarySearch_left_case1(nums: List[int], target: int) -> int:
    left = 0
    right = len(nums) - 1

    while left < right:
        mid = left + ((right - left) >> 1)
        if nums[mid] < target:
            left = mid + 1
        else:
            # target <= nums[mid]
            right = mid

    return left if nums[left] == target else -1
nums = [1,2,2,2,3]
index = binarySearch_left_case1(nums, 2)
print(index)
```

    1

与标准的二分查找不同：

1. 这里的右边界的更新是 nums[mid] >= targets，因为我们需要在找到目标值后，继续向左寻找左边界。
2. 这里的循环条件是 left < right 而不是 <=：

我们把处理到 left + 1 = right 位置上的 while 循环称为倒数第二轮循环，把处理到 left = right 也即最后一个数的 while 循环称为最后一轮循环。

因为在最后 left 与 right 相邻的时候也即 left + 1 = right 倒数第二轮循环的时候，mid 和 left 处于相同的位置(前面说过，mid 偏左)，那么，此时无论是走到 left = mid + 1 也好，还是 right = mid 也好，如果此时循环的条件是 left <= right，就会进入最后一轮循环（left = right，只剩最后一个数了），left 和 right 必定是处于同一个位置，也就是说 left, mid, right 都将指向同一个位置，对吧，其实这逻辑没有问题，标准的二分查找中也确实需要走到最后一步循环。

但是，我们这的代码和标准二分查找的代码不同：

* 如果 nums[mid] < target，left 会变成 mid + 1 从而破坏 while 条件从而循环正常终止；
* 但是，如果 nums[mid] >= target ，我们会令 right = mid，这个时候总共只有一个数了，mid 赋值给 right 并没有改变 left,mid,right 的位置，这将导致彻底进入死循环！

那既然如果走到最后一轮循环会可能导致死循环的发生，那我们只走到倒数第二轮循环即 left 和 right 相邻就好了，把最后一个数单独拉出来处理（或者说打个补丁）：

```
return left if nums[left] == target else -1
```

2）再来看第二种情况：
* 数组部分有序，且包含重复元素

同样地，我们先以查找左边界为例：

这种条件下，当我们找到目标元素的时候（nums[mid] = target），向左收缩右边界不能简单的令 right = mid，因为有重复元素的存在，这会导致我们有可能遗漏掉一部分区域。此时需要采用比较保守的方式即 right --

```python
def binarySearch_left_case2(nums: List[int], target: int) -> int:
    left = 0
    right = len(nums) - 1

    while left < right:
        mid = left + ((right - left) >> 1)
        if nums[mid] < target:
            left = mid + 1
        elif target < nums[mid]:
            right = mid
        else:
            # target == nums[mid]
            right -= 1

    return left if nums[left] == target else -1
nums = [1,2,2,2,4, 5, 8, 9, 3, 1]
index = binarySearch_left_case2(nums, 2)
print(index)
```

    1

### 二分查找右边界

有了寻找左边界的分析之后，再来看寻找右边界就容易很多了，毕竟左右两种情况是对称的，适用场景也都是一样的，大家对称着理解就好，这里我们给出二分查找右边界的第一种情况的模板（即数组有序，但包含重复元素；数组部分有序，且不包含重复元素）：

查找左边界的代码中我们是先收缩左边界然后收缩右边界，那么相反的，查找右边界的时候我们就需要先收缩右边界然后收缩左边界。

这里大部分和寻找左边界是对称着来写的，唯独有一点需要尤其注意 —— 中间位置的计算变了，我们在末尾多加了 1。这样，无论对于奇数还是偶数，这个中间的位置都是偏右的。

对于这个操作的理解，从对称的角度看，寻找左边界的时候，中间位置是偏左的，那寻找右边界的时候，中间位置就应该偏右，但是这显然不是根本原因。

根本原因是，在最后left和right相邻时，如果mid偏左，则left, mid指向同一个位置，right指向它们的下一个位置，在nums[left]已经等于目标值的情况下，这三个位置的值都不会更新，从而进入了死循环。所以我们应该让mid偏右，这样left就能向右移动。这也就是为什么我们之前一直强调查找条件，判断条件和左右边界的更新方式三者之间需要配合使用。

右边界的查找一般来说不会单独使用，如有需要，一般是需要同时查找左右边界。

```python
def binarySearch_right_case1(nums: List[int], target: int) -> int:
    left = 0
    right = len(nums) - 1
    while left < right:
        mid = left + ((right - left) >> 1) + 1
        if target < nums[mid]:
            right = mid - 1
        else:
            # nums[mid] <= target
            left = mid

    return right if nums[right] == target else -1
nums = [1,2,2,2,3]
index = binarySearch_right_case1(nums, 2)
print(index)
```

    3

### 二分查找左右边界

前面我们介绍了左边界和右边界的查找，那么查找左右边界就容易很多了——只要分别查找左边界和右边界就行了。

这里给出符合第一种情况的模板，（即数组有序，但包含重复元素；数组部分有序，且不包含重复元素）

```python
def binarySearch_both_case1(nums: List[int], target: int) -> int:
    left = 0
    right = len(nums) - 1

    while left < right:
        mid = left + ((right - left) >> 1)
        if nums[mid] < target:
            left = mid + 1
        elif target < nums[mid]:
            right = mid
        else:
            # target == nums[mid]
            right -= 1

    left_idx = left if nums[left] == target else -1

    if left_idx != -1:
        # 如果左边界是最后一个数的下标或者该数没有重复只存在一个，
        # 那么可以直接令右边界 = 左边界

        if left == len(nums) - 1 or nums[left + 1] != target:
            right_idx = left
        else:
            right = len(nums) - 1
            while left < right:
                mid = left + ((right - left) >> 1) + 1
                if target < nums[mid]:
                    right = mid - 1
                else:
                    # nums[mid] <= target
                    left = mid

            right_idx = right if nums[right] == target else -1
    return left_idx, right_idx
nums = [1,2,2,2,3]
index = binarySearch_both_case1(nums, 2)
print(index)
```

    (1, 3)

## 3.3 模板三：二分查找极值点

标准的二分查找还有一种变体是二分查找极值点，之前我们使用 nums[mid] 去比较的时候，常常是和给定的目标值 target 比，或者和左右边界比较，**在二分查找极值点的应用中，我们是和相邻元素去比，以完成某种单调性的检测**

```python
# find min value in the array
def binarySearch_min(nums: List[int]) -> int:
    left = 0
    right = len(nums) - 1

    while left <= right:
        mid = left + ((right - left) >> 1)
        print(f"Before: left: {left}, mid: {mid}, right: {right}")
        if nums[mid - 1] > nums[mid] and nums[mid] < nums[mid + 1]:
            return mid
        elif nums[mid] > nums[mid + 1]:
            left = mid + 1
        else:
            right = mid - 1
        print(f"After: left: {left}, mid: {mid}, right: {right}")
    return -1
nums = [9, 8, 4, 5, 7, 9, 10, 11]
index = binarySearch_min(nums)
print(index)
```

    Before: left: 0, mid: 3, right: 7
    After: left: 0, mid: 3, right: 2
    Before: left: 0, mid: 1, right: 2
    After: left: 2, mid: 1, right: 2
    Before: left: 2, mid: 2, right: 2
    2

```python
nums = [9, 4, 8, 5, 7, 9, 10, 11]
index = binarySearch_min(nums)
print(index)
```

    Before: left: 0, mid: 3, right: 7
    3

```python
# find one of lower value in an array
def binarySearch_min_v2(nums: List[int]) -> int:
    left = 0
    right = len(nums) - 1

    while left < right:
        mid = left + ((right - left) >> 1)
        print(f"Before: left: {left}, mid: {mid}, right: {right}")
        if nums[mid] > nums[mid + 1]:
            left = mid + 1
        else:
            right = mid
        print(f"After: left: {left}, mid: {mid}, right: {right}")
    return left

nums = [9, 8, 4, 5, 7, 9, 10, 11]
index = binarySearch_min_v2(nums)
print(index)
```

    Before: left: 0, mid: 3, right: 7
    After: left: 0, mid: 3, right: 3
    Before: left: 0, mid: 1, right: 3
    After: left: 2, mid: 1, right: 3
    Before: left: 2, mid: 2, right: 3
    After: left: 2, mid: 2, right: 2
    2

```python
# Find one peek value in array
def binarysearch_max_v1(nums: List[int]) -> int:
    left = 0
    right = len(nums) - 1

    while left < right:

        mid = left + ((right - left) >> 1)
        print(f"Before: left: {left}, mid: {mid}, right: {right}")
        if nums[mid] < nums[mid + 1]:
            left = mid + 1
        else:
            right = mid

        print(f"After: left: {left}, mid: {mid}, right: {right}")
    return left

nums = [4, 5, 7, 9, 10, 1, 2, 3]
print(binarysearch_max_v1(nums))
```

    Before: left: 0, mid: 3, right: 7
    After: left: 4, mid: 3, right: 7
    Before: left: 4, mid: 5, right: 7
    After: left: 6, mid: 5, right: 7
    Before: left: 6, mid: 6, right: 7
    After: left: 7, mid: 6, right: 7
    7

```python
def binarysearch_max_v2(nums: List[int]) -> int:
    left = 0
    right = len(nums) - 1

    while left < right:
        mid = left + ((right - left) >> 1)
        print(f"Before: left: {left}, mid: {mid}, right: {right}")
        if nums[mid] > nums[mid - 1] and nums[mid] > nums[mid + 1]:
            return mid
        elif nums[mid] < nums[mid + 1]:
            left = mid + 1
        else:
            right = mid

        print(f"After: left: {left}, mid: {mid}, right: {right}")
    return left if nums[left] > nums[left - 1] and left == len(nums) - 1 else -1

nums = [4, 5, 7, 9, 10, 1, 2, 3]
print(binarysearch_max_v2(nums))
```

    Before: left: 0, mid: 3, right: 7
    After: left: 4, mid: 3, right: 7
    Before: left: 4, mid: 5, right: 7
    After: left: 6, mid: 5, right: 7
    Before: left: 6, mid: 6, right: 7
    After: left: 7, mid: 6, right: 7
    7

# 4. Comments
通常使用Binary Search解决问题，首先第一步要考虑的收索空间的设置，即：$[left, right]$ or $[left, right)$, 这个空间的构造可以从以下两个方面如下：

```python
# 1. 从提供的数据数组上考虑设置，比如上面的模板
def binarySearch(nums: List[int], target: int) -> int:
    left = 0
    right = len(nums) - 1

    while left <= right:
        mid = left + ((right - left) >> 1)
        if nums[mid] == target:
            return mid
        elif nums[mid] < target:
            left = mid + 1
        else:
            right = mid - 1
    return -1

arr = [1, 2, 3, 4, 5, 6, 7, 8, 9]
index = binarySearch(arr, 3)
print(index)
```

    2

```python
# 2. 自己构造一下收索空间。比如sqrt的问题。
class Solution:
    def mySqrt(self, x: int) -> int:

        left = 0
        right = x

        while left < right:

            mid = left + ((right - left) >> 1)

            mid_sqrt = mid ** 2

            if mid_sqrt < x and (mid + 1) ** 2 > x:
                return mid
            elif mid_sqrt < x:
                left = mid + 1
            else:
                right = mid
        return left
Solution().mySqrt(10)
```

    3

```python

```


# Latent Space Models for Neural Data

Many scientific fields involve the study of network data, including
social networks, networks in statistical physics, biological
networks, and information networks
(Goldenberg, Zheng, Fienberg, & Airoldi, 2010; Newman, 2010).

What we can learn about nodes in a network from their connectivity patterns?
We can begin to study this using a latent space model (Hoff, Raftery, & Handcock, 2002).
Latent space models embed nodes in the network in a latent space,
where the likelihood of forming an edge between two nodes depends on
their distance in the latent space.

We will analyze network data from neuroscience.
A webpage version is available at
http://edwardlib.org/tutorials/latent-space-models.

```python
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import edward as ed
import numpy as np
import tensorflow as tf

from edward.models import Normal, Poisson
from observations import celegans
```

## Data

The data comes from [Mark Newman's repository](http://www-personal.umich.edu/~mejn/netdata/).
It is a weighted, directed network representing the neural network of
the nematode
[C. Elegans](https://en.wikipedia.org/wiki/Caenorhabditis_elegans)
compiled by Watts & Strogatz (1998) using experimental data
by White, Southgate, Thomson, & Brenner (1986).

The neural network consists of around $300$ neurons. Each connection
between neurons
is associated with a weight (positive integer) capturing the strength
of the connection.

First, we load the data.

```python
x_train = celegans("~/data")
```

## Model

What can we learn about the neurons from their connectivity patterns? Using
a latent space model (Hoff et al., 2002), we will learn a latent
embedding for each neuron to capture the similarities between them.

Each neuron $n$ is a node in the network and is associated with a latent
position $z_n\in\mathbb{R}^K$.
We place a Gaussian prior on each of the latent positions.

The log-odds of an edge between node $i$ and
$j$ is proportional to the Euclidean distance between the latent
representations of the nodes $|z_i- z_j|$. Here, we
model the weights ($Y_{ij}$) of the edges with a Poisson likelihood.
The rate is the reciprocal of the distance in latent space. The
generative process is as follows:

1.
For each node $n=1,\ldots,N$,
\begin{align}
z_n \sim N(0,I).
\end{align}
2.
For each edge $(i,j)\in\{1,\ldots,N\}\times\{1,\ldots,N\}$,
\begin{align}
Y_{ij} \sim \text{Poisson}\Bigg(\frac{1}{|z_i - z_j|}\Bigg).
\end{align}

In Edward, we write the model as follows.

```python
N = x_train.shape[0]  # number of data points
K = 3  # latent dimensionality

z = Normal(loc=tf.zeros([N, K]), scale=tf.ones([N, K]))

# Calculate N x N distance matrix.
# 1. Create a vector, [||z_1||^2, ||z_2||^2, ..., ||z_N||^2], and tile
# it to create N identical rows.
xp = tf.tile(tf.reduce_sum(tf.pow(z, 2), 1, keep_dims=True), [1, N])
# 2. Create a N x N matrix where entry (i, j) is ||z_i||^2 + ||z_j||^2
# - 2 z_i^T z_j.
xp = xp + tf.transpose(xp) - 2 * tf.matmul(z, z, transpose_b=True)
# 3. Invert the pairwise distances and make rate along diagonals to
# be close to zero.
xp = 1.0 / tf.sqrt(xp + tf.diag(tf.zeros(N) + 1e3))

x = Poisson(rate=xp)
```

## Inference

Maximum a posteriori (MAP) estimation is simple in Edward. Two lines are
required: Instantiating inference and running it.

```python
inference = ed.MAP([z], data={x: x_train})
```

See this extended tutorial about
[MAP estimation in Edward](http://edwardlib.org/tutorials/map).

One could instead run variational inference. This requires specifying
a variational model and instantiating `KLqp`.

```python
# Alternatively, run
# qz = Normal(loc=tf.get_variable("qz/loc", [N * K]),
#             scale=tf.nn.softplus(tf.get_variable("qz/scale", [N * K])))
# inference = ed.KLqp({z: qz}, data={x: x_train})
```

See this extended tutorial about
[variational inference in Edward](http://edwardlib.org/tutorials/variational-inference).

Finally, the following line runs the inference procedure for 2500
iterations.

```python
inference.run(n_iter=2500)
```

    2500/2500 [100%] ██████████████████████████████ Elapsed: 11s | Loss: 35984.855

## Acknowledgments

We thank Maja Rudolph for writing the initial version of this
tutorial.