```python
%run ./_setup.ipynb
```

```python
# function for nonlinear least squares fitting
def fixed_line(t, slope):
    return slope*t
```

# 3. Calculation of Inflow

To calculate the freshwater inflow to Galveston Bay, we use a control volume of the Bay and apply conservation of volume:

$\begin{align}
\partial V/\partial t &= Q_{fw} + Q_{ch} ~~~~~~~~(1) \\
\Rightarrow Q_{fw}(t) &= \partial V/\partial t - Q_{ch}(t)
\end{align}$

where $V(t)$ is the volume of Galveston Bay, $Q_{fw}(t)$ is the freshwater input to Galveston Bay, and $Q_{ch}(t)$ is the flux in Galveston Bay channel ($Q_{ch}$ is positive into the bay and negative on ebb tide out of the bay.

## 3.1 Bay Volume

Assuming that any change in x/y area of the bay due to flooding is limited (and therefore the surface area and bathymetry are constant), we have the following:

$\begin{align}
\partial V/\partial t &\approx A_s \partial \eta/\partial t, ~~~~~~~~(2)
\end{align}$

where $A_s$ is the surface area of Galveston Bay and $\eta$ is the sea level height.
We use the sea level at station 8771013 in the middle of the main bay to represent the sea level in both the main and east bays, and at stations 8771486 and 8771450 in the west bay to represent that area.

We have three possible surface area values to account for uncertainty. These values were calculated and saved in `0_process_bathymetry.ipynb`. Volume is calculated in the `calc_harvey` function below so that all parameter combinations can be used. Some definitions are set up here for later use.

Now calculate multiple options for the volume of the Bay, $V(t)$, and change in volume of the Bay, $dV/dt$. A rolling mean of 3 hours has been applied to $dV/dt$.

```python
# Sea surface stations used
eta_main = '8771013'  # main bay and East bay together
eta_west = '8771486'  # West bay
```

## 3.2 Channel Flux

### 3.2.1 ADCP data

We have data from an upward-facing ADCP from NOAA PORTS station `g06010` that spans much of the water column. We choose three ways to extrapolate the velocity data to be continuous from seabed from sea surface: piecewise linear, 2nd order polynomial, and 3rd order polynomial. This is done in the "Gather Data" notebook.

```python
# Calculate the standared deviation and mean across along-channel velocity signals to use later
std_uch = df[['smean','smean_2nd','smean_3rd']].std(axis=1)
mean_uch = df[['smean','smean_2nd','smean_3rd']].mean(axis=1)
```

### 3.2.2 Salinity conservation

Conservation of salinity in Galveston Bay gives:

$
\begin{align}
d (S V)/d t &= S_{in}Q_{fw} + S_{ch} Q_{ch} \\
d (S V)/d t &=  S_{ch} Q_{ch} \\
\Rightarrow Q_{ch} &= \frac{1}{S_{ch}} \frac{d (S V)}{d t}  ~~~~~~~~~~~~~~~~~~ (4)
\end{align}
$

where the input salinity, $S_{in}$, is assumed to be fresh so $S_{in}=0$ and $S_{ch}$ is the channel salinity.

We have data from TWDB station BOLI to represent $S_{ch}$. Rayson et al 2017 finds that the volume-averaged salinity in the bay is reasonably represented by three TWDB stations: $S(t) = \frac{1}{3} TRIN + \frac{1}{3} MIDG + \frac{1}{3} BOLI$; we use this result.

The derivative term is calculated as follows, where the volume is approximated as the surface area times the average depth plus a representative sea surface height:

$
\begin{align}
    \frac{\partial (S V)}{\partial t} &= A_s \frac{\partial [S(H+\eta)]} {\partial t}.~~~~~~~~~~~~~~~~~~ (5)
\end{align}
$

We calculate this term for both of the bays (main and East together, and West Bay separately).

```python
# Calculate volume averaged salinity
Sch = df['BOLI: Salinity']  # rename channel salinity
```

Solve for $Q_{ch}$ from conservation of salinity, Equation 4. Also solve for $Q_{fw}$, the freshwater flux, from Equation 1 using this $Q_{ch}$. It does not apply during the flooding when $S_{ch}=0$, but it does give results at other times.

The following is to find a range of values for $A_{eq}$ and metrics for when they should be valid.

```python
start = pd.Timestamp('2017-8-1')
stop = pd.Timestamp('2017-10-1')
```

```python
## finding equivalent area and then final result for different inputs, Section 3.2.3 ##
def calc_harvey(smean, midgweight, Amain, Awest, Hmain, Hwest, Vconstant, date):

    ## Calculate things ##
    # salinity
    S = ((1-0.33-midgweight)*df['TRIN: Salinity'] + midgweight*df['MIDG: Salinity'] + 0.33*Sch)  # volume averaged salinity

    if Vconstant:
        V = Amain * Hmain + Awest * Hwest
        dVdt = 0
    else:  # V variable
        V = (Amain * (Hmain + df['%s: Water Level [m, MHW]' % eta_main])) + (Awest * (Hwest + df['%s: Water Level [m, MHW]' % eta_west]))
        dVdt = (V.diff()/dt).rolling(center=True, window=4*3).mean()[start:stop]

    dSVdt = (S*V).diff()/dt
    Qch = (dSVdt/Sch)
#     Qr = Qch + dVdt
    ####

    ## Prep for more calculations ##
    enddate = date + pd.Timedelta('%i days' % ndays)

    Ut = df[smean][date:enddate].copy()
    Qt = Qch.loc[date:enddate].copy()
    Scht = Sch[date:enddate].copy()
    St = S[date:enddate].copy()

    indf = Ut > 0  # into the bay, flood tide
    inde = Ut < 0  # out of the bay, ebb tide
    ####

    ## flood ##
    Uf = Ut[indf].copy()
    Qf = Qt[indf].copy()
    Stf = St[indf].copy()
    Schtf = Scht[indf].copy()
    ind = (~Uf.isnull()) & (~Qf.isnull())
    Uf = Uf[ind]; Qf = Qf[ind]; Stf = Stf[ind]; Schtf = Schtf[ind]
    ####

    ## ebb ##
    Ue = Ut[inde].copy()
    Qe = Qt[inde].copy()
    Ste = St[inde].copy()
    Schte = Scht[inde].copy()
    ind = (~Ue.isnull()) & (~Qe.isnull())
    Ue = Ue[ind]; Qe = Qe[ind]; Schte = Schte[ind]; Ste = Ste[ind]
    ####

    npts = (len(Ue) + len(Uf))/len(Ut)  # ratio pts included

    Qcht = df[smean].copy()
    Uchf = df[smean][df[smean] > 0]; Uche = df[smean][df[smean] < 0]

    # may not work bc missing data
    # straight line
    if not Ue.empty and not Qe.empty and not Uf.empty and not Qf.empty:
        fite = curve_fit(fixed_line, Ue, Qe)
        fitf = curve_fit(fixed_line, Uf, Qf)

        ## Equation 3 after finding equivalent area for this loop ##
        Qcht[df[smean] > 0] = Uchf*fitf[0]
        Qcht[df[smean] < 0] = Uche*fite[0]

        area_ebb = fite[0][0]
        area_flood = fitf[0][0]

        # this is now the improved flux calculation which applies during the storm too
        Qchnew = Qcht.copy()[start:stop]

        # Section 3.3: inflow rate calculation
        Qrfin = (-Qchnew + dVdt)

        # Inflow volume calculation
        # save Harvey inflow
        harvey = Qrfin[date_outflow:date_endoutflow].sum()*dt/1000**3

        # save Harvey inflow WITHOUT dVdt term for comparison
        harvey_nodVdt = -Qchnew[date_outflow:date_endoutflow].sum()*dt/1000**3

        Qch = Qch[start:stop]

    else:
        harvey = np.nan
        Qrfin = np.nan
        harvey_nodVdt = np.nan
        area_ebb = np.nan
        area_flood = np.nan
        dVdt = np.nan
        Qchnew = np.nan
        Qch = np.nan

    return npts, harvey, Qrfin, harvey_nodVdt, area_ebb, area_flood, dVdt, Qchnew, Qch

```

```python
# Things to vary or set otherwise
smeans = ['smean', 'smean_2nd', 'smean_3rd']

# data time periods
ndays = 27  # 28 days total since first day is included
dates_use = dates[::4*24*7*2]

# weights for salinity
# uniform distribution of 10 points between 1/6 and 1/2
midgweights = np.linspace(1/3-1/6, 1/3+1/6, 10)

# From Gaussian distribution in 0_process_bathy.ipynb, chose ratio of
# 2 to 2 to 1 to 1 for the 4 bathymetries, respectively
arr = bathy['Amain'].values*1000**2
Amains = [arr[0], arr[0], arr[1], arr[2]]  # m^2
arr = bathy['Awest'].values
Awests = [arr[0], arr[0], arr[1], arr[2]]  # m^2
arr = bathy['Hmain'].values
Hmains = [arr[0], arr[0], arr[1], arr[2]]  # m^2
arr = bathy['Hwest'].values
Hwests = [arr[0], arr[0], arr[1], arr[2]]  # m^2

# Have volume calculated either without (assume eta=0) or with SSH data
Vconstants = [True, False]
```

```python
%%time
ncalcs = len(smeans)*len(dates_use)*midgweights.size*len(Amains)*len(Vconstants)

cols = ['npts','harvey inflow [km^3]','harvey inflow - no dVdt [km^3]',
        'effective area (ebb)','effective area (flood)', 'midgweight',
                    'Amain', 'Awest', 'Hmain', 'Hwest', 'Vconstant']
index = pd.date_range(start, stop, freq='15T')

i = 0
tsarr = np.zeros((index.size,ncalcs))
dVdtarr = np.zeros((index.size,ncalcs))
Qcharr = np.zeros((index.size,ncalcs))
Qchsaltarr = np.zeros((index.size,ncalcs))
smeanlist = []
dateslist = []
harvarr = np.zeros((ncalcs,len(cols)))  # number of calculations by columns to save

for smean in smeans:
    for date in dates_use:
        for midgweight in midgweights:
            for Amain, Awest, Hmain, Hwest in zip(Amains, Awests, Hmains, Hwests):
                for Vconstant in Vconstants:

                    a, b, c, d, e, f, g, h, k = calc_harvey(smean, midgweight, Amain, Awest, Hmain, Hwest,
                                                             Vconstant, date)

                    # save time series into arrays during loop. Too slow to put directly into dataframe.
                    tsarr[:,i] = c
                    dVdtarr[:,i] = g
                    Qcharr[:,i] = h
                    Qchsaltarr[:,i] = k

                    harvarr[i,:] = np.array([a,b,d,e,f, midgweight, Amain, Awest, Hmain, Hwest, Vconstant])

                    smeanlist.append(smean)
                    dateslist.append(date)

                    i += 1

# reconstitute
ts = pd.DataFrame(tsarr, index=index)
dVdt = pd.DataFrame(dVdtarr, index=index)
Qch = pd.DataFrame(Qcharr, index=index)
Qchsalt = pd.DataFrame(Qchsaltarr, index=index)
harv = pd.DataFrame(harvarr, index=np.arange(ncalcs), columns=cols)
harv['smean'] = smeanlist
harv['date'] = dateslist

```

    CPU times: user 19min 15s, sys: 3min 12s, total: 22min 28s
    Wall time: 19min 26s

Calculate r$^2$ for correlation between U$_{ch}$ and Q$_{ch}$ estimated from salinity conservation.

```python
from sklearn.linear_model import LinearRegression
lr = LinearRegression()

r2 = []
for smean in smeans:
    for date in dates:
        enddate = date + pd.Timedelta('%i days' % ndays)

        Uch = df[smean][date:enddate].copy()

        for midgweight in midgweights:

            S = ((1-0.33-midgweight)*df['TRIN: Salinity'] + midgweight*df['MIDG: Salinity'] + 0.33*Sch)  # volume averaged salinity

            for Amain, Awest, Hmain, Hwest in zip(Amains, Awests, Hmains, Hwests):
                V = (Amain * (Hmain + df['%s: Water Level [m, MHW]' % eta_main])) + (Awest * (Hwest + df['%s: Water Level [m, MHW]' % eta_west]))

                dSVdt = (S*V).diff()/dt
                Qcht = (dSVdt/Sch)[date:enddate]

                innan = (~np.isnan(Uch)) & (~np.isnan(Qcht))
                X = Uch[innan][:,np.newaxis]
                y = Qcht[innan][:,np.newaxis]
                lr.fit(X,y)
                r2.append(lr.score(X,y))

len(r2)
```

    2160

```python
r2 = np.array(r2)
r2.mean(), np.median(r2), r2.std()
```

    (0.37309247152100544, 0.3854656001613981, 0.054428304724295475)

```python
# find index for the final set of parameters analyzed in the loop above, to have a single example to show
irun = np.where((harv.date == date) & (harv.smean == smean) & (harv.midgweight == midgweight) & (harv.Amain == Amain) & (harv.Awest == Awest))[0][0]

# areas found for example
Ae, Af = harv.loc[irun]['effective area (ebb)'], harv.loc[irun]['effective area (flood)']
```

```python
import matplotlib.gridspec as gridspec

fig = plt.figure(constrained_layout=True, figsize=(15,10))
gs = fig.add_gridspec(2, 2)
ax1 = fig.add_subplot(gs[0, :])
ax1.plot(Qch[irun]['2017-8-1':'2017-8-24'], 'k', lw=2)
ax1.plot(Qchsalt[irun]['2017-8-1':'2017-8-24'], 'k', alpha=0.5, lw=1)
ax1.tick_params(labelsize=12)
ax1.set_ylabel('Channel flux [m$^3$/s]', fontsize=14)
ax1.set_xlim('2017-7-31 12:00','2017-8-25 12:00')
ax1.legend(['Combined channel flux estimate', 'Channel flux estimate from salt conservation'],
          loc='lower left', fontsize=12, framealpha=0)

ax2 = fig.add_subplot(gs[1, 0])
ax2.plot(Uch, Qcht, 'o', color='0.2')
ax2.set_xlabel('Along-channel velocity at g06010 [m/s]', fontsize=14)
ax2.set_ylabel('Channel flux from salt conservation [m$^3$/s]', fontsize=14)
ax2.set_title('Example relationship starting ' + Uch.index[0].strftime('%Y-%m-%d'), fontsize=14)
ax2.tick_params(labelsize=12)
xlims, ylims = ax2.get_xlim(), ax2.get_ylim()
# ax2.plot(Uch[Uch<0], Qcht[Uch<0], 'x')
# ax2.hlines(0, *xlims)
ax2.vlines(0, *ylims, linestyle='--', lw=1)
ax2.plot([0,Uch.min()], [0, Ae*Uch.min()], '0.7', lw=3, ls='-')
ax2.plot([0,Uch.max()], [0, Af*Uch.max()], '0.7', lw=3, ls='-')
ax2.text(0.02, 0.02, 'ebb', transform=ax2.transAxes, fontsize=12)
ax2.text(0.9, 0.02, 'flood', transform=ax2.transAxes, fontsize=12)
ax2.text(0.4, 0.9, 'r$^2$=%.2f' % r2[-1], transform=ax2.transAxes, fontsize=12)

ax3 = fig.add_subplot(gs[1, 1])
ax3.hist(r2, bins=8, color='0.2');
ax3.set_xlabel('r$^2$ value', fontsize=14)
ax3.set_ylabel('Number of instances', fontsize=14)
ax3.tick_params(labelsize=12)
ylims = ax3.get_ylim()
med, std = np.median(r2), r2.std()
ax3.vlines(med, *ylims)
ax3.vlines([med-std, med+std], *ylims, ls='--')

for ax, char in zip([ax1, ax2, ax3],['A','B','C']):
    ax.text(0.025, 0.9, char, transform=ax.transAxes, fontsize=16)

fig.savefig('../figures/example.pdf', bbox_inches='tight')
```

Run through all the combinations of inputs to get a best estimate for inflow and uncertainty.

```python
nptsmin = 0.8

# save full results and make harv only have time-changing volume calculations
# harvsave = harv.copy()
# harv = harv[harv['Vconstant'] == 0]

inpts = harv['npts'] > nptsmin
iharvey = (harv['date'] > pd.Timestamp('2017-7-25')) & (harv['date'] < pd.Timestamp('2017-9-24'))
ibad2018 = harv['date'] > pd.Timestamp('2018-5-1')
ivolume = harv['Vconstant'] == 0  # want these --- volume not constant
iuse = inpts & ~iharvey & ~ibad2018 & ivolume
key = 'harvey inflow [km^3]'
ax = (harv[key][iuse]).hist(bins=15, color='0.7')
ax.set_xlabel('Harvey inflow [km$^3$]')
ax.set_ylabel('Number of instances')

Vfw = harv[key][iuse].median()
Vfwstd = harv[key][iuse].std()
# Vfw_nodVdt = harv['harvey inflow - no dVdt [km^3]'][iuse].median()

print('Median of Harvey inflow estimates: %2.1f' % Vfw)
print('Mean of Harvey inflow estimates: %2.1f' % harv[key][iuse].mean())  # for reference
print('Standard deviation of Harvey inflow estimates: %2.1f' % Vfwstd)
print('Number of valid Harvey inflow estimates: %2i' % len(harv[iuse]))
print('Number of total Harvey inflow estimates: %2i' % len(harv))
# print('Median of Harvey inflow estimates with no dV/dt included: %2.1f' % Vfw_nodVdt)
# print('%% difference in inflow when dVdt excluded: %2.1f' % ((Vfw_nodVdt-Vfw)/Vfw*100))

plt.savefig('../figures/histogram.pdf', bbox_inches='tight')
```

### Uncertainty

How much uncertainty do the input parameters introduce?

```python
harv2 = harv[iuse]

## uncertainty due to salinity station weights
ncalcs = len(smeans)*len(harv2['date'].unique())*len(Amains)
std_midgweights = np.zeros((ncalcs,1))
i=0

for smean in smeans:
    for date in harv2['date'].unique():  # just use dates that result in good numbers
        for Amain, Awest in zip(Amains, Awests):

                inds = (harv2['Amain'] == Amain) & (harv2['Awest'] == Awest) & (harv2['smean'] == smean) \
                        & (harv2['date'] == date)
                std_midgweights[i] = harv2[inds][key].std()
                i += 1

## uncertainty due to ADCP profile extrapolation method
ncalcs = midgweights.size*len(harv2['date'].unique())*len(Amains)
std_smeans = np.zeros((ncalcs,1))
i=0

for midgweight in midgweights:
    for date in harv2['date'].unique():  # just use dates that result in good numbers
        for Amain, Awest in zip(Amains, Awests):

                inds = (harv2['Amain'] == Amain) & (harv2['Awest'] == Awest) & (harv2['midgweight'] == midgweight) \
                        & (harv2['date'] == date)
                std_smeans[i] = harv2[inds][key].std()
                i += 1

## uncertainty due to dates used
ncalcs = midgweights.size*len(smeans)*len(Amains)
std_dates = np.zeros((ncalcs,1))
i=0

for midgweight in midgweights:
    for smean in smeans:  # just use dates that result in good numbers
        for Amain, Awest in zip(Amains, Awests):

                inds = (harv2['Amain'] == Amain) & (harv2['Awest'] == Awest) & (harv2['midgweight'] == midgweight) \
                        & (harv2['smean'] == smean)
                std_dates[i] = harv2[inds][key].std()
                i += 1

## uncertainty due to surface area
ncalcs = midgweights.size*len(harv2['date'].unique())*len(smeans)
std_areas = np.zeros((ncalcs,1))
i=0

for midgweight in midgweights:
    for date in harv2['date'].unique():  # just use dates that result in good numbers
        for smean in smeans:

                inds = (harv2['smean'] == smean) & (harv2['midgweight'] == midgweight) & (harv2['date'] == date)
                std_areas[i] = harv2[inds][key].std()
                i += 1

## uncertainty due to volume being constant or not
# not using constant volume in calculations except for here
inpts = harv['npts'] > nptsmin
iharvey = (harv['date'] > pd.Timestamp('2017-7-25')) & (harv['date'] < pd.Timestamp('2017-9-24'))
ibad2018 = harv['date'] > pd.Timestamp('2018-5-1')
iuse3 = inpts & ~iharvey & ~ibad2018
harv3 = harv[iuse3]
ncalcs = midgweights.size*len(smeans)*len(Amains)*len(harv3['date'].unique())
std_Vconstant = np.zeros((ncalcs,1))
i=0

for midgweight in midgweights:
    for date in harv3['date'].unique():  # just use dates that result in good numbers
        for smean in smeans:  # just use dates that result in good numbers
            for Amain, Awest in zip(Amains, Awests):

                inds = (harv3['Amain'] == Amain) & (harv3['Awest'] == Awest) & (harv3['midgweight'] == midgweight) \
                         & (harv3['date'] == date) & (harv3['smean'] == smean)
                std_Vconstant[i] = harv3[inds][key].std()
                i += 1

```

Histograms of the distribution of standard deviation for each instance of the calculation of total inflow. The median of each histogram is printed and included in the paper.

```python
fig, ax = plt.subplots(1,1)
ax.hist(std_dates, alpha=0.5, density=True)
ax.hist(std_areas, alpha=0.5, density=True)
ax.hist(std_midgweights, alpha=0.5, density=True)
ax.hist(std_smeans, alpha=0.5, density=True)
ax.hist(std_Vconstant, alpha=0.95, density=True)

ax.set_xlabel('Standard deviation due to input parameter')
ax.set_ylabel('Number of instances')

print(np.median(std_dates), np.median(std_areas), np.median(std_midgweights),
                    np.median(std_smeans), np.nanmedian(std_Vconstant))

```

Read in TWDB inflow estimates

```python
# read in TWDB prediction of ungauged/modeled inflow by TWDB subbasin
Files = glob('%s/galv_ws_modelled_runoff/*' % (baset))
twdb = pd.DataFrame()
convert = 0.0142764  # 1 acre-foot per day = 0.0142764 m^3/s
for File in Files:
    dft = pd.read_table(File, delim_whitespace=True, parse_dates=[[0,1,2]], index_col=0)*convert
    dft.columns = [dft.columns[0][:5] + ' [m^3/s]'] # rename
    twdb = pd.concat([twdb, dft], axis=1)
```

```python
# read in TWDB prediction of gauged inflow by TWDB subbasin TWDB prediction of
Files = glob('%s/gauged_flow/*' % (baset))
for File in Files:
    dft = pd.read_table(File, delim_whitespace=True, parse_dates=[[0,1,2]], index_col=0)*convert
    dft.columns = [dft.columns[0].split('_')[0] + ' [m^3/s]'] # rename
    if not dft.columns[0] in twdb.columns:
        twdb = pd.concat([twdb, dft], axis=1)
```

```python
dstart = '2017-8-21'; dend = '2017-9-11'
tsmed = (ts.loc[:,iuse]).median(axis=1).rolling(window=24*4*1, center=True).mean()
tsstd = (ts.loc[:,iuse]).std(axis=1).rolling(window=24*4*1, center=True).mean()

fig, axes = plt.subplots(2, 1, sharex=True, figsize=(14,12))
fig.subplots_adjust(hspace=0.05)

ax = axes[0]

## freshwater inflow

# present calculation: median and uncertainty
tsmed[dstart:dend].plot(color='k', lw=3, label='$Q_{fw}$ median', ax=ax)
ax.fill_between(df[dstart:dend].index, (tsmed - tsstd)[dstart:dend], (tsmed + tsstd)[dstart:dend],
                color='k', alpha=0.1)#, label='$Q_{fw}$ uncertainty')

# TWDB inflow estimate
twdb.sum(axis=1)[dstart:dend].plot(color='k', lw=3, marker='.', markersize=15, linestyle='', label='TWDB', ax=ax)
# twdb['gauged'][dstart:dend].plot(color='k', lw=3, alpha=0.3, linestyle='--', label='TWDB gauged', ax=ax)

# (twdb.sum(axis=1)['2017-8-25':'2017-9-11'] + bayrain/86400).plot(color='k', lw=3, marker='x', markersize=15, ax=ax)#, linestyle='', label='TWDB', ax=ax)

# rain inflow into watershed
# newdates = pd.date_range('2017-8-25','2017-8-31', freq='1D')
# ax.scatter(newdates, rain/86400, color='k', marker='+', label='Rain into watershed')

# dVdt
col = '0.7'
dVdtmed = dVdt.loc[:,iuse].median(axis=1).rolling(window=24*4*1, center=True).mean()
dVdtstd = dVdt.loc[:,iuse].std(axis=1).rolling(window=24*4*1, center=True).mean()
dVdtmed[dstart:dend].plot(color=col, lw=3, label='dV/dt Median', ax=ax)
ax.fill_between(df[dstart:dend].index, (dVdtmed - dVdtstd)[dstart:dend], (dVdtmed + dVdtstd)[dstart:dend],
                color=col, alpha=0.1)#, label='dV/dt Uncertainty')

# Calculate channel flux median time series to compare with ADCP approach below
col = '0.4'
Qchmed = Qch.loc[:,iuse].median(axis=1).rolling(window=24*4*1, center=True).mean()
Qchstd = Qch.loc[:,iuse].std(axis=1).rolling(window=24*4*1, center=True).mean()
Qchmed[dstart:dend].plot(color=col, lw=3, label='Q$_{ch}$', ax=ax)
ax.fill_between(df[dstart:dend].index, (Qchmed - Qchstd)[dstart:dend], (Qchmed + Qchstd)[dstart:dend],
                color=col, alpha=0.1)#, label='Uncertainty')

# Calculate channel flux using NOAA PORTS signal and cross-sectional area of the entrance channel calculated in
# 0_process_bathymetry notebook
Acs = bathy.attrs['Acs0'] + bathy.attrs['Dcs']*df['8771341: Water Level [m, MHW]']
Qch_adcp = df['g06010: Along [cm/s] (water)']/100*Acs
Qch_adcp.rolling(window=24*4*1, center=True).mean()[dstart:dend].plot(ax=ax, color=col, lw=2, ls='--', label='Q$_{ch}$ from ADCP')

# # Qch from salinity conservation
# Qchsaltmed = Qchsalt.loc[:,iuse].median(axis=1).rolling(window=24*4*1, center=True).mean()
# Qchsaltmed[dstart:dend].plot(color=col, lw=3, linestyle=':', ax=ax)

ax.set_ylabel('Flux terms [m$^3$s$^{-1}$]', fontsize=14)
ax.text(0.05, 0.90, 'Inflow', transform=ax.transAxes, fontsize=12)
ax.text(0.05, 0.05, 'Outflow', transform=ax.transAxes, fontsize=12)
ax.legend(loc='lower right')

# Cumulative inflow calculation
ax2 = axes[1]
(tsmed[date_outflow:dend].cumsum()*dt/1000**3).plot(color='k', lw=2, ax=ax2, label='Cumulative sum of Q$_{fw}$')

# rain time series
newdates = pd.date_range('2017-8-25','2017-9-11', freq='1D')
ax2.scatter(newdates, np.cumsum(radar.loc['all','rain [m^3]'])/1000**3, marker='+', color='k', label='Cumulative sum of rainfall in watershed')

# twdb cumulative sum
(twdb.sum(axis=1)[date_outflow:dend].cumsum()*86400/1000**3).plot(color='k', lw=3, marker='.', markersize=15, linestyle='', label='TWDB', ax=ax2)

ax2.set_ylabel('Freshwater inflow volume, $V_{fw}$ [km$^3$]', fontsize=14)
ax2.legend(loc='lower right')

ax.set_ylabel('Flow rate [m$^3$s$^{-1}$]', fontsize=14)
ax2.set_xlabel('');
ax.legend(loc='lower right')

ax.set_xlim(dstart, dend)

# add grids
ax.set_yticks(np.arange(-50000, 70000, 10000), minor=True)
# ax.yaxis.grid(True, linewidth=1.5, alpha=0.5, which='major')
ax.yaxis.grid(True, linewidth=0.7, alpha=0.5, which='both')
# ax.yaxis.grid(True, linewidth=1, alpha=0.5, axis='minor')
ax2.yaxis.grid(True, linewidth=0.7, alpha=0.5)
# # minor = mpl.dates.HourLocator(byhour=np.arange(0,24,6))
# # ax.xaxis.set_minor_locator(minor)
# major = mpl.dates.HourLocator(byhour=np.arange(0,24,24))
# ax.xaxis.set_major_locator(major)
# ax.xaxis.set_major_formatter(mpl.dates.DateFormatter('%b %d'))

for ax in axes:
    ylims = ax.get_ylim()
    ax.vlines(date_outflow, *ylims, linewidth=2, alpha=0.5, color='0.5', linestyle='--')
    ax.vlines(date_endoutflow, *ylims, linewidth=2, alpha=0.5, color='0.5', linestyle='--')
    ax.hlines(0, dstart, dend, lw=1, alpha=0.5)
    ax.tick_params(labelsize=12)

plt.savefig('../figures/inflow_comp.pdf', bbox_inches='tight')
```

Simple version of plot

```python
dstart = '2017-8-21'; dend = '2017-9-11'
tsmed = (ts.loc[:,iuse]).median(axis=1).rolling(window=24*4*1, center=True).mean()
tsstd = (ts.loc[:,iuse]).std(axis=1).rolling(window=24*4*1, center=True).mean()

fig, ax = plt.subplots(1, 1, sharex=True, figsize=(14,6))

# present calculation: median and uncertainty
tsmed[dstart:dend].plot(color='k', lw=3, label='$Q_{fw}$ median', ax=ax)
ax.fill_between(df[dstart:dend].index, (tsmed - tsstd)[dstart:dend], (tsmed + tsstd)[dstart:dend],
                color='k', alpha=0.1)#, label='$Q_{fw}$ uncertainty')

# TWDB inflow estimate
twdb.sum(axis=1)[dstart:dend].plot(color='k', lw=3, marker='.', markersize=15, linestyle='', label='TWDB', ax=ax)

ax.set_ylabel('Flux terms [m$^3$s$^{-1}$]', fontsize=14)
ax.text(0.05, 0.90, 'Inflow', transform=ax.transAxes, fontsize=14)
ax.text(0.05, 0.05, 'Outflow', transform=ax.transAxes, fontsize=14)
ax.legend(loc='lower right')
ax.set_ylabel('Flow rate [m$^3$s$^{-1}$]', fontsize=18)
ax.set_xlabel('', fontsize=18)
ax.legend(loc='lower right', fontsize=14)

# add grids
ax.set_yticks(np.arange(-50000, 70000, 10000), minor=True)
ax.yaxis.grid(True, linewidth=0.7, alpha=0.5, which='both')
ax.set_xlim(dstart, dend)

# for ax in axes:
ax.set_ylim(-20000, 50000)
ylims = ax.get_ylim()
ax.vlines(date_outflow, *ylims, linewidth=2, alpha=0.5, color='0.5', linestyle='--')
ax.vlines(date_endoutflow, *ylims, linewidth=2, alpha=0.5, color='0.5', linestyle='--')
ax.hlines(0, dstart, dend, lw=1, alpha=0.5)
ax.tick_params(labelsize=16, which='both')

plt.savefig('../figures/inflow_comp_presentation.pdf', bbox_inches='tight', dpi=100)
plt.savefig('../figures/inflow_comp_presentation.png', bbox_inches='tight', dpi=100)
```

For the approximate cross-sectional area at the entrance channel, what speed would the water exiting the channel have to be traveling, uniformly across the channel?

```python
# they are very similar
(Qchmed/Acs)[dstart:dend].plot(color='k', linewidth=2, label='Q$_{ch}$')
(df['g06010: Along [cm/s] (water)']/100)[dstart:dend].plot(figsize=(15,4))

```

Compare effective flood/ebb areas with the calculated cross-sectional area

```python
fig, axes = plt.subplots(1, 2, sharey=True, sharex=True, figsize=(14,4))
key = 'effective area (ebb)'
harv[key][iuse].hist(ax=axes[0], bins=20)
ylims = axes[0].get_ylim()
axes[0].vlines(Acs[date_outflow:date_endoutflow].mean(), 0, ylims[1])
key = 'effective area (flood)'
harv[key][iuse].hist(ax=axes[1], bins=20)
axes[1].vlines(Acs[date_outflow:date_endoutflow].mean(), 0, ylims[1])

```

## 4. Discussion

### 4.1 Comparisons

### TWDB inflow estimate

```python
# integral over flooding time period
VTWDB = (twdb.sum(axis=1)[date_outflow:date_endoutflow]*86400).sum()/1000**3
print('TWDB freshwater volume inflow: %2.1f km^3' % VTWDB)

# VTWDB_gauged = (twdb['gauged'][date_outflow:date_endoutflow]*86400).sum()/1000**3
# VTWDB_ungauged = (twdb['modelled'][date_outflow:date_endoutflow]*86400).sum()/1000**3
# print('TWDB ungauged freshwater volume inflow: %2.1f km^3' % (VTWDB_ungauged))
# print('TWDB gauged freshwater volume inflow: %2.1f km^3' % (VTWDB_gauged))

per = ((Vfw - VTWDB)/VTWDB)*100
print('The present study predicts an inflow that is %2.0f percent more than the TWDB model.' % per)

per = ((Vfw - VTWDB)/Vfw)*100
print('The TWDB model inflow prediction is %2.0f percent lower than the prediction in the present study.' % per)
```

    TWDB freshwater volume inflow: 13.1 km^3
    The present study predicts an inflow that is 27 percent more than the TWDB model.
    The TWDB model inflow prediction is 21 percent lower than the prediction in the present study.

TWDB inflow estimate doesn't include direct rainfall onto bay, which from radar rainfall integrates over the storm period to about 1 km^3.

## Compare with ungauged subbasins

```python
secondsperday = 86400
colorungauged = '0.5'
markerfillcolorun = '0.5'
colorgauged = '0.5'
markersize = 15
markerg = '.'
markerug = '.'

fig, axes = plt.subplots(1, 2, figsize=(14,6))

# ungauged
ax = axes[0]
ungauged = radar[radar['kind'] == 'ungauged'].index
rain_radar = radar.loc[ungauged, 'rain total [m^3]']/1000**3
rain_twdb = [(twdb['%s [m^3/s]' % str(basin).zfill(5)]['2017-8-25':'2017-9-9'].sum()*secondsperday)/1000**3 for basin in ungauged]
ax.plot(rain_radar, rain_twdb, marker=markerug, ls='', markersize=markersize, alpha=0.5, color=colorungauged, mfc=markerfillcolorun)
ax.set_title('Ungauged subbasins', fontsize=14)
xpts = np.linspace(0, 0.8)
ax.plot(xpts, xpts, 'k', label=None, linestyle=':')
ax.set_ylabel('TWDB estimate [km$^3$]', fontsize=14)
ax.set_xlabel('Radar rain [km$^3$]', fontsize=14)
ax.tick_params(labelsize=12)
ax.axis('tight')
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)

# gauged
bax = axes[1]

# for Lake Houston
lkhouston = radar[radar['kind'] == 'Lake Houston (gauged)'].index
rain_radar = (radar.loc[lkhouston, 'rain total [m^3]']/1000**3).sum()
rain_twdb = (twdb['lkhouston [m^3/s]']['2017-8-25':'2017-9-9'].sum()*secondsperday)/1000**3
plkh, = bax.plot(rain_radar, rain_twdb, marker='^', ls='', markersize=markersize, alpha=0.5, color=colorgauged)

# for Trinity
trinity = radar[radar['kind'] == 'gauged (large)'].index
rain_radar = (radar.loc[trinity, 'rain total [m^3]']/1000**3).sum()
rain_twdb = (twdb['trinity [m^3/s]']['2017-8-25':'2017-9-9'].sum()*secondsperday)/1000**3
ptrin, = bax.plot(rain_radar, rain_twdb, marker='*', ls='', markersize=markersize, alpha=0.5, color=colorgauged)

gauged = radar[radar['kind'] == 'gauged'].index
gauged = gauged.drop(['10074','10073'])  # plot separately
rain_radar = radar.loc[gauged, 'rain total [m^3]']/1000**3
rain_twdb = [(twdb['%s [m^3/s]' % str(basin).zfill(5)]['2017-8-25':'2017-9-9'].sum()*secondsperday)/1000**3 for basin in gauged]
bax.plot(rain_radar, rain_twdb, marker=markerg, ls='', markersize=markersize, alpha=0.5, color=colorgauged)
# example
rain_radar = radar.loc['10074', 'rain total [m^3]']/1000**3
rain_twdb = (twdb['%s [m^3/s]' % str(10074).zfill(5)]['2017-8-25':'2017-9-9'].sum()*secondsperday)/1000**3
p3, = bax.plot(rain_radar, rain_twdb, marker='s', ls='', markersize=10, alpha=0.5, color=colorgauged)
# example
rain_radar = radar.loc['10073', 'rain total [m^3]']/1000**3
rain_twdb = (twdb['%s [m^3/s]' % str(10073).zfill(5)]['2017-8-25':'2017-9-9'].sum()*secondsperday)/1000**3
p4 = bax.plot(rain_radar, rain_twdb, marker='s', ls='', markersize=10, alpha=0.5, color=colorgauged)

xpts = np.linspace(0, 5)
bax.plot(xpts, xpts, 'k', label=None, linestyle=':')
bax.set_xlabel('Radar rain [km$^3$]', fontsize=14)
bax.set_title('Gauged subbasins', fontsize=14)
bax.spines['right'].set_visible(False)
bax.spines['top'].set_visible(False)

bax.tick_params(labelsize=12)

bax.legend(handles=[plkh,ptrin,p3], labels=['Lake Houston', 'Trinity','Flood Example'], loc='lower right')

fig.savefig('../figures/twdb_radar_comp.pdf', bbox_inches='tight')
fig.savefig('../figures/twdb_radar_comp.png', bbox_inches='tight')

```

### Galveston Bay volume

```python
Vall = (np.nansum(bathy['main'][:,:,0]*bathy.attrs['dx']*bathy.attrs['dy']) + np.nansum(bathy['west'][:,:,0]*bathy.attrs['dx']*bathy.attrs['dy']))/1000**3

print('The flood water could fill the bay approximately %1.0f times' % abs(Vfw/Vall))
```

    The flood water could fill the bay approximately 5 times


# Frequentist Inference Case Study - Part B

## Learning objectives

Welcome to Part B of the Frequentist inference case study! The purpose of this case study is to help you apply the concepts associated with Frequentist inference in Python. In particular, you'll practice writing Python code to apply the following statistical concepts:
* the _z_-statistic
* the _t_-statistic
* the difference and relationship between the two
* the Central Limit Theorem, including its assumptions and consequences
* how to estimate the population mean and standard deviation from a sample
* the concept of a sampling distribution of a test statistic, particularly for the mean
* how to combine these concepts to calculate a confidence interval

In the previous notebook, we used only data from a known normal distribution. **You'll now tackle real data, rather than simulated data, and answer some relevant real-world business problems using the data.**

## Hospital medical charges

Imagine that a hospital has hired you as their data scientist. An administrator is working on the hospital's business operations plan and needs you to help them answer some business questions.

In this assignment notebook, you're going to use frequentist statistical inference on a data sample to answer the questions:
* has the hospital's revenue stream fallen below a key threshold?
* are patients with insurance really charged different amounts than those without?

Answering that last question with a frequentist approach makes some assumptions, and requires some knowledge, about the two groups.

We are going to use some data on medical charges obtained from [Kaggle](https://www.kaggle.com/easonlai/sample-insurance-claim-prediction-dataset).

For the purposes of this exercise, assume the observations are the result of random sampling from our single hospital. Recall that in the previous assignment, we introduced the Central Limit Theorem (CLT), and its consequence that the distributions of sample statistics approach a normal distribution as $n$ increases. The amazing thing about this is that it applies to the sampling distributions of statistics that have been calculated from even highly non-normal distributions of data! Recall, also, that hypothesis testing is very much based on making inferences about such sample statistics. You're going to rely heavily on the CLT to apply frequentist (parametric) tests to answer the questions in this notebook.

```python
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import t, ttest_ind
medical = pd.read_csv('data/insurance2.csv')
```

```python
medical.shape
```

    (1338, 8)

```python
medical.head()
```

      age
      sex
      bmi
      children
      smoker
      region
      charges
      insuranceclaim

      0
      19
      0
      27.900
      0
      1
      3
      16884.92400
      1

      1
      18
      1
      33.770
      1
      0
      2
      1725.55230
      1

      2
      28
      1
      33.000
      3
      0
      2
      4449.46200
      0

      3
      33
      1
      22.705
      0
      0
      1
      21984.47061
      0

      4
      32
      1
      28.880
      0
      0
      1
      3866.85520
      1

__Q1:__ Plot the histogram of charges and calculate the mean and standard deviation. Comment on the appropriateness of these statistics for the data.

__A:__

```python
charges_mean = np.mean(medical.charges)
charges_std = np.std(medical.charges)

print(f"The mean of charges is: {charges_mean:,.2f}")
print(f"The standard deviation of charges is: {charges_std:,.2f}")

```

    The mean of charges is: 13,270.42
    The standard deviation of charges is: 12,105.48

It looks odd that the standard deviation is so high given the mean.
We would be in negative charges (which doesn't make sense) after 1 standard deviation.
The data must be skewed heavily to the right (a few large charges with a lot of near zero charges).

```python
_ = plt.hist(medical.charges)
_ = plt.xlabel('charges ($)')
_ = plt.ylabel('number of people')
_ = plt.title('Distribution of charges at hospital')
_ = plt.axvline(charges_mean, color='r')
_ = plt.axvline(charges_mean+charges_std, color='r', linestyle='--')
_ = plt.axvline(charges_mean-charges_std, color='r', linestyle='--')
_ = plt.axvline(charges_mean+charges_std*2, color='r', linestyle='-.')
```

__Q2:__ The administrator is concerned that the actual average charge has fallen below 12,000, threatening the hospital's operational model. On the assumption that these data represent a random sample of charges, how would you justify that these data allow you to answer that question? And what would be the most appropriate frequentist test, of the ones discussed so far, to apply?

__A:__ There are 1,338 charges in this data set.
Therefore, we should have a large enough sample size to easily calculate various statistics accurately for the hospital's population.
We will use a t test because the data is skewed to the right (it is not normally distributed).

__Q3:__ Given the nature of the administrator's concern, what is the appropriate confidence interval in this case? A ***one-sided*** or ***two-sided*** interval? (Refresh your understanding of this concept on p. 399 of the *AoS*). Calculate the critical value and the relevant 95% confidence interval for the mean, and comment on whether the administrator should be concerned.

__A:__ A 95% confidence interval is the standard one to use unless a different one is requested.
In this case, a one-sided test makes sense because of the skew of the data.
The calculations below show that 12,000 is lower than the 95% confidence interval (12,621.20).
Therefore, we do NOT reject the null hypothesis and the administrator doesn't need to worry that the actual charges have fallen this low.

```python
sample_size =  medical.charges.count()
print(f"The sample size is {sample_size}.")
```

    The sample size is 1338.

```python
def calculate_confidence_interval_and_critical_value(sample_size, confidence_percentage=0.95):
    """ Steps to calculate a Confidence Interval For a Sample
    1) Subtract 1 from your sample size.
    2) Subtract the confidence level from 1 and then divide by two.
    3) Lookup answers from step 1 and 2 in t-distribution table or calculate them.
    4) Divide sample standard deviation by the square root of the sample size.
    5) Multiple step 3 and 4.
    6) Subtract step 5 from the sample mean for the lower end of the range.
    7) Add step 5 from the sample mean for the upper end of the range.
    :param
    sample_size: from the population
    operation: + or -. Which direction from the mean we're analyzing
    confidence_percentage: for the hypothesis test. Defaults to 0.95 for 95%.
    :return: critical value, confidence interval
    """

    lower_range = (1 - confidence_percentage) / 2
    upper_range = 1 - lower_range
    first_steps = t(sample_size-1).ppf([lower_range, upper_range])
    step_4 = charges_std / np.sqrt(sample_size)
    step_5 = first_steps * step_4
    last_steps = charges_mean + step_5

    return first_steps, last_steps
```

```python
critical_value, conf_interval = calculate_confidence_interval_and_critical_value(sample_size)
print(f"The critical value range is from {critical_value[0]:.2f} to {critical_value[1]:.2f}.")
print(f"The 95% confidence interval range for the median is from {conf_interval[0]:,.2f} to {conf_interval[1]:,.2f}.")
```

    The critical value range is from -1.96 to 1.96.
    The 95% confidence interval range for the median is from 12,621.20 to 13,919.65.

```python
# Even changing the confidence interval to 99% still shows that the charges haven't fallen to 12,000.
critical_value, conf_interval = calculate_confidence_interval_and_critical_value(sample_size, 0.99)
print(f"The critical value range is from {critical_value[0]:.2f} to {critical_value[1]:.2f}.")
print(f"The 99% confidence interval range for the median is from {conf_interval[0]:,.2f} to {conf_interval[1]:,.2f}.")
```

    The critical value range is from -2.58 to 2.58.
    The 99% confidence interval range for the median is from 12,416.75 to 14,124.10.

The administrator then wants to know whether people with insurance really are charged a different amount to those without.

__Q4:__ State the null and alternative hypothesis here. Use the _t_-test for the difference between means, where the pooled standard deviation of the two groups is given by:
\begin{equation}
s_p = \sqrt{\frac{(n_0 - 1)s^2_0 + (n_1 - 1)s^2_1}{n_0 + n_1 - 2}}
\end{equation}

and the *t*-test statistic is then given by:

\begin{equation}
t = \frac{\bar{x}_0 - \bar{x}_1}{s_p \sqrt{1/n_0 + 1/n_1}}.
\end{equation}

(If you need some reminding of the general definition of ***t-statistic***, check out the definition on p. 404 of *AoS*).

What assumption about the variances of the two groups are we making here?

__A:__
That they are the same.

__Q5:__ Perform this hypothesis test both manually, using the above formulae, and then using the appropriate function from [scipy.stats](https://docs.scipy.org/doc/scipy/reference/stats.html#statistical-tests) (hint, you're looking for a function to perform a _t_-test on two independent samples). For the manual approach, calculate the value of the test statistic and then its probability (the p-value). Verify you get the same results from both.

__A:__
#### Manual Calculations

```python
print("Patients With Insurance")
claim = medical.insuranceclaim
has_insurance = medical.charges[claim == 1]
print(f"The total is {has_insurance.count()}.")
insurance_mean = np.mean(has_insurance)
print(f"The mean of the charges is {insurance_mean:,.2f}.")
insurance_std = np.std(has_insurance)
print(f"The standard deviation of the charges is {insurance_std:,.2f}.")
```

    Patients With Insurance
    The total is 783.
    The mean of the charges is 16,423.93.
    The standard deviation of the charges is 14,036.96.

```python
print("Patients Without Insurance")
no_insurance = medical.charges[claim == 0]
print(f"The total is {no_insurance.count()}.")
no_insurance_mean = np.mean(no_insurance)
print(f"The mean of the charges is {no_insurance_mean:,.2f}.")
no_insurance_std = np.std(no_insurance)
print(f"The standard deviation of the charges is {no_insurance_std:,.2f}.")
```

    Patients Without Insurance
    The total is 555.
    The mean of the charges is 8,821.42.
    The standard deviation of the charges is 6,440.70.

```python
n0, n1 = len(has_insurance), len(no_insurance)
x0, x1 = insurance_mean, no_insurance_mean
s0, s1 = insurance_std, no_insurance_std

# First calculation above
sp = np.sqrt(((n0 - 1) * s0 ** 2 + (n1 - 1) * s1 ** 2) / (n0 + n1 - 2))
# Second calculation above
t = (x0 - x1) / (sp * np.sqrt(1 / n0 + 1 / n1))

print(f"The pooled standard deviation is {sp:,.2f}.")
print(f"The t-test statistic is {t:.2f}.")
```

    The pooled standard deviation is 11,512.28.
    The t-test statistic is 11.90.

From these statistics it appears that people with insurance are charged a different amount.
In this case, a significantly higher amount.

#### Calculations using [scipy.stats.ttest_ind](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.ttest_ind.html) function

```python
statistic, pvalue = ttest_ind(has_insurance, no_insurance)
print(f"The t-test statistic is {statistic:.2f}.")
print(f"The pvalue is {pvalue}.")
```

    The t-test statistic is 11.89.
    The pvalue is 4.461230231620717e-31.

Congratulations! Hopefully you got the exact same numerical results. This shows that you correctly calculated the numbers by hand. Secondly, you used the correct function and saw that it's much easier to use. All you need to do is pass your data to it.

__Q6:__ Conceptual question: look through the documentation for statistical test functions in scipy.stats. You'll see the above _t_-test for a sample, but can you see an equivalent one for performing a *z*-test from a sample? Comment on your answer.

__A:__ The closest thing in scipy is for a z-score at [scipy.stats.zscore](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.zscore.html).
- The signature is: scipy.stats.zscore(a, axis=0, ddof=0, nan_policy='propagate').

From [Investopedia](https://www.investopedia.com/terms/z/z-test.asp), some key takeaways:
- Z-test is a statistical test to determine whether two population means are different when the variances are known and the sample size is large.
- Z-test is a hypothesis test in which the z-statistic follows a normal distribution.
- A z-statistic, or z-score, is a number representing the result from the z-test.
- Z-tests are closely related to t-tests, but t-tests are best performed when an experiment has a small sample size.
- Z-tests assume the standard deviation is known, while t-tests assume it is unknown.

## Learning outcomes

Having completed this project notebook, you now have good hands-on experience:
* using the central limit theorem to help you apply frequentist techniques to answer questions that pertain to very non-normally distributed data from the real world
* performing inference using such data to answer business questions
* forming a hypothesis and framing the null and alternative hypotheses
* testing this using a _t_-test

Frequentist Inference Case Study - Part B (2).ipynb


```python
# https://arxiv.org/pdf/quant-ph/0104030.pdf
# ^^^ Need to be able to prepare arbitrary state!
import numpy as np
import cirq
from functools import reduce
from sympy import *
```

```python
# https://arxiv.org/pdf/quant-ph/9503016.pdf
```

bottom of pg 663[https://rdo.psu.ac.th/sjstweb/journal/27-3/18mathices.pdf]

## Roots of diagonalizable matrices

In this section, we consider an nth root of a diagonalizable matrix.

- Theorem 2.1: Let A be an $m\times m$ complex matrix. If A is diagonalizable, then A has an nth root, for anypositive integer n.

Proof:

Let $A$ be a diagonalizable matrix, i.e., there exists a non-singular matrix S such that $A = SDS^{-1}$where $D=[d_{ij}]_{m\times m}$ is a diagonal matrix.

Let $D^{\frac{1}{n}}=[d_{ij}^{\frac{1}{n}}]_{m \times m}$, where $d_{ij}^{\frac{1}{n}}$ is an n-th root of $d_{ij}$.

So $A = S (D^{\frac{1}{n}})^{n} S^{-1} = (SD^{\frac{1}{n}}S)^{n}$ Therefore an n-th root of A exists.

https://math.stackexchange.com/questions/1168438/the-nth-root-of-the-2x2-square-matrix

```python
X=Matrix(cirq.X._unitary_())
X
```

```python

X = Matrix([[0,1],[1j,0]])

from sympy.physics.quantum import Dagger
Dagger(X)
```

```python
np.array(X)
```

```python
from numpy import linalg as LA
from sympy import *
from sympy.physics.quantum import Dagger

# SYMPY calculation gets exact diagonal!!! (note matrices are Hermitian)

class Build_V_Gate():

    # V^{n} = U

    def __init__(self,U, n_power):
        self.U=U
        self.n = n_power

        self.D = None
        self.V = None
        self.V_dag = None

    def _diagonalise_U(self):

        # find diagonal matrix:
        U_matrix = Matrix(self.U)
        self.S, self.D = U_matrix.diagonalize()
        self.S_inv = self.S**-1
        # where U = S D S^{-1}

        if not np.allclose(np.array(self.S*(self.D*self.S_inv), complex), self.U):
            raise ValueError('U != SDS-1')

    def Get_V_gate_matrices(self):

        if self.D is None:
            self._diagonalise_U()

#         D_nth_root = np.power(self.D, 1/self.n)
        D_nth_root = self.D**(1/self.n)

#         self.V = np.array(self.S,complex).dot(np.array(D_nth_root,complex)).dot(np.array(self.S_inv,complex))
#         self.V_dag = self.V.conj().transpose()
        self.V = self.S * D_nth_root * self.S_inv
        self.V_dag = Dagger(self.V)

        if not np.allclose(reduce(np.matmul, [np.array(self.V, complex) for _ in range(self.n)]), self.U, atol=1e-10):
            raise ValueError('U != V^{}'.format(self.n))

        return np.array(self.V, complex), np.array(self.V_dag, complex)

n_root=2
mat = cirq.X._unitary_()
aa = Build_V_Gate(mat, n_root)
V, V_dag = aa.Get_V_gate_matrices()
reduce(np.matmul, [V for _ in range(n_root)])
```

```python
# from numpy import linalg as LA
# # NUMPY VERSION... NOT as good!

# class Build_V_Gate():

#     # V^{n} = U

#     def __init__(self,U, n_power):
#         self.U=U
#         self.n = n_power

#         self.D = None
#         self.V = None
#         self.V_dag = None

#     def _diagonalise_U(self):

#         val,vec = np.linalg.eig(self.U)

#         #sorting
#         idx = val.argsort()[::-1]
#         val_sorted = val[idx]
#         vec_sorted = vec[:,idx]

#         # find diagonal matrix:
#         vec_sorted_inv = np.linalg.inv(vec_sorted)
#         self.D = vec_sorted_inv.dot(self.U.dot(vec_sorted))

#         self.S=vec_sorted
#         self.S_inv = vec_sorted_inv
#         # where U = S D S^{-1}

#         if not np.allclose(self.S.dot(self.D).dot(self.S_inv), self.U):
#             raise ValueError('U != SDS-1')

#     def Get_V_gate_matrices(self):

#         if self.D is None:
#             self._diagonalise_U()

#         D_nth_root = np.power(self.D, 1/self.n)
# #         D_nth_root = np.sqrt(self.D)

#         self.V = self.S.dot(D_nth_root).dot(self.S_inv)
#         self.V_dag = self.V.conj().transpose()

#         if not np.allclose(reduce(np.matmul, [self.V for _ in range(self.n)]), self.U, atol=1e-1):
#             raise ValueError('U != V^{}'.format(self.n))

#         return self.V, self.V_dag

# mat = cirq.X._unitary_()
# aa = Build_V_Gate(mat, 2)
# V, V_dag = aa.Get_V_gate_matrices()

# np.around(V.dot(V), 3)
```

```python
# aa = Build_V_Gate(mat, 4)
# V, V_dag = aa.Get_V_gate_matrices()
# np.around(((V.dot(V)).dot(V)).dot(V), 3)
```

```python

```

```python
class My_V_gate(cirq.SingleQubitGate):
    """
    Description

    Args:
        theta (float): angle to rotate by in radians.
        number_control_qubits (int): number of control qubits
    """

    def __init__(self, V, V_dag, dagger_gate = False):
        self.V = V
        self.V_dag = V_dag
        self.dagger_gate = dagger_gate
    def _unitary_(self):
        if self.dagger_gate:
            return self.V_dag
        else:
            return self.V

    def num_qubits(self):
        return 1

    def _circuit_diagram_info_(self,args):
        if self.dagger_gate:
            return 'V^{†}'
        else:
            return 'V'

    def __str__(self):
        if self.dagger_gate:
            return 'V^{†}'
        else:
            return 'V'

    def __repr__(self):
        return self.__str__()
```

```python
n_root=4
mat = cirq.X._unitary_()
aa = Build_V_Gate(mat, n_root)
V, V_dag = aa.Get_V_gate_matrices()

GATE = My_V_gate(V, V_dag, dagger_gate=True)

circuit = GATE.on(cirq.LineQubit(2))
cirq.Circuit(circuit)
```

```python

```

```python
def int_to_Gray(num, n_qubits):
    # https://en.wikipedia.org/wiki/Gray_code

    # print(np.binary_repr(num, n_qubits)) # standard binary form!

    # The operator >> is shift right. The operator ^ is exclusive or
    gray_int = num^(num>>1)

    return np.binary_repr(gray_int,n_qubits)

### example... note that grey code reversed as indexing from left to right: [0,1,-->, N-1]
for i in range(2**3):
    print(int_to_Gray(i, 3)[::-1])

int_to_Gray(6, 4)
```

```python
def check_binary_str_parity(binary_str):
    """
    Returns 0 for EVEN parity
    Returns 1 for ODD parity
    """
    parity = sum(map(int,binary_str))%2

    return parity

check_binary_str_parity('0101')
```

```python
# NOTE pg 17 of Elementary gates for quantum information
class n_control_U(cirq.Gate):
    """
    """

    def __init__(self, V, V_dag, list_of_control_qubits, list_control_vals, U_qubit):
        self.V = V
        self.V_dag = V_dag

        if len(list_of_control_qubits)!=len(list_control_vals):
            raise ValueError('incorrect qubit control bits or incorrect number of control qubits')

        self.list_of_control_qubits = list_of_control_qubits
        self.list_control_vals = list_control_vals
        self.U_qubit = U_qubit

        self.n_ancilla=len(list_of_control_qubits)

    def flip_control_to_zero(self):
        for index, control_qubit in enumerate(self.list_of_control_qubits):
            if self.list_control_vals[index]==0:
                yield cirq.X.on(control_qubit)

    def _get_gray_control_lists(self):

        grey_cntrl_bit_lists=[]
        n_ancilla = len(self.list_of_control_qubits)
        for grey_index in range(1, 2**n_ancilla):

            gray_control_str = int_to_Gray(grey_index, n_ancilla)[::-1] # note reversing order
            control_list = list(map(int,gray_control_str))
            parity = check_binary_str_parity(gray_control_str)

            grey_cntrl_bit_lists.append((control_list, parity))
        return grey_cntrl_bit_lists

    def _decompose_(self, qubits):

        ## flip if controlled on zero
        X_flip = self.flip_control_to_zero()
        yield X_flip

        ## perform controlled gate
        n_ancilla = len(self.list_of_control_qubits)

        grey_control_lists = self._get_gray_control_lists()

        for control_index, binary_control_tuple in enumerate(grey_control_lists):

            binary_control_seq, parity = binary_control_tuple
            control_indices = np.where(np.array(binary_control_seq)==1)[0]
            control_qubit = control_indices[-1]

            if parity==1:
                gate = self.V.controlled(num_controls=1, control_values=[1]).on(self.list_of_control_qubits[control_qubit], self.U_qubit)
#                 gate= 'V'
            else:
                gate = self.V_dag.controlled(num_controls=1, control_values=[1]).on(self.list_of_control_qubits[control_qubit], self.U_qubit)
#                 gate= 'V_dagg'

            if control_index==0:
                yield gate
#                 print(gate, control_qubit)
            else:
                for c_index in range(len(control_indices[:-1])):
                    yield cirq.CNOT(self.list_of_control_qubits[control_indices[c_index]], self.list_of_control_qubits[control_indices[c_index+1]])
#                     print('CNOT', control_indices[c_index], control_indices[c_index+1])
#                 print(gate, control_qubit)
                yield gate
                for c_index in list(range(len(control_indices[:-1])))[::-1]:
#                     print('CNOT', control_indices[c_index], control_indices[c_index+1])
                    yield cirq.CNOT(self.list_of_control_qubits[control_indices[c_index]], self.list_of_control_qubits[control_indices[c_index+1]])

        ## unflip if controlled on zero
        X_flip = self.flip_control_to_zero()
        yield X_flip

    def _circuit_diagram_info_(self, args):

#         return cirq.CircuitDiagramInfo(
#             wire_symbols=tuple([*['@' for _ in range(len(self.list_of_control_qubits))],'U']),exponent=1)
        return cirq.protocols.CircuitDiagramInfo(
            wire_symbols=tuple([*['@' if bit==1 else '(0)' for bit in self.list_control_vals],'U']),
            exponent=1)

    def num_qubits(self):
        return len(self.list_of_control_qubits) + 1 #(+1 for U_qubit)

```

```python
n_control_qubits=6
n_power = 2**(n_control_qubits-2)

```

```python
cirq.X._unitary_()
```

```python

cirq.LineQubit.range(3)
```

```python
### setup V gate ##
n_control_qubits=3
n_power = 2**(n_control_qubits-1)

theta= np.pi/2
# U_GATE_MATRIX = np.array([
#                     [np.cos(theta), np.sin(theta)],
#                     [np.sin(theta), -1* np.cos(theta)]
#                 ])

U_GATE_MATRIX =  cirq.X._unitary_()

get_v_gate_obj = Build_V_Gate(U_GATE_MATRIX, n_power)
V, V_dag = get_v_gate_obj.Get_V_gate_matrices()

V_gate_DAGGER = My_V_gate(V, V_dag, dagger_gate=True)
V_gate = My_V_gate(V, V_dag, dagger_gate=False)

circuit = V_gate_DAGGER.on(cirq.LineQubit(2))
cirq.Circuit(circuit)

## setup n-control-U ###
list_of_control_qubits = cirq.LineQubit.range(3)
list_control_vals=[0,1,0]
U_qubit = cirq.LineQubit(3)

xx = n_control_U(V_gate, V_gate_DAGGER, list_of_control_qubits, list_control_vals, U_qubit)

Q_circuit = cirq.Circuit(cirq.decompose_once(
        (xx(*cirq.LineQubit.range(xx.num_qubits())))))

print(cirq.Circuit((xx(*cirq.LineQubit.range(xx.num_qubits())))))

Q_circuit
```

```python

```

```python
op = cirq.X.controlled(num_controls=3, control_values=[0, 1, 0]).on(*[cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(2)], cirq.LineQubit(3))
print(cirq.Circuit(op))
np.isclose(cirq.Circuit(op).unitary(), Q_circuit.unitary(), atol=1e-9)
```

```python
### setup V gate ##
n_control_qubits=2
n_power = 2**(n_control_qubits-1)

theta= np.pi/2

U_GATE_MATRIX =  cirq.X._unitary_()

get_v_gate_obj = Build_V_Gate(U_GATE_MATRIX, n_power)
V, V_dag = get_v_gate_obj.Get_V_gate_matrices()

V_gate_DAGGER = My_V_gate(V, V_dag, dagger_gate=True)
V_gate = My_V_gate(V, V_dag, dagger_gate=False)

circuit = V_gate_DAGGER.on(cirq.LineQubit(2))
cirq.Circuit(circuit)

## setup n-control-U ###
list_of_control_qubits = cirq.LineQubit.range(2)
list_control_vals=[0,1]
U_qubit = cirq.LineQubit(2)

xx = n_control_U(V_gate, V_gate_DAGGER, list_of_control_qubits, list_control_vals, U_qubit)

Q_circuit = cirq.Circuit(cirq.decompose_once(
        (xx(*cirq.LineQubit.range(xx.num_qubits())))))

print(cirq.Circuit((xx(*cirq.LineQubit.range(xx.num_qubits())))))

Q_circuit
```

```python
op = cirq.X.controlled(num_controls=2, control_values=[0, 1]).on(*[cirq.LineQubit(0), cirq.LineQubit(1)], cirq.LineQubit(2))
print(cirq.Circuit(op))
np.isclose(cirq.Circuit(op).unitary(), Q_circuit.unitary(), atol=1e-6)
```

```python

```

```python

```

```python

```

```python

```

```python
cirq.X.controlled(num_controls=2, control_values=[0,1]).on(
                            *list(cirq.LineQubit.range(3)))
```

```python
op = cirq.X.controlled(num_controls=3, control_values=[0, 0, 1]).on(*[cirq.LineQubit(1), cirq.LineQubit(2), cirq.LineQubit(3)], cirq.LineQubit(4))
print(cirq.Circuit(op))

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python
def int_to_Gray(num, n_qubits):
    # https://en.wikipedia.org/wiki/Gray_code

    # print(np.binary_repr(num, n_qubits)) # standard binary form!

    # The operator >> is shift right. The operator ^ is exclusive or
    gray_int = num^(num>>1)

    return np.binary_repr(gray_int,n_qubits)

### example... note that grey code reversed as indexing from left to right: [0,1,-->, N-1]
for i in range(2**3):
    print(int_to_Gray(i, 3)[::-1])

int_to_Gray(6, 4)
```

    000
    100
    110
    010
    011
    111
    101
    001

    '0101'

```python
def check_binary_str_parity(binary_str):
    """
    Returns 0 for EVEN parity
    Returns 1 for ODD parity
    """
    parity = sum(map(int,binary_str))%2

    return parity

check_binary_str_parity('0101')
```

    0

```python

```

```python
class My_U_Gate(cirq.SingleQubitGate):
    """
    Description

    Args:
        theta (float): angle to rotate by in radians.
        number_control_qubits (int): number of control qubits
    """

    def __init__(self, theta):
        self.theta = theta
    def _unitary_(self):
        Unitary_Matrix = np.array([
                    [np.cos(self.theta), np.sin(self.theta)],
                    [np.sin(self.theta), -1* np.cos(self.theta)]
                ])
        return Unitary_Matrix
    def num_qubits(self):
        return 1

    def _circuit_diagram_info_(self,args):
        # return cirq.CircuitDiagramInfo(
        #     wire_symbols=tuple([*['@' for _ in range(self.num_control_qubits-1)],' U = {} rad '.format(self.theta.__round__(4))]),exponent=1)
        return ' U = {} rad '.format(self.theta.__round__(4))

    def __str__(self):
        return ' U = {} rad '.format(self.theta.__round__(4))

    def __repr__(self):
        return ' U_arb_state_prep'
```

```python
class My_V_gate(cirq.SingleQubitGate):
    """
    Description

    Args:
        theta (float): angle to rotate by in radians.
        number_control_qubits (int): number of control qubits
    """

    def __init__(self, V_mat, V_dag_mat, dagger_gate = False):
        self.V_mat = V_mat
        self.V_dag_mat = V_dag_mat
        self.dagger_gate = dagger_gate
    def _unitary_(self):
        if self.dagger_gate:
            return self.V_dag_mat
        else:
            return self.V_mat

    def num_qubits(self):
        return 1

    def _circuit_diagram_info_(self,args):
        if self.dagger_gate:
            return 'V^{†}'
        else:
            return 'V'

    def __str__(self):
        if self.dagger_gate:
            return 'V^{†}'
        else:
            return 'V'

    def __repr__(self):
        return self.__str__()
```

```python
from sympy import *
from sympy.physics.quantum import Dagger

# NOTE pg 17 of Elementary gates for quantum information
class n_control_U(cirq.Gate):
    """
    """

    def __init__(self, list_of_control_qubits, list_control_vals, U_qubit, U_cirq_gate, n_control_qubits):
        self.U_qubit = U_qubit
        self.U_cirq_gate = U_cirq_gate

        if len(list_of_control_qubits)!=len(list_control_vals):
            raise ValueError('incorrect qubit control bits or incorrect number of control qubits')

        self.list_of_control_qubits = list_of_control_qubits
        self.list_control_vals = list_control_vals

        self.n_ancilla=len(list_of_control_qubits)
        self.D = None
        self.n_root = 2**(n_control_qubits-1)
        self.n_control_qubits = n_control_qubits

        self.V_mat = None
        self.V_dag_mat = None

    def _diagonalise_U(self):

        # find diagonal matrix:
        U_matrix = Matrix(self.U_cirq_gate._unitary_())
        self.S, self.D = U_matrix.diagonalize()
        self.S_inv = self.S**-1
        # where U = S D S^{-1}

        if not np.allclose(np.array(self.S*(self.D*self.S_inv), complex), self.U_cirq_gate._unitary_()):
            raise ValueError('U != SDS-1')

    def Get_V_gate_matrices(self, check=True):

        if self.D is None:
            self._diagonalise_U()
        D_nth_root = self.D**(1/self.n_root)

        V_mat = self.S * D_nth_root * self.S_inv
        V_dag_mat = Dagger(V_mat)

        self.V_mat = np.array(V_mat, complex)
        self.V_dag_mat = np.array(V_dag_mat, complex)

        if check:
            V_power_n = reduce(np.matmul, [self.V_mat for _ in range(self.n_root)])
            if not np.allclose(V_power_n, self.U_cirq_gate._unitary_()):
                raise ValueError('V^{n} != U')

    def flip_control_to_zero(self):
        for index, control_qubit in enumerate(self.list_of_control_qubits):
            if self.list_control_vals[index]==0:
                yield cirq.X.on(control_qubit)

    def _get_gray_control_lists(self):

        grey_cntrl_bit_lists=[]
        n_ancilla = len(self.list_of_control_qubits)
        for grey_index in range(1, 2**n_ancilla):

            gray_control_str = int_to_Gray(grey_index, n_ancilla)[::-1] # note reversing order
            control_list = list(map(int,gray_control_str))
            parity = check_binary_str_parity(gray_control_str)

            grey_cntrl_bit_lists.append((control_list, parity))
        return grey_cntrl_bit_lists

    def _decompose_(self, qubits):
        if (self.V_mat is None) or (self.V_dag_mat is None):
            self.Get_V_gate_matrices()

        V_gate_DAGGER = My_V_gate(self.V_mat, self.V_dag_mat, dagger_gate=True)
        V_gate = My_V_gate(self.V_mat, self.V_dag_mat, dagger_gate=False)

        ## flip if controlled on zero
        X_flip = self.flip_control_to_zero()
        yield X_flip

        ## perform controlled gate
        n_ancilla = len(self.list_of_control_qubits)

        grey_control_lists = self._get_gray_control_lists()

        for control_index, binary_control_tuple in enumerate(grey_control_lists):

            binary_control_seq, parity = binary_control_tuple
            control_indices = np.where(np.array(binary_control_seq)==1)[0]
            control_qubit = control_indices[-1]

            if parity==1:
                gate = V_gate.controlled(num_controls=1, control_values=[1]).on(self.list_of_control_qubits[control_qubit], self.U_qubit)
            else:
                gate = V_gate_DAGGER.controlled(num_controls=1, control_values=[1]).on(self.list_of_control_qubits[control_qubit], self.U_qubit)

            if control_index==0:
                yield gate
            else:
                for c_index in range(len(control_indices[:-1])):
                    yield cirq.CNOT(self.list_of_control_qubits[control_indices[c_index]], self.list_of_control_qubits[control_indices[c_index+1]])
                yield gate
                for c_index in list(range(len(control_indices[:-1])))[::-1]:
                    yield cirq.CNOT(self.list_of_control_qubits[control_indices[c_index]], self.list_of_control_qubits[control_indices[c_index+1]])

        ## unflip if controlled on zero
        X_flip = self.flip_control_to_zero()
        yield X_flip

    def _circuit_diagram_info_(self, args):
#         return cirq.protocols.CircuitDiagramInfo(
#             wire_symbols=tuple([*['@' if bit==1 else '(0)' for bit in self.list_control_vals],'U']),
#             exponent=1)
        return cirq.protocols.CircuitDiagramInfo(
            wire_symbols=tuple([*['@' if bit==1 else '(0)' for bit in self.list_control_vals],self.U_cirq_gate.__str__()]),
            exponent=1)

    def num_qubits(self):
        return len(self.list_of_control_qubits) + 1 #(+1 for U_qubit)

    def check_Gate_gate_decomposition(self, tolerance=1e-9):
        """
        function compares single and two qubit gate construction of n-controlled-U
        against perfect n-controlled-U gate

        tolerance is how close unitary matrices are required

        """

        # decomposed into single and two qubit gates
        decomposed = self._decompose_(None)
        n_controlled_U_quantum_Circuit = cirq.Circuit(decomposed)

#         print(n_controlled_U_quantum_Circuit)

        # perfect gate
        perfect_circuit_obj = self.U_cirq_gate.controlled(num_controls=self.n_control_qubits, control_values=self.list_control_vals).on(
                            *self.list_of_control_qubits, self.U_qubit)

        perfect_circuit = cirq.Circuit(perfect_circuit_obj)

#         print(perfect_circuit)

        if not np.allclose(n_controlled_U_quantum_Circuit.unitary(), perfect_circuit.unitary(), atol=tolerance):
            raise ValueError('V^{n} != U')
        else:
#             print('Correct decomposition')
            return True

```

```python

```

```python

```

```python
## setup
n_control_qubits=2
theta= np.pi/4
U_gate = My_U_Gate(theta)

list_of_control_qubits = cirq.LineQubit.range(2)
list_control_vals=[0,1]
U_qubit = cirq.LineQubit(2)

xx = n_control_U(list_of_control_qubits, list_control_vals, U_qubit, U_gate, n_control_qubits)

Q_circuit = cirq.Circuit(cirq.decompose_once(
        (xx(*cirq.LineQubit.range(xx.num_qubits())))))

# NOT decomposing:
print(cirq.Circuit((xx(*cirq.LineQubit.range(xx.num_qubits())))))

# decomposing
Q_circuit
```

    0: ───(0)────────────────
          │
    1: ───@──────────────────
          │
    2: ─── U = 0.7854 rad ───

<pre style="overflow: auto; white-space: pre;">0: ───X───@───@───────────@───X───
          │   │           │
1: ───────┼───X───@───────X───@───
          │       │           │
2: ───────V───────V^{†}───────V───</pre>

```python
xx.check_Gate_gate_decomposition(tolerance=1e-15)
```

    True

```python
U_single_qubit = My_U_Gate(theta)
perfect_circuit_obj = U_single_qubit.controlled(num_controls=n_control_qubits, control_values=list_control_vals).on(
                            *list_of_control_qubits, U_qubit)
perfect_circuit = cirq.Circuit(perfect_circuit_obj)
perfect_circuit
```

<pre style="overflow: auto; white-space: pre;">0: ───(0)────────────────
      │
1: ───@──────────────────
      │
2: ─── U = 0.7854 rad ───</pre>

```python
np.allclose(Q_circuit.unitary(), perfect_circuit.unitary(), atol=1e-6)
```

    True

```python

```

quchem_examples/Multi control gate decomposition.ipynb


## Simulating overdispersion using the `Spp` class

In the present notebook, which is based on ex07-Spp, we show how overdispersion can be simulated using the `pyross.stochastic.Spp` class. We consider the same SIR model as used in ex07-Spp, and compare several options for the overdispersion.

#### Overdispersion

Without overdispersion, during one infection event precisely one individual gets infected. With overdispersion, on average $\lambda $ individuals get infected. The value $\lambda $ is set via the parameter "overdispersion", as will be illustrated further below.

With overdispersion, the number $K$ of individuals that get infected in one infection event is a random variable.
The overdispersion mode is set by the parameter "overdispersion_mode", and there are currently two distributions available in PyRoss:

##### 1. Poisson-distributed overdispersion (overdispersion_mode = 0, default)

Here, the number $K$ of newly infected individuals in one infection event is distributed according to a (shifted) Poisson distribution,
\begin{equation}
    K = X + 1,
\end{equation}
where $X$ is distributed according to
\begin{equation}
P(X=n) = (\lambda-1)^{n} \frac{e^{-(\lambda-1)}}{n!}, \qquad n = \{0,1,2, ...\}.
\end{equation}
Note that $\langle K \rangle = \lambda$, and that $K \geq 1$, so that at least one person is infected in an infection event. By default, $\lambda = 1$, which corresponds to no overdispersion.

##### 2. Delta-peak overdispersion (overdispersion_mode = 1)

In this mode, the input $\lambda$ is rounded to next smallest integer, and $K$ is distributed according to a Kronecker delta,
\begin{equation}
P(K=k) = \delta_{k,\lambda},
\end{equation}
where $\lambda$ is understood to be an integer.

```python
%matplotlib inline
import numpy as np
import pyross
import matplotlib.pyplot as plt
#from matplotlib import rc; rc('text', usetex=True)
```

### The SIR model

Below you will find the model-specification dictionary for the SIR model with some constant injection into the S class:

```python
model_spec = {
    "classes" : ["S", "I"],

    "S" : {
        "constant"  : [ ["k"] ],
        "infection" : [ ["I", "-beta", "overdispersion"] ]
    },

    "I" : {
        "linear"    : [ ["I", "-gamma"] ],
        "infection" : [ ["I", "beta", "overdispersion"] ]
    }
}

```

## Example 1: Small population

We now consider an SIR model with a rather small population, and compare

1. Gillespie simulations using Poisson-distributed overdispersion, with
2. Gillespie simulations using Delta-peak-distributed overdispersion.

```python
overdispersion = 10

parameters = {
    'beta' : 0.1,
    'gamma' : 0.1,
    'k' : 1,
    'overdispersion': overdispersion,
    'overdispersion_mode': 0, # Poisson distributed
}
```

We start by plotting a single example trajectory for Poisson-distributed overdispersion:

```python
M = 3
Ni = 1000*np.ones(M)
N = np.sum(Ni)

# Initial conditions as an array
x0 = np.array([
    999, 1000, 1000, # S
    1,   0,    0,    # I
])

# Or initial conditions as a dictionary
I0 = [20, 20, 20]
S0 = [n-20 for n in Ni]

x0 = {
    'S' : S0,
    'I' : I0
}

CM = np.array([
    [1,   0.5, 0.1],
    [0.5, 1,   0.5],
    [0.1, 0.5, 1  ]
], dtype=float)

def contactMatrix(t):
    return CM

# duration of simulation and data file
Tf = 160;  Nf=Tf+1;

model = pyross.stochastic.Spp(model_spec, parameters, M, Ni)
# simulate model
data = model.simulate(x0, contactMatrix, Tf, Nf)
```

```python
# plot the data and obtain the epidemic curve
S = np.sum(model.model_class_data('S', data), axis=1)
I = np.sum(model.model_class_data('I', data), axis=1)
R = np.sum(model.model_class_data('R', data), axis=1)
t = data['t']

fig = plt.figure(num=None, figsize=(10, 8), dpi=80, facecolor='w', edgecolor='k')
plt.rcParams.update({'font.size': 22})

plt.fill_between(t, 0, S/N, color="#348ABD", alpha=0.3)
plt.plot(t, S, '-', color="#348ABD", label='$S$', lw=4)

plt.fill_between(t, 0, I/N, color='#A60628', alpha=0.3)
plt.plot(t, I, '-', color='#A60628', label='$I$', lw=4)

plt.fill_between(t, 0, R/N, color="dimgrey", alpha=0.3)
plt.plot(t, R, '-', color="dimgrey", label='$R$', lw=4)

plt.legend(fontsize=26); plt.grid()
plt.autoscale(enable=True, axis='x', tight=True)
plt.ylabel('Compartment value')
plt.xlabel('Days');
```

Now we run 100 simulations of both Poisson-distributed and Delta-peak-distributed overdispersion, and compare the resulting mean trajectory and corresponding standard deviation.

```python
Ns = 100

trajectories = np.zeros( [Ns,3,Nf] ,
                        dtype=float)

for i in range(Ns):
    print(i,end='\r')
    model = pyross.stochastic.Spp(model_spec, parameters, M, Ni)
    data = model.simulate(x0.copy(), contactMatrix, Tf, Nf)
    S = np.sum(model.model_class_data('S', data), axis=1)
    I = np.sum(model.model_class_data('I', data), axis=1)
    R = np.sum(model.model_class_data('R', data), axis=1)
    trajectories[i,0] = S.copy()
    trajectories[i,1] = I.copy()
    trajectories[i,2] = R.copy()
    #print(R)

mean_traj = np.mean(trajectories,axis=0)
std_traj = np.sqrt( np.var(trajectories,axis=0) )
```

    99

```python

parameters_2 = {
    'beta' : 0.1,
    'gamma' : 0.1,
    'k' : 1,
    'overdispersion':overdispersion,
    'overdispersion_mode': 1, # Delta-peak
}

trajectories_2 = np.zeros( [Ns,M,Nf] ,
                        dtype=float)

for i in range(Ns):
    print(i,end='\r')
    model_2 = pyross.stochastic.Spp(model_spec, parameters_2, M, Ni)
    data = model_2.simulate(x0.copy(), contactMatrix, Tf, Nf,
                             method='tau-leaping',
                           nc=2*overdispersion,
                           epsilon = 0.03/overdispersion)
    S = np.sum(model_2.model_class_data('S', data), axis=1)
    I = np.sum(model_2.model_class_data('I', data), axis=1)
    R = np.sum(model_2.model_class_data('R', data), axis=1)
    trajectories_2[i,0] = S.copy()
    trajectories_2[i,1] = I.copy()
    trajectories_2[i,2] = R.copy()

mean_traj_2 = np.mean(trajectories_2,axis=0)
std_traj_2 = np.sqrt( np.var(trajectories_2,axis=0) )
```

    99

```python
fig,ax = plt.subplots(1,1,figsize=(15,8))

colors=['dodgerblue','orange','limegreen']
colors_2=['cyan','blue','red']
labels=['S','I','R']

plot_var = True

for i,e in enumerate(mean_traj):
    if plot_var:
        ax.plot(t,(e - std_traj[i])/N,
                lw=0.5,
                color=colors[i])
        ax.plot(t,(e + std_traj[i])/N,
                lw=0.5,
                color=colors[i])
    ax.plot(t,e/N,
            lw=2,
            color=colors[i],
           label='Poisson, ' + str(labels[i]))

for i,e in enumerate(mean_traj_2):
    if plot_var:
        ax.plot(t,(e - std_traj_2[i])/N,
                lw=1,
                ls='dotted',
                color=colors_2[i])
        ax.plot(t,(e + std_traj_2[i])/N,
                lw=1,
                ls='dotted',
                color=colors_2[i])
    ax.plot(t,e/N,ls='--',
            lw=2,
            color=colors_2[i],
            label='Delta-peak, ' + str(labels[i]))

ax.set_xlim(np.min(t),np.max(t))
ax.set_xlabel(r'Days')
ax.set_ylabel(r'Compartment fraction')
ax.legend(fontsize=15)
plt.show()
plt.close(fig)
```

## Example 2: Large population and large overdispersion

We now consider an SIR model with a larger population, and compare

1. Gillespie simulations using Delta-peak-distributed overdispersion, with
2. Tau-leaping simulations using Poisson-distributed overdispersion.

Note that for tau-leaping, both the parameters "nc" and "epsilon" need to be adapted, to account for the fact that overdispersion can lead to faster variation of populations.

We start with an example trajectory from a Gillespie simulation using Delta-peak-distributed overdispersion:

```python
M = 3
Ni = 100000*np.ones(M)
N = np.sum(Ni)

# Or initial conditions as a dictionary
I0 = [2000, 2000, 2000]
S0 = [n-2000 for n in Ni]

x0 = {
    'S' : S0,
    'I' : I0
}

CM = np.array([
    [1,   0.5, 0.1],
    [0.5, 1,   0.5],
    [0.1, 0.5, 1  ]
], dtype=float)

overdispersion = 300

parameters = {
    'beta' : 0.1,
    'gamma' : 0.1,
    'k' : 1,
    'overdispersion': overdispersion,
    'overdispersion_mode': 1,
}

def contactMatrix(t):
    return CM

# duration of simulation and data file
Tf = 50;  Nf=Tf+1;

model = pyross.stochastic.Spp(model_spec, parameters, M, Ni)
# simulate model
data = model.simulate(x0, contactMatrix, Tf, Nf)
```

```python
# plot the data and obtain the epidemic curve
S = np.sum(model.model_class_data('S', data), axis=1)
I = np.sum(model.model_class_data('I', data), axis=1)
R = np.sum(model.model_class_data('R', data), axis=1)
t = data['t']

fig = plt.figure(num=None, figsize=(10, 8), dpi=80, facecolor='w', edgecolor='k')
plt.rcParams.update({'font.size': 22})

plt.fill_between(t, 0, S/N, color="#348ABD", alpha=0.3)
plt.plot(t, S, '-', color="#348ABD", label='$S$', lw=4)

plt.fill_between(t, 0, I/N, color='#A60628', alpha=0.3)
plt.plot(t, I, '-', color='#A60628', label='$I$', lw=4)

plt.fill_between(t, 0, R/N, color="dimgrey", alpha=0.3)
plt.plot(t, R, '-', color="dimgrey", label='$R$', lw=4)

plt.legend(fontsize=26); plt.grid()
plt.autoscale(enable=True, axis='x', tight=True)
plt.ylabel('Compartment value')
plt.xlabel('Days');
```

```python
Ns = 100

trajectories = np.zeros( [Ns,3,Nf] ,
                        dtype=float)

for i in range(Ns):
    print(i,end='\r')
    model = pyross.stochastic.Spp(model_spec, parameters, M, Ni)
    data = model.simulate(x0.copy(), contactMatrix, Tf, Nf)
    S = np.sum(model.model_class_data('S', data), axis=1)
    I = np.sum(model.model_class_data('I', data), axis=1)
    R = np.sum(model.model_class_data('R', data), axis=1)
    trajectories[i,0] = S.copy()
    trajectories[i,1] = I.copy()
    trajectories[i,2] = R.copy()

mean_traj = np.mean(trajectories,axis=0)
std_traj = np.sqrt( np.var(trajectories,axis=0) )
```

    99

```python
parameters_2 = {
    'beta' : 0.1,
    'gamma' : 0.1,
    'k' : 1,
    'overdispersion': overdispersion,
    'overdispersion_mode': 0,
}

trajectories_2 = np.zeros( [Ns,M,Nf] ,
                        dtype=float)

for i in range(Ns):
    print(i,end='\r')
    model_2 = pyross.stochastic.Spp(model_spec, parameters_2, M, Ni,
                                   )
    data = model_2.simulate(x0.copy(), contactMatrix, Tf, Nf,
                            method='tau-leaping',
                           nc=min( 3*overdispersion, 30 ), # definitely switch to gillespie when a
                                                        # population is comparable to the overdispersion
                           epsilon = 0.03/overdispersion)  # decrease the timestep by the overdispersion
    S = np.sum(model_2.model_class_data('S', data), axis=1)
    I = np.sum(model_2.model_class_data('I', data), axis=1)
    R = np.sum(model_2.model_class_data('R', data), axis=1)
    trajectories_2[i,0] = S.copy()
    trajectories_2[i,1] = I.copy()
    trajectories_2[i,2] = R.copy()

mean_traj_2 = np.mean(trajectories_2,axis=0)
std_traj_2 = np.sqrt( np.var(trajectories_2,axis=0) )
```

    99

```python
fig,ax = plt.subplots(1,1,figsize=(15,8))

colors=['dodgerblue','orange','limegreen']
colors_2=['cyan','blue','red']
labels=['S','I','R']

plot_var = True

for i,e in enumerate(mean_traj):
    if plot_var:
        ax.plot(t,(e - std_traj[i])/N,
                lw=0.5,
                color=colors[i])
        ax.plot(t,(e + std_traj[i])/N,
                lw=0.5,
                color=colors[i])
    ax.plot(t,e/N,
            lw=2,
            color=colors[i],
           label='Poisson, ' + str(labels[i]))

for i,e in enumerate(mean_traj_2):
    if plot_var:
        ax.plot(t,(e - std_traj_2[i])/N,
                lw=1,
                ls='dotted',
                color=colors_2[i])
        ax.plot(t,(e + std_traj_2[i])/N,
                lw=1,
                ls='dotted',
                color=colors_2[i])
    ax.plot(t,e/N,ls='--',
            lw=2,
            color=colors_2[i],
            label='Delta-peak, ' + str(labels[i]))

ax.set_xlim(np.min(t),np.max(t))
ax.set_xlabel(r'Days')
ax.set_ylabel(r'Compartment fraction')
ax.legend(fontsize=15)
plt.show()
plt.close(fig)
```

```python

```


# Realization of Non-Recursive Filters

*This jupyter notebook is part of a [collection of notebooks](../index.ipynb) on various topics of Digital Signal Processing. Please direct questions and suggestions to [Sascha.Spors@uni-rostock.de](mailto:Sascha.Spors@uni-rostock.de).*

## Segmented Convolution

In many applications one of the signals of a convolution is much longer than the other. For instance when filtering a speech signal $x_L[k]$ of length $L$ with a room impulse response $h_N[k]$ of length $N \ll L$. In such cases the [fast convolution](fast_convolution.ipynb), as introduced before, does not bring a benefit since both signals have to be zero-padded to a total length of at least $N+L-1$. Applying the fast convolution may then even be impossible in terms of memory requirements or overall delay. The filtering of a signal which is captured in real-time is also not possible by the fast convolution.

In order to overcome these limitations, various techniques have been developed that perform the filtering on limited portions of the signals. These portions are known as partitions, segments or blocks. The respective algorithms are termed as *segmented* or *block-based* algorithms. The following section introduces two techniques for the segmented convolution of signals. The basic concept of these is to divide the convolution $y[k] = x_L[k] * h_N[k]$ into multiple convolutions operating on (overlapping) segments of the signal $x_L[k]$.

### Overlap-Add Algorithm

The [overlap-add algorithm](https://en.wikipedia.org/wiki/Overlap%E2%80%93add_method) is based on splitting the signal $x_L[k]$ into non-overlapping segments $x_p[k]$ of length $P$

\begin{equation}
x_L[k] = \sum_{p = 0}^{L/P - 1} x_p[k - p \cdot P]
\end{equation}

where the segments $x_p[k]$ are defined as

\begin{equation}
x_p[k] = \begin{cases} x_L[k + p \cdot P] & \text{ for } k=0,1,\dots,P-1 \\ 0 & \text{ otherwise} \end{cases}
\end{equation}

Note that $x_L[k]$ might have to be zero-padded so that its total length is a multiple of the segment length $P$. Introducing the segmentation of $x_L[k]$ into the convolution yields

\begin{align}
y[k] &= x_L[k] * h_N[k] \\
&= \sum_{p = 0}^{L/P - 1} x_p[k - p \cdot P] * h_N[k] \\
&= \sum_{p = 0}^{L/P - 1} y_p[k - p \cdot P]
\end{align}

where $y_p[k] = x_p[k] * h_N[k]$. This result states that the convolution of $x_L[k] * h_N[k]$ can be split into a series of convolutions $y_p[k]$ operating on the samples of one segment only. The length of $y_p[k]$ is $N+P-1$. The result of the overall convolution is given by summing up the results from the segments shifted by multiples of the segment length $P$. This can be interpreted as an overlapped superposition of the results from the segments, as illustrated in the following diagram

The overall procedure is denoted by the name *overlap-add* technique. The convolutions $y_p[k] = x_p[k] * h_N[k]$ can be realized efficiently by the [fast convolution](fast_convolution.ipynb) using zero-padding and fast Fourier transformations (FFTs) of length $M \geq P+N-1$.

A drawback of the overlap-add technique is that the next input segment is required to compute the result for the actual segment of the output. For real-time applications this introduces an algorithmic delay of one segment.

#### Example

The following example illustrates the overlap-add algorithm by showing the (convolved) segments and the overall result.

```python
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig

L = 64  # length of input signal
N = 8  # length of impulse response
P = 16  # length of segments

# generate input signal
x = sig.triang(L)
# generate impulse response
h = sig.triang(N)

# overlap-add convolution
xp = np.zeros((L//P, P))
yp = np.zeros((L//P, N+P-1))
y = np.zeros(L+P-1)
for p in range(L//P):
    xp[p, :] = x[p*P:(p+1)*P]
    yp[p, :] = np.convolve(xp[p, :], h, mode='full')
    y[p*P:(p+1)*P+N-1] += yp[p, :]
y = y[0:N+L]

# plot signals
plt.figure(figsize=(10, 2))

plt.subplot(121)
plt.stem(x)
for n in np.arange(L//P)[::2]:
    plt.axvspan(n*P, (n+1)*P-1, facecolor='g', alpha=0.5)
plt.title(r'Signal $x[k]$ and segments')
plt.xlabel(r'$k$')
plt.ylabel(r'$x[k]$')
plt.axis([0, L, 0, 1])

plt.subplot(122)
plt.stem(h)
plt.title(r'Impulse response $h[k]$')
plt.xlabel(r'$k$')
plt.ylabel(r'$h[k]$')
plt.axis([0, L, 0, 1])

for p in np.arange(L//P):
    plt.figure(figsize=(10, 2))

    plt.stem(np.concatenate((np.zeros(p*P), yp[p, :])))
    plt.title(r'Result of segment $p=%d$' % (p))
    plt.xlabel(r'$k$')
    plt.ylabel(r'$y_%d[k - %d P]$' % (p, p))
    plt.axis([0, L+P, 0, 4])

plt.figure(figsize=(10, 2))
plt.stem(y)
plt.title(r'Result $y[k] = x[k] * h[k]$ by overlap-add algorithm')
plt.xlabel(r'$k$')
plt.ylabel(r'$y[k]$')
plt.axis([0, L+P, 0, 4])
```

**Exercises**

* Change the length `N` of the impulse response and the length `P` of the segments. What changes?
* What influence have these two lengths on the numerical complexity of the overlap-add algorithm?

Solution: The parameters `N` and `P` influence the overlap in the output and the total number of segments. The number of overlapping samples of two consecutive output segments $y_p[k]$ and $y_{p+1}[k]$ is given as $N-1$, and the total number of segments as $\frac{L}{P}$. The segmented convolution requires $\frac{L}{P}$ linear convolutions of length $P+N-1$ each. The numerical complexity is mainly determined by the overall number of multiplications which is given as $\frac{L}{P} (P+N-1)^2$. For fixed $L$ and $N$, the optimum segment length is computed by finding the minimum in terms of multiplications. It is given as $P=N-1$.

### Overlap-Save Algorithm

The [overlap-save](https://en.wikipedia.org/wiki/Overlap%E2%80%93save_method) algorithm, also known as *overlap-discard algorithm*, follows a different strategy as the overlap-add technique introduced above. It is based on an overlapping segmentation of the input $x_L[k]$ and application of the periodic convolution for the individual segments.

Lets take a closer look at the result of the periodic convolution $x_p[k] \circledast_P h_N[k]$, where $x_p[k]$ denotes a segment of length $P$ of the input signal and $h_N[k]$ the impulse response of length $N$. The result of a linear convolution $x_p[k]* h_N[k]$ would be of length $P + N -1$. The result of the periodic convolution of period $P$ for $P > N$ would suffer from a circular shift (time aliasing) and superposition of the last $N-1$ samples to the beginning. Hence, the first $N-1$ samples are not equal to the result of the linear convolution. However, the remaining $P- N + 1$ do so.

This motivates to split the input signal $x_L[k]$ into overlapping segments of length $P$ where the $p$-th segment overlaps its preceding $(p-1)$-th segment by $N-1$ samples

\begin{equation}
x_p[k] = \begin{cases}
x_L[k + p \cdot (P-N+1) - (N-1)] & \text{ for } k=0,1, \dots, P-1 \\
0 & \text{ otherwise}
\end{cases}
\end{equation}

The part of the circular convolution $x_p[k] \circledast_P h_N[k]$ of one segment $x_p[k]$ with the impulse response $h_N[k]$ that is equal to the linear convolution of both is given as

\begin{equation}
y_p[k] = \begin{cases}
x_p[k] \circledast_P h_N[k] & \text{ for } k=N-1, N, \dots, P-1 \\
0 & \text{ otherwise}
\end{cases}
\end{equation}

The output $y[k]$ is simply the concatenation of the $y_p[k]$

\begin{equation}
y[k] = \sum_{p=0}^{\frac{L}{P-N+1} - 1} y_p[k - p \cdot (P-N+1) + (N-1)]
\end{equation}

The overlap-save algorithm is illustrated in the following diagram

For the first segment $x_0[k]$, $N-1$ zeros have to be appended to the beginning of the input signal $x_L[k]$ for the overlapped segmentation. From the result of the periodic convolution $x_p[k] \circledast_P h_N[k]$ the first $N-1$ samples are discarded, the remaining $P - N + 1$ are copied to the output $y[k]$. This is indicated by the alternative notation *overlap-discard* used for the technique. The periodic convolution can be realized efficiently by a FFT/IFFT of length $P$.

#### Example

The following example illustrates the overlap-save algorithm by showing the results of the periodic convolutions of the segments. The discarded parts are indicated by the red background.

```python
L = 64  # length of input signal
N = 8  # length of impulse response
P = 24  # length of segments

# generate input signal
x = sig.triang(L)
# generate impulse response
h = sig.triang(N)

# overlap-save convolution
nseg = (L+N-1)//(P-N+1) + 1
x = np.concatenate((np.zeros(N-1), x, np.zeros(P)))
xp = np.zeros((nseg, P))
yp = np.zeros((nseg, P))
y = np.zeros(nseg*(P-N+1))

for p in range(nseg):
    xp[p, :] = x[p*(P-N+1):p*(P-N+1)+P]
    yp[p, :] = np.fft.irfft(np.fft.rfft(xp[p, :]) * np.fft.rfft(h, P))
    y[p*(P-N+1):p*(P-N+1)+P-N+1] = yp[p, N-1:]
y = y[0:N+L]

plt.figure(figsize=(10, 2))

plt.subplot(121)
plt.stem(x[N-1:])
plt.title(r'Signal $x[k]$')
plt.xlabel(r'$k$')
plt.ylabel(r'$x[k]$')
plt.axis([0, L, 0, 1])

plt.subplot(122)
plt.stem(h)
plt.title(r'Impulse response $h[k]$')
plt.xlabel(r'$k$')
plt.ylabel(r'$h[k]$')
plt.axis([0, L, 0, 1])

for p in np.arange(nseg):
    plt.figure(figsize=(10, 2))
    plt.stem(yp[p, :])
    plt.axvspan(0, N-1+.5, facecolor='r', alpha=0.5)
    plt.title(
        r'Result of periodic convolution of $x_%d[k]$ and $h_N[k]$' % (p))
    plt.xlabel(r'$k$')
    plt.axis([0, L+P, 0, 4])

plt.figure(figsize=(10, 2))
plt.stem(y)
plt.title(r'Result $y[k] = x[k] * h[k]$')
plt.xlabel(r'$k$')
plt.ylabel(r'$y[k]$')
plt.axis([0, L+P, 0, 4])
```

**Exercise**

* Change the length `N` of the impulse response and the length `P` of the segments. What changes?
* How many samples of the output signal $y[k]$ are computed per segment for a particular choice of these two values?
* What would be a good choice for the segment length `P` with respect to the length `N` of the impulse response?

Solution: Decreasing the segment length $P$ or increasing the length of the impulse response $N$ decreases the number of valid output samples per segment which is given as $P-N+1$. The computation of $L$ output samples requires $\frac{L}{P-N+1}$ cyclic convolutions of length $P$ each. Regarding the total number of multiplications, an optimal choice for the segment length is $P = 2 N - 2$.

### Practical Aspects and Extensions

* For both the overlap-add and overlap-save algorithm the length $P$ of the segments influences the lengths of the convolutions, FFTs and the number of output samples per segment. The segment length is often chosen as

    * $P=N$ for overlap-add and
    * $P = 2 N$ for overlap-save.

For both algorithms this requires FFTs of length $2 N$ to compute $P$ output samples. The overlap-add algorithm requires $P$ additional additions per segment in comparison to overlap-save.

* For real-valued signals $x_L[k]$ and impulse responses $h_N[k]$ real-valued FFTs lower the computational complexity significantly. As alternative, the $2 N$ samples in the FFT can be distributed into the real and complex part of a FFT of length $N$ [[Zölzer](../index.ipynb#Literature)].

* The impulse response can be changed in each segment in order to simulate time-variant linear systems. This is often combined with an overlapping computation of the output in order to avoid artifacts due to instationarities.

* For long impulse responses $h_N[k]$ or low-delay applications, algorithms have been developed which base on an additional segmentation of the impulse response. This is known as *partitioned convolution*.

**Copyright**

This notebook is provided as [Open Educational Resource](https://en.wikipedia.org/wiki/Open_educational_resources). Feel free to use the notebook for your own purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT). Please attribute the work as follows: *Sascha Spors, Digital Signal Processing - Lecture notes featuring computational examples*.


Project 0: Inaugural project

Labor Supply Problem

Following labor supply problem is given:

$$
c^*,l^* = log(c) - v \frac{l^{1+\frac{1}{\epsilon}}}{1+\frac{1}{\epsilon}}
\\
x = m + wl - [\tau_0wl+\tau_1 \max(wl-\kappa,0)]\\

c \in [0,x]\\
l \in [0,1]\\

$$

Where:
c is consumption,
l is labor supply,
m is cash-on-hand,
w is the wage rate,
$$t_0$$ is the standard labor income tax
$$t_1$$ is the top bracket labor income tax,
k is the cut-off of top labor income bracket
x is total resources
v scales disutility of labor
E is the Frisch elasticity of labor supply
utility is monotonically increasing in consumption, which implies $$c^* = x$$

Question 1

```python
# All used packages are imported

import numpy as np
import sympy as sm
from scipy import optimize

```

```python
t0 = sm.symbols('t_0')
t1 = sm.symbols('t_1')
```

```python

```

```python
m = 1 #cash-on-hand
v = 10 #disutility of labor
e = 0.3 #elasticity of labor supply
t0 = 0.4 #standard labor income tax
t1 = 0.1 #top bracket labor income tax
k = 0.4 #cut-off for top labor income tax

# Defining utility

def utility(c,v,l,e):
    u = np.log(c) - v*(l**(1+1/e)/(1+1/e))
    return u

# Defining constraint

def constraint(m,w,l,t0,t1,k):
    x = m + w*l - (t0*w*l + t1*np.max(w*l-k,0))
    return x

def consumption(l,w,e,v,t0,t1,k):
    c = constraint(m,w,l,t0,t1,k)
    return -utility(c,v,l,e)

def optimizer(w,e,v,t0,t1,k,m):
    res = optimize.minimize_scalar(
    consumption, method='bounded',
    bounds=(0,1), args=(w,e,v,t0,t1,k))

    labor_star = res.x
    cons_star = constraint(m,w,labor_star,t0,t1,k)
    utility_star = utility(cons_star,v,labor_star,e)

    return labor_star,cons_star,utility_star

labor_star = optimizer(0.5,e,v,t0,t1,k,m)[0]
cons_star = optimizer(0.5,e,v,t0,t1,k,m)[1]
u_star = optimizer(0.5,e,v,t0,t1,k,m)[2]

print('labour supply is:' + str(labor_star))
print('consumption is:' + str(cons_star))
print('utility:' +  str(u_star))
```

    labour supply is:0.31961536193545265
    consumption is:1.119903840483863
    utility:0.09677772523865749

Question 2

```python
import matplotlib.pyplot as plt

plt.style.use('grayscale')

# Plot l_star and c_star with w going from 0.5 to 1.5
# The definitions are defined - the used packages is defined above
N = 10000
w_vector = np.linspace(0.5,1.5,num=N)
c_optimal = np.empty(N)
l_optimal = np.empty(N)

# a loop is generated to test the range of W

for i, w in enumerate(w_vector):
    optimization = optimizer(w,e,v,t0,t1,k,m)
    l_optimal[i]=optimization[0]
    c_optimal[i]=optimization[1]

fig = plt.figure(figsize=(10,4))

# Left plot
axis_left = fig.add_subplot(1,2,1)
axis_left.plot(w_vector,l_optimal)
axis_left.set_title('Optimal labor supply given w')
axis_left.set_xlabel('$w$')
axis_left.set_ylabel('$l$')
axis_left.grid(True)

# Right plot
axis_right = fig.add_subplot(1,2,2)
axis_right.plot(w_vector,c_optimal)
axis_right.set_title('Optimal consumption given w')
axis_right.set_xlabel('$w1$')
axis_right.set_ylabel('$c$')
axis_right.grid(True)

plt.show

```

Question 3

```python
# Calculate the tax revenue

tax_revenue = np.sum( t0 * w_vector * l_optimal + t1 * np.max( w_vector * l_optimal - k ,0 ))
print('Total tax revenue is: ' + str(tax_revenue))

```

    Total tax revenue is: 1775.3896759006836

Question 4

```python
# How does the tax revenue change when e = 0.1?
# New epsilon is defined
e_new = 0.1
l_optimal_e_new = np.empty(N)

# Same loop is used as above but only a new labor
# supply is calculated as consumption isn't included
# in the tax revenue formula
for i, w in enumerate(w_vector):
    optimization = optimizer(w,e_new,v,t0,t1,k,m)
    l_optimal_e_new[i]=optimization[0]

# then the new tax revenue can be calculated
tax_revenue_e_new = np.sum( t0 * w_vector * l_optimal_e_new + t1 * np.max( w_vector * l_optimal_e_new - k ,0))
print('New total tax revenue: '+str(tax_revenue_e_new))

# Thus the difference in tax revenue can be calucalted as
print('The difference in tax revenue is: '+ str(tax_revenue_e_new-tax_revenue))
```

    New total tax revenue: 3578.900497991557
    The difference in tax revenue is: 1803.5108220908735

Question 5

```python
# Optimize the tax

# Same optimization formula as above
def tax_optimize(t0,t1,k):
    tax_optimal = optimize.minimize_scalar(tax_revenue , method='bounded' , x=[0.1,0.1,0.1])
    t0_optimal = tax_optimal.x
    t1_optimal = tax_optimal.x
    k_optimal = tax_optimal.x
    return t0_optimal, t1_optimal, k_optimal

t0_optimal = tax_optimize(t0,t1,k)[0]
t1_optimal = tax_optimize(t0,t1,k)[1]
k_optimal = tax_optimize(t0,t1,k)[2]

print('Optimal t0 is: ' + str(t0_optimal))
```