# Aiyagari (1994) in Continuous Time

#### By [SeHyoun Ahn](http://www.princeton.edu/~sehyouna/) and [Benjamin Moll](http://www.princeton.edu/~moll/)

The material in this notebook is based on [Achdou et al. (2015) "Heterogeneous Agent Models in Continuous Time"](http://www.princeton.edu/~moll/HACT.pdf) and follows closely the material in the paper's [online Appendix](http://www.princeton.edu/~moll/HACTproject/HACT_Numerical_Appendix.pdf). Additional codes (mainly in MATLAB) can be found [here](http://www.princeton.edu/~moll/HACTproject.htm). The code and its performance should be compared to [QuantEcon's discrete-time version of the Aiyagari model](http://quant-econ.net/py/aiyagari.html). The structure of the code in the present notebook is purposely kept as close as possible to that of the discrete-time code. The continuous-time code is considerably faster.

We begin by importing some packages

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from scipy import sparse
from scipy.sparse.linalg import spsolve
```

## Aiyagari Economy
The economy can be represented by the following system of equations which we aim to solve numerically:
$$
\begin{align*}
\rho v_1(a) &=  \max_c \ u(c) + v_1'(a)(wz_1 + ra - c) + \lambda_1(v_2(a) - v_1(a))\\
\rho v_2(a) &=  \max_c \ u(c) + v_2'(a)(wz_2 + ra - c) + \lambda_2(v_1(a) - v_2(a))\\
0 &= - \frac{d}{da}[s_1(a)g_1(a)] - \lambda_1 g_1(a) + \lambda_2 g_2(a)\\
0 &= - \frac{d}{da}[s_2(a)g_2(a)] - \lambda_2 g_2(a) + \lambda_1 g_1(a)\\
1 &= \int_{\underline{a}}^\infty g_1(a)da + \int_{\underline{a}}^\infty g_2(a)da\\
K &= \int_{\underline{a}}^\infty a g_1(a)da + \int_{\underline{a}}^\infty a g_2(a)da\\
r &= \alpha K^{\alpha-1} - \delta, \quad w=(1-\alpha)K^\alpha
\end{align*}
$$

where $z_1$ < $z_2$ and $s_j(a)=wz_j + ra -c_j(a)$ and $c_j(a) = (u')^{-1}(v_j(a))$ are optimal savings and consumption. Finally, there is a state constraint $a\geq \underline{a}$. The first order condition $u'(c_j(\underline{a}))=v'_j(\underline{a})$ still holds at the borrowing constraint. However, in order to respect the constraint we need $s_j(\underline{a}) = z_j + ra - c_j(\underline{a}) \geq 0$. Combining this with the FOC, the state constraint motivates a boundary condition
$$
\begin{align} v_j'(\underline{a}) \geq u'(z_j + r \underline{a}), \quad j=1,2 \end{align}
$$

We use a finite difference method, an in particular an "implicit upwind scheme." The details are explained in the paper's [online Appendix](http://www.princeton.edu/~moll/HACTproject/HACT_Numerical_Appendix.pdf). We here provide a brief summary. We approximate the functions $(v_1,v_2,g_1,g_2)$ at $I$ discrete points in the space dimension, $a_i,i=1,...,I$. We use equispaced grids, denote by $\Delta a$ the distance between grid points, and use the short-hand notation $v_{i,j} \equiv v_j(a_i)$ and so on. The derivative $v_{i,j}'=v_j'(a_i)$ is approximated with either a forward or a backward difference approximation
$$
\begin{align*}
v_j'(a_i) \approx \frac{v_{i+1,j} - v_{i,j}}{\Delta a} \equiv v_{i,j,F}' \\
v_j'(a_i) \approx \frac{v_{i-1,j} - v_{i,j}}{\Delta a} \equiv v_{i,j,B}'
\end{align*}
$$
An upwind scheme means that we approximate the derivative $v_j'(a_i)$ with a forward difference approximation whenever the drift of the state variable is positive and the backward difference approximation whenever it is negative. An implicit scheme is a particular way of iterating on the value function.

In a nutshell, the discretized HJB equation can be written as
$$\rho v = u + \mathbf{A} v$$
where $\mathbf{A}$ is $N \times N$ transition matrix with $N = 2 \times I$ and where $I$ is the number of wealth grid points. The matrix $\mathbf{A}$ depends on $v$, i.e. this is a nonlinear problem and we therefore need to iterate (this is where the implicit scheme comes in). The matrix $\mathbf{A}$ has the interpretation of a Poisson transition matrix (or "intensity matrix") on the discretized state space $(a_i,z_j)$.

Similarly, one can show that the discretized Kolmogorov Forward equation is
$$0 = \mathbf{A}^T g$$
which is an eigenvalue problem. That is, the discretized stationary distribution $g$ is the eigenvector corresponding to a zero eigenvalue of the transpose of the Poisson transition matrix $\mathbf{A}$. The transpose comes from the fact that the differential operator in the KF equation is the "adjoint" of the operator in the HJB equation. And an "adjoint" is the infinite-dimensional analogue of matrix transpose.

The matrix $\mathbf{A}$ is found from the discretized HJB equation. Skipping a number of steps it can be written as
$$
\begin{align*}
&\frac{v^{n+1}_{i,j} - v^{n}_{i,j}}{\Delta} + \rho v_{i,j}^{n+1} = u(c_{i,j}^n) + v_{i-1,j}^{n+1}x_{i,j}   +  v^{n+1}_{i,j} y_{i,j} + v_{i+1,j}^{n+1} z_{i,j} + v_{i,-j}^{n+1}\lambda_j \quad \mbox{where}\\
&x_{i,j} = -\frac{(s^n_{i,j,B})^-}{\Delta a},\\
&y_{i,j} = - \frac{(s^n_{i,j,F})^+}{\Delta a}  +  \frac{ (s^n_{i,j,B})^-}{\Delta a} - \lambda_j,\\
&z_{i,j} = \frac{(s^n_{i,j,F})^+}{\Delta a}
\end{align*}
$$
where $s^n_{i,j,F}$ and $s^n_{i,j,B}$ are the discretized optimal household savings at grid points $(a_i,z_j)$ using forward and backward approximations, and where for any number $x$, the notation $x^+$ means "the positive part of $x$", i.e. $x^+ = \max\{x,0\}$ and analogously $x^{-} = \min\{x,0\}$. This part is what makes it an upwind scheme.

This is a system of $2 \times I$ linear equations which can be written in matrix notation as:
\begin{equation}\frac{1}{\Delta}(v^{n+1} - v^n) + \rho v^{n+1} = u^n + \mathbf{A}^n v^{n+1} \end{equation}
where
$$\mathbf{A}^n = \left[\begin{matrix}y_{1,1} & z_{1,1} & 0 & \cdots & 0 & \lambda_1 & 0 & 0 & \cdots & 0 \\ x_{2,1} & y_{2,1} & z_{2,1} & 0 & \cdots & 0 & \lambda_1 & 0 & 0 & \cdots \\ 0 & x_{3,1} & y_{3,1} & z_{3,1} & 0 & \cdots & 0 & \lambda_1 & 0 & 0 \\ \vdots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \vdots \\ 0 & \ddots & \ddots & x_{I,1} & y_{I,1} & 0 & 0 & 0 & 0 & \lambda_1\\  \lambda_2 & 0 & 0 & 0 & 0 & y_{1,2} & z_{1,2} & 0 & 0 & 0 \\ 0 & \lambda_2 & 0 & 0 & 0 & x_{2,2} & y_{2,2} & z_{2,2} & 0 & 0  \\ 0 & 0 & \lambda_2 & 0 & 0 & 0 & x_{3,2}  & y_{3,2} & z_{3,2} & 0 \\ 0 & 0 & \ddots & \ddots & \ddots & \ddots & \ddots  & \ddots & \ddots & \ddots \\ 0 & \cdots & \cdots & 0 & \lambda_2 & 0 & \cdots  & 0 & x_{I,2} & y_{I,2}  \end{matrix}\right], \quad u^n = \left[\begin{matrix} u(c_{1,1}^n)\\ \vdots \\ \vdots \\ u(c_{I,1}^n)\\  u(c_{1,2}^n) \\ \vdots \\ \vdots \\ u(c_{I,2}^n)\end{matrix}\right]$$
This system can in turn be written as
$$
\begin{align*}\mathbf{B}^n v^{n+1} = b^n, \quad \quad \mathbf{B}^n = \left(\frac{1}{\Delta} + \rho\right)\mathbf{I} - \mathbf{A}^n, \quad  b^n = u^n + \frac{1}{\Delta}v^n \end{align*}
$$
This system of linear equations can be solved very efficiently using sparse matrix routines. In Python this is implemented with the function "spsolve()."

#### Summary of Algorithm
First consider the algorithm for solving the HJB equations. Guess $v^0_{i,j},i=1,...,I,j=1,2$ and for $n=0,1,2,...$ follow
1. Compute $(v^n_{i,j})'$ using the current guess of the value function and the upwind scheme (forward difference when drift is positive, backward difference when drift is negative)
2. Compute $c^n$ from $c_{i,j}^n = (u')^{-1}[(v_{i,j}^n)']$
3. Find $v^{n+1}$ by solving the system of linear equations involving the matrix $\mathbf{A}$ described above (implicit scheme)
4. If $v^{n+1}$ is close enough to $v^n$: stop. Otherwise, go to step 1.

After solving the HJB equations, solving the Kolmogorov Forward equation: simply solve the eigenvalue problem $0=\mathbf{A}^T g$ described above. That is, once the HJB equation is solved, we basically get the Kolmogorov Forward equation "for free."

#### Overview of Code
Given this overview of the algorithm, we now briefly describe the code. As in the discrete-time version, we define a "household" class. Household objects contain all the data relevant to solving a household's decision problem. The household class  contain:
* economic parameters (e.g., w, r)
* utility paramters (e.g., discount factor $\rho$)
* asset and skill level represented on a grid

The household's decision problem is solved by invoking the function solve_bellman() which solves the HJB equation given the relevant parameters saving the value-function as v.

We also include the stationary wealth distribution of households in the household object. This is a natural thing to do since this stationary distribution is computed using the household's decision rule. Hence, after computing the decision problem of households, the stationary distribution can be found by invoking compute_stationary_distribution()

The definition of the household class is given below.

```python
class Household(object):
    def __init__(self,
                 dep=0.05,   ###ADDED DEPRECIATION
                 r=0.03,         # interest rate
                 w=1,            # wages
                 rho=0.04,       # discount factor
                 a_min=1e-10,    # minimum asset amount
                 pi=[[-0.33, 0.33], [0.33, -0.33]],  # poisson Jumps
                 z_vals=[1.0, 2.0],                # exogenous income states
                 a_max=40,
                 a_size=1000,    # number of asset grid points
                 delta=1000.0):
        # Initialize values, and set up grids over a and z
        self.r, self.w, self.rho, self.dep = r, w, rho, dep
        self.a_min, self.a_max, self.a_size = a_min, a_max, a_size
        self.da = (self.a_max-self.a_min)/(self.a_size-1)
        self.k = 10
        self.pi = np.asarray(pi)
        self.z_vals = np.asarray(z_vals)
        self.z_size = len(z_vals)
        self.a_vals = np.linspace(self.a_min, self.a_max, self.a_size)
        self.n = self.a_size * self.z_size
        self.delta = delta
        ###### ADDED TO MATCH LABOR SUPPLY IN .m
        self.z_ave = (self.z_vals[0]*self.pi[0, 1] +
                      self.z_vals[1]*self.pi[1, 0]) / \
            (self.pi[0, 1] + self.pi[1, 0])

        # Initial Guess of Value Function
        self.v = np.log(np.tile(self.a_vals,(self.z_size,1))*self.r
                        +self.w*np.tile(self.z_vals,(self.a_size,1)).transpose())/self.rho

        # Build skill_transition, the matrix summarizing transitions due to the Poisson income shocks
        # This is analogous to the Q matrix in the discrete time version of the QuantEcon Aiyagari model
        self.z_transition = sparse.kron(self.pi,sparse.eye(self.a_size), format="csr")

        # Preallocation
        self.v_old = np.zeros((self.z_size,self.a_size))
        self.g = np.zeros((self.z_size,self.a_size))
        self.dv = np.zeros((self.z_size,self.a_size-1))
        self.cf = np.zeros((self.z_size,self.a_size-1))
        self.c0 = np.zeros((self.z_size,self.a_size))
        self.ssf = np.zeros((self.z_size,self.a_size))
        self.ssb = np.zeros((self.z_size,self.a_size))
        self.is_forward = np.zeros((self.z_size,self.a_size),'bool')
        self.is_backward = np.zeros((self.z_size,self.a_size),'bool')
        self.diag_helper = np.zeros((self.z_size,self.a_size))
        self.A = self.z_transition.copy()
        self.B = self.z_transition.copy()
        self.AT = self.z_transition.copy()

    def set_prices(self, r, w):
        """
        Resets prices
        Calling the method will resolves the Bellman Equation.

        Parameters:
        -----------------
        r : Interest rate
        w : wage
        """
        self.r, self.w = r, w
        self.solve_bellman()

    def reinitialize_v(self):
        """
        Reinitializes the value function if the value function
        became NaN
        """
        self.v = np.log(np.tile(self.a_vals,(self.z_size,1))*self.r
                        +self.w*np.tile(self.z_vals,(self.a_size,1)).transpose())/self.rho

    def solve_bellman(self,maxiter=100,crit=1e-6):
        """
        This function solves the decision problem with the given parameters

        Parameters:
        -----------------
        maxiter : maximum number of iteration before haulting value function iteration

        crit : convergence metric, stops if value function does not change more than crit
        """
        dist=100.0
        for i in range(maxiter):
            # compute saving and consumption implied by current guess for value function, using upwind method
            self.dv = (self.v[:,1:]-self.v[:,:-1])/self.da
            self.cf = 1.0/self.dv
            self.c0 = np.tile(self.a_vals,(self.z_size,1))*self.r \
                            +self.w*np.tile(self.z_vals,(self.a_size,1)).transpose()

            # computes savings with forward forward difference and backward difference
            self.ssf[:,:-1] = self.c0[:,:-1]-self.cf
            self.ssb[:,1:] = self.c0[:,1:]-self.cf
            # Note that the boundary conditions are handled implicitly as ssf will be zero at a_max and ssb at a_min
            self.is_forward = self.ssf>0
            self.is_backward = self.ssb<0
            # Update consumption based on forward or backward difference based on direction of drift
            self.c0[:,:-1] += (self.cf-self.c0[:,:-1])*self.is_forward[:,:-1]
            self.c0[:,1:] += (self.cf-self.c0[:,1:])*self.is_backward[:,1:]
            ######
            # UNCOMMENT FOR DEBUGGING
            #plt.plot(self.a_vals, self.c0.transpose())
            #plt.show()
            self.c0 = np.log(self.c0)

            # Build the matrix A that summarizes the evolution of the process for (a,z)
            # This is a Poisson transition matrix (aka intensity matrix) with rows adding up to zero
            self.A = self.z_transition.copy()
            self.diag_helper = (-self.ssf*self.is_forward/self.da \
                               + self.ssb*self.is_backward/self.da).reshape(self.n)
            self.A += sparse.spdiags(self.diag_helper,0,self.n,self.n)
            self.diag_helper = (-self.ssb*self.is_backward/self.da).reshape(self.n)
            self.A += sparse.spdiags(self.diag_helper[1:],-1,self.n,self.n)
            self.diag_helper = (self.ssf*self.is_forward/self.da).reshape(self.n)
            self.A += sparse.spdiags(np.hstack((0,self.diag_helper)),1,self.n,self.n)
            # Solve the system of linear equations corresponding to implicit finite difference scheme
            self.B = sparse.eye(self.n)*(1/self.delta + self.rho) - self.A
            self.b = self.c0.reshape(self.n,1) + self.v.reshape(self.n,1)/self.delta
            self.v_old = self.v.copy()
            self.v = spsolve(self.B,self.b).reshape(self.z_size,self.a_size)

            # Compute convergence metric and stop if it satisfies the convergence criterion
            dist = np.amax(np.absolute(self.v_old-self.v).reshape(self.n))
            if dist < crit:
                break

    def compute_stationary_distribution(self):
        """
        Solves for the stationary distribution given household decision rules

        Output:
        Capital level from the stationary distribution
        """
        self.AT = self.A.transpose().tocsr()

        # The discretized Kolmogorov Forward equation AT*g=0 is an eigenvalue problem
        # AT is singular because one of the equation is the distribution adding
        # up to 1. Here we solve the eigenvalue problem by setting g(1,1)=0.1
        # and the equation is solved relative to that value.
        # Alternatively, one could use a routine for solving eigenvalue problems.
        b = np.zeros((self.n,1))
        b[0] = 0.1
        self.AT.data[1:self.AT.indptr[1]] = 0
        self.AT.data[0] = 1.0
        self.AT.indices[0] = 0
        self.AT.eliminate_zeros()
        self.g = spsolve(self.AT,b).reshape(self.z_size,self.a_size)

        # Since g was solved taking one of g(1,1) as given, g needs to be
        # renormalized to add up to 1
        self.g = self.g/np.sum(self.g)
        return np.sum(self.g*(np.tile(self.a_vals,(self.z_size,1))))
```

For example, if interest rate is 0.05 and wage is 1, we can initialize a household, solve it decision problem, and find the stationary distribution by running

```python
lam = 0.11
PI = [[-lam, lam], [lam, -lam]]
am=Household(rho=0.05, r=0.02,w=1, pi=PI)
am.solve_bellman()
am.compute_stationary_distribution()
```

    0.69274641340853271

Once, a household object is created, the object can be reused to solve a different problem by changing parameters. For example, if the interest rate were to change to 0.03 and wage to 0.9, the new problem can be solved by setting parameters directly by

```python
am.r=0.02
am.w=0.9
```

and running

```python
am.solve_bellman()
am.compute_stationary_distribution()
```

    0.62323720534758664

Alternately, a helper function can be created to automate this process. For example, set_prices(r,w) function resets parameters, and solves the decision problem again. For example, you can run the following to the same effect.

```python
am.set_prices(r=0.03,w=0.9)
am.compute_stationary_distribution()
```

    1.129833308836365

Given the household class, solving for the steady state becomes really simple. To find the equilibrium, we solve for the capital level that is consistent with household's decisions.

Before solving for the actual steady state, we can visualize the capital demand and capital supply.

```python
A = 0.1
##### NOT NECESSARY
#N = 0.05
alpha = 0.33

def r_to_w(am, r):
    return A*(1 - alpha) * \
        (alpha*A / (am.dep + r))**(alpha / (1 - alpha))

def rd(am, K):
    return A*alpha*(am.z_ave / K)**(1 - alpha) - am.dep

def prices_to_capital_stock(am, r):
    """
    Map prices to the induced level of capital stock.

    Parameters:
    ----------
    am : Household
        An instance of the Household class
    r : float
        The interest rate
    """
    w = r_to_w(am, r)

    # Set new prices and solve the Bellman equation
    am.set_prices(r, w)

    # Compute the stationary distribution and capital
    return am.compute_stationary_distribution()
```

We make a grid of interest rate points, and plot the resulting capital.

```python
num_points = 20
r_vals = np.linspace(0.02, 0.048, num_points)

# Compute supply of capital
k_vals = np.empty(num_points)
for i, r in enumerate(r_vals):
    k_vals[i] = prices_to_capital_stock(am,r)

# Plot supply and demand of capital
fig, ax = plt.subplots(figsize=(11, 8))
ax.plot(k_vals, r_vals, lw=2, alpha=0.6, label='supply of capital')
ax.plot(k_vals, rd(am, k_vals), lw=2, alpha=0.6, label='demand for capital')
ax.grid()
ax.set_xlabel('capital')
ax.set_ylabel('interest rate')
ax.legend(loc='upper right')
plt.show()
```

Finally, the equilibrium interest rate can be found by using the bisection method, and we can see the equilibrium distribution of assets.

```python
# Set parameters for bisection method
crit = 1e-6
r_min = 0.02
r_max = 0.05
r = 0.03

# Bisection loop
for i in range(100):
    am.set_prices(r,r_to_w(am, r))
    r_new = rd(am, am.compute_stationary_distribution())
    if np.absolute(r_new-r)<crit:
        break
    elif r_new > r:
        r_min = r
        r = (r_max+r_min)/2.
    else:
        r_max = r
        r = (r_max+r_min)/2.

# Plot stationary distribution at the equilibrium
fig, ax = plt.subplots(figsize=(11, 8))
n=50                      # Determine the max asset level to show in the plot
ax.plot(am.a_vals[0:n], am.g[0,0:n], lw=2, alpha=0.6, label='low income')
ax.plot(am.a_vals[0:n], am.g[1,0:n], lw=2, alpha=0.6, label='high income')
ax.grid()
ax.set_xlabel('asset position')
ax.set_ylabel('distribution')
ax.legend(loc='upper right')
plt.show()
```

```python
am.r
```

    0.04605979919433595


```python
from IPython.display import display, HTML
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"
#InteractiveShell.ast_node_interactivity = "last"
```

```python
# Load inputs
import pandas as pd
df = pd.read_pickle("bau.pickle")

# Already have OAT, no need for wx.
#wx = pd.read_pickle("wxhydro2018.pickle")
#df = dfIn.merge(wx,how='outer',left_index=True,right_index=True)
#df = pd.DataFrame(df[df['OAT']<16]['OAT'].dropna())
#df.describe()
```

# Introduction

### Programmable Stat
My house is occupied 24/7, but everyone sleeps on the same schedule, so we could consider implementing a night setback program.  NSB in a fancy commercial BAS builds a lookup table so it know's each zone's thermal response speed to temperature perturbations, but most residential programmable stats only let you modify the RTSP on a schedule.

The hard way to estimate savings from a reduction in interior temperature is to model the heating system's behaviour.  It's hard because it requires a solid heat transfer model of the entire house, including ventilation, to carefully determine the heat loss, solving the differential equations to account for the exponential decay in transfer rates as temperatures approach steady state, and then trying to figure out how the heating system will compensate.

But fortunately, we already know how much heat is coming through the walls, and for a first-order approximation, heat loss is proportional to the delta-T between inside and outside. Our savings will be the time-weighted change in delta-T that we can produce with scheduling.

In my house, if I turn the heat off overnight with the OAT at around 10C, it will cool from 22 to 19 in about 8 hours.  It takes about three hours to warm back up from 18 no matter what the OAT.

Applying Newton's law of cooling $T(t)=OAT+(RT-OAT)e^{-kt}$, we find

\begin{align*}
T(8)=19 & =10+12e^{-8k}\\
ln(\frac{9}{12}) & =-8k\\
k & =0.036\\
\end{align*}

Heating dominates over losses, so the warmup rate is quite linear:

$$WARMUP\:RATE = 0.5\frac{^{\circ}C}{hr}$$

Bedtime's around 10PM and reveille around 7AM, so let's use a NSB schedule of 8PM-4AM.  We can tolerate things cooling off a bit just before bedtime, and the 4AM start guarantees that we'll have hit our "occupied" RTSP by 7AM on all but very cold nights.

```python
# RTSP = 22 Occupied, 18 Night
#df = pd.DataFrame(bau['20170901':'20180301']['OAT'].dropna())
df.loc[(df.index.hour>=4)&(df.index.hour<20),'nrtsp']=22
df.loc[(df.index.hour<4)|(df.index.hour>=20),'nrtsp']=18
```

```python
import numpy as np
WARMRATE=0.5
# First we'll calculate the hourly RT for NSB: nrt
df['nrt']=0.0
rt=22
newtonK=0.036
for i,row in df.iterrows():
    rtsp = row.nrtsp
    delta = float(rtsp-rt)
    #print("D: %f"%delta)
    if delta > 0:
        # increase temperature by delta to a max of 1
        if delta > 1: delta = WARMRATE
    elif delta < 0:
        # newton's law of cooling for T=+1
        rt1 = row.OAT + (rt-row.OAT)*np.exp(-newtonK)
        CDRATE = rt1-rt
        #print("CD: %f"%CDRATE)
        if delta < CDRATE: delta = CDRATE
    rt = rt + delta
    df.nrt.at[i]= rt
    #print("%d rtsp=%f rt=%f delta=%f oat=%f" % (i.hour,rtsp,rt,delta,row.OAT))

```

```python
# Now calculate the sum of delta-T for BAU vs NSB (heating only)
df['bauhd']=df.rt-df.OAT
df.loc[df['bauhd']<0,'bauhd']=0
df['nsbhd']=df.nrt-df.OAT
df.loc[df['nsbhd']<0,'nsbhd']=0
#project savings
heatDTsaved = (df.bauhd.sum() - df.nsbhd.sum())
heatfracsaved = heatDTsaved/df.bauhd.sum()
print("Heating savings: %.2f%%"%(100*heatfracsaved))
#df.nsbhd.describe()
#df.bauhd.describe()
```

    Heating savings: 7.61%

```python
df.nsbhd.describe()
df.bauhd.describe()
df.nsbhd.sum() / df.bauhd.sum()
```

    count    8760.000000
    mean       10.333060
    std         4.850989
    min         0.000000
    25%         6.171439
    50%        10.998450
    75%        14.100000
    max        23.900000
    Name: nsbhd, dtype: float64

    count    8760.000000
    mean       11.183767
    std         5.257502
    min         0.000000
    25%         6.500000
    50%        12.000000
    75%        15.500000
    max        25.000000
    Name: bauhd, dtype: float64

    0.9239337241165206

The following chart illustrates how much more quickly things cool down when it's cold out.  The green squiggle shows the sum of reduced $\Delta T$ values (RT vs OAT) over two periods of three days, as compared to a constant RT of 22C.

```python
import matplotlib.pyplot as plt
fig, axes = plt.subplots(nrows=2, ncols=1,figsize=(7,12),sharey=True)
fig.suptitle("Typical NSB HDH Savings")

dd=pd.DataFrame(df['2013-05-10':'2013-05-13'])
dd['SavHDH'] = (dd.rt-dd.nrt).cumsum()
a=0

tit="Warm OAT ~ %d°C"%dd.OAT.mean()
ax=dd[['nrt','OAT','SavHDH']].plot(ax=axes[a],kind='line',title=tit,secondary_y='SavHDH')
dd[['nrtsp']].plot(ax=axes[a],kind='line',drawstyle="steps-post",alpha=.3)
lines = ax.get_lines() + ax.right_ax.get_lines() # Need to add "lines" from right ax to label it.
#ax.legend(lines,['RT','OAT','SavHDH','RTSP'])
ax.set_ylabel('Degrees C')

dd=pd.DataFrame(df['2013-12-6':'2013-12-9'])
dd['SavHDH'] = (dd.rt-dd.nrt).cumsum()
a=1
tit="Cold OAT ~ %d°C"%dd.OAT.mean()
ax=dd[['nrt','OAT','SavHDH']].plot(ax=axes[a],kind='line',title=tit,secondary_y='SavHDH')
dd[['nrtsp']].plot(ax=axes[a],kind='line',drawstyle="steps-post",alpha=.3)
lines = ax.get_lines() + ax.right_ax.get_lines() # Need to add "lines" from right ax to label it.
#ax.legend(lines,['RT','OAT','SavHDH','RTSP'])
ax.right_ax.set_ylabel('Sum of saved degree-hours')

plt.show()

```

## NSB Hourly Consumption
The average annual savings number is nice to know, but what we *really* want to know is the new hourly consumption -- otherwise we have no way of knowing how the NSB will interact with other measures.

We can get fairly close with the assumptions we've already made:

  - Cooldown: hourly heating load is 0
  - Steady state: same as BAU, but with savings proportional to the reduction in $\Delta T$ vs OAT.
  - Total heating savings $BAU \rightarrow NSB$ are known from the approximation above
  - We can now solve for the sum of heating loads during warmup periods
    - distribute this sum between all warmup periods
    - pro rate such that hours with a fractional warmup get a fractional share.

\begin{align}
\sum H_{NSB}	=\sum H_{BAU}\times(1-Savings\sim6.5\%)\\
H_{NSB}	=H_{cooldown}+H_{SS}+H_{warmup}\\
\begin{aligned}H_{cool} &=0\\
H_{SS} & =H_{BAU}\times\left(1-\frac{\Delta T_{NSB}}{\Delta T_{BAU}}\right)\\
\end{aligned}\\
\sum H_{warmup}	=\sum H_{BAU}(0.935)-\sum H_{SS}
\end{align}

Distribute as follows:

\begin{align}
\sum Hours_{warm}=\sum f(t) \begin{cases}
1 & \Delta T(t)\ge1C\\
\Delta T(t) & \Delta T(t)<1C
\end{cases}\\
H(t)=\frac{\sum H}{\sum Hours} \times f(t)\\
\end{align}

```python
def saveDT(row):
    RTN=row.nrt
    RTB=row.rt
    OAT=row.OAT
    # Computes the fractional reduction in delta T for NSB as compared to BAU
    DTN=RTN-OAT
    DTB=RTB-OAT
    if(DTB<=0): return 0 # There are no savings possible if the BAU DT indicates no heating.
    if(DTN<=0): DTN=0 # Best we can do is turn off the heat and save 100%
    return (DTB-DTN)/DTB

# Categorize periods and calculate new heating load for each period
h = df[(df.hload!=0)].dropna()
# Set default values for cooling
h['cooldown']= h.nrt > h.nrtsp
h.loc[h.cooldown,'savefrac'] = 1
h.loc[h.cooldown,'nhload'] = 0

#compute savings fraction and hload for the Steady State periods
# I think there's an error here.
h['ss']= h.nrt == h.nrtsp
#h['savefrac'] =
h.loc[h.ss,'savefrac']=h[h.ss].apply(saveDT,axis=1)
h.loc[h.ss,'nhload']=h[h.ss]['hload'] * (1-h[h.ss]['savefrac'])

#now get ready for the warmup rows
h['warmup']= h.nrt < h.nrtsp
```

```python

```

```python

```

```python
# calculate the expected total NSB heating load, and the number of hours over which it should be distributed
hw = pd.DataFrame(h[h.warmup])

# Count the whole and fractional hours where warmup occurs
hw['whrs'] = hw.nrtsp - hw.nrt # This is correct where DT<=1
hw.loc[hw.whrs>1,'whrs'] = 1 # Because whrs is 1 where DT>1
whrs = hw.whrs.sum()

# Total Warmup heating is expected NSB total, less expected NSB at SS.
expectedHnsb = h.hload.sum() * (1-heatfracsaved)
expectedHss = h[h.ss].hload.sum()
expectHwarmup = expectedHnsb - expectedHss

# Now distribute the warmup load
HwarmPerWhr = expectHwarmup / whrs
h.loc[hw.index,'nhload'] = (hw.whrs * HwarmPerWhr)

# There is an error here:
h.nhload.sum() / h.hload.sum()

```

    0.9188191831627831

```python
# These should match
#expectedHss
#h.nhload[h.ss].sum()

# So should these
#expectHwarmup
#h.nhload[h.warmup].sum()

h.nhload.sum()
```

    6369.993951481103

```python
h.hload.sum()
```

    6932.8047

```python

```

```python

```


Power Law Distributions: Zipf's Law
=======================

(Note: Be sure to either download the data that goes with this notebook or uncomment/edit the lines below needed to get your own data before executing the code in this notebook.)

It turns out that "power law" distributions are ubitquitous in nature, see:

1 [http://arxiv.org/abs/cond-mat/0412004](http://arxiv.org/abs/cond-mat/0412004)

2 [https://arxiv.org/abs/1812.09662](https://arxiv.org/abs/1812.09662)

In the next feww cells of this notebook I measure the number of files of different sizes on my hard disk. Then I proceed to quantify (in the simples possible way) the distribution of file sizes in a certain range as a power law distribution.

```python
import numpy as np
import matplotlib.pyplot as pl

import pandas as pd                    # pandas is the "python data analysis package"
import os                              # operating system services
from scipy.optimize import curve_fit   # non-linear curve fitting
```

```python
def getFileSizes(folder, minsize=2000, maxsize=10**9):
    """
    return a list of file sizes in a folder
    """
    sizes = []
    for (path, dirs, files) in os.walk(folder):
        for file in files:
            try:
                filename = os.path.join(path, file)
                size=os.path.getsize(filename)
                if size>minsize and size<maxsize:
                    sizes.append(size)
            except OSError:
                pass
    return sizes
```

```python
if 0: # set this to "1" to actaully collect data, but edit the paths to match your computer.
    minsize=10**4  # the minimum file size to consider
    maxsize=10**10 # the maximum file size to consider
    sizes = getFileSizes('/Users/steve/Desktop',minsize,maxsize) # <--- edit this to match some folder with a lot of files
    print(len(sizes), 'files between', minsize, 'and', maxsize)
```

```python

if 0: # set this to "1" to write the file
    df=pd.DataFrame({'file_sizes':sizes})
    df.to_csv('myData.csv', sep=',')
else: # read the file
    df=pd.read_csv('myData.csv')
    df.head()
    sizes = df.file_sizes
```

```python
ns,bins,patches = pl.hist(np.log(np.array(sizes)+1),40)
```

```python
xvals=(bins[:-1]+bins[1:])/2.0
counts=ns+0.5
yvals=np.log(counts)
pl.title("distribution of file sizes")
pl.ylabel("log(count)")
pl.xlabel("log(size)")
pl.plot(xvals,yvals,'b.')
```

Uncertainties in the Parameters
-------------------------------

There are a variety of techniques for estimating the uncertainty in the parameters you compute. The approach we'll take is one of the more "brute force" approaches, but it's pretty robust. We're going to fabricate data that has the same statistics as our original data (do the degree that we know and/or care). Then we'll fit the fabricated data many times (1000 times in this case) and we'll look at the statistics of the fit parameters produced. From this we can tell how well our fit has done at narrowing down the parameters we computed.

```python
#
# Just to see same data fit using the non-linear scipy package "curve_fit":
#

def fLinear(x, m, b):
    return m*x + b

m0 = (yvals[-1]-yvals[0])/(xvals[-1]-xvals[0])  # estimate slope
b0 = yvals[0] - m0*xvals[0]                     # estimate intercept
sigma=1.0/np.sqrt(np.abs(counts))

popt, pcov = curve_fit(fLinear, xvals, yvals, p0=(m0, b0), sigma=sigma)

m=popt[0]          # slope
dm=np.sqrt(pcov[0,0]) # sqrt(variance(slope))
b=popt[1]          # int
db=np.sqrt(pcov[1,1]) # sqrt(variance(int))
ystar=fLinear(xvals, m, b)

N=1000    # number of samples to take
mList=[]  # keep track of monte-carlo fit parameters
bList=[]

for i in range(N):
    """
    Generate mc data with the same statistical properties as the real data.
    Repeat the fit for each set, and record the parameters.
    """
    mcY = ystar + sigma*np.random.normal(size=len(xvals))  # generate fabricated data to fit
    mcpopt, mcpcov = curve_fit(fLinear, xvals, mcY, p0=(m,b), sigma=sigma)
    mList.append(mcpopt[0])  # store the fit paramters for the fab-data
    bList.append(mcpopt[1])
```

```python
# check slope histogram

pl.hist(mList)
```

```python
# check intercept histogram

pl.hist(bList)
```

```python
#
# Compute the statistics of the mc-results
#

marr=np.array(mList)
barr=np.array(bList)

mAvg = marr.mean()
bAvg = barr.mean()
sigM = marr.std()
sigB = barr.std()

pl.errorbar(xvals, yvals, sigma, fmt='r.')
pl.title("File Sizes on My HD")
pl.ylabel("log(count)")
pl.xlabel("log(size)")
pl.plot(xvals,ystar, 'g-')
print("Slope=", m, '+/-', sigM, "(", m-sigM,",",m+sigM, ")")
print("Intercept=", b, '+/-', sigB, "(", b-db,",",b+sigB, ")")

print() # print a blank link
print("Compare to cov-matrix for fun:")
print("sqsrt(pcov[0,0]) (should be sigma m)", np.sqrt(pcov[0,0]))
print("sqrt(pcov[1,1]) (should be sigma b)", np.sqrt(pcov[1,1]))
```

Project 7
=========

See this reference for other systems that exhibit power law behavior:

<http://arxiv.org/abs/cond-mat/0412004>

Or test your intuition and see if it's right. Many system exhibit this behavior!

Get some data, and test your idea. Use the matrix technique to do the fitting.

Matrix Approach
--------------------

# WARNING: THE REST OF THIS NOTEBOOK IS OPTIONAL. YOU DO NOT NEED TO READ FURTHER UNLESS YOU ARE CURIOUS/INTERSTED IN THE INNER MACHINERY OF CURVE FITTING

The remainder of this notebook details an approach to parameter estimation
using matrices. If you're interested, you're welcome to use this approach to
fit your data. It has the (minor) advantage that there are no "black boxes"
like `curve_fit`. The underlying mathematics is open for your inspection! However
it's also pretty mathematically "heavy" so, please use `curve_fit` if you prefer.

Fitting Data with Matrices
===========================

We can see that we have something that looks like it might be a power law. How can we tell? We'll use "linear least squares" analysis.

When fitting some experimental data to a model, there are typically
a set of paramters $(\alpha_1, \alpha_2, \ldots, \alpha_n)$ that
need to be found so that the function $\chi^2$ is a minimum.
We can think of the paramters as an n-vector called $\alpha$. The
data consist of a set of paired numbers $(x_i, y_i)$ and we can
define $\chi^2$ as

\begin{equation}
\chi^2 = \sum_{i=1}^{m} {(y_i - y^\star(x_i,\alpha))^2\over \sigma_i^2},
\end{equation}

where $\sigma_i$ is the uncertainty in $y_i$, and $y^\star$ is the
model value corresponding to $x_i$ (and the value of the parameters
$\alpha$). If the model $y^\star(x,\vec\alpha)$ is *linear* then we
can write $y^\star$ as:

\begin{equation}
y^\star (x) = \alpha_1 X_1(x)+\alpha_2 X_2(x)+\ldots+\alpha_n X_n(x).
\end{equation}

or

\begin{equation}
y^\star(x) = \sum_{j=1}^{n} \alpha_j X_j(x),
\end{equation}

where the $X_j$ are a set of functions that are to be superposed
with coefficients $\alpha_j$ to form the model. A common example
of this would be a straight line fit $y^\star(x) = mx+b$ where the
functions are $x$ and 1 for the paramters $m$ and $b$ respectively.

Notice that, since the $x$'s are at fixed measured values, this whole thing can be
cast into a single matrix equation. If we define the following $M$ matrix:

\begin{equation}
M_{ij} = X_j(x_i)
\end{equation}

and if we think of $y^\star$ as an $m$-vector of model
values evaluated at the $m$ data points $x_i$ and $y$ as an
experimental $m$-vector $y$, then the model vector
can be evaluated by matrix multiplication as:

\begin{equation}
y^\star = M \alpha
\end{equation}

and the difference vector between the model and the data becomes simply

\begin{equation}
y - y^\star = y - M\alpha
\end{equation}

Finally we can write $\chi^2$ as

\begin{equation}
\chi^2 = (y - M\alpha)^T S (y - M\alpha) \,\,\,\,\,\,\,\,(1)
\end{equation}

where $S$ is a diagonal matrix of $1/\sigma_i^2$ and
the $T$ on the left side means "transpose".

Now it is a simple matter to find the values of the paramters that
minimize $\chi^2$. We require that $\nabla_{\alpha}\chi^2 = 0$.
If you play around with this you can convince yourself that this means:

\begin{equation}
M^T S(y - M\alpha) = 0
\end{equation}

and therefore:

\begin{equation}
\alpha = (M^T SM)^{-1}M^T S y.  \,\,\,\,\,\,\,\,(2)
\end{equation}

```python
#
# define the function values (these are the columns of the "M" matrix
#
# generate arrays that contain "functions" of x that are combined to form the model
#

def doFit(funcs, xvals, yvals, sigma):
    """
    doFit uses matrices to solve for the linear least squares
    fit parameters based on the list of 'funcs', and the x-y values
    and their corresponding uncertainties (sigma).

    returns the array of model parameters
    """
    #
    # define the S matrix in terms of uncertainties
    #

    S=np.diag((1.0/sigma)**2)

    #
    # define the model matrix
    #

    M = np.array(funcs).T
    MtM = M.T.dot(S.dot(M))
    MtMInv = np.linalg.inv(MtM)
    MtY = M.T.dot(S.dot(yvals))
    alpha = MtMInv.dot(MtY)      # see eq 2 above.

    #
    # solve for the parameters, and return
    #
    ystar = M.dot(alpha)

    return (alpha, ystar, MtMInv)

fx1 = np.ones(len(xvals))   # X_1(x) = 1
fx2 = xvals              # X_2(x) = x

sigma=1.0/np.sqrt(np.abs(counts))

alpha, ystar, fcov = doFit([fx1, fx2], xvals, yvals, sigma)

mFit = alpha[1]
bFit = alpha[0]

pl.errorbar(xvals, yvals, sigma, fmt='r.')
pl.title("File Sizes on My HD")
pl.ylabel("log(count)")
pl.xlabel("log(size)")
pl.plot(xvals,ystar, 'g-')

```

```python
#
# Now do monte-carlo data fabrication and fit analysis
#

N=1000    # number of samples to take
mList=[]  # keep track of monte-carlo fit parameters
bList=[]

for i in range(N):
    """
    Generate mc data with the same statistical properties as the real data.
    Repeat the fit for each set, and record the parameters.
    """
    mcY = ystar + sigma*np.random.normal(size=len(xvals))
    mcAlpha,mcYstar,mccov = doFit([fx1,fx2],xvals,mcY,sigma)   # repeatedly fit mc data
    mList.append(mcAlpha[1])
    bList.append(mcAlpha[0])

#
# Compute the statistics of the mc-results
#

marr=np.array(mList)
barr=np.array(bList)

mAvg = marr.sum()/N
bAvg = barr.sum()/N
delM = marr-mAvg
delB = barr-bAvg
sigM = np.sqrt((delM*delM).sum()/(N-1))  # sigM is the std-deviation of the m values
sigB = np.sqrt((delB*delB).sum()/(N-1))  # sigB is the std-deviation of the b values

#
# plot the fit
#
print("Slope=", mFit, '+/-', sigM, "(", mFit - sigM,",",mFit + sigM, ")")
print("Intercept=", bFit, '+/-', sigB, "(", bFit - sigB,",",bFit + sigB, ")")

#
# Just for fun print out the diagonals of the covariance matrix.
#
print # print a blank link
print("Compare to cov-matrix for fun:")
print("sqsrt(cov[1,1]) (should be sigma m)", np.sqrt(fcov[1,1]))
print("sqrt(cov[0,0]) (should be sigma b)", np.sqrt(fcov[0,0]))

```

    Slope= -0.40101297210387443 +/- 0.0069164777066884975 ( -0.40792944981056295 , -0.3940964943971859 )
    Intercept= 10.251110164592234 +/- 0.07947579449122216 ( 10.171634370101012 , 10.330585959083455 )
    Compare to cov-matrix for fun:
    sqsrt(cov[1,1]) (should be sigma m) 0.006919937608028965
    sqrt(cov[0,0]) (should be sigma b) 0.08007908396499791

```python
pl.hist(marr) # look at the variation in "m" values
print("m average:", mAvg)
print("m sigma:",sigM)
```

```python
pl.hist(barr)  # look at the variation in "b" values
print("b average:", bAvg)
print("b sigma:",sigB)
```

```python

```

P07-Parameter Estimation (Zipf's Law).ipynb


```python
%matplotlib inline
%config InlineBackend.figure_formats = ['svg']
import numpy as np
import matplotlib.pyplot as plt
```

#  Pulses and Waveforms

A waveform is simply a time varying signal. When a pulse operation occurs, a digital to analog converter produces a signal by combinding the following data:

* a time varying complex function $ u(t) $ which describes the baseband waveform
* a *frequency* $f$, in Hertz,
* a *phase* $\theta$, in radians,
* a unitless *scale* $\alpha$.

The signal generated from this has the mathematical form
\begin{equation}
x(t) = \text{Re}[\alpha u(t)e^{i (\theta + 2 \pi f t)}],
\end{equation}
where $\text{Re}$ denotes the real part of a complex number.

In accordance with the usual conventions, we will refer to the real part of $u(t)$ as the **in-phase** component, and the imaginary part as the **quadrature** component. A general complex number $z = x + iy$ will sometimes be referred to as an **IQ-value**, with $I = x$ and $Q = y$.

Quil-T provides the programmer with pulse and capture-level control, both by allowing for custom baseband waveforms, as well as allowing for direct control over the run-time frequency, phase, and scale.

## Waveform Templates and References

There are two ways to specify a Quil-T waveform:
- by using a pre-existing template, which has a shape dictated by certain parameters
- by referencing a custom waveform definition.

Consider the following Quil-T code:
```
PULSE 0 "rf" flat(duration: 1e-6, iq: 0.5 + 0.5*i)
```
This denotes a pulse operation, on frame `0 "rf"`, with a total duration of one microsecond and with baseband waveform given by the `flat` template, (corresponding to $u(t) = 1/2 + i/2$ for the duration of the signal). The resulting signal produced depends also on the phase, frequency, and scaling factor associated with frame `0 "rf"`. These may be set explicitly, as in
```
SET-SCALE 0 "rf" 1.0
SET-FREQUENCY 0 "rf" 5e9
SET-PHASE 0 "rf" 0.0
```

On the other hand, a custom waveform may be defined and used. This is done using the `DEFWAVEFORM` form,
```
DEFWAVEFORM my_waveform:
    0.01, 0.01+0.01*i, ...
```
the body of which consists of a list of complex numbers. Such a custom waveform may be referenced directly by name, as in
```
PULSE 0 "rf" my_waveform.
```

The precise meaning of this depends on the *sample rate* of the associated frame: this indicates the number of samples per second consumed by the underlying digital-to-analog converter. For example, suppose that the definition of `0 "rf"` looked like this
```
DEFFRAME 0 "rf":
    SAMPLE-RATE: 1000000000.0
    INITIAL-FREQUENCY: 4807541957.13474
    DIRECTION: "tx"
```
and `my_waveform` has a definition consisting of complex numbers $z_1, \ldots, z_N$. Then, letting $r=10^9$ denote the sample rate, the resulting pulse has total duration $\left \lceil{\frac{N}{r}}\right \rceil$, corresponding to the baseband waveform $ u(t) = z_{\left \lfloor{tr}\right \rfloor}. $ Here $\left \lceil{x}\right \rceil $ and $\left \lfloor{x}\right \rfloor$ denote the ceiling and floor of $x$, respectively.

As before, this baseband waveform is combined with the frame's scale, frequency, and phase during digital to analog conversion.

## A Catalog of Template Waveforms

Rigetti provides a number of templates by default. These include

- `flat`, corresponding to simple rectangular waveforms
- `gaussian`, for a Gaussian waveform
- `drag_gaussian`, for a Gaussian waveform modified by the Derivative Removal by Adiabatic Gate (DRAG) technique
- `hrm_gaussian`, for a DRAG Gaussian waveform with second-order corrections
- `erf_square`, for a flat waveform with smooth edges derived from the Gaussian error function
- `boxcar_kernel`, for a flat waveform which is normalized to integrate to 1

Each of these waveforms has a corresponding definition in `pyquil.quiltwaveforms`. In addition to providing documentation on the meaning of each of the waveform parameters, this module also contains routines for generating samples for each template.

Below we look at each individidually, discussing the meaning of various parameters and plotting the real part of the waveform envelope.

```python
from pyquil.quilatom import TemplateWaveform

def plot_waveform(wf: TemplateWaveform, sample_rate: float):
    """ Plot a template waveform by sampling at the specified sample rate. """
    samples = wf.samples(sample_rate)
    times = np.arange(len(samples))/sample_rate

    print(wf)
    plt.plot(times, samples.real)
    plt.show()
```

### flat

A flat waveform is simple: it represents a constant signal for a certain duration. There are two required parameters:

* `duration`: the length of the waveform, in seconds
* `iq`: a complex number

```python
from pyquil.quiltwaveforms import FlatWaveform

plot_waveform(FlatWaveform(duration=1e-6, iq=1.0), sample_rate=1e9)
```

    flat(duration: 1e-06, iq: 1.0)

#### Scale, Phase, Detuning

In addition to the parameters specific to each template waveform, there are also a few generic parameters. One of these we have met: each template waveform has a required `duration` argument, indicating the length of the waveform in seconds.

The other arguments are *optional*, and are used to modulate the basic shape. These are

* scale $\alpha$, which has the effect of scaling the baseband $u(t) \mapsto \alpha u(t)$
* phase $\theta$, in radians, which has the effect of a phase shift on the baseband $u(t) \mapsto e^{i \theta} u(t)$
* detuning $f_d$, in Hertz, which has the effect of $u(t) \mapsto e^{2 \pi i f_d} u(t)$

These may be provided as arguments to any template waveform, e.g.
```
PULSE 0 "rf" flat(duration: 1e-8, iq: 1.0, scale: 0.3, phase: 1.570796, detuning: 1e8)
```

Below we consider this by way of the PyQuil bindings.

```python
plot_waveform(FlatWaveform(duration=1e-6, iq=1.0, detuning=1e7), sample_rate=1e9)
```

    flat(duration: 1e-06, iq: 1.0, detuning: 10000000.0)

### gaussian

Several of the template waveforms provided are derived from a standard (unnormalized) Gaussian. Here we have
\begin{equation}
u(t) = \text{exp}\Big(-\frac{(t-t_0)^2}{2 \sigma^2}\Big),
\end{equation}
where $t_0$ denotes the center of the Gaussian, and $\sigma$ is the usual standard deviation.

By Quil-T convention, this is parameterized by the Gaussian's [full width at half maximum](https://en.wikipedia.org/wiki/Full_width_at_half_maximum) (FWHM), which is defined to be

\begin{equation}
\text{FWHM} = 2 \sqrt{2 \ln(2)} \sigma.
\end{equation}

As with all Quil-T waveforms, a Quil-T `gaussian` has a finite duration, and thus corresponds to a truncation of a true Gaussian.

In short, the parameters are:

* `duration`: the duration of the waveform, in seconds. The Gaussian will be truncated to $[0, \text{duration}]$
* `t0`: the center of the Gaussian, in seconds
* `fwhm`: the full width half maximum of the Gaussian, in seconds

```python
from pyquil.quiltwaveforms import GaussianWaveform

plot_waveform(GaussianWaveform(duration=1e-6, t0=5e-7, fwhm=4e-7), sample_rate=1e9)
```

    gaussian(duration: 1e-06, fwhm: 4e-07, t0: 5e-07)

### drag_gaussian

The `drag_gaussian` waveform extends the basic Gaussian with an additional correction factor (cf. https://arxiv.org/abs/1809.04919 and references therein). The shape is given by
\begin{equation}
u(t) = \Big(1 + i\frac{\alpha}{2 \pi \eta \sigma^2}\Big)\text{exp}\Big(-\frac{(t-t_0)^2}{2 \sigma^2}\Big),
\end{equation}
where $\eta$ is the anharmonicity constant, in Hz, and $\alpha$ is a dimensionless shape parameter.

As before, rather than providing $\sigma$ explicity, `drag_gaussian` takes a full-width half max parameter.

In summary, the required arguments to `drag_gaussian` are:

* `duration`: the duration of the waveform, in seconds. The Gaussian will be truncated to $[0, \text{duration}]$
* `t0`: the center of the Gaussian, in seconds
* `fwhm`: the full width half maximum of the Gaussian, in seconds
* `anh`: the anharmonicity constant $\eta$, in Hertz
* `alpha`: dimensionless shape parameter $\alpha$.

```python
from pyquil.quiltwaveforms import DragGaussianWaveform

plot_waveform(
    DragGaussianWaveform(duration=1e-6, t0=5e-7, fwhm=4e-7, anh=1.1, alpha=1.0),
    sample_rate=1e9
)
```

    drag_gaussian(duration: 1e-06, fwhm: 4e-07, t0: 5e-07, anh: 1.1, alpha: 1.0)

Of course, we only plotted the *real part* of the waveform above. The imaginary part is relevant when the baseband waveform is converted to a passband waveform, as we can demonstrate below via the optional `detuning` argument.

```python
from pyquil.quiltwaveforms import DragGaussianWaveform

plot_waveform(
    DragGaussianWaveform(duration=1e-6, t0=5e-7, fwhm=4e-7, anh=1.1, alpha=1.0, detuning=1e7),
    sample_rate=1e9
)
```

    drag_gaussian(duration: 1e-06, fwhm: 4e-07, t0: 5e-07, anh: 1.1, alpha: 1.0, detuning: 10000000.0)

### hrm_gaussian

The `hrm_gaussian` waveform is a variant of `drag_gaussian` which incorporates a higher order term. The shape is given by
\begin{equation}
u(t) = \Big(1 - H_2 \frac{(t-t_0)^2}{2 \sigma^2} + i\frac{\alpha (t-t_0)}{2 \pi \eta \sigma^2} \big(1 - H_2 \big(\frac{(t-t_0)^2}{2\sigma^2} - 1\big)\Big)\text{exp}\Big(-\frac{(t-t_0)^2}{2 \sigma^2}\Big),
\end{equation}
where $\alpha$ is a dimensionless DRAG parameter, $\eta$ is the anharmonicity constant, and $H_2$ is a second order correction coefficient (cf. [1]). Note that when $H_2$ equals 0 this reduces to an ordinary `drag_gaussian`.

The required arguments to `hrm_gaussian` are:

* `duration`: the duration of the waveform, in seconds. The Gaussian will be truncated to $[0, \text{duration}]$
* `t0`: the center of the Gaussian, in seconds
* `fwhm`: the full width half maximum of the Gaussian, in seconds
* `anh`: the anharmonicity constant $\eta$, in Hertz
* `alpha`: dimensionless shape parameter $\alpha$
* `second_order_hrm_coeff`: the constant $H_2$.

[1] Warren, W. S. (1984). Effects of arbitrary laser or NMR pulse shapes on population inversion and coherence. The Journal of Chemical Physics, 81(12), 5437–5448. doi:10.1063/1.447644

```python
from pyquil.quiltwaveforms import HrmGaussianWaveform

plot_waveform(
    HrmGaussianWaveform(duration=1e-6, t0=5e-7, fwhm=4e-7, anh=1.1, alpha=1.0, second_order_hrm_coeff=0.5),
    sample_rate=1e9
)
```

    hrm_gaussian(duration: 1e-06, fwhm: 4e-07, t0: 5e-07, anh: 1.1, alpha: 1.0, second_order_hrm_coeff: 0.5)

### erf_square

The `erf_square` waveform is a variant of `flat_waveform` with the boundary discontinuities smoothed via the [error function](https://en.wikipedia.org/wiki/Error_function) (erf), and additional zero-padding.

The required arguments are:

* `duration`: the duration of the nonzero part of the waveform, in seconds
* `risetime`: width of each of the rise and fall sections of the pulse, in seconds
* `pad_left`: amount of zero-padding to add to the left of the pulse, in seconds
* `pad_right`: amount of zero-padding to add to the right of the pulse, in seconds

**NOTE**: The total duration of the waveform is `duration + pad_left + pad_right`; the total duration of the *support* (nonzero entries) is `duration`.

```python
from pyquil.quiltwaveforms import ErfSquareWaveform

plot_waveform(
    ErfSquareWaveform(duration=1e-6, risetime=1e-7, pad_left=1e-7, pad_right=1e-7),
    sample_rate=1e9
)
```

    erf_square(duration: 1e-06, risetime: 1e-07, pad_left: 1e-07, pad_right: 1e-07)

## Capture and Kernels

In Quil-T, waveforms are used in two places:

* in a `PULSE` operation, to specify what signal is generated
* in a `CAPTURE` operation, to specify how to resolve a signal into a number

The `CAPTURE` operation involves reading in a signal on a suitable signal line, and integrating it with respect to a *kernel*.

Mathematically, there is a real-valued signal $ s(t), $ corresponding to a reading on the signal line. This is combined with a complex-valued "baseband" kernel $k(t)$ to get a resulting **IQ value** $z$ by
\begin{equation}
z = \int_{t_\text{min}}^{t_\text{max}} k(t) s(t) e^{-2 \pi i f t} \, dt,
\end{equation}
where $f$ denotes the frequency of the associated capture frame.

In Quil-T, this integrating kernel is specified by a waveform, with the usual convention that is is scaled to satisfy $\int k(t) \, dt = 1.$ The most common example is the `boxcar_kernel`, which corresponds to a flat pulse scaled to satisfy this condition.

### boxcar_kernel

Because of the normalization condition, the `boxcar_kernel` requires only a `duration` argument for its construction.

```python
from pyquil.quiltwaveforms import BoxcarAveragerKernel

plot_waveform(
    BoxcarAveragerKernel(duration=1e-6),
    sample_rate=1e9
)
```

    boxcar_kernel(duration: 1e-06)

The samples should sum to (roughly) one.

```python
assert np.isclose(
    np.sum(BoxcarAveragerKernel(duration=1e-6).samples(1e9)),
    1.0
)
```

**Note**: The reference implementations of these waveforms makes use of Python's double precision floating point arithmetic. The *actual* implementation involves a certain amount of hardware dependence. For example, Rigetti's waveform generation hardware currently makes use of 16 bit fixed point arithmetic.

## Compile Time versus Run Time

When thinking about parameters in Quil-T, there are two times to consider:
* the time of compilation, when a Quil-T program is translated to a binary format executable on Rigetti hardware
* the time at which the program is run.

**All template parameters must be resolved at compile time.** The Quil-T compiler depends on being able to determine the size and contents of waveforms in advane. In other words, the following is not allowed:
```
DECLARE theta REAL
PULSE 0 "rf" flat(duration: 1e-8, iq: 1.0, phase: theta)
```
However, the following program is valid, and does something equivalent
```
DECLARE theta REAL
SHIFT-PHASE 0 "rf" theta
PULSE 0 "rf" flat(duration: 1e-8, iq: 1.0)
SHIFT-PHASE 0 "rf" -theta
```
(Why the second `SHIFT-PHASE`? To leave the frame `0 "rf"` as it was before we started!)

Explicit control over the *run-time* phase, frequency, or scale requires the use of one of the following instructions.

* `SET-SCALE`
* `SET-PHASE`
* `SHIFT-PHASE`
* `SET-FREQUENCY`
* `SHIFT-FREQUENCY`

all of which support run-time parameter arguments.

```python

```


# ATW relaxation notebook

```python
import numpy as np
import matplotlib.pyplot as plt
import warnings
from copy import deepcopy

# Global constants
f = 1e-4      # [s-1]
g = 9.81      # [m s-2]

%matplotlib inline
plt.rcParams['font.size'] = 14
warnings.simplefilter('ignore')
```

### Analytical solution

Start with the linearized, steady state shallow water equations with linear friction and longshore windstress. Assume cross-shore geostrophic balance.

\begin{align}
f\mathbf{k}\times\mathbf{u} & = -g\nabla\eta + \frac{1}{h}\left(\tau_y - \mu v\right)\hat{\jmath} \tag{1a} \\
0 & = \nabla\cdot h\mathbf{u} \tag{1b}
\end{align}

Taking the curl of (1a) and solving for $\eta$ gives the the Arrested Topography Wave (ATW) of Csanady (1978 *JPO*). I have oriented the problem to $x\to-\infty$ offshore such that $\frac{\partial h}{\partial x} = -s$.

$$\frac{\partial^2\eta}{\partial x^2} - \frac{1}{\kappa}\frac{\partial\eta}{\partial y} = 0, \hspace{0.5cm} \frac{1}{\kappa} = \frac{fs}{\mu}\tag{2}$$

The coastal boundary condition (obtained from 1a) requires $u \to 0$ and $h \to 0$

$$\frac{\partial\eta}{\partial x}(0, y) = \frac{\tau_yf}{\mu g} = q_0 \tag{3}$$

Equation (2) is analogous to a constant heat flux boundary condition. The solution is given by Carslaw and Jaeger 1959 (p. 112)

$$\eta(x, y) = \frac{\kappa q_0y}{L} + q_0L\left\{\frac{3(x + L)^2 - L^2}{6L^2} - \frac{2}{\pi^2}\sum_{n=1}^\infty\frac{(-1)^n}{n^2}\exp\left(\frac{-\kappa n^2\pi^2y}{L^2}\right)\cos\left(\frac{n\pi(x + L)}{L}\right)\right\} \tag{4}$$

which, as $y\to\infty$, reduces to

$$\eta(x, y) = \frac{\kappa q_0y}{L} + q_0L\frac{3(x + L)^2 - L^2}{6L^2} \tag{5}$$

Calculate $\eta$ according to equation (5)

```python
def calc_eta(x, y, L, kappa, q_0):
    """Calculate eta according to equation 5
    """

    return kappa * q_0 * y / L + q_0 * L * (3 * (x + L)**2 - L**2) / (6 * L**2)
```

Find $\eta$ given problem parameters

```python
# Constants
L = 1e3        # Slope width [m]
tau_y = -1e-4  # Kinematic wind stress [m2 s-2]
mu = 1e-2      # Linear friction coefficient [s-1]
s = 1          # Shelf slope [dimensionless]

# Terms (heat equation analogues)
kappa = mu / (f * s)        # 'Diffusivity' of eta
q_0 = tau_y * f / (mu * g)  # 'Flux' of eta through boundary
print(q_0)
print(kappa)

# Coordinates
dL = L * 1e-2
xi = np.arange(-L, 0, dL)  # x-direction, second coordinate
yi = np.arange(0, L, dL)   # y-direction, first coordinate
x, y = np.meshgrid(xi, yi)

# Solution
eta = calc_eta(x, y, L, kappa, q_0)
```

    -1.0193679918450561e-07
    100.0

Plot $\eta$ solution

```python
# Plot eta
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
ax.contour(xi/L, yi/L, eta, colors='k')
for tick in np.arange(0, 1, 0.015):
    ax.plot([0, 0.005], [tick, tick+0.005], 'k-', clip_on=False)
ax.set_xlabel('$\longleftarrow$ $X/L$')
ax.set_ylabel('$\longleftarrow$ $Y/L$')
ax.xaxis.set_ticks([-1, 0])
ax.yaxis.set_ticks([0, 1])
ax.xaxis.set_ticklabels(['$-L$', 0])
ax.yaxis.set_ticklabels([0, '$L$'])
ax.tick_params(direction='out', pad=8)
ax.set_xlim([0, -1])
ax.set_ylim([1, 0])
ax.text(0.02, 0.05, 'Low $\eta$', transform=ax.transAxes)
ax.text(0.85, 0.9, 'High $\eta$', transform=ax.transAxes)
ax.text(0.03, 0.46, '$\\tau_y$', transform=ax.transAxes)
ax.arrow(0.04, 0.5, 0, 0.1, transform=ax.transAxes, head_width=0.01, facecolor='k')
ax.set_title('Cross-shelf bottom slope (ATW) solution')

plt.show()
```

### Relaxation solution

Three schemes:

Centered difference

$$r_{i, j}^{(n)} = \frac{\eta_{i, j+1}^{(n)} - \eta_{i, j-1}^{(n)}}{2\Delta y} - \kappa\frac{\eta_{i+1, j}^{(n)} - 2\eta_{i, j}^{(n)} + \eta_{i-1, j}^{(n)}}{\Delta x^2}$$

$$\eta_{i, j}^{(n+1)} = \eta_{i, j}^{(n)} - \frac{\mu\Delta x^2}{2\kappa}r_{i, j}^{(n)}$$

Upstream Euler

$$r_{i, j}^{(n)} = \frac{\eta_{i, j+1}^{(n)} - \eta_{i, j}^{(n)}}{\Delta y} - \kappa\frac{\eta_{i+1, j}^{(n)} - 2\eta_{i, j}^{(n)} + \eta_{i-1, j}^{(n)}}{\Delta x^2}$$

$$\eta_{i, j}^{(n+1)} = \eta_{i, j}^{(n)} - \frac{\mu}{\left(\frac{2\kappa}{\Delta x} - 1\right)}r_{i, j}^{(n)}$$

Downstream Euler

$$r_{i, j}^{(n)} = \frac{\eta_{i, j}^{(n)} - \eta_{i, j-1}^{(n)}}{\Delta y} - \kappa\frac{\eta_{i+1, j}^{(n)} - 2\eta_{i, j}^{(n)} + \eta_{i-1, j}^{(n)}}{\Delta x^2}$$

$$\eta_{i, j}^{(n+1)} = \eta_{i, j}^{(n)} - \frac{\mu}{\left(\frac{2\kappa}{\Delta x} + 1\right)}r_{i, j}^{(n)}$$

Only the downstream Euler is stable.

Find $\eta$ by relaxation

```python
# Find phi by relaxation

# Parameters
M = eta.shape[0] # matrix size
mu = 1           # SOR convergence parameter
TOL = 1e-4       # Convergence tolerance
dissipation = 0.

# Allocate arrays
# Solution
eta_soln = calc_eta(x, y, L, kappa, q_0)
eta_next = np.copy(eta_soln)            # getting the boundary conditions correct here
eta = np.zeros(eta.shape)               # start from zero to show this is working
res = np.zeros(eta.shape)

# Make figure array
fig, axs = plt.subplots(1, 3, figsize=(17, 6))

N=100
# Relaxation loop
for n in range(N):
    for i in range(1, M-1): # Longshore step
        for j in range(2, M-1): # Cross-shore step : start from 2 to preserve gradient boundary condition
            #Downstream Euler : note the switch in i and j, i is dy, j is dx
            res[i, j] = (eta[i, j] - eta[i-1, j]) / (dL) - kappa * (eta[i, j+1] - 2 * eta[i, j] + eta[i, j-1]) / dL**2
            eta_next[i, j] = eta[i, j] - mu / (2 * kappa / dL + 1) * res[i, j]
            eta = eta_next  # move this into the loop for faster convergence
    if dL**2 * np.max(abs(res)) / np.max(abs(eta)) < TOL and n > 5:           # n > 5 just because I'm starting from 0.
        print('done', n)
        break

# Plot results
mesh = axs[0].pcolormesh(xi/L, yi/L, eta_soln)
fig.colorbar(mesh, ax=axs[0])
mesh = axs[1].contour(xi/L, yi/L, res, colors='b')
fig.colorbar(mesh, ax=axs[1])
mesh=axs[2].pcolormesh(xi/L, yi/L, eta_next)
fig.colorbar(mesh, ax=axs[2])
for ax in axs:
    ax.set_xlim([0, -1])
    ax.set_ylim([1, 0])

```

```python

```


```python
import numpy as np
import matplotlib.pyplot as plt
from numpy import sin, cos, exp

from scipy.integrate import solve_ivp

from typing import Callable

import sim_utils

# Fix Python 3's weird rounding function
# https://stackoverflow.com/a/44888699/538379
round2=lambda x,y=None: round(x+1e-15,y)
```

```python
def reference_bank_angle(t, X, params):
    """Generates the reference bank angle profile"""
    v = X[2]
    if v >= 3500:
        return np.deg2rad(75);
    elif v <= 1500:
        return np.deg2rad(50);
    else:
        return np.deg2rad(50 + (75-50)*(v-1500)/(3500-1500))

V = np.linspace(1000, 4000, 301)

bank_angle_deg = np.degrees([reference_bank_angle(0, [0, 0, v, 0], {}) for v in V], dtype=np.float64)

plt.plot(V, bank_angle_deg)
plt.xlabel('V [m/s]')
plt.ylabel('Bank Angle [deg]')
plt.grid(True)
plt.title('Reference Bank Angle')
plt.show()
```

```python
def traj_eom(t: float,
             state: np.array,
             params: dict,
             bank_angle_fn: Callable[[float, np.array, dict], float]
            ):
    h, s, v, gam = state
    u = bank_angle_fn(t, state, params)

    rho0 = params['rho0']
    H = params['H']
    beta = params['beta']   # m/(Cd * Aref)
    LD = params['LD']
    R_m = params['R_m']
    g = params['g']

    v2 = v*v
    rho = rho0 * exp(-h/H)
    D_m = rho * v2 / (2 * beta)  # Drag Acceleration (D/m)
    r = R_m + h
    return np.array([v * sin(gam),       # dh/dt
                     v * cos(gam),       # ds/dt
                     -D_m - g*sin(gam),  # dV/dt
                     (v2 * cos(gam)/r + D_m*LD*cos(u) - g*cos(gam))/v] # dgam/dt
                   )
```

```python
# Initial conditions
h0 = 120e3; # Entry altitude
V0 = 5500;  # Entry velocity
gamma0_deg = -14.5; # Entry flight path angle
s0 = 0

# Model params
params = {'H': 11.1e3,
              'rho0': 0.020, # kg/m^3
              'beta': 120,
              'LD': 0.24,
              'R_m': 3380e3,
              'g': 3.73}

# Terminal velocity
v_f = 600

gamma0 = np.deg2rad(gamma0_deg)
X0 = np.array([h0, s0, V0, gamma0])
t0 = 0
tf = 500.
tspan = np.linspace(t0, tf, 101)
```

```python
ref_traj = sim_utils.simulate_entry_trajectory(traj_eom, t0, tf, X0, 2, v_f, params, reference_bank_angle, tspan)

plt.plot(ref_traj.X[:,2]/1e3, ref_traj.X[:,0]/1e3)
plt.xlabel('V [km/s]')
plt.ylabel('h [km]')
plt.grid(True)
plt.show()
```

# Computing Co-States and Gains

Once the reference trajectory has been propagated, additional information is required for formulating the guidance law. These are the costates $\lambda_h$, $\lambda_s$, $\lambda_v$, $\lambda_\gamma$ and $\lambda_u$. Of these,  $\lambda_s$ was found to have a constant value of $1$ throughout the trajectory.
The remaining costates have these terminal boundary conditions associated with them:

$$
\begin{align}
\lambda_h(t_f) &= -\cot{\gamma(t_f)} \\
\lambda_v(t_f) &= 0 \\
\lambda_\gamma(t_f) &= 0 \\
\lambda_u(t_f) &= 0 \\
\end{align}
$$

Starting at these terminal boundary conditions, we can integrate the equations in reverse to obtain the time history of these co-states.

```python
def traj_eom_with_costates(t: float,
             state: np.array,
             params: dict,
             bank_angle_fn: Callable[[float, np.array, dict], float]
            ):
    lamS = 1
    h, s, V, gam, lamH, lamV, lamGAM, lamU = state

    u = bank_angle_fn(t, state, params)

    rho0 = params['rho0']
    H = params['H']
    beta = params['beta']
    LD = params['LD']
    R_m = params['R_m']
    g = params['g']

    r = R_m + h

    v = V
    V2 = V*V
    rho = rho0 * exp(-h/H)
    D_m = rho * V2 / (2 * beta)  # Drag Acceleration (D/m)

#     lamHDot = D_m*LD*lamGAM*cos(u)/(H*v) - D_m*lamV/H + lamGAM*v*cos(gam)/r**2
#     lamVDot = D_m*LD*lamGAM*cos(u)/v**2 - LD*lamGAM*rho*cos(u)/beta - g*lamGAM*cos(gam)/v**2 - lamGAM*cos(gam)/r \
#               - lamH*sin(gam) \
#               - lamS*cos(gam) \
#               + lamV*rho*v/beta
#     lamGAMDot = g*lamV*cos(gam) - lamGAM*(g*sin(gam) - v**2*sin(gam)/r)/v - lamH*v*cos(gam) + lamS*v*sin(gam)

    lamHdot = D_m*LD*lamGAM*cos(u)/(H*v) - D_m*lamV/H + lamGAM*v*cos(gam)/r**2
    lamVdot = D_m*LD*lamGAM*cos(u)/v**2 - LD*lamGAM*rho*cos(u)/beta - g*lamGAM*cos(gam)/v**2 - lamGAM*cos(gam)/r - lamH*sin(gam) - lamS*cos(gam) + lamV*rho*v/beta
    lamGAMdot = -g*lamGAM*sin(gam)/v + g*lamV*cos(gam) + lamGAM*v*sin(gam)/r - lamH*v*cos(gam) + lamS*v*sin(gam)

#     lamUdot = -LD*lamGAM*rho*v*sin(u)/(2*beta)
    lamUdot = LD*lamGAM*rho*v*sin(u)/(2*beta)
    return np.array([V * sin(gam),       # dh/dt
                     V * cos(gam),       # ds/dt
                     -D_m - g*sin(gam),  # dV/dt
                     (V2 * cos(gam)/r + D_m*LD*cos(u) - g*cos(gam))/V, # dgam/dt

                     lamHdot,
                     lamVdot,
                     lamGAMdot,
                     lamUdot]
                   )
```

```python
ref_tf = ref_traj.t[-1]
ref_tspan_rev = ref_traj.t[::-1] # Reverse the time span
Xf = np.copy(ref_traj.X[-1,:])

# Ensure monotonic decreasing V
def V_event(t,X,p,_):
    return X[3] - 5500

V_event.direction = 1
V_event.terminal = True

X_and_lam0 = np.concatenate((Xf, [-1/np.tan(Xf[3]), 0, 0, 0]))
output = solve_ivp(traj_eom_with_costates, # lambda t,X,p,u: -traj_eom_with_costates(t,X,p,u),
                   [ref_tf, 0],
                   X_and_lam0,
                   t_eval=ref_traj.t[::-1],
                   rtol=1e-8,
                   events=V_event,
                   args=(params, reference_bank_angle))
lam = output.y.T[:,4:][::-1]
X_and_lam = output.y.T[::-1]

np.set_printoptions(suppress=True)
```

```python
class ApolloReferenceData:
    def __init__(self, X_and_lam: np.array, u: np.array, tspan: np.array, params: dict):
        """
        X_and_lam: [h, s, v, gam, lamH, lamV, lamGAM, lamU] - 8 x n matrix
        tspan: 1 x n vector
        """
        self.X_and_lam = X_and_lam
        self.tspan = tspan
        self.params = params
        self.u = u

        assert len(X_and_lam.shape) == 2 and X_and_lam.shape[0] > 1, "Need at least two rows of data"
        self.num_rows = X_and_lam.shape[0]

        self.delta_v = abs(X_and_lam[1,2] - X_and_lam[0,2])
        assert self.delta_v > 0, "Reference trajectory has repeated velocites in different rows"

        self.start_v = X_and_lam[0,2]

        F1, F2, F3, D_m, hdot_ref = self._compute_gains_and_ref()
        F3[-1] = F3[-2]   # Account for F3=0 at t=tf
        # Stack the columns as follows:
        # [t, h, s, v, gam, F1, F2, F3, D/m]
        self.data = np.column_stack((tspan, X_and_lam[:,:4], F1, F2, F3, D_m, hdot_ref))

    def _compute_gains_and_ref(self):
        h = self.X_and_lam[:,0]
        v = self.X_and_lam[:,2]
        gam = self.X_and_lam[:,3]

        lamH = self.X_and_lam[:,4]
        lamGAM = self.X_and_lam[:,6]
        lamU = self.X_and_lam[:,7]

        rho0 = self.params['rho0']
        H = self.params['H']
        beta = self.params['beta']   # m/(Cd * Aref)

        v2 = v*v
        rho = rho0 * exp(-h/H)
        D_m = rho * v2 / (2 * beta)  # Drag Acceleration (D/m)
        hdot = v * sin(gam)

        F1 = H * lamH/D_m
        F2 = lamGAM/(v * np.cos(gam))
        F3 = lamU
        return F1, F2, F3, D_m, hdot

    def get_row_by_velocity(self, v: float):
        """
        Returns data row closest to given velocity
        """
        all_v = self.data[:,3]
        dist_to_v = np.abs(all_v - v)
        index = min(dist_to_v) == dist_to_v
        return self.data[index,:][0]

    def save(self, filename: str):
        """Saves the reference trajectory data to a file"""
        np.savez(filename, X_and_lam=self.X_and_lam, u=self.u, tspan=self.tspan, params=self.params)

    @staticmethod
    def load(filename: str):
        """Initializes a new ApolloReferenceData from a saved data file"""
        npzdata = np.load(filename, allow_pickle=True)
        X_and_lam = npzdata.get('X_and_lam')
        u = npzdata.get('u')
        tspan = npzdata.get('tspan')
        params = npzdata.get('params').item()
        return ApolloReferenceData(X_and_lam, u, tspan, params)
```

```python
# Test loading and saving of data
apollo_ref = ApolloReferenceData(X_and_lam, ref_traj.u, ref_traj.t, params)
apollo_ref.save('apollo_data_vref.npz')

# Load data back and check that it matches the original
ref = ApolloReferenceData.load('apollo_data_vref.npz')
assert np.allclose(ref.data, apollo_ref.data)
```

```python

```