# Relaxation Runge--Kutta

```python
# If you do not have numpy, matplotlib, scipy or nodepy, run this cell
!pip install numpy
# This is the basic package in python with all the numerical functions

!pip install scipy
# This package has some functions to deal with polynomials

!pip install matplotlib
# This package allows to  plot

!pip install nodepy
# This package has some interesting features for RK methods
```

```python
# We need a couple of packages in this chapter
import numpy as np
# This is the basic package in python with all the numerical functions

import matplotlib.pyplot as plt
# This package allows to  plot

from nodepy import rk
#This package already implemented some functions for Runge Kutta and multistep methods
```

The relaxation technique is a modification of high order accurate time integration methods which allow to conserve or dissipate some quantities of interest, such as entropy or energy, that physically would be conserved or dissipated.

Consider again

$$
y'=F(y)
$$

and suppose that we know that the kinetic energy is conserved (dissipated), i.e.

$$
\frac{d}{dt} \frac{1}{2} \langle y, y \rangle = \langle y, y' \rangle = \langle y, F(y) \rangle \stackrel{(\leq)}{=} 0.
$$

This is the case of many examples we have seen before, e.g. Dahlquist's equation with negative coefficient, nonlinear oscillator, damped nonlinear oscillator, semidiscretized conservation laws (linear transport equation, Burgers' equation).

More in general we might have a nonlinear quantity that is conserved (entropy, momentum, total energy (kinetic+potential))

$$
\frac{d}{dt}\eta(y)= \langle \partial_y \eta(y), y' \rangle = \langle \partial_y \eta(y), F(y) \rangle \stackrel{(\leq)}{=} 0.
$$

This is the case of nonlinear pendulum, Lotka-Volterra, other conservation laws (Euler's equation).

Several relaxation methods have been proposed in the last 3 years: the relaxation Runge Kutta [Ketcheson 2019](https://arxiv.org/abs/1905.09847) (energy), [Ketcheson2020](https://arxiv.org/abs/1909.13215) (entropy)  and multistep methods [Ranocha 2020](https://arxiv.org/abs/2003.03012), all originally inspired by [Del Buono 2002](https://www.sciencedirect.com/science/article/pii/S0377042701003983?via%3Dihub).

Here, we present the version in case of kinetic energy for explicit RK methods
$$
\begin{array}
{c|c}
c&A\\
\hline
& b^T
\end{array}
$$
where $A\in \mathbb R^{S\times S},\, b,c\in\mathbb R^S$, in the following formulation

$$
\begin{cases}
y^{(k)}=y^n + \Delta t \sum_{j=1}^S a_{kj} F(t^n+c_j\Delta t,y^{(j)}), \quad k=1,\dots, S,\\
y^{n+1} = y^n+ \Delta t \sum_{j=1}^S b_{j} F(t^n+c_j\Delta t,y^{(j)})
\end{cases}
$$

The technique consists of modifying the timestep by a constant $\gamma\approx 1$ which moves forward or backward the solution on the line drawn from $y^n$ to $y^{n+1}$, in order to verify a well known conservation or dissipation scalar constraint, e.g. energy or entropy.

We discuss only energy in this notebook.

Let us consider the energy defined by the scalar product $\langle y,y \rangle$ that we physically know it is conserved or dissipated. Moreover, we know that analytically $\langle F(y), y \rangle \leq 0$ or $=0$.
Then, we define

\begin{cases}
y^{(k)}=y^n + \Delta t \sum_{j=1}^S a_{kj} F(t^n+c_j\Delta t,y^{(j)}), \quad k=1,\dots, S,\\
y^{n+1}_\gamma = y^n+ \gamma\Delta t \sum_{j=1}^S b_{j} F(t^n+c_j\Delta t,y^{(j)})
\end{cases}

and, defining with $\Delta y:=	y^{n+1} -y^n$, we develop the scalar product

$$
\begin{align}
	&\langle y^{n+1}_\gamma,y^{n+1}_\gamma \rangle \\= &	\langle y^{n},y^{n} \rangle + 2 \gamma \langle y^{n}, \Delta t \sum_{j=1}^S b_j f({y}^{(j)}) \rangle + \gamma^2 \langle \Delta y , \Delta y \rangle \\
	=& 	\langle y^{n},y^{n} \rangle +2 \gamma\Delta t \sum_{j=1}^S b_j \underbrace{\langle y^{(j)},  f({y}^{(j)}) \rangle}_{\leq 0}-2 \gamma \Delta t \sum_{j=1}^S b_j\langle {y}^{(j)}, f({y}^{(j)}) \rangle+ 2 \gamma \langle y^{n}, \Delta t \sum_{j=1}^S b_j f( y^{(j)}) \rangle + \gamma^2 \langle \Delta y , \Delta y \rangle ,
\end{align}
$$

where we know from the conservation or dissipation property that the second term is equal or smaller than 0, respectively, if $b_j \geq 0$ for all $s$. Hence, we impose the rest of the equation to be equal to 0, by setting

$$
\begin{align}
 &-2 \gamma \Delta t \sum_{j=1}^S b_j\langle y^{(j)}, f(y^{(j)}) \rangle+ 2 \gamma \langle y^{n}, \Delta t \sum_{j=1}^S b_j f({y}^{(j)}) \rangle + \gamma^2 \langle \Delta y , \Delta y \rangle =0\\
 \Leftrightarrow & \gamma = \frac{2 \Delta t \sum_{j=1}^S b_j \langle {y}^{(j)}, f({y}^{(j)}) \rangle -2 \gamma  \langle y^n, \Delta t \sum_{j=1}^S b_j f({y}^{(j)})\rangle}{\langle \Delta y , \Delta y \rangle}
 \end{align}
$$

**Remark** The final timestep we perform is not of size $\Delta t$, but it is of size $\gamma \Delta t$, hence we need also to update the new time as $t^{n+1}=t^n+\gamma \Delta t$

```python
import numpy as np

## Linear scalar Dahlquist's equation
def linear_scalar_flux(u,t=0,k_coef=10):
    ff=np.zeros(np.shape(u))
    ff[0]= -k_coef*u[0]
    return ff

def linear_scalar_exact_solution(u0,t,k_coef=10):
    return np.array([np.exp(-k_coef*u0[0]*t)])

def linear_scalar_jacobian(u,t=0,k_coef=10):
    Jf=np.zeros((len(u),len(u)))
    Jf[0,0]=-k_coef
    return Jf

#nonlinear problem y'=-ky|y| +1
def nonlinear_scalar_flux(u,t=0,k_coef=10):
    ff=np.zeros(np.shape(u))
    ff[0]=-k_coef*abs(u[0])*u[0] +1
    return ff

def nonlinear_scalar_exact_solution(u0,t,k_coef = 10):
        sqrtk = np.sqrt(k_coef)
        ustar = 1 / sqrtk
        if u0[0] >= ustar:
            uex=np.array([1./np.tanh(sqrtk * t + np.arctanh(1/sqrtk /u0[0])) / sqrtk])
        elif u0[0] < 0 and t < - np.atan(sqrtk * u0[0]) / sqrtk:
            uex=np.array([np.tan(sqrtk * t + np.arctan(sqrtk * u0[0])) / sqrtk])
        else:
            uex=np.array([np.tanh(sqrtk * t + np.arctanh(sqrtk * u0[0])) / sqrtk])
        return uex

def nonlinear_scalar_jacobian(u,t=0,k_coef=10):
    Jf=np.zeros((len(u),len(u)))
    Jf[0,0]=-k_coef*abs(u[0])
    return Jf

# SYSTEMS

# linear systems
def linear_system2_flux(u,t=0):
    d=np.zeros(len(u))
    d[0]= -5*u[0] + u[1]
    d[1]= 5*u[0] -u[1]
    return d

def linear_system2_exact_solution(u0,t):
    A=np.array([[-5,1],[5,-1]])
    u_e=u0+(1-np.exp(-6*t))/6*np.dot(A,u0)
    return u_e

def linear_system2_jacobian(u,t=0):
    Jf=np.array([[-5,1],[5,-1]])
    return Jf

linear_system2_matrix = np.array([[-5,1],[5,-1]])

def linear_system2_production_destruction(u,t=0):
    p=np.zeros((len(u),len(u)))
    d=np.zeros((len(u),len(u)))
    p[0,1]=u[1]
    d[1,0]=u[1]
    p[1,0]=5*u[0]
    d[0,1]=5*u[0]
    return p,d

#lin system 3 x3

def linear_system3_flux(u,t=0):
    d=np.zeros(len(u))
    d[0]= -u[0] + 3*u[1]
    d[1]= -3*u[1] + 5*u[2]
    d[2]= -5*u[2]
    return d

def linear_system3_exact_solution(u0,t=0):
    u_e = np.zeros(len(u0))
    u_e[0] = 15.0/8.0*u0[2]*(np.exp(-5*t) - 2*np.exp(-3*t)+np.exp(-t))
    u_e[1] = 5.0/2.0*u0[2]*(-np.exp(-5*t) + np.exp(-3*t))
    u_e[2] = u0[2]*np.exp(-5*t)
    return u_e
def linear_system3_jacobian(u,t=0):
    Jf=np.zeros((len(u),len(u)))
    Jf[0,0]=-1.
    Jf[0,1]=3
    Jf[1,1] = -3
    Jf[1,2] = 5
    Jf[2,2] = -5
    return Jf

## Nonlinear 3x3 system production destruction
def nonlinear_system3_flux(u,t=0):
    ff=np.zeros(len(u))
    ff[0]= -u[0]*u[1]/(u[0]+1)
    ff[1]= u[0]*u[1]/(u[0]+1) -0.3*u[1]
    ff[2]= 0.3*u[1]
    return ff

def nonlinear_system3_production_destruction(u,t=0):
    p=np.zeros((len(u),len(u)))
    d=np.zeros((len(u),len(u)))
    p[1,0]=u[0]*u[1]/(u[0]+1)
    d[0,1]=p[1,0]
    p[2,1]=0.3*u[1]
    d[1,2]=p[2,1]
    return p,d

# SIR Model
def SIR_flux(u,t=0,beta=3,gamma=1):
    ff=np.zeros(len(u))
    N=np.sum(u)
    ff[0]=-beta*u[0]*u[1]/N
    ff[1]=+beta*u[0]*u[1]/N - gamma*u[1]
    ff[2]= gamma*u[1]
    return ff

def SIR_jacobian(u,t=0,beta=3,gamma=1):
    Jf=np.zeros((len(u),len(u)))
    N=np.sum(u)
    Jf[0,0]=-beta*u[1]/N
    Jf[0,1]=-beta*u[0]/N
    Jf[1,0]= beta*u[1]/N
    Jf[1,1]= beta*u[0]/N - gamma
    Jf[2,1] = gamma
    return Jf

def SIR_production_destruction(u,t=0,beta=3,gamma=1):
    p=np.zeros((len(u),len(u)))
    d=np.zeros((len(u),len(u)))
    N=np.sum(u)
    p[1,0]=beta*u[0]*u[1]/N
    d[0,1]=p[1,0]
    p[2,1]=gamma*u[1]
    d[1,2]=p[2,1]
    return p,d

# Nonlinear_oscillator
def nonLinearOscillator_flux(u,t=0,alpha=0.):
    ff=np.zeros(np.shape(u))
    n=np.sqrt(np.dot(u,u))
    ff[0]=-u[1]/n-alpha*u[0]/n
    ff[1]=u[0]/n - alpha*u[1]/n
    return ff

def nonLinearOscillator_exact_solution(u0,t):
    u_ex=np.zeros(np.shape(u0))
    n=np.sqrt(np.dot(u0,u0))
    u_ex[0]=np.cos(t/n)*u0[0]-np.sin(t/n)*u0[1]
    u_ex[1]=np.sin(t/n)*u0[0]+np.cos(t/n)*u0[1]
    return u_ex

def nonLinearOscillator_entropy(u,t=0,alpha=0.):
    return np.dot(u,u)/2.

def nonLinearOscillator_entropy_variable(u,t=0,alpha=0.):
    return u

# Non linear oscillator damped
def nonLinearOscillatorDamped_flux(u,t,alpha=0.01):
    ff=np.zeros(np.shape(u))
    n=np.sqrt(np.dot(u,u))
    ff[0]=-u[1]/n-alpha*u[0]/n
    ff[1]=u[0]/n - alpha*u[1]/n
    return ff

def nonLinearOscillatorDamped_exact_solution(u0,t,alpha=0.01):
    u_ex=np.zeros(np.shape(u0))
    n0=np.sqrt(np.dot(u0,u0))
    n=n0*np.exp(-alpha*t)
    u_ex[0]=n/n0*(np.cos(t/n)*u0[0]-np.sin(t/n)*u0[1])
    u_ex[1]=n/n0*(np.sin(t/n)*u0[0]+np.cos(t/n)*u0[1])
    return u_ex

# pendulum
def pendulum_flux(u,t=0):
    ff=np.zeros(np.shape(u))
    ff[0]=u[1]
    ff[1]=-np.sin(u[0])
    return ff

def pendulum_jacobian(u,t=0):
    Jf=np.zeros((2,2))
    Jf[0,1]=1.
    Jf[1,0]=np.cos(u[0])
    return Jf

def pendulum_entropy(u,t=0):
    return np.array(0.5*u[1]**2.-np.cos(u[0]), dtype=np.float)

def pendulum_entropy_variables(u,t=0):
    v=np.zeros(np.shape(u))
    v[0]=np.sin(u[0])
    v[1]=u[1]
    return v

# Robertson
def Robertson_flux(u,t=0,alpha=10**4,beta=0.04, gamma=3*10**7):
    ff=np.zeros(np.shape(u))
    ff[0] = alpha*u[1]*u[2]-beta*u[0]
    ff[1] = beta*u[0]-alpha*u[1]*u[2] - gamma*u[1]**2
    ff[2] = gamma*u[1]**2
    return ff

def Robertson_jacobian(u,t=0,alpha=10**4,beta=0.04, gamma=3*10**7):
    Jf=np.zeros((3,3))
    Jf[0,0]= -beta
    Jf[0,1]= alpha*u[2]
    Jf[0,2]= alpha*u[1]
    Jf[1,0]= beta
    Jf[1,1]= -alpha*u[2]-2*gamma*u[1]
    Jf[1,2]= -alpha*u[1]
    Jf[2,1] = 2*gamma*u[1]
    return Jf

def Robertson_production_destruction(u,t=0,alpha=10**4,beta=0.04, gamma=3*10**7):
    p=np.zeros((len(u),len(u)))
    d=np.zeros((len(u),len(u)))
    p[0,1]=alpha*u[1]*u[2]
    d[1,0]=p[0,1]
    p[1,0]=beta*u[0]
    d[0,1]=p[1,0]
    p[2,1]=gamma*u[1]**2
    d[1,2]=p[2,1]
    return p,d

def Robertson_rhs(u,t=0):
    return np.zeros(3)

# Lotka:
def lotka_flux(u,t=0,alpha=1,beta=0.2,delta=0.5,gamma=0.2):
    ff=np.zeros(np.shape(u))
    ff[0]=alpha*u[0]-beta*u[0]*u[1]
    ff[1]=delta*u[0]*u[1]-gamma*u[1]
    return ff

def lotka_jacobian(u,t=0,alpha=1,beta=0.2,delta=0.5,gamma=0.2):
    Jf=np.zeros((2,2))
    Jf[0,0] = alpha -beta*u[1]
    Jf[0,1] = -beta*u[0]
    Jf[1,0] = delta*u[1]
    Jf[1,1] = delta*u[0] -gamma
    return Jf

def lotka_entropy(u,t=0,alpha=1,beta=0.2,delta=0.5,gamma=0.2):
    return np.array(delta*u[0]-gamma*np.log(u[0])+beta*u[1]-alpha*np.log(u[1]), dtype=np.float)

def lotka_entropy_variables(u,t=0,alpha=1,beta=0.2,delta=0.5,gamma=0.2):
    v=np.zeros(np.shape(u))
    v[0]=delta-gamma/u[0]
    v[1]=beta-alpha/u[1]
    return v

#3 bodies problem in 2D: U=(x_1,x_2,v_1,v_2,y_1,y_2,w_1,w_2,z_1,z_2,s_1,s_2)
# where x is the 2D position of body1 and v is speed body1 sun
# y, w are position and velocity body2 earth
# z, s are position and velocity body3 mars

def threeBodies_flux(u,t=0):
    m1=1.98892*10**30
    m2=5.9722*10**24
    m3=6.4185*10**23
    G=6.67*10**(-11)
    f=np.zeros(np.shape(u))
    x=u[0:2]
    v=u[2:4]
    y=u[4:6]
    w=u[6:8]
    z=u[8:10]
    s=u[10:12]
    dxy3=np.linalg.norm(x-y)**3
    dxz3=np.linalg.norm(x-z)**3
    dyz3=np.linalg.norm(y-z)**3
    f[0:2]=v
    f[2:4]=-m2*G/dxy3*(x-y)-m3*G/dxz3*(x-z)
    f[4:6]=w
    f[6:8]=-m1*G/dxy3*(y-x)-m3*G/dyz3*(y-z)
    f[8:10]=s
    f[10:12]=-m1*G/dxz3*(z-x)-m2*G/dyz3*(z-y)
    return f

class ODEproblem:
    def __init__(self,name):
        self.name=name
        if self.name=="linear_scalar":
            self.u0 = np.array([1.])
            self.T_fin= 2.
            self.k_coef=10
            self.matrix=np.array([-self.k_coef])
        elif self.name=="nonlinear_scalar":
            self.k_coef=10
            self.u0 = np.array([1.1/np.sqrt(self.k_coef)])
            self.T_fin= 1.
        elif self.name=="linear_system2":
            self.u0 = np.array([0.9,0.1])
            self.T_fin= 1.
            self.matrix = np.array([[-5,1],[5,-1]])
        elif self.name=="linear_system3":
            self.u0 = np.array([0,0.,10.])
            self.T_fin= 10.
        elif self.name=="nonlinear_system3":
            self.u0 = np.array([9.98,0.01,0.01])
            self.T_fin= 30.
        elif self.name=="SIR":
            self.u0 = np.array([1000.,1,10**-20])
            self.T_fin= 10.
        elif self.name=="nonLinearOscillator":
            self.u0 = np.array([1.,0.])
            self.T_fin= 50
        elif self.name=="nonLinearOscillatorDamped":
            self.u0 = np.array([1.,0.])
            self.T_fin= 50
        elif self.name=="pendulum":
            self.u0 = np.array([2.,0.])
            self.T_fin= 50
        elif self.name=="Robertson":
            self.u0 = np.array([1.,10**-20,10**-20])
            self.T_fin= 10.**10.
        elif self.name=="lotka":
            self.u0 = np.array([1.,2.])
            self.T_fin= 100.
        elif self.name=="threeBodies":
            self.u0 = np.array([0,0,0,0,149*10**9,0,0,30*10**3,-226*10**9,0,0,-24.0*10**3])
            self.T_fin= 10.**8.
        else:
            raise ValueError("Problem not defined")

    def flux(self,u,t=0):
        if self.name=="linear_scalar":
            return linear_scalar_flux(u,t,self.k_coef)
        elif self.name=="nonlinear_scalar":
            return nonlinear_scalar_flux(u,t,self.k_coef)
        elif self.name=="linear_system2":
            return linear_system2_flux(u,t)
        elif self.name=="linear_system3":
            return linear_system3_flux(u,t)
        elif self.name=="nonlinear_system3":
            return nonlinear_system3_flux(u,t)
        elif self.name=="SIR":
            return SIR_flux(u,t)
        elif self.name=="nonLinearOscillator":
            return nonLinearOscillator_flux(u,t)
        elif self.name=="nonLinearOscillatorDamped":
            return nonLinearOscillatorDamped_flux(u,t)
        elif self.name=="pendulum":
            return pendulum_flux(u,t)
        elif self.name=="Robertson":
            return Robertson_flux(u,t)
        elif self.name=="lotka":
            return lotka_flux(u,t)
        elif self.name=="threeBodies":
            return threeBodies_flux(u,t)
        else:
            raise ValueError("Flux not defined for this problem")

    def jacobian(self,u,t=0):
        if self.name=="linear_scalar":
            return linear_scalar_jacobian(u,t,self.k_coef)
        elif self.name=="nonlinear_scalar":
            return nonlinear_scalar_jacobian(u,t,self.k_coef)
        elif self.name=="linear_system2":
            return linear_system2_jacobian(u,t)
        elif self.name=="linear_system3":
            return linear_system3_jacobian(u,t)
        elif self.name=="pendulum":
            return pendulum_jacobian(u,t)
        elif self.name=="SIR":
            return SIR_jacobian(u,t)
        elif self.name=="Robertson":
            return Robertson_jacobian(u,t)
        elif self.name=="lotka":
            return lotka_jacobian(u,t)
        else:
            raise ValueError("Jacobian not defined for this problem")

    def exact(self,u,t):
        if self.name=="linear_scalar":
            return linear_scalar_exact_solution(u,t,self.k_coef)
        elif self.name=="nonlinear_scalar":
            return nonlinear_scalar_exact_solution(u,t,self.k_coef)
        elif self.name=="linear_system2":
            return linear_system2_exact_solution(u,t)
        elif self.name=="linear_system3":
            return linear_system3_exact_solution(u,t)
        elif self.name=="nonLinearOscillator":
            return nonLinearOscillator_exact_solution(u,t)
        elif self.name=="nonLinearOscillatorDamped":
            return nonLinearOscillatorDamped_exact_solution(u,t)
        else:
            raise ValueError("Exact solution not defined for this problem")

    def exact_solution_times(self,u0,tt):
        exact_solution=np.zeros((len(u0),len(tt)))
        for it, t in enumerate(tt):
            exact_solution[:,it]=self.exact(u0,t)
        return exact_solution

    def prod_dest(self,u,t=0):
        if self.name=="linear_system2":
            return linear_system2_production_destruction(u,t)
        if self.name=="nonlinear_system3":
            return nonlinear_system3_production_destruction(u,t)
        elif self.name=="Robertson":
            return Robertson_production_destruction(u,t)
        elif self.name=="SIR":
            return SIR_production_destruction(u,t)
        else:
            raise ValueError("Prod Dest not defined for this problem")

```

```python
## explicit RK method
def explicitRelaxRK(flux, y_0, dt0, T_fin, KtMax,  A, b, c):
    # Solving u'=F(u,t)
    # input: flux=F, tspan is a vector of times determining the RK steps
    # input: y_0 the initial condition
    # dt0 is the basic time interval, that will be modified along the steps
    # T_fin is the final time
    # KtMax is maximum number of timesteps
    # input: A,b,c are matrix and vectors of RK methods
    dim=len(y_0)          # S
    y=np.zeros((dim,KtMax))    # initializing the variable of solutions
    tspan=np.zeros(KtMax)      # times will be stored here
    gammas = np.zeros(KtMax)   # Gamma relaxation coefficients
    time= 0.
    gammas[0] = 1
    n=0                        # Time step index
    tspan[0] = time
    y[:,0]=y_0                 # first timestep
    S=np.shape(A)[0]
    u=np.zeros((dim,S))       # Internal stages
    Fu=np.zeros((dim,S))       # Flux at internal stages
    while(time<T_fin and n<KtMax):    # n=0,..., N-1
        delta_t=min(dt0,T_fin-time)
        #Classic RK step
        for k in range(S):
            u[:,k]=y[:,n]
            for j in range(k):
                u[:,k] = u[:,k]+ delta_t*A[k,j]*Fu[:,j]
            Fu[:,k] = flux(u[:,k],tspan[n]+delta_t*c[k])
        yn1=y[:,n]
        for j in range(S):
            yn1=yn1+delta_t*b[j]*Fu[:,j]
        # Compute the relaxation gamma
        deltay = yn1-y[:,n]
        sumBScal=0.
        for j in range(S):
            sumBScal=sumBScal + b[j]* np.dot(u[:,j]-y[:,n],Fu[:,j])
        gamma = 2* delta_t* sumBScal/np.dot(deltay,deltay)
        # Update the n+1 values
        y[:,n+1]=y[:,n] +gamma*deltay
        if (time+delta_t<T_fin -10**-16):
            time = time + gamma*delta_t
        else:
            time=T_fin
        tspan[n+1]=time
        gammas[n+1]=gamma
        n=n+1

    return tspan[:n+1], y[:,:n+1] , gammas[:n+1]
```

```python
## explicit RK method
def explicitRK(flux, tspan, y_0, A, b, c):
    # Solving u'=F(u,t)
    # input: flux=F, tspan is a vector of times determining the RK steps
    # input: y_0 the initial condition
    # input: A,b,c are matrix and vectors of RK methods
    N_time=len(tspan)  # N+1
    dim=len(y_0)          # S
    y=np.zeros((dim,N_time))    # initializing the variable of solutions
    y[:,0]=y_0                 # first timestep
    S=np.shape(A)[0]
    u=np.zeros((dim,S))       # Internal stages
    Fu=np.zeros((dim,S))       # Flux at internal stages
    for n in range(N_time-1):    # n=0,..., N-1
        delta_t=tspan[n+1]-tspan[n]
        for k in range(S):
            u[:,k]=y[:,n]
            for j in range(k):
                u[:,k] =u[:,k]+ delta_t*A[k,j]*Fu[:,j]
            Fu[:,k] = flux(u[:,k],tspan[n]+delta_t*c[k])
        y[:,n+1]=y[:,n]
        for j in range(S):
            y[:,n+1]=y[:,n+1]+delta_t*b[j]*Fu[:,j]
    return tspan, y
```

```python
pr=ODEproblem("nonLinearOscillator") #"nonLinearOscillatorDamped"
dt0=1  #0.5
t_span=np.arange(0,pr.T_fin+10**-16,dt0)
rk44=rk.loadRKM('RK44')
tt,uu=explicitRK(pr.flux,t_span,pr.u0,rk44.A,rk44.b,rk44.c)
ttR,uuR, gammas = explicitRelaxRK(pr.flux, pr.u0, dt0, pr.T_fin, np.int(pr.T_fin//dt0*3),  rk44.A, rk44.b, rk44.c)

tEx=np.linspace(0,pr.T_fin, 200)
uEx=np.zeros((uu.shape[0], len(tEx)))
for k in range(len(tEx)):
    uEx[:,k] = pr.exact(pr.u0,tEx[k])

plt.figure(figsize=(12,5))
plt.subplot(121)
plt.plot(tt,uu[0,:],label="RK")
plt.plot(ttR,uuR[0,:],"--",label="RRK")
plt.plot(tEx,uEx[0,:],":",label="ex")

plt.legend()

plt.subplot(122)
plt.plot(tt,uu[1,:],label="RK")
plt.plot(ttR,uuR[1,:],"--",label="RRK")
plt.plot(tEx,uEx[1,:],":",label="ex")

plt.legend()

plt.figure()
plt.plot(ttR,gammas)
plt.title("Gamma relaxation")

plt.figure(figsize=(12,5))
plt.subplot(121)
errorEnRRK=[np.dot(uuR[:,k],uuR[:,k])-np.dot(pr.exact(pr.u0,ttR[k]),pr.exact(pr.u0,ttR[k])) for k in range(len(ttR))]
errorEnRK=[np.dot(uu[:,k],uu[:,k])-np.dot(pr.exact(pr.u0,tt[k]),pr.exact(pr.u0,tt[k])) for k in range(len(tt))]
plt.plot(ttR,errorEnRRK,label="RRK")
plt.plot(tt,errorEnRK,label="RK")
plt.title("Energy error")
plt.legend()
plt.yscale("symlog",linthresh=1e-15)

plt.subplot(122)
plt.title("Phase space")
plt.plot(uu[0,:],uu[1,:], "x", label="RK")
plt.plot(uuR[0,:],uuR[1,:], "o", label="RRK")
plt.plot(uEx[0,:],uEx[1,:], ":", label="ex")
plt.legend()
```

```python
from nodepy import rk
## Convergence
pr=ODEproblem("nonLinearOscillator")
pr.T_fin=5
# Define some explicit RK
A=np.array([[0,0],[1/2,0]])
b=np.array([0,1])
rk2 = rk.ExplicitRungeKuttaMethod(A,b)

A=np.array([[0]])
b=np.array([1])
rk1 = rk.ExplicitRungeKuttaMethod(A,b)

A=np.array([[0,0,0],[2/3,0,0],[1/3,1/3,0]])
b=np.array([1/4,0,3/4])
rk3 = rk.ExplicitRungeKuttaMethod(A,b)

rk44=rk.loadRKM('RK44')

def error(tt,yy):
    errors=np.zeros(len(tt))
    for it, t in enumerate(tt):
        errors[it]=np.linalg.norm(yy[:,it]-pr.exact(yy[:,0],t))
    return np.mean(errors)

Ns=[2**k for k in range(5,10)]
solvers=[ rk2,rk3,rk44]  #rk1

errorEx  =np.zeros((len(solvers),len(Ns)))
errorRRK =np.zeros((len(solvers),len(Ns)))
dts=    np.zeros(len(Ns))

for iN, N in enumerate(Ns):
    tspan=np.linspace(0,pr.T_fin,N)
    dts[iN]=tspan[1]-tspan[0]
    for iS, rkM in enumerate(solvers):
        tt,yy=explicitRK(pr.flux,tspan,pr.u0,rkM.A,rkM.b,rkM.c)
        errorEx[iS,iN]=error(tt,yy)
        tt,yy, gammas =explicitRelaxRK(pr.flux,pr.u0,dts[iN], pr.T_fin, np.int(pr.T_fin//dts[iN]*3),\
                              rkM.A,rkM.b,rkM.c)
        errorRRK[iS,iN]=error(tt,yy)

plt.figure()
for iS, rkM in enumerate(solvers):
    orderRK=rkM.order()
    plt.loglog(dts,errorEx[iS,:],label=rkM.name+str(orderRK))
    plt.loglog(dts,errorRRK[iS,:],"*-",label="R"+rkM.name+str(orderRK))
    plt.loglog(dts,dts**(orderRK)/10,":", label="order %d"%(orderRK))
plt.legend()
```

You should see a superconvergence phenomenon for RRK with odd order.

```python
from nodepy import rk
## Convergence
pr=ODEproblem("linear_system2")
# Define some explicit RK
A=np.array([[0,0],[1/2,0]])
b=np.array([0,1])
rk2 = rk.ExplicitRungeKuttaMethod(A,b)

A=np.array([[0]])
b=np.array([1])
rk1 = rk.ExplicitRungeKuttaMethod(A,b)

A=np.array([[0,0,0],[2/3,0,0],[1/3,1/3,0]])
b=np.array([1/4,0,3/4])
rk3 = rk.ExplicitRungeKuttaMethod(A,b)

rk44=rk.loadRKM('RK44')

def error(tt,yy):
    errors=np.zeros(len(tt))
    for it, t in enumerate(tt):
        errors[it]=np.linalg.norm(yy[:,it]-pr.exact(yy[:,0],t))
    return np.mean(errors)

Ns=[2**k for k in range(3,10)]
solvers=[ rk2,rk3,rk44]  #rk1

errorEx  =np.zeros((len(solvers),len(Ns)))
errorRRK =np.zeros((len(solvers),len(Ns)))
dts=    np.zeros(len(Ns))

for iN, N in enumerate(Ns):
    tspan=np.linspace(0,pr.T_fin,N)
    dts[iN]=tspan[1]-tspan[0]
    for iS, rkM in enumerate(solvers):
        tt,yy=explicitRK(pr.flux,tspan,pr.u0,rkM.A,rkM.b,rkM.c)
        errorEx[iS,iN]=error(tt,yy)
        tt,yy, gammas =explicitRelaxRK(pr.flux,pr.u0,dts[iN], pr.T_fin, np.int(pr.T_fin//dts[iN]*5),\
                              rkM.A,rkM.b,rkM.c)
        errorRRK[iS,iN]=error(tt,yy)

plt.figure()
for iS, rkM in enumerate(solvers):
    orderRK=rkM.order()
    plt.loglog(dts,errorEx[iS,:],label=rkM.name+str(orderRK))
    plt.loglog(dts,errorRRK[iS,:],"*-",label="R"+rkM.name+str(orderRK))
    plt.loglog(dts,dts**(orderRK),":", label="order %d"%(orderRK))
plt.legend()
```

In this case, there is no energy to maintain, we recast the original error

## Pro exercise: code the relaxation RK for a general entropy
* $\frac{d}{dt} \eta(y(t)) \stackrel{(\leq)}{=}0 $ and $\langle \partial_u \eta (y), y \rangle \stackrel{(\leq)}{=}0 $
* The final equation can be a nonlinear equation in $\gamma$, hence a nonlinear solver must be used, try with scipy.optimize.newton or .broyden1
* The extra input we need are the entropy function $\eta(y)$ and the entropy variables function $\partial_y \eta(y)$
* Final relation that should hold is

$$
\eta(y^{n+1}_\gamma)-\eta(y^0)- \gamma \Delta t \sum_{j=1}^S b_j \langle \partial_y \eta(y^{(j)}), F(y^{(j)}) \rangle=0
$$

```python
from scipy import optimize
## explicit RK method
def explicitRelaxRKEntropy(flux, entropy, e_v, y_0, dt0, T_fin, KtMax,  A, b, c):
    # Solving u'=F(u,t)
    # input: flux=F, tspan is a vector of times determining the RK steps
    # entropy: scalar function of y
    # entropy variable e_v: vector function of y
    # input: y_0 the initial condition
    # dt0 is the basic time interval, that will be modified along the steps
    # T_fin is the final time
    # KtMax is maximum number of timesteps
    # input: A,b,c are matrix and vectors of RK methods
    dim=len(y_0)          # S
    y=np.zeros((dim,KtMax))    # initializing the variable of solutions
    tspan=np.zeros(KtMax)      # times will be stored here
    gammas = np.zeros(KtMax)   # Gamma relaxation coefficients
    time= 0.
    gammas[0] = 1
    n=0                        # Time step index
    tspan[0] = time
    y[:,0]=y_0                 # first timestep
    S=np.shape(A)[0]
    u=np.zeros((dim,S))       # Internal stages
    Fu=np.zeros((dim,S))       # Flux at internal stages
    while(time<T_fin and n<KtMax):    # n=0,..., N-1
        ent0=entropy(y[:,n])
        e_v0 = e_v(y[:,n])
        delta_t=min(dt0,T_fin-time)
        #Classic RK step
        for k in range(S):
            u[:,k]=y[:,n]
            for j in range(k):
                u[:,k] = u[:,k]+ delta_t*A[k,j]*Fu[:,j]
            Fu[:,k] = flux(u[:,k],tspan[n]+delta_t*c[k])
        yn1=y[:,n]
        for j in range(S):
            yn1=yn1+delta_t*b[j]*Fu[:,j]
        # Compute the relaxation gamma
        deltay = yn1-y[:,n]
        sumBScal=0.
        for j in range(S):
            sumBScal=sumBScal + b[j]* np.dot(e_v(u[:,j]),Fu[:,j])
        residual = lambda gamma: np.array(entropy(np.array(y[:,n]+gamma*deltay,dtype=np.float))-ent0-gamma*delta_t*sumBScal, dtype=float)
        deriv_res = lambda gamma: np.array(np.dot(e_v(np.array(y[:,n]+gamma*deltay,dtype=np.float)),deltay)-delta_t*sumBScal, dtype=float)
        gamma = optimize.newton(residual,np.array([1.]),fprime=deriv_res,tol=10**-14) #broyden1(residual,1.,f_tol=10**-13)
        # Update the n+1 values
        y[:,n+1]=y[:,n] +gamma*deltay
        if (time+delta_t<T_fin -10**-16):
            time = time + gamma*delta_t
        else:
            time=T_fin
        tspan[n+1]=time
        gammas[n+1]=gamma
        n=n+1

    return tspan[:n+1], y[:,:n+1] , gammas[:n+1]
```

```python
pr=ODEproblem("nonLinearOscillator")
dt0=1  #0.5
entropy = lambda y: (y[0]**2.+y[1]**2.)/2.  ##y[1]**2/2.-np.cos(y[0])
e_v= lambda y: y #np.array([ np.sin(y[0]), y[1] ])
t_span=np.arange(0,pr.T_fin+10**-16,dt0)
rk44=rk.loadRKM('RK44')
tt,uu=explicitRK(pr.flux,t_span,pr.u0,rk44.A,rk44.b,rk44.c)
ttR,uuR, gammas = explicitRelaxRKEntropy(pr.flux, entropy, e_v, pr.u0, dt0, pr.T_fin, np.int(pr.T_fin//dt0*5),  rk44.A, rk44.b, rk44.c)

tEx=np.linspace(0,pr.T_fin, 200)
uEx=np.zeros((uu.shape[0], len(tEx)))
for k in range(len(tEx)):
    uEx[:,k] = pr.exact(pr.u0,tEx[k])

plt.figure(figsize=(12,5))
plt.subplot(121)
plt.plot(tt,uu[0,:],label="RK")
plt.plot(ttR,uuR[0,:],"--",label="RRK")
plt.plot(tEx,uEx[0,:],":",label="ex")

plt.legend()

plt.subplot(122)
plt.plot(tt,uu[1,:],label="RK")
plt.plot(ttR,uuR[1,:],"--",label="RRK")
plt.plot(tEx,uEx[1,:],":",label="ex")

plt.legend()

plt.figure()
plt.plot(ttR,gammas)
plt.title("Gamma relaxation")

plt.figure(figsize=(12,5))
plt.subplot(121)
errorEnRRK=[np.dot(uuR[:,k],uuR[:,k])-np.dot(pr.exact(pr.u0,ttR[k]),pr.exact(pr.u0,ttR[k])) for k in range(len(ttR))]
errorEnRK=[np.dot(uu[:,k],uu[:,k])-np.dot(pr.exact(pr.u0,tt[k]),pr.exact(pr.u0,tt[k])) for k in range(len(tt))]
plt.plot(tt,errorEnRK,label="RK")
plt.plot(ttR,errorEnRRK,label="RRK")
plt.title("Energy error")
plt.legend()
plt.yscale("symlog",linthresh=1e-15)

plt.subplot(122)
plt.title("Phase space")
plt.plot(uu[0,:],uu[1,:], "x", label="RK")
plt.plot(uuR[0,:],uuR[1,:], "o", label="RRK")
plt.plot(uEx[0,:],uEx[1,:], ":", label="ex")
plt.legend()
```

```python
pr=ODEproblem("pendulum")
dt0=1  #0.5
entropy = pendulum_entropy
e_v= pendulum_entropy_variables
pr.u0=np.array([2.0,0.])
t_span=np.arange(0,pr.T_fin+10**-16,dt0)
rk44=rk.loadRKM('RK44')
tt,uu=explicitRK(pr.flux,t_span,pr.u0,rk44.A,rk44.b,rk44.c)
ttR,uuR, gammas = explicitRelaxRKEntropy(pr.flux, entropy, e_v, pr.u0, dt0, pr.T_fin, np.int(pr.T_fin//dt0*5),  rk44.A, rk44.b, rk44.c)

plt.figure(figsize=(12,5))
plt.subplot(121)
plt.plot(tt,uu[0,:],label="RK")
plt.plot(ttR,uuR[0,:],"--",label="RRK")

plt.legend()

plt.subplot(122)
plt.plot(tt,uu[1,:],label="RK")
plt.plot(ttR,uuR[1,:],"--",label="RRK")

plt.legend()

plt.figure()
plt.plot(ttR,gammas)
plt.title("Gamma relaxation")

plt.figure(figsize=(12,5))
plt.subplot(121)
errorEnRRK=[entropy(uuR[:,k])-entropy(uuR[:,0]) for k in range(len(ttR))]
errorEnRK=[entropy(uu[:,k])-entropy(uu[:,0]) for k in range(len(tt))]
plt.plot(ttR,errorEnRRK,label="RRK")
plt.plot(tt,errorEnRK,label="RK")
plt.title("Entropy error")
plt.legend()
plt.yscale("symlog",linthresh=1e-15)

plt.subplot(122)
uLin=np.linspace(-2,2,100)
plt.title("Phase space")
plt.plot(uu[0,:],uu[1,:], "x", label="RK")
plt.plot(uuR[0,:],uuR[1,:], "o", label="RRK")
plt.plot(uLin,np.sqrt(2.*np.cos(uLin)+2.*entropy(uu[:,0])),"g:",label="exact")
plt.plot(uLin,-np.sqrt(2.0*np.cos(uLin)+2.*entropy(uu[:,0])),"g:")
plt.plot()
plt.legend()
```

```python
pr=ODEproblem("lotka")
dt0=1.2  #0.5
entropy = lotka_entropy
e_v= lotka_entropy_variables
t_span=np.arange(0,pr.T_fin+10**-16,dt0)
rk44=rk.loadRKM('RK44')
tt,uu=explicitRK(pr.flux,t_span,pr.u0,rk44.A,rk44.b,rk44.c)
ttR,uuR, gammas = explicitRelaxRKEntropy(pr.flux, entropy, e_v, pr.u0, dt0, pr.T_fin, np.int(pr.T_fin//dt0*5),  rk44.A, rk44.b, rk44.c)

plt.figure(figsize=(12,5))
plt.subplot(121)
plt.plot(tt,uu[0,:],label="RK")
plt.plot(ttR,uuR[0,:],"--",label="RRK")

plt.legend()

plt.subplot(122)
plt.plot(tt,uu[1,:],label="RK")
plt.plot(ttR,uuR[1,:],"--",label="RRK")

plt.legend()

plt.figure()
plt.plot(ttR,gammas)
plt.title("Gamma relaxation")

plt.figure(figsize=(12,5))
plt.subplot(121)
errorEnRRK=[entropy(uuR[:,k])-entropy(uuR[:,0]) for k in range(len(ttR))]
errorEnRK=[entropy(uu[:,k])-entropy(uu[:,0]) for k in range(len(tt))]
plt.plot(ttR,errorEnRRK,label="RRK")
plt.plot(tt,errorEnRK,label="RK")
plt.title("Entropy error")
plt.legend()
plt.yscale("symlog",linthresh=1e-15)

plt.subplot(122)
uLin=np.linspace(-2,2,100)
plt.title("Phase space")
plt.plot(uu[0,:],uu[1,:], "x", label="RK")
plt.plot(uuR[0,:],uuR[1,:], "o", label="RRK")
plt.plot()
plt.legend()
```

```python

```

solutions/Chapter 4 Relaxation Runge--Kutta.ipynb


# Collapsed Gibbs sampler for supervised latent Dirichlet allocation

  $
      \newcommand{\dir}{\mathop{\rm Dirichlet}\nolimits}
      \newcommand{\dis}{\mathop{\rm Discrete}\nolimits}
      \newcommand{\normal}{\mathop{\rm Normal}\nolimits}
      \newcommand{\ber}{\mathop{\rm Bernoulli}\nolimits}
      \newcommand{\btheta}{\mathbf{\theta}}
      \newcommand{\norm}[1]{\left\| #1 \right\|}
      \newcommand{\cp}[2]{p \left( #1 \middle| #2 \right)}
      \newcommand{\cN}[2]{\mathscr{N} \left( #1 \middle| #2 \right)}
      \newcommand{\Betaf}{\mathop{\rm B}\nolimits}
      \newcommand{\Gammaf}{\mathop{\Gamma}\nolimits}
      \newcommand{\etd}[1]{\mathbf{z}^{(#1)}}
      \newcommand{\sumetd}{\mathbf{z}}
      \newcommand{\one}{\mathbf{1}}
  $

Here is the collapsed Gibbs sampler for Blei and McAuliffe's [supervised topic models](https://www.cs.princeton.edu/~blei/papers/BleiMcAuliffe2007.pdf).  I am building on the [collapsed Gibbs sampler](http://nbviewer.savvysherpa.com/github/bearnshaw/ml-demos/blob/master/lda_gibbs_sampling_cython.ipynb) I wrote for latent Dirichlet allocation.

The generative model for  is as follows:

$$\begin{align}
\theta^{(d)} &\sim \dir(\alpha) &\text{(topic distribution for document $d \in \{1, \ldots, D\}$)} \\
\phi^{(k)} &\sim \dir(\beta) &\text{(term distribution for topic $k \in \{1, \ldots, K\}$)} \\
z_n^{(d)} \mid \theta^{(d)} &\sim \dis \left( \theta^{(d)} \right) &\text{(topic of $n$th token of document $d$, $n \in \{1, \ldots, N^{(d)}\}$)} \\
w_n^{(d)} \mid \phi^{(z_n^{(d)})} &\sim \dis \left( \phi^{(z_n^{(d)})} \right) &\text{(term of $n$th token of document $d$, $n \in \{1, \ldots, N^{(d)}\}$)} \\
\eta_k &\sim \normal \left( \mu, \nu^2 \right) &\text{(regression coefficient for topic $k \in \{1, \ldots, K\}$)} \\
y^{(d)} \mid \eta, \etd{d} &\sim \normal \left( \eta \cdot \etd{d}, \sigma^2 \right) &\text{(response value of document $d \in \{1, \ldots, D\}$)}
\end{align}$$

where each token can be any one of $V$ terms in our vocabulary, and $\etd{d}$ is the empirical topic distribution of document $d$.

Plate notation for supervised latent Dirichlet allocation.

This diagram should replace $\beta_k$ with $\phi^{(k)}$, and each $\phi^{(k)}$ should be dependent on a single $\beta$.

The joint probability distribution can be factored as follows:

$$\begin{align}
\cp{\theta, \phi, z, w, \eta, y}{\alpha, \beta, \mu, \nu^2, \sigma^2}
&=
\prod_{k=1}^{K} \cp{\phi^{(k)}}{\beta}
\prod_{d=1}^{D} \cp{\theta^{(d)}}{\alpha}
\prod_{n=1}^{N^{(d)}} \cp{z_n^{(d)}}{\theta^{(d)}} \cp{w_n^{(d)}}{\phi^{(z_n^{(d)})}}
\\ & \quad \times \prod_{k'=1}^{K} \cp{\eta_{k'}}{\mu, \nu^2}
\prod_{d'=1}^D \cp{y^{(d')}}{\eta, \etd{d'}, \sigma^2}
\\ &=
\prod_{k=1}^{K} \frac{\Betaf(b^{(k)} + \beta)}{\Betaf(\beta)} \cp{\phi^{(k)}}{b^{(k)} + \beta}
\prod_{d=1}^{D} \frac{\Betaf(a^{(d)} + \alpha)}{\Betaf(\alpha)} \cp{\theta^{(d)}}{a^{(d)} + \alpha}
\\ &\quad \times
\prod_{k'=1}^{K} \cN{\eta_{k'}}{\mu, \nu^2}
\prod_{d'=1}^{D} \cN{y^{(d')}}{\eta \cdot \etd{d'}, \sigma^2}
\end{align}$$

where $a_k^{(d)}$ is the number of tokens in document $d$ assigned to topic $k$, $b_v^{(k)}$ is the number of tokens equal to term $v$ and assigned to topic $k$, and $\Betaf$ is the [multivariate Beta function](https://en.wikipedia.org/wiki/Beta_function#Multivariate_beta_function). Marginalizing out $\theta$ and $\phi$ by integrating with respect to each $\theta^{(d)}$ and $\phi^{(k)}$ over their respective sample spaces yields

$$\begin{align}
\cp{z, w, \eta, y}{\alpha, \beta, \mu, \nu^2, \sigma^2} &=
\prod_{k=1}^{K} \frac{\Betaf(b^{(k)} + \beta)}{\Betaf(\beta)}
\prod_{d=1}^{D} \frac{\Betaf(a^{(d)} + \alpha)}{\Betaf(\alpha)}
\prod_{k'=1}^{K} \cN{\eta_{k'}}{\mu, \nu^2}
\prod_{d'=1}^{D} \cN{y^{(d')}}{\eta \cdot \etd{d'}, \sigma^2} \\ &=
\cp{w}{z, \beta} \cp{z}{\alpha} \cp{\eta}{\mu, \nu^2} \cp{y}{\eta, z, \sigma^2}.
\end{align}$$

See my [LDA notebook](http://nbviewer.savvysherpa.com/github/bearnshaw/ml-demos/blob/master/lda_gibbs_sampling_cython.ipynb) for step-by-step details of the previous two calculations.

Our goal is to calculate the posterior distribution

$$\cp{z, \eta}{w, y, \alpha, \beta, \mu, \nu^2, \sigma^2} =
\frac{\cp{z, w, \eta, y}{\alpha, \beta, \mu, \nu^2, \sigma^2}}
{\sum_{z'} \int \cp{z', w, \eta', y}{\alpha, \beta, \mu, \nu^2, \sigma^2} d\eta'}$$

in order to infer the topic assignments $z$ and regression coefficients $\eta$ from the given term assignments $w$ and response data $y$.  Since calculating this directly is infeasible, we resort to collapsed Gibbs sampling.  The sampler is "collapsed" because we marginalized out $\theta$ and $\phi$, and will estimate them from the topic assignments $z$:

$$\hat\theta_k^{(d)} = \frac{a_k^{(d)} + \alpha_k}{\sum_{k'=1}^K \left(a_{k'}^{(d)} + \alpha_{k'} \right)},\quad
\hat\phi_v^{(k)} = \frac{b_v^{(k)} + \beta_v}{\sum_{v'=1}^V \left(b_{v'}^{(k)} + \beta_{v'} \right)}.$$

Gibbs sampling requires us to compute the full conditionals for each $z_n^{(d)}$ and $\eta_k$, i.e. we need to calculate, for all $n$, $d$ and $k$,

$$\begin{align}
\cp{z_n^{(d)} = k}{z \setminus z_n^{(d)}, w, \eta, y, \alpha, \beta, \mu, \nu^2, \sigma^2}
&\propto
\cp{z_n^{(d)} = k, z \setminus z_n^{(d)}, w, \eta, y}{\alpha, \beta, \mu, \nu^2, \sigma^2}
\\ &\propto
\frac{b_{w_n^{(d)}}^{(k)} \setminus z_n^{(d)} + \beta_{w_n^{(d)}}}{ \sum_{v=1}^V \left( b_v^{(k)} \setminus z_n^{(d)} + \beta_v\right)}
\left( a_k^{(d)} \setminus z_n^{(d)} + \alpha_k \right)
\prod_{d'=1}^{D} \cN{y^{(d')}}{\eta \cdot \etd{d'}, \sigma^2}
\\ &\propto
\frac{b_{w_n^{(d)}}^{(k)} \setminus z_n^{(d)} + \beta_{w_n^{(d)}}}{ \sum_{v=1}^V \left( b_v^{(k)} \setminus z_n^{(d)} + \beta_v\right)}
\left( a_k^{(d)} \setminus z_n^{(d)} + \alpha_k \right)
\exp \left( \frac{1}{2 \sigma^2} \frac{\eta_k}{N^{(d)}} \left( 2 \left[ y^{(d)} - \eta \cdot \left( \etd{d} \setminus z_n^{(d)} \right) \right] - \frac{\eta_k}{N^{(d)}} \right) \right)
\end{align}$$

where the "set-minus" notation $\cdot \setminus z_n^{(d)}$ denotes the variable the notation is applied to with the entry $z_n^{(d)}$ removed (again, see my [LDA notebook](http://nbviewer.savvysherpa.com/github/bearnshaw/ml-demos/blob/master/lda_gibbs_sampling_cython.ipynb) for details).  This final proportionality is true since

$$\begin{align}
\prod_{d'=1}^{D} \cN{y^{(d')}}{\eta \cdot \etd{d'}, \sigma^2}
&\propto
\prod_{d'=1}^{D} \exp \left( -\frac{ \left( y^{(d')} - \eta \cdot \etd{d'} \right)^2 }{2 \sigma^2} \right)
\\ &\propto
\prod_{d'=1}^{D} \exp \left( \frac{ 2 y^{(d')} \eta \cdot \etd{d'} - \left( \eta \cdot \etd{d'} \right)^2 }{2 \sigma^2} \right)
\\ &=
\prod_{d'=1}^{D} \exp \left( \frac{ 2 y^{(d')} \left( \eta \cdot \left( \etd{d'} \setminus z_n^{(d)} \right) + \delta_{d, d'} \frac{\eta_k}{N^{(d)}} \right) - \left( \eta \cdot \left( \etd{d'} \setminus z_n^{(d)} \right) + \delta_{d, d'} \frac{\eta_k}{N^{(d)}} \right)^2 }{2 \sigma^2} \right)
\\ &=
\prod_{d'=1}^{D} \exp \left( \frac{ 2 y^{(d')} \eta \cdot \left( \etd{d'} \setminus z_n^{(d)} \right) - \left( \eta \cdot \left( \etd{d'} \setminus z_n^{(d)} \right) \right)^2 }{2 \sigma^2} \right)
\exp \left( \frac{1}{2 \sigma^2} \frac{\eta_k}{N^{(d)}} \left( 2 \left[ y^{(d)} - \eta \cdot \left( \etd{d} \setminus z_n^{(d)} \right) \right] - \frac{\eta_k}{N^{(d)}} \right) \right)
\\ &\propto
\exp \left( \frac{1}{2 \sigma^2} \frac{\eta_k}{N^{(d)}} \left( 2 \left[ y^{(d)} - \eta \cdot \left( \etd{d} \setminus z_n^{(d)} \right) \right] - \frac{\eta_k}{N^{(d)}} \right) \right)
\end{align}$$

where $\delta_{d, d'}$ is the [Kronecker delta](https://en.wikipedia.org/wiki/Kronecker_delta).

We also need to calculate the full conditional for $\eta$.  In order to do this, let $Z = (\etd{1} \cdots \etd{D})$ be the matrix whose columns are the empirical topic distributions $\etd{d}$, let $I$ be the identity matrix and $\one$ be the vector of ones, and note that

$$\prod_{k=1}^{K} \cN{\eta_{k}}{\mu, \nu^2} = \cN{\eta}{\mu \one, \nu^2 I}$$
$$\prod_{d=1}^{D} \cN{y^{(d)}}{\eta \cdot \etd{d'}, \sigma^2} = \cN{y}{Z^T \eta, \sigma^2 I}.$$

Therefore

$$\begin{align}
\cp{\eta}{z, w, y, \alpha, \beta, \mu, \nu^2, \sigma^2}
&\propto
\cp{z, w, \eta, y}{\alpha, \beta, \mu, \nu^2, \sigma^2}
\\ &\propto
\cN{\eta}{\mu \one, \nu^2 I} \cN{y}{Z^T \eta, \sigma^2 I}
\\ &\propto
\cN{\eta}{\Sigma \left( \nu^{-2} \mu \one + \sigma^{-2} Zy \right), \Sigma}
\end{align}$$

where $\Sigma^{-1} = \nu^{-2} I + \sigma^{-2} ZZ^T$ (see Section 9.3 of [Kevin Murphy's notes](http://www.cs.ubc.ca/~murphyk/Teaching/CS340-Fall07/reading/gauss.pdf) for a derivation of Bayes rule for linear Gaussian systems).  It is interesting to consider the mean and variance of $\eta$ in the two variance regimes $\sigma \gg \nu$ and $\sigma \ll \nu$.  If $\sigma \gg \nu$, then

$$\Sigma^{-1} = \nu^{-2} \left( I + \left( \frac{\nu}{\sigma} \right)^2 ZZ^T \right)
\approx \nu^{-2} I$$

which implies that the covariance structure of $\eta$ is $\Sigma \approx \nu^2 I$ and the mean of $\eta$ is

$$\Sigma \left( \nu^{-2} \mu \one + \sigma^{-2} Zy \right)
\approx \mu \one + \left( \frac{\nu}{\sigma} \right)^2 Zy
\approx \mu \one,$$

thus $\eta$ is approximately distributed according to its prior distribution.  On the other hand, if $\sigma \ll \nu$, then

$$\Sigma^{-1} = \sigma^{-2} \left( \left( \frac{\sigma}{\nu} \right)^2 I + ZZ^T \right)
\approx \sigma^{-2} ZZ^T.$$

Notice that $ZZ^T$ is almost surely positive definite, and hence almost surely invertible.  Therefore $\Sigma \approx \sigma^2 (ZZ^T)^{-1}$ and

$$\Sigma \left( \nu^{-2} \mu \one + \sigma^{-2} Zy \right)
\approx \left( \frac{\sigma}{\nu} \right)^2 \mu (ZZ^T)^{-1} \one + (ZZ^T)^{-1} Zy
\approx (ZZ^T)^{-1} Zy,$$

thus $\eta$ is approximately distributed as the least-squares solution of $y = Z^T \eta$.

## Graphical test

```python
%matplotlib inline

from modules.helpers import plot_images
from functools import partial
from sklearn.metrics import (mean_squared_error)
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np

imshow = partial(plt.imshow, cmap='gray', interpolation='nearest', aspect='auto')
rmse = lambda y_true, y_pred: np.sqrt(mean_squared_error(y_true, y_pred))
sns.set(style='white')
```

### Generate topics

We assume a vocabulary of 25 terms, and create ten "topics", where each topic assigns exactly 5 consecutive terms equal probability.

```python
V = 25
K = 10
N = 100
D = 1000
```

```python
topics = []
topic_base = np.concatenate((np.ones((1, 5)) * 0.2, np.zeros((4, 5))), axis=0).ravel()
for i in range(5):
    topics.append(np.roll(topic_base, i * 5))
topic_base = np.concatenate((np.ones((5, 1)) * 0.2, np.zeros((5, 4))), axis=1).ravel()
for i in range(5):
    topics.append(np.roll(topic_base, i))
topics = np.array(topics)
plt.figure(figsize=(10, 5))
plot_images(plt, topics, (5, 5), layout=(2, 5), figsize=(10, 5))
```

### Generate documents from topics

We generate 1,000 documents from these 10 topics by sampling 1,000 topic distributions, one for each document, from a Dirichlet distribution with parameter $\alpha = (1, \ldots, 1)$.

```python
alpha = np.ones(K)
np.random.seed(42)
thetas = np.random.dirichlet(alpha, size=D)
topic_assignments = np.array([np.random.choice(range(K), size=100, p=theta)
                              for theta in thetas])
word_assignments = np.array([[np.random.choice(range(V), size=1, p=topics[topic_assignments[d, n]])[0]
                              for n in range(N)] for d in range(D)])
doc_term_matrix = np.array([np.histogram(word_assignments[d], bins=V, range=(0, V - 1))[0] for d in range(D)])
imshow(doc_term_matrix)
```

### Generate responses

```python
# choose parameter values
nu2 = 10
sigma2 = 1
np.random.seed(42)
eta = np.random.normal(scale=nu2, size=K)
y = [np.dot(eta, thetas[i]) for i in range(D)] + np.random.normal(scale=sigma2, size=D)
# plot histogram of responses
print(eta)
_ = plt.hist(y, bins=20)
```

### Estimate parameters

```python
from slda.topic_models import SLDA
```

```python
_K = 10
_alpha = alpha
_beta = np.repeat(0.01, V)
_mu = 0
_nu2 = nu2
_sigma2 = sigma2
n_iter = 500
slda = SLDA(_K, _alpha, _beta, _mu, _nu2, _sigma2, n_iter, seed=42)
```

```python
%%time
slda.fit(doc_term_matrix, y)
```

    CPU times: user 12.7 s, sys: 23.3 ms, total: 12.8 s
    Wall time: 12.8 s

```python
plot_images(plt, slda.phi, (5, 5), (2, 5), figsize=(10, 5))
print(slda.phi)
print(np.sum(slda.phi, axis=0))
print(np.sum(slda.phi, axis=1))
```

```python
topic_order = [1, 2, 4, 3, 9, 0, 7, 5, 6, 8]
plot_images(plt, slda.phi[topic_order], (5, 5), (2, 5), figsize=(10, 5))
```

```python
imshow(slda.theta)
```

```python
plt.plot(slda.loglikelihoods)
```

```python
plt.plot(np.diff(slda.loglikelihoods)[-100:])
```

```python
burn_in = max(n_iter - 100, int(n_iter / 2))
slda.loglikelihoods[burn_in:].mean()
```

    29182.323844355378

```python
eta_pred = slda.eta[burn_in:].mean(axis=0)
print(eta)
print(eta_pred[topic_order])
np.linalg.norm(eta - eta_pred[topic_order])
```

    [  4.96714153  -1.38264301   6.47688538  15.23029856  -2.34153375
      -2.34136957  15.79212816   7.67434729  -4.69474386   5.42560044]
    [  4.44236917  -1.37428922   5.69171791  13.92661777  -1.90987367
      -1.35671636  15.39498326   7.42637763  -2.77220459   4.94725851]

    2.8091957908971406

### Predict response of test documents

Create 1,000 test documents using the same generative process as our training documents, and compute their actual responses.

```python
np.random.seed(42^2)
thetas_test = np.random.dirichlet(np.ones(K), size=D)
topic_assignments_test = np.array([np.random.choice(range(K), size=100, p=theta)
                                   for theta in thetas_test])
word_assignments_test = np.array([[np.random.choice(range(V), size=1, p=topics[topic_assignments_test[d, n]])[0]
                                   for n in range(N)] for d in range(D)])
doc_term_matrix_test = np.array([np.histogram(word_assignments_test[d], bins=V, range=(0, V - 1))[0] for d in range(D)])
y_test = [np.dot(eta, thetas_test[i]) for i in range(D)]
imshow(doc_term_matrix_test)
```

Estimate their topic distributions using the trained model, then calculate the predicted responses using the mean of our samples of $\eta$ after burn-in as an estimate for $\eta$.

```python
thetas_test_slda = slda.transform(doc_term_matrix_test)
y_slda = [np.dot(eta_pred, thetas_test_slda[i]) for i in range(D)]
```

Measure the goodness of our prediction using [root mean square error](https://en.wikipedia.org/wiki/Root-mean-square_deviation).

```python
rmse(y_test, y_slda)
```

    0.82113166468454046

### Two-step learning: learn topics, then learn regression

```python
from slda.topic_models import LDA
```

```python
lda = LDA(_K, _alpha, _beta, n_iter, seed=42)
```

```python
%%time
lda.fit(doc_term_matrix)
```

    2016-06-10 13:13:01.718094 start iterations
    2016-06-10 13:13:01.810971 0:00:00.092877 elapsed, iter   10, LL -518559.0163, 8.97% change from last
    2016-06-10 13:13:01.886179 0:00:00.168085 elapsed, iter   20, LL -468960.0555, 9.56% change from last
    2016-06-10 13:13:01.956784 0:00:00.238690 elapsed, iter   30, LL -445434.0461, 5.02% change from last
    2016-06-10 13:13:02.026697 0:00:00.308603 elapsed, iter   40, LL -429314.6963, 3.62% change from last
    2016-06-10 13:13:02.099862 0:00:00.381768 elapsed, iter   50, LL -419699.1884, 2.24% change from last
    2016-06-10 13:13:02.167735 0:00:00.449641 elapsed, iter   60, LL -415442.5080, 1.01% change from last
    2016-06-10 13:13:02.234883 0:00:00.516789 elapsed, iter   70, LL -408896.4141, 1.58% change from last
    2016-06-10 13:13:02.301891 0:00:00.583797 elapsed, iter   80, LL -402998.4074, 1.44% change from last
    2016-06-10 13:13:02.368526 0:00:00.650432 elapsed, iter   90, LL -398554.0889, 1.10% change from last
    2016-06-10 13:13:02.437271 0:00:00.719177 elapsed, iter  100, LL -396059.2952, 0.63% change from last
    2016-06-10 13:13:02.507638 0:00:00.789544 elapsed, iter  110, LL -393946.6660, 0.53% change from last
    2016-06-10 13:13:02.573177 0:00:00.855083 elapsed, iter  120, LL -393350.3096, 0.15% change from last
    2016-06-10 13:13:02.638212 0:00:00.920118 elapsed, iter  130, LL -391582.5349, 0.45% change from last
    2016-06-10 13:13:02.703620 0:00:00.985526 elapsed, iter  140, LL -389979.6403, 0.41% change from last
    2016-06-10 13:13:02.769137 0:00:01.051043 elapsed, iter  150, LL -390480.8703, -0.13% change from last
    2016-06-10 13:13:02.839140 0:00:01.121046 elapsed, iter  160, LL -389499.9872, 0.25% change from last
    2016-06-10 13:13:02.904864 0:00:01.186770 elapsed, iter  170, LL -389235.6565, 0.07% change from last
    2016-06-10 13:13:02.969981 0:00:01.251887 elapsed, iter  180, LL -388759.9274, 0.12% change from last
    2016-06-10 13:13:03.034976 0:00:01.316882 elapsed, iter  190, LL -388820.9941, -0.02% change from last
    2016-06-10 13:13:03.100102 0:00:01.382008 elapsed, iter  200, LL -388260.5027, 0.14% change from last
    2016-06-10 13:13:03.170388 0:00:01.452294 elapsed, iter  210, LL -388249.9451, 0.00% change from last
    2016-06-10 13:13:03.236702 0:00:01.518608 elapsed, iter  220, LL -388140.9833, 0.03% change from last
    2016-06-10 13:13:03.301888 0:00:01.583794 elapsed, iter  230, LL -388708.6353, -0.15% change from last
    2016-06-10 13:13:03.368221 0:00:01.650127 elapsed, iter  240, LL -388470.8166, 0.06% change from last
    2016-06-10 13:13:03.433243 0:00:01.715149 elapsed, iter  250, LL -388586.6879, -0.03% change from last
    2016-06-10 13:13:03.503369 0:00:01.785275 elapsed, iter  260, LL -388759.3606, -0.04% change from last
    2016-06-10 13:13:03.569231 0:00:01.851137 elapsed, iter  270, LL -388112.9995, 0.17% change from last
    2016-06-10 13:13:03.634155 0:00:01.916061 elapsed, iter  280, LL -388325.5728, -0.05% change from last
    2016-06-10 13:13:03.699221 0:00:01.981127 elapsed, iter  290, LL -388194.6193, 0.03% change from last
    2016-06-10 13:13:03.764164 0:00:02.046070 elapsed, iter  300, LL -388022.2504, 0.04% change from last
    2016-06-10 13:13:03.829746 0:00:02.111652 elapsed, iter  310, LL -388035.5122, -0.00% change from last
    2016-06-10 13:13:03.901004 0:00:02.182910 elapsed, iter  320, LL -387834.7560, 0.05% change from last
    2016-06-10 13:13:03.973892 0:00:02.255798 elapsed, iter  330, LL -387896.1408, -0.02% change from last
    2016-06-10 13:13:04.045441 0:00:02.327347 elapsed, iter  340, LL -387689.7851, 0.05% change from last
    2016-06-10 13:13:04.115698 0:00:02.397604 elapsed, iter  350, LL -387410.1707, 0.07% change from last
    2016-06-10 13:13:04.186569 0:00:02.468475 elapsed, iter  360, LL -387194.6374, 0.06% change from last
    2016-06-10 13:13:04.257784 0:00:02.539690 elapsed, iter  370, LL -387956.5177, -0.20% change from last
    2016-06-10 13:13:04.328896 0:00:02.610802 elapsed, iter  380, LL -387593.8931, 0.09% change from last
    2016-06-10 13:13:04.399921 0:00:02.681827 elapsed, iter  390, LL -387841.1371, -0.06% change from last
    2016-06-10 13:13:04.470958 0:00:02.752864 elapsed, iter  400, LL -388194.7104, -0.09% change from last
    2016-06-10 13:13:04.541473 0:00:02.823379 elapsed, iter  410, LL -387557.6239, 0.16% change from last
    2016-06-10 13:13:04.612267 0:00:02.894173 elapsed, iter  420, LL -387654.8058, -0.03% change from last
    2016-06-10 13:13:04.682486 0:00:02.964392 elapsed, iter  430, LL -388271.5004, -0.16% change from last
    2016-06-10 13:13:04.753022 0:00:03.034928 elapsed, iter  440, LL -387805.9252, 0.12% change from last
    2016-06-10 13:13:04.823710 0:00:03.105616 elapsed, iter  450, LL -388152.3984, -0.09% change from last
    2016-06-10 13:13:04.894584 0:00:03.176490 elapsed, iter  460, LL -388539.1905, -0.10% change from last
    2016-06-10 13:13:04.965270 0:00:03.247176 elapsed, iter  470, LL -387818.0986, 0.19% change from last
    2016-06-10 13:13:05.035921 0:00:03.317827 elapsed, iter  480, LL -387694.9035, 0.03% change from last
    2016-06-10 13:13:05.106130 0:00:03.388036 elapsed, iter  490, LL -387909.3207, -0.06% change from last
    CPU times: user 3.45 s, sys: 15.4 ms, total: 3.46 s
    Wall time: 3.46 s

```python
plot_images(plt, lda.phi, (5, 5), (2, 5), figsize=(10, 5))
```

```python
topic_order_lda = [1, 2, 0, 3, 9, 4, 7, 5, 6, 8]
plot_images(plt, lda.phi[topic_order_lda], (5, 5), (2, 5), figsize=(10, 5))
```

```python
imshow(lda.theta)
```

```python
plt.plot(lda.loglikelihoods)
```

```python
thetas_test_lda = lda.transform(doc_term_matrix_test)
```

#### Unregularized linear regression

- train linear regression on training data
- calculate response on test data
- measure the goodness of prediction using root mean square error.

```python
from sklearn.linear_model import LinearRegression
```

```python
lr = LinearRegression(fit_intercept=False)
lr.fit(lda.theta, y)
y_lr = lr.predict(thetas_test_lda)
rmse(y_test, y_lr)
```

    0.82647226636711646

```python
print(eta)
print(lr.coef_[topic_order_lda])
np.linalg.norm(eta - lr.coef_[topic_order_lda])
```

    [  4.96714153  -1.38264301   6.47688538  15.23029856  -2.34153375
      -2.34136957  15.79212816   7.67434729  -4.69474386   5.42560044]
    [  4.29349447  -1.35662767   5.86690601  13.56743486  -2.225693    -1.3096669
      15.68197752   7.74802135  -2.83164145   4.58107028]

    2.978496259038288

#### L2-regularized linear regression

- train ridge regression on training data
- calculate response on test data
- measure the goodness of prediction using root mean square error.

```python
from sklearn.linear_model import Ridge
```

```python
lrl2 = Ridge(alpha=1., fit_intercept=False)
lrl2.fit(lda.theta, y)
y_lrl2 = lrl2.predict(thetas_test_lda)
rmse(y_test, y_lrl2)
```

    0.86174658080698563

```python
print(eta)
print(lrl2.coef_[topic_order_lda])
np.linalg.norm(eta - lrl2.coef_[topic_order_lda])
```

    [  4.96714153  -1.38264301   6.47688538  15.23029856  -2.34153375
      -2.34136957  15.79212816   7.67434729  -4.69474386   5.42560044]
    [  4.28631637  -0.95278145   5.67416731  12.66995942  -1.48722133
      -0.7396339   14.37412703   7.40026967  -2.0793953    4.49485813]

    4.5755413457461787

#### Gradient boosted regression trees

- train gradient boosted regressor on training data
- calculate response on test data
- measure the goodness of prediction using root mean square error.

```python
from sklearn.ensemble import GradientBoostingRegressor
```

```python
gbr = GradientBoostingRegressor()
gbr.fit(lda.theta, y)
y_gbr = gbr.predict(thetas_test_lda)
rmse(y_test, y_gbr)
```

    0.98198962236756349

### Conclusion

SLDA is slightly better than unregularized linear regression, and better than ridge regression or gradient boosted regression trees.  The similar performance to unregularized linear regression is likely due to the fact that this test was set up as an exact problem - all parameters used in training, except $\beta$ (because we hand-picked the topics) and $\eta$ (because that was the one parameter we wanted to learn), were those used to generate documents, and in SLDA the likelihood of $y$ introduces a regularization on $\eta$ in the full-conditional for $z$, while unregularized linear regression does not enforce such a penalty.

## Test with fewer topics

We now redo the previous test, but this time with fewer topics, in order to determine whether

1. the supervised portion of SLDA will produce topics different from those produced by LDA, and
1. prediction with SLDA is improved over LDA-and-a-regression.

Because we will no longer be solving an exact problem (the number of topics, and hence $\alpha$, will both be different from the document generation process), we expect SLDA to do better than LDA-and-a-regression, including unregularized linear regression.

```python
_K = 5
_alpha = np.repeat(1. / _K, _K)
_beta = np.repeat(0.01, V)
_mu = 0
_nu2 = nu2
_sigma2 = sigma2
n_iter = 500
slda1 = SLDA(_K, _alpha, _beta, _mu, _nu2, _sigma2, n_iter, seed=42)
```

```python
%%time
slda1.fit(doc_term_matrix, y)
```

    2016-06-10 13:13:08.127915 start iterations
    2016-06-10 13:13:08.283610 0:00:00.155695 elapsed, iter   10, LL -48424.3678, 57.47% change from last
    2016-06-10 13:13:08.426567 0:00:00.298652 elapsed, iter   20, LL -26669.8841, 44.92% change from last
    2016-06-10 13:13:08.565257 0:00:00.437342 elapsed, iter   30, LL -23012.4319, 13.71% change from last
    2016-06-10 13:13:08.697874 0:00:00.569959 elapsed, iter   40, LL -22168.0622, 3.67% change from last
    2016-06-10 13:13:08.833184 0:00:00.705269 elapsed, iter   50, LL -20919.7416, 5.63% change from last
    2016-06-10 13:13:08.971179 0:00:00.843264 elapsed, iter   60, LL -20442.9833, 2.28% change from last
    2016-06-10 13:13:09.113854 0:00:00.985939 elapsed, iter   70, LL -21271.4782, -4.05% change from last
    2016-06-10 13:13:09.262225 0:00:01.134310 elapsed, iter   80, LL -21511.5328, -1.13% change from last
    2016-06-10 13:13:09.404185 0:00:01.276270 elapsed, iter   90, LL -21498.6877, 0.06% change from last
    2016-06-10 13:13:09.546102 0:00:01.418187 elapsed, iter  100, LL -21450.9848, 0.22% change from last
    2016-06-10 13:13:09.689745 0:00:01.561830 elapsed, iter  110, LL -21264.0617, 0.87% change from last
    2016-06-10 13:13:09.832946 0:00:01.705031 elapsed, iter  120, LL -21124.3302, 0.66% change from last
    2016-06-10 13:13:09.972801 0:00:01.844886 elapsed, iter  130, LL -21409.3146, -1.35% change from last
    2016-06-10 13:13:10.114312 0:00:01.986397 elapsed, iter  140, LL -20748.2025, 3.09% change from last
    2016-06-10 13:13:10.254950 0:00:02.127035 elapsed, iter  150, LL -21784.6926, -5.00% change from last
    2016-06-10 13:13:10.399699 0:00:02.271784 elapsed, iter  160, LL -21593.1937, 0.88% change from last
    2016-06-10 13:13:10.543634 0:00:02.415719 elapsed, iter  170, LL -21755.2391, -0.75% change from last
    2016-06-10 13:13:10.679132 0:00:02.551217 elapsed, iter  180, LL -21838.7155, -0.38% change from last
    2016-06-10 13:13:10.812856 0:00:02.684941 elapsed, iter  190, LL -22390.0410, -2.52% change from last
    2016-06-10 13:13:10.951874 0:00:02.823959 elapsed, iter  200, LL -21844.0314, 2.44% change from last
    2016-06-10 13:13:11.088078 0:00:02.960163 elapsed, iter  210, LL -22356.2309, -2.34% change from last
    2016-06-10 13:13:11.221326 0:00:03.093411 elapsed, iter  220, LL -22476.3578, -0.54% change from last
    2016-06-10 13:13:11.357571 0:00:03.229656 elapsed, iter  230, LL -21854.6413, 2.77% change from last
    2016-06-10 13:13:11.494444 0:00:03.366529 elapsed, iter  240, LL -22209.4309, -1.62% change from last
    2016-06-10 13:13:11.635861 0:00:03.507946 elapsed, iter  250, LL -22552.7551, -1.55% change from last
    2016-06-10 13:13:11.771394 0:00:03.643479 elapsed, iter  260, LL -22492.8254, 0.27% change from last
    2016-06-10 13:13:11.904591 0:00:03.776676 elapsed, iter  270, LL -21363.6117, 5.02% change from last
    2016-06-10 13:13:12.039422 0:00:03.911507 elapsed, iter  280, LL -21755.4643, -1.83% change from last
    2016-06-10 13:13:12.174180 0:00:04.046265 elapsed, iter  290, LL -21026.2013, 3.35% change from last
    2016-06-10 13:13:12.310638 0:00:04.182723 elapsed, iter  300, LL -21019.0059, 0.03% change from last
    2016-06-10 13:13:12.445832 0:00:04.317917 elapsed, iter  310, LL -21956.2779, -4.46% change from last
    2016-06-10 13:13:12.581865 0:00:04.453950 elapsed, iter  320, LL -21733.3422, 1.02% change from last
    2016-06-10 13:13:12.718181 0:00:04.590266 elapsed, iter  330, LL -21173.1530, 2.58% change from last
    2016-06-10 13:13:12.858645 0:00:04.730730 elapsed, iter  340, LL -21356.0593, -0.86% change from last
    2016-06-10 13:13:12.994323 0:00:04.866408 elapsed, iter  350, LL -20630.9341, 3.40% change from last
    2016-06-10 13:13:13.128332 0:00:05.000417 elapsed, iter  360, LL -20467.3912, 0.79% change from last
    2016-06-10 13:13:13.263090 0:00:05.135175 elapsed, iter  370, LL -21047.2518, -2.83% change from last
    2016-06-10 13:13:13.402083 0:00:05.274168 elapsed, iter  380, LL -21253.7650, -0.98% change from last
    2016-06-10 13:13:13.539833 0:00:05.411918 elapsed, iter  390, LL -21504.8653, -1.18% change from last
    2016-06-10 13:13:13.677880 0:00:05.549965 elapsed, iter  400, LL -21712.5750, -0.97% change from last
    2016-06-10 13:13:13.811928 0:00:05.684013 elapsed, iter  410, LL -22065.4088, -1.63% change from last
    2016-06-10 13:13:13.948051 0:00:05.820136 elapsed, iter  420, LL -21642.4839, 1.92% change from last
    2016-06-10 13:13:14.082922 0:00:05.955007 elapsed, iter  430, LL -21622.6513, 0.09% change from last
    2016-06-10 13:13:14.224415 0:00:06.096500 elapsed, iter  440, LL -22010.7755, -1.79% change from last
    2016-06-10 13:13:14.359776 0:00:06.231861 elapsed, iter  450, LL -22175.4719, -0.75% change from last
    2016-06-10 13:13:14.495657 0:00:06.367742 elapsed, iter  460, LL -22189.0535, -0.06% change from last
    2016-06-10 13:13:14.630771 0:00:06.502856 elapsed, iter  470, LL -21675.9985, 2.31% change from last
    2016-06-10 13:13:14.766175 0:00:06.638260 elapsed, iter  480, LL -21900.0566, -1.03% change from last
    2016-06-10 13:13:14.900684 0:00:06.772769 elapsed, iter  490, LL -21701.0455, 0.91% change from last
    CPU times: user 6.88 s, sys: 33.5 ms, total: 6.91 s
    Wall time: 6.91 s

```python
plot_images(plt, slda1.phi, (5, 5), (1, 5), figsize=(10, 5))
```

```python
imshow(slda1.theta)
```

```python
plt.plot(slda1.loglikelihoods)
```

```python
burn_in1 = max(n_iter - 100, int(n_iter / 2))
slda1.loglikelihoods[burn_in1:].mean()
```

    -22055.005897817184

```python
eta_pred1 = slda1.eta[burn_in1:].mean(axis=0)
eta_pred1
```

    array([  2.22220264,  -0.15788853,   9.55428965,  10.2525672 ,   3.22556876])

```python
thetas_test_slda1 = slda1.transform(doc_term_matrix_test)
y_slda1 = [np.dot(eta_pred1, thetas_test_slda1[i]) for i in range(D)]
```

```python
rmse(y_test, y_slda1)
```

    1.2455951548532411

```python
lda1 = LDA(_K, _alpha, _beta, n_iter, seed=42)
```

```python
%%time
lda1.fit(doc_term_matrix)
```

    2016-06-10 13:13:16.087473 start iterations
    2016-06-10 13:13:16.143708 0:00:00.056235 elapsed, iter   10, LL -404994.8904, 13.25% change from last
    2016-06-10 13:13:16.188322 0:00:00.100849 elapsed, iter   20, LL -380416.0984, 6.07% change from last
    2016-06-10 13:13:16.232188 0:00:00.144715 elapsed, iter   30, LL -374718.7031, 1.50% change from last
    2016-06-10 13:13:16.275514 0:00:00.188041 elapsed, iter   40, LL -372662.5340, 0.55% change from last
    2016-06-10 13:13:16.318583 0:00:00.231110 elapsed, iter   50, LL -372592.3679, 0.02% change from last
    2016-06-10 13:13:16.361522 0:00:00.274049 elapsed, iter   60, LL -371993.9718, 0.16% change from last
    2016-06-10 13:13:16.404200 0:00:00.316727 elapsed, iter   70, LL -370991.9632, 0.27% change from last
    2016-06-10 13:13:16.446894 0:00:00.359421 elapsed, iter   80, LL -372006.8433, -0.27% change from last
    2016-06-10 13:13:16.490452 0:00:00.402979 elapsed, iter   90, LL -372236.5623, -0.06% change from last
    2016-06-10 13:13:16.532751 0:00:00.445278 elapsed, iter  100, LL -373132.6808, -0.24% change from last
    2016-06-10 13:13:16.575374 0:00:00.487901 elapsed, iter  110, LL -372307.9178, 0.22% change from last
    2016-06-10 13:13:16.618148 0:00:00.530675 elapsed, iter  120, LL -373064.4300, -0.20% change from last
    2016-06-10 13:13:16.661584 0:00:00.574111 elapsed, iter  130, LL -372896.9690, 0.04% change from last
    2016-06-10 13:13:16.704663 0:00:00.617190 elapsed, iter  140, LL -373666.8725, -0.21% change from last
    2016-06-10 13:13:16.747776 0:00:00.660303 elapsed, iter  150, LL -375130.4420, -0.39% change from last
    2016-06-10 13:13:16.791435 0:00:00.703962 elapsed, iter  160, LL -375266.7223, -0.04% change from last
    2016-06-10 13:13:16.834670 0:00:00.747197 elapsed, iter  170, LL -375293.3698, -0.01% change from last
    2016-06-10 13:13:16.878672 0:00:00.791199 elapsed, iter  180, LL -376238.5505, -0.25% change from last
    2016-06-10 13:13:16.922925 0:00:00.835452 elapsed, iter  190, LL -375997.5441, 0.06% change from last
    2016-06-10 13:13:16.966738 0:00:00.879265 elapsed, iter  200, LL -375560.6490, 0.12% change from last
    2016-06-10 13:13:17.011130 0:00:00.923657 elapsed, iter  210, LL -375438.6514, 0.03% change from last
    2016-06-10 13:13:17.053203 0:00:00.965730 elapsed, iter  220, LL -375317.4478, 0.03% change from last
    2016-06-10 13:13:17.100620 0:00:01.013147 elapsed, iter  230, LL -375611.4260, -0.08% change from last
    2016-06-10 13:13:17.142928 0:00:01.055455 elapsed, iter  240, LL -375084.7172, 0.14% change from last
    2016-06-10 13:13:17.186384 0:00:01.098911 elapsed, iter  250, LL -374881.3455, 0.05% change from last
    2016-06-10 13:13:17.230083 0:00:01.142610 elapsed, iter  260, LL -375142.8976, -0.07% change from last
    2016-06-10 13:13:17.275077 0:00:01.187604 elapsed, iter  270, LL -374933.0420, 0.06% change from last
    2016-06-10 13:13:17.317606 0:00:01.230133 elapsed, iter  280, LL -374389.2267, 0.15% change from last
    2016-06-10 13:13:17.359826 0:00:01.272353 elapsed, iter  290, LL -373889.6761, 0.13% change from last
    2016-06-10 13:13:17.406300 0:00:01.318827 elapsed, iter  300, LL -373999.4793, -0.03% change from last
    2016-06-10 13:13:17.454818 0:00:01.367345 elapsed, iter  310, LL -374055.1543, -0.01% change from last
    2016-06-10 13:13:17.498775 0:00:01.411302 elapsed, iter  320, LL -374575.8085, -0.14% change from last
    2016-06-10 13:13:17.543644 0:00:01.456171 elapsed, iter  330, LL -375653.2470, -0.29% change from last
    2016-06-10 13:13:17.587968 0:00:01.500495 elapsed, iter  340, LL -376213.0731, -0.15% change from last
    2016-06-10 13:13:17.630953 0:00:01.543480 elapsed, iter  350, LL -375100.3852, 0.30% change from last
    2016-06-10 13:13:17.675820 0:00:01.588347 elapsed, iter  360, LL -374749.8943, 0.09% change from last
    2016-06-10 13:13:17.719917 0:00:01.632444 elapsed, iter  370, LL -374970.1822, -0.06% change from last
    2016-06-10 13:13:17.765062 0:00:01.677589 elapsed, iter  380, LL -375792.9562, -0.22% change from last
    2016-06-10 13:13:17.809414 0:00:01.721941 elapsed, iter  390, LL -375580.5170, 0.06% change from last
    2016-06-10 13:13:17.853399 0:00:01.765926 elapsed, iter  400, LL -376051.9297, -0.13% change from last
    2016-06-10 13:13:17.896977 0:00:01.809504 elapsed, iter  410, LL -375554.2968, 0.13% change from last
    2016-06-10 13:13:17.940666 0:00:01.853193 elapsed, iter  420, LL -375327.1425, 0.06% change from last
    2016-06-10 13:13:17.984309 0:00:01.896836 elapsed, iter  430, LL -374650.8399, 0.18% change from last
    2016-06-10 13:13:18.028397 0:00:01.940924 elapsed, iter  440, LL -375878.8372, -0.33% change from last
    2016-06-10 13:13:18.072662 0:00:01.985189 elapsed, iter  450, LL -376199.3532, -0.09% change from last
    2016-06-10 13:13:18.116349 0:00:02.028876 elapsed, iter  460, LL -375946.0416, 0.07% change from last
    2016-06-10 13:13:18.160254 0:00:02.072781 elapsed, iter  470, LL -376240.8013, -0.08% change from last
    2016-06-10 13:13:18.203976 0:00:02.116503 elapsed, iter  480, LL -376057.7271, 0.05% change from last
    2016-06-10 13:13:18.247017 0:00:02.159544 elapsed, iter  490, LL -376733.2730, -0.18% change from last
    CPU times: user 2.2 s, sys: 13.7 ms, total: 2.21 s
    Wall time: 2.22 s

```python
plot_images(plt, lda1.phi, (5, 5), (1, 5), figsize=(10, 5))
```

We plot the SLDA topics again and note that they are indeed different!

```python
plot_images(plt, slda1.phi, (5, 5), (1, 5), figsize=(10, 5))
```

```python
imshow(lda1.theta)
```

```python
plt.plot(lda1.loglikelihoods)
```

```python
thetas_test_lda1 = lda1.transform(doc_term_matrix_test)
```

#### Unregularized linear regression

```python
lr1 = LinearRegression(fit_intercept=False)
lr1.fit(lda1.theta, y)
y_lr1 = lr1.predict(thetas_test_lda1)
rmse(y_test, y_lr1)
```

    1.7255065570452643

#### L2-regularized linear regression

```python
lrl21 = Ridge(alpha=1., fit_intercept=False)
lrl21.fit(lda1.theta, y)
y_lrl21 = lrl21.predict(thetas_test_lda1)
rmse(y_test, y_lrl21)
```

    1.7266314501798159

#### Gradient boosted regression trees

```python
gbr1 = GradientBoostingRegressor()
gbr1.fit(lda1.theta, y)
y_gbr1 = gbr1.predict(thetas_test_lda1)
rmse(y_test, y_gbr1)
```

    1.790809644071969

#### Unregularized linear regression with SLDA topics

```python
lr1_0 = LinearRegression(fit_intercept=False)
lr1_0.fit(slda1.theta, y)
y_lr1_0 = lr1_0.predict(thetas_test_slda1)
rmse(y_test, y_lr1_0)
```

    1.2424204767358347

#### L2-regularized linear regression with SLDA topics

```python
lrl21_0 = Ridge(alpha=1., fit_intercept=False)
lrl21_0.fit(slda1.theta, y)
y_lrl21_0 = lrl21_0.predict(thetas_test_slda1)
rmse(y_test, y_lrl21_0)
```

    1.2497459065224035

#### Gradient boosted regression trees with SLDA topics

```python
gbr1_0 = GradientBoostingRegressor()
gbr1_0.fit(slda1.theta, y)
y_gbr1_0 = gbr1_0.predict(thetas_test_slda1)
rmse(y_test, y_gbr1_0)
```

    1.3036177236719277

### Conclusion for test with fewer topics

SLDA performs at least 27.5% better in root mean square error than any of the LDA-and-a-regression methods, and similar to the SLDA-and-a-linear-regression methods.


# Brief introduction to JAX

murphyk@gmail.com, Last update: 2020-12-31.

[JAX](https://github.com/google/jax) is a  version of NumPy that runs fast on CPU, GPU and TPU, by compiling down to XLA. It also has an excellent automatic differentiation library, extending the earlier [autograd](https://github.com/hips/autograd) package.

The JAX interface is almost identical to NumPy (by design), but with some small differences, and additional features, some of which we explain below. (More detail can be find in the official documentation.)

```
# Standard Python libraries
from __future__ import absolute_import, division, print_function, unicode_literals

import os
import time
import numpy as np
np.set_printoptions(precision=3)
import glob
import matplotlib.pyplot as plt
import PIL
import imageio

from IPython import display
%matplotlib inline

import sklearn

```

```

# Load JAX
import jax
import jax.numpy as jnp
from jax import random
#import jax.numpy as np
#import numpy as onp # original numpy
from jax import grad, hessian, jit, vmap
print("jax version {}".format(jax.__version__))

```

    jax version 0.2.7

# Random number generation

One of the biggest differences from NumPy is the way Jax treates pseudo random number generation (PRNG).
This is because Jax does not maintain any global state, i.e., it is purely functional.
This design "provides reproducible results invariant to compilation boundaries and backends,
while also maximizing performance by enabling vectorized generation and parallelization across random calls"
(to quote [the official page](https://github.com/google/jax#a-brief-tour)).

Thus, whenever we do anything stochastic, we need to give it a fresh RNG key. We can do this by splitting the existing key into pieces. We can do this indefinitely, as shown below.

```
import jax.random as random

key = random.PRNGKey(0)
print(random.normal(key, shape=(3,)))  # [ 1.81608593 -0.48262325  0.33988902]
print(random.normal(key, shape=(3,)))  # [ 1.81608593 -0.48262325  0.33988902]  ## identical results

# To make a new key, we split the current key into two pieces.
key, subkey = random.split(key)
print(random.normal(subkey, shape=(3,)))  # [ 1.1378783  -1.22095478 -0.59153646]

# We can continue to split off new pieces from the global key.
key, subkey = random.split(key)
print(random.normal(subkey, shape=(3,)))  # [-0.06607265  0.16676566  1.17800343]

# We can always use original numpy if we like (although this may interfere with the deterministic behavior of jax)
np.random.seed(42)
print(np.random.randn(3))
```

    [ 1.816 -0.483  0.34 ]
    [ 1.816 -0.483  0.34 ]
    [ 1.138 -1.221 -0.592]
    [-0.066  0.167  1.178]
    [ 0.497 -0.138  0.648]

# GPU magic

```
# Check if GPU is available
!nvidia-smi

```

    Thu Dec 31 17:51:58 2020
    +-----------------------------------------------------------------------------+
    | NVIDIA-SMI 460.27.04    Driver Version: 418.67       CUDA Version: 10.1     |
    |-------------------------------+----------------------+----------------------+
    | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
    | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
    |                               |                      |               MIG M. |
    |===============================+======================+======================|
    |   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
    | N/A   37C    P0    27W /  70W |  13613MiB / 15079MiB |      0%      Default |
    |                               |                      |                 ERR! |
    +-------------------------------+----------------------+----------------------+

    +-----------------------------------------------------------------------------+
    | Processes:                                                                  |
    |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
    |        ID   ID                                                   Usage      |
    |=============================================================================|
    |  No running processes found                                                 |
    +-----------------------------------------------------------------------------+

```

# Check if JAX is using GPU
print("jax backend {}".format(jax.lib.xla_bridge.get_backend().platform))
```

    jax backend gpu

Let's see how JAX can speed up things like matrix-matrix multiplication using a GPU.

First the numpy/CPU version.

```
# Standard CPU

size = 1000
x = np.random.normal(size=(size, size)).astype(np.float32)
print(type(x))
%timeit -o np.dot(x, x.T)

```

    <class 'numpy.ndarray'>
    100 loops, best of 3: 17.7 ms per loop

    <TimeitResult : 100 loops, best of 3: 17.7 ms per loop>

```
res = _ # get result of last cell
time_cpu = res.best
print(time_cpu)
```

    0.01767426187999945

Now the GPU version. We added that block_until_ready because JAX uses [asynchronous execution](https://jax.readthedocs.io/en/latest/async_dispatch.html) by default.

```
# GPU version
x = jax.random.normal(key, (size, size), dtype=jnp.float32)
print(type(x))
%timeit -o jnp.dot(x, x.T).block_until_ready()
```

    <class 'jax.interpreters.xla._DeviceArray'>
    The slowest run took 7.90 times longer than the fastest. This could mean that an intermediate result is being cached.
    1000 loops, best of 3: 719 µs per loop

    <TimeitResult : 1000 loops, best of 3: 719 µs per loop>

```
res = _
time_gpu = res.best
print('GPU time {:0.6f}, CPU time {:0.6f}, speedup {:0.6f}'.format(
    time_gpu, time_cpu, time_cpu/time_gpu))
```

    GPU time 0.000719, CPU time 0.017674, speedup 24.593663

We can move numpy arrays to the GPU for speed. The result will be transferred back to CPU for printing, saving, etc.

```
from jax import device_put

x = np.random.normal(size=(size, size)).astype(np.float32)
print(type(x))
%timeit np.dot(x, x.T)

x = device_put(x)
print(type(x))
%timeit jnp.dot(x, x.T).block_until_ready()
```

    <class 'numpy.ndarray'>
    100 loops, best of 3: 18.3 ms per loop
    <class 'jax.interpreters.xla._DeviceArray'>
    1000 loops, best of 3: 863 µs per loop

# Vmap

To illustrate vmap, consider a binary logistic regression model.

```
def sigmoid(x): return 0.5 * (jnp.tanh(x / 2.) + 1)

def predict_single(w, x):
    return sigmoid(jnp.dot(w, x)) # <(D) , (D)> = (1) # inner product

def predict_batch(w, X):
    return sigmoid(jnp.dot(X, w)) # (N,D) * (D,1) = (N,1) # matrix-vector multiply

D = 2
N = 3

#np.random.state(42)
#w = np.random.randn(D)
#X = np.random.randn(N, D)
#y = np.random.randint(0, 2, N)

w = jax.random.normal(key, shape=(D,))
X = jax.random.normal(key, shape=(N,D))
y = jax.random.choice(key, 2, shape=(N,)) # uniform binary labels

print(X)
print(y)

# We can apply predict_batch to a matrix of data, but we cannot apply predict_single in this way
# because the order of the arguments to np.dot is incorrect.

p1 = predict_batch(w, X)
print(p1)
try:
    p2 = predict_single(w, X)
except:
    print('cannot apply to batch')
```

    [[-1.455  0.973]
     [-0.217  0.691]
     [-1.011  0.401]]
    [0 0 1]
    [0.938 0.773 0.814]
    cannot apply to batch

To avoid having to think about batch shape, it is often easier to write a function that works on single
input vectors. We can then apply this in a loop.

```
p3 = [predict_single(w, x) for x in X]
assert np.allclose(p1, p3)
```

Unfortunately, mapping down a list is slow.
Fortunately, JAX provides `vmap`, which has the same effect, but can be parallelized.

We first apply the `predict_single` function to its first arugment, w, to get a function that only
depends on x. We then vectorize this, and map the resulting modified function along rows (dimension 0)
of the data matrix.

```
from functools import partial

predict_single_w = partial(predict_single, w)
predict_batch_w = vmap(predict_single_w)
p4 = predict_batch_w(X)
assert np.allclose(p1, p4)

# More concise
p5 = vmap(predict_single, in_axes=(None, 0))(w, X)
assert np.allclose(p1, p5)

p6 = vmap(partial(predict_single, w))(X)
assert np.allclose(p1, p6)

```

# Autograd

In this section, we illustrate automatic differentation using JAX.

## Simple convex functions

```
from jax import grad, hessian, jacfwd, jacrev, vmap, jit
```

Linear function: multi-input, scalar output.

$$
\begin{align}
f(x; a) &= a^T x\\
\nabla_x f(x;a) &= a
\end{align}
$$

```
# We construct a single output linear function.
# In this case, the Jacobian and gradient are the same.
def fun1d(x):
    return jnp.dot(a, x)[0]

Din = 3; Dout = 1;
a = jax.random.normal(key, shape=(Dout, Din))
x = jax.random.normal(key, shape=(Din,))
g = grad(fun1d)(x)
assert np.allclose(g, a)
J = jacrev(fun1d)(x)
assert np.allclose(J, g)
```

Linear function: multi-input, multi-output.

$$
\begin{align}
f(x;A) &= A x \\
\nabla_x f(x;A) &= A
\end{align}
$$

```
# We construct a multi-output linear function.
# We check forward and reverse mode give same Jacobians.

def fun(x):
    return jnp.dot(A, x)

Din = 3; Dout = 4;
A = jax.random.normal(key, shape=(Dout, Din))
x = jax.random.normal(key, shape=(Din,))
Jf = jacfwd(fun)(x)
Jr = jacrev(fun)(x)
assert np.allclose(Jf, Jr)
assert np.allclose(Jf, A)
```

Quadratic form.

$$
\begin{align}
f(x;A) &= x^T A x \\
\nabla_x f(x;A) &= (A+A^T) x \\
\nabla^2 x^2 f(x;A) &= A + A^T
\end{align}
$$

```

D = 4
A = jax.random.normal(key, shape=(D,D))
x = jax.random.normal(key, shape=(D,))

quadfun = lambda x: jnp.dot(x, jnp.dot(A, x))

J = jacfwd(quadfun)(x)
assert np.allclose(J, jnp.dot(A+A.T, x))

H1 = hessian(quadfun)(x)
assert np.allclose(H1, A+A.T)

def my_hessian(fun):
  return jacfwd(jacrev(fun))

H2 = my_hessian(quadfun)(x)
assert np.allclose(H1, H2)
```

Chain rule applied to sigmoid function.

$$
\begin{align}
f(x;w) &=\sigma(w^T x) \\
\nabla_w f(x;w) &= \sigma'(w^T x) x \\
\sigma(a) &= \sigma(a) * (1-\sigma(a))
\end{align}
$$

```

D = 4
w = jax.random.normal(key, shape=(D,))
x = jax.random.normal(key, shape=(D,))
y = 0

def sigmoid(x): return 0.5 * (jnp.tanh(x / 2.) + 1)
def mu(w): return sigmoid(jnp.dot(w,x))
def deriv_mu(w): return mu(w) * (1-mu(w)) * x
deriv_mu_jax =  grad(mu)

print(deriv_mu(w))
print(deriv_mu_jax(w))

assert np.allclose(deriv_mu(w), deriv_mu_jax(w), atol=1e-3)

```

    [0.002 0.003 0.004 0.002]
    [0.002 0.003 0.004 0.002]

## Binary logistic regression

```

# negative log likelihood
def loss(weights, inputs, targets):
    preds = predict_batch(weights, inputs)
    logprobs = jnp.log(preds) * targets + jnp.log(1 - preds) * (1 - targets)
    return -jnp.sum(logprobs)

D = 2
N = 3
w = jax.random.normal(key, shape=(D,))
X = jax.random.normal(key, shape=(N,D))
y = jax.random.choice(key, 2, shape=(N,)) # uniform binary labels

print(loss(w, X, y))

# Gradient function
grad_fun = grad(loss)

# Gradient of each example in the batch - 2 different ways
grad_fun_w = partial(grad_fun, w)
grads = vmap(grad_fun_w)(X,y)
print(grads)
assert grads.shape == (N,D)

grads2 = vmap(grad_fun, in_axes=(None, 0, 0))(w, X, y)
assert np.allclose(grads, grads2)

# Gradient for entire batch
grad_sum = jnp.sum(grads, axis=0)
assert grad_sum.shape == (D,)
print(grad_sum)
```

    4.468591
    [[-1.365  0.913]
     [-0.168  0.534]
     [ 0.188 -0.075]]
    [-1.345  1.373]

```
# Textbook implementation of gradient
def NLL_grad(weights, batch):
    X, y = batch
    N = X.shape[0]
    mu = predict_batch(weights, X)
    g = jnp.sum(jnp.dot(jnp.diag(mu - y), X), axis=0)
    return g

grad_sum_batch = NLL_grad(w, (X,y))
print(grad_sum_batch)
assert np.allclose(grad_sum, grad_sum_batch)
```

    [-1.345  1.373]

```
# We can also compute Hessians, as we illustrate below.
from jax import hessian

hessian_fun = hessian(loss)

# Hessian on one example
H0 = hessian_fun(w, X[0,:], y[0])
print('Hessian(example 0)\n{}'.format(H0))

# Hessian for batch
Hbatch = vmap(hessian_fun, in_axes=(None, 0, 0))(w, X, y)
print('Hbatch shape {}'.format(Hbatch.shape))

Hbatch_sum = jnp.sum(Hbatch, axis=0)
print('Hbatch sum\n {}'.format(Hbatch_sum))
```

    Hessian(example 0)
    [[ 0.123 -0.082]
     [-0.082  0.055]]
    Hbatch shape (3, 2, 2)
    Hbatch sum
     [[ 0.286 -0.17 ]
     [-0.17   0.163]]

```
# Textbook implementation of Hessian

def NLL_hessian(weights, batch):
  X, y = batch
  mu = predict_batch(weights, X)
  S = jnp.diag(mu * (1-mu))
  H = jnp.dot(jnp.dot(X.T, S), X)
  return H

H2 = NLL_hessian(w, (X,y) )
assert np.allclose(Hbatch_sum, H2, atol=1e-2)
```

# JIT (just in time compilation)

In this section, we illustrate how to use the Jax JIT compiler to make code go faster (even on a CPU). However, it does not work on arbitrary Python code, as we explain below.

```
grad_fun_jit = jit(grad_fun) # speedup gradient function
grads_jit = vmap(partial(grad_fun_jit, w))(X,y)
assert np.allclose(grads, grads_jit)

```

```
# We can apply JIT to non ML applications as well.

def slow_f(x):
  # Element-wise ops see a large benefit from fusion
  return x * x + x * 2.0

x = jnp.ones((5000, 5000))
%timeit  slow_f(x)

fast_f = jit(slow_f)
%timeit fast_f(x)

assert np.allclose(slow_f(x), fast_f(x))
```

    100 loops, best of 3: 5 ms per loop
    The slowest run took 8.36 times longer than the fastest. This could mean that an intermediate result is being cached.
    1000 loops, best of 3: 1.37 ms per loop

We can also add the `@jit` decorator in front of a function.

```
@jit
def faster_f(x):
  return x * x + x * 2.0
%timeit faster_f(x)
assert np.allclose(faster_f(x), fast_f(x))
```

    The slowest run took 12.84 times longer than the fastest. This could mean that an intermediate result is being cached.
    1000 loops, best of 3: 1.37 ms per loop

## Static argnum

Note that JIT compilation requires that the control flow through the function  can be determined by the shape (but not concrete value) of its inputs. The function below violates this, since when x<3, it takes one branch, whereas when x>= 3, it takes the other.

```
@jit
def f(x):
  if x < 3:
    return 3. * x ** 2
  else:
    return -4 * x

# This will fail!
try:
  print(f(2))
except Exception as e:
  print("ERROR:", e)

```

    ERROR: Abstract tracer value encountered where concrete value is expected.

    The problem arose with the `bool` function.

    While tracing the function f at <ipython-input-66-3da05647f18e>:1, this concrete value was not available in Python because it depends on the value of the arguments to f at <ipython-input-66-3da05647f18e>:1 at flattened positions [0], and the computation of these values is being staged out (that is, delayed rather than executed eagerly).

    You can use transformation parameters such as `static_argnums` for `jit` to avoid tracing particular arguments of transformed functions, though at the cost of more recompiles.

    See https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error for more information.

    Encountered tracer value: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=0/1)>

We can fix this by telling JAX to trace the control flow through the function using concrete values of some of its arguments. JAX will then compile different versions, depending on the input values. See below for an example.

```
def f(x):
  if x < 3:
    return 3. * x ** 2
  else:
    return -4 * x

f2 = jit(f, static_argnums=(0,))

print(f2(5))
```

    -20

Unfortunately, the static argnum method fails with vmap, which passes in different inputs.

```

xs = jnp.arange(5)
try:
  ys = vmap(f)(xs)
  print('used vmap')
except:
  ys = jnp.array([f(x) for x in xs])
  print('did not use vmap')
print(ys)

```

    did not use vmap
    [  0.   3.  12. -12. -16.]

## Side effects

There are a few other subtleties. If your function has global side-effects, JAX's tracer can cause weird things to happen. A common gotcha is trying to print arrays inside jit'd functions:

```
def f(x):
  print(x)
  y = 2 * x
  print(y)
  return y
y1 = f(2)
print(y1)

print('jit version follows')
@jit
def f(x):
  print(x)
  y = 2 * x
  print(y)
  return y
y2 = f(2)
print(y2)

print('call jitted function a second time')
y2 = f(2)
print(y2)
```

    2
    4
    4
    jit version follows
    Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
    Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>
    4
    call jitted function a second time
    4

# Looping constructs

For loops in Python are slow, even when JIT-compiled. However, there are built-in primitives for loops that are fast, as we illustrate below.

## For loops.

The semantics of the for loop function in JAX is as follows:
```
def fori_loop(lower, upper, body_fun, init_val):
  val = init_val
  for i in range(lower, upper):
    val = body_fun(i, val)
  return val
```
We see that ```val``` is used to accumulate the results across iterations.

Below is an example.

```
# sum from 1 to N = N*(N+1)/2

def sum_exact(N):
  return int(N*(N+1)/2)

def sum_slow(N):
  s = 0
  for i in range(1,N+1):
    s += i
  return s

N = 10

assert sum_slow(N) == sum_exact(N)

def sum_fast(N):
  s = jax.lax.fori_loop(1, N+1, lambda i,partial_sum: i+partial_sum, 0)
  return s

assert sum_fast(N) == sum_exact(N)
```

```
N = 1000
%timeit sum_slow(N)
%timeit sum_fast(N)
```

    10000 loops, best of 3: 44.1 µs per loop
    10 loops, best of 3: 41 ms per loop

```
N = 100000
%timeit sum_slow(N)
%timeit sum_fast(N)
```

    100 loops, best of 3: 5.04 ms per loop
    1 loop, best of 3: 2.88 s per loop

```
# Let's do more compute per step of the for loop

D = 10
X = jax.random.normal(key, shape=(D,D))

def sum_slow(N):
  s = jnp.zeros_like(X)
  for i in range(1,N+1):
    s += jnp.dot(X, X)
  return s

def sum_fast(N):
  s = jnp.zeros_like(X)
  s = jax.lax.fori_loop(1, N+1, lambda i,s: s+jnp.dot(X,X), s)
  return s

N = 10
assert np.allclose(sum_fast(N), sum_slow(N))
```

```
N = 1000
%timeit sum_slow(N)
%timeit sum_fast(N)
```

    1 loop, best of 3: 482 ms per loop
    10 loops, best of 3: 46.3 ms per loop

## While loops

Here is the semantics of the JAX while loop

```
def while_loop(cond_fun, body_fun, init_val):
  val = init_val
  while cond_fun(val):
    val = body_fun(val)
  return val
```

Below is an example.

```

def sum_slow_while(N):
  s = 0
  i = 0
  while (i <= N):
    s += i
    i += 1
  return s

def sum_fast_while(N):
  init_val = (0,0)
  def cond_fun(val):
    s,i = val
    return i<=N
  def body_fun(val):
    s,i = val
    s += i
    i += 1
    return (s,i)
  val = jax.lax.while_loop(cond_fun, body_fun, init_val)
  s2 = val[0]
  return s2

N = 10
assert sum_slow_while(N) == sum_exact(N)
assert sum_slow_while(N) == sum_fast_while(N)
```

# Mutation of arrays

Since JAX is functional, you cannot mutate arrays in place,
since this makes program analysis and transformation very difficult. JAX requires a pure functional expression of a numerical program.
Instead, JAX offers the functional update functions: `index_update`, `index_add`, `index_min`, `index_max`, and the `index` helper. These are illustrated below.

Note: If the input values of `index_update` aren't reused, jit-compiled code will perform these operations in-place, rather than making a copy.

```
# You cannot assign directly to elements of an array.

A = jnp.zeros((3,3), dtype=np.float32)

# In place update of JAX's array will yield an error!
try:
  A[1, :] = 1.0
except:
  print('must use index_update')
```

    must use index_update

```
from jax.ops import index, index_add, index_update

D = 3
A = 2*jnp.ones((D,D))
print("original array:")
print(A)

A2 = index_update(A, index[1, :], 42.0) # A[1,:] = 42
print("original array:")
print(A) # unchanged
print("new array:")
print(A2)

A3 = A.at[1,:].set(42.0) # A3=np.copy(A),  A3[1,:] = 42
print("original array:")
print(A) # unchanged
print("new array:")
print(A3)

A4 = A.at[1,:].mul(42.0) # A4=np.copy(A),  A4[1,:] *= 42
print("original array:")
print(A) # unchanged
print("new array:")
print(A4)

```

    original array:
    [[2. 2. 2.]
     [2. 2. 2.]
     [2. 2. 2.]]
    original array:
    [[2. 2. 2.]
     [2. 2. 2.]
     [2. 2. 2.]]
    new array:
    [[ 2.  2.  2.]
     [42. 42. 42.]
     [ 2.  2.  2.]]
    original array:
    [[2. 2. 2.]
     [2. 2. 2.]
     [2. 2. 2.]]
    new array:
    [[ 2.  2.  2.]
     [42. 42. 42.]
     [ 2.  2.  2.]]
    original array:
    [[2. 2. 2.]
     [2. 2. 2.]
     [2. 2. 2.]]
    new array:
    [[ 2.  2.  2.]
     [84. 84. 84.]
     [ 2.  2.  2.]]

# Implicitly casting lists to vectors

You cannot treat a list of numbers as a vector. Instead you must explicitly create the vector using the np.array() constructor.

```
# You cannot treat a list of numbers as a vector.
try:
  S = jnp.diag([1.0, 2.0, 3.0])
except:
  print('must convert indices to np.array')
```

    must convert indices to np.array

```
# Instead you should explicitly construct the vector.

S = jnp.diag(jnp.array([1.0, 2.0, 3.0]))
```

# JAX neural net libraries

JAX is a purely functional library, which differs from Tensorflow and
Pytorch, which are stateful. The main advantages of functional programming
are that  we can safely transform the code, and/or run it in parallel, without worrying about
global state changing behind the scenes. The main disadvantage is that code (especially DNNs) can be harder to write.
To simplify the task, various DNN libraries have been designed, as we list below. In this book, we use Flax.

|Name|Description|
|----|----|
|[Stax](https://github.com/google/jax/blob/master/jax/experimental/stax.py)|Barebones library for specifying DNNs|
|[Flax](https://github.com/google/flax)|Library for specifying and training DNNs|
|[Haiku](https://github.com/deepmind/dm-haiku)|Library for specifying DNNs, similar to Sonnet|
|[Trax](https://github.com/google/trax)|Library for specifying and training DNNs, with a focus on sequence models|
|[Objax](https://github.com/google/objax)|Stateful (object-oriented) DNN framework, similar to PyTorch, not compatible with other JAX libraries|

# Other JAX  libraries

There are many other useful JAX libraries, most of which are purely functional, and therefore compose nicely.

|Name|Description|
|----|----|
|[NumPyro](https://github.com/pyro-ppl/numpyro)|Library for (deep) probabilistic modeling|
|[Optax](https://github.com/deepmind/optax)|Library for defining gradient-based optimizers|
|[RLax](https://github.com/deepmind/rlax)|Library for reinforcement learning|
|[Chex](https://github.com/deepmind/chex)|Library for debugging and developing reliable JAX code|

```

```


```python
import numpy as np
import pandas as pd
import linearsolve as ls
import matplotlib.pyplot as plt
plt.style.use('classic')
%matplotlib inline
```

# Discussion: Week 7

## Exercise: The Baseline RBC Model without Labor

The equilibrium conditions for the RBC model without labor are:

\begin{align}
\frac{1}{C_t} & = \beta E_t \left[\frac{\alpha A_{t+1}K_{t+1}^{\alpha-1} +1-\delta }{C_{t+1}}\right]\\
K_{t+1} & = I_t + (1-\delta) K_t\\
Y_t & = A_t K_t^{\alpha}\\
Y_t & = C_t + I_t\\
\log A_{t+1} & = \rho \log A_t + \epsilon_{t+1}
\end{align}

where $\epsilon_{t+1} \sim \mathcal{N}(0,\sigma^2)$.

The objective is use `linearsolve` to create a stochastic simulation of the RBC model. USe the following parameter values for the simulation:

| $\rho$ | $\sigma$ | $\beta$ | $\alpha$ | $\delta $ |
|--------|----------|---------|----------|-----------|
| 0.75   | 0.006    | 0.99    | 0.35     |  0.025    |

### Model Preparation

Before proceding, let's recast the model in the form required for `linearsolve`. Write the model with all variables moved to the lefthand side of the equations and dropping the expecations operator $E_t$ and the exogenous shock $\epsilon_{t+1}$:

\begin{align}
0 & = \beta\left[\frac{\alpha A_{t+1}K_{t+1}^{\alpha-1} +1-\delta }{C_{t+1}}\right] - \frac{1}{C_t}\\
0 & = A_t K_t^{\alpha} - Y_t\\
0 & = I_t + (1-\delta) K_t - K_{t+1}\\
0 & = C_t + I_t - Y_t\\
0 & = \rho \log A_t - \log A_{t+1}
\end{align}

Remember, capital and TFP are called *state variables* because they're $t+1$ values are predetermined. Output, consumption, and investment are called a *costate* or *control* variables. Note that the model as 5 equations in 5 endogenous variables.

### Initialization, Approximation, and Solution

The next several cells initialize the model in `linearsolve` and then approximate and solve it.

```python
# Create a variable called 'parameters' that stores the model parameter values in a Pandas Series. CELL PROVIDED
parameters = pd.Series()
parameters['rho'] = .75
parameters['sigma_squared'] = 0.006**2
parameters['beta'] = 0.99
parameters['alpha'] = 0.35
parameters['delta'] = 0.025

# Create variable called 'var_names' that stores the variable names in a list with state variables ordered first
var_names = ['a','k','y','c','i']

# Create variable called 'shock_names' that stores an exogenous shock name for each state variable.
shock_names = ['e_a','e_k']

# Define a function that evaluates the equilibrium conditions of the model solved for zero.
def equilibrium_equations(variables_forward,variables_current,parameters):

    # Parameters. PROVIDED
    p = parameters

    # Current variables. PROVIDED
    cur = variables_current

    # Forward variables. PROVIDED
    fwd = variables_forward

    # Euler equation
    euler_equation = p.beta*(p.alpha*fwd.a*fwd.k**(p.alpha-1)+1-p.delta)/fwd.c - 1/cur.c

    # Production function
    production_function = cur.a*cur.k**p.alpha - cur.y

    # Capital evolution
    capital_evolution = cur.i + (1 - p.delta)*cur.k - fwd.k

    # Market clearing
    market_clearing = cur.c+cur.i - cur.y

    # Exogenous tfp
    tfp_process = p.rho*np.log(cur.a) - np.log(fwd.a)

    # Stack equilibrium conditions into a numpy array
    return np.array([
        euler_equation,
        production_function,
        capital_evolution,
        market_clearing,
        tfp_process
        ])

# Initialize the model into a variable named 'rbc_model'
rbc_model = ls.model(equations = equilibrium_equations,
                 n_states=2,
                 var_names=var_names,
                 shock_names=shock_names,
                 parameters=parameters)

# Compute the steady state numerically using .compute_ss() method of rbc_model
guess = [1,4,1,1,1]
rbc_model.compute_ss(guess)

# Find the log-linear approximation around the non-stochastic steady state and solve using .approximate_and_solve() method of rbc_model
rbc_model.approximate_and_solve()
```

### Stochastic Simulation

Compute a 201 period stochastic simulation of the model. Set the seed for the simulation to 126.

Recall that `linearsolve` handles the model *as if* there is an exogenous shock for every state variable. In this case, that means two exogenous shocks $\epsilon_t^a$ and $\epsilon^k_t$. For the stochastic simulation, we need to specify the *covariance matrix* for the two shocks. The covariance matrix has the variance of each shock in the diagonal elements and covariances on the off-diagonal elements. Since there is no shock to capital in our model,

\begin{align}
\text{Covariance matrix} & = \left[\begin{array}{cc}\sigma^2 & 0\\ 0 & 0\end{array}\right]
\end{align}

```python
# Compute the stochastic simulation using the .stoch_sim() method of rbc_model
rbc_model.stoch_sim(T=201,seed=126,cov_mat = [[parameters['sigma_squared'],0],[0,0]])
```

```python
# Print the first 5 rows of rbc_model.simulated
print(rbc_model.simulated.head())
```

            e_a  e_k         a         k         y         c         i
    0  0.001484  0.0  0.003862  0.003977  0.005254  0.002906  0.012323
    1  0.001427  0.0  0.004323  0.004186  0.005788  0.003091  0.013910
    2  0.003155  0.0  0.006397  0.004429  0.007947  0.003499  0.021342
    3  0.011305  0.0  0.016103  0.004852  0.017801  0.004973  0.056433
    4 -0.001778  0.0  0.010299  0.006141  0.012449  0.005031  0.034786

```python
# Print the first 5 rows of simulated output
print(rbc_model.simulated['y'].head())
```

    0    0.005254
    1    0.005788
    2    0.007947
    3    0.017801
    4    0.012449
    Name: y, dtype: float64

## Exercise: Analyze Simulation Results

On a single axis, plot the simulated values for output, consumption and investment.

```python
# Plot simulated results
fig = plt.figure(figsize=(12,4))
ax = fig.add_subplot(1,1,1)
rbc_model.simulated[['y','c','i']].plot(ax=ax,lw=3,grid=True,alpha=0.75)
ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
```

**Questions**

1. What is the exogenous *cause* of the fluctuations in output, consumption and investment in your simulation?
2. Looking at the plot of your simulated data, which quantity — output, consumption or investment — has the smallest fluctuations from the steady state? What feature of the model leads to this?
3. Which quantity — output, consumption or investment — has the largest fluctuations from the steady state? What feature of the model leads to this?

**Answers**

1. Exogenous shocks to TFP. That is, exogenous shocks to the economy's production function.<!-- answer -->
2. Consumption. The household's Euler equation implies that the household wants to *smooth* consumption over time.<!-- answer -->
3. Investment. Again, the Euler equation. The household uses investment fluctuations to smooth consumption. When output is high, the household raises investment in order to accumulate capital and save some of the high production for later. Similarly, when output is low, the household reduces investment in order to deplete capital and therefore borrow against future output.<!-- answer -->