## Algorithm for computing natural cubic splines
Cubic splines have polynomial pieces of the form ${\displaystyle P_{i}(x)=a_{i}+b_{i}(x-x_{i})+c_{i}(x-x_{i})^{2}+d_{i}(x-x_{i})^{3}.}$ Given ${\displaystyle k+1}$ coordinates ${\displaystyle (x_{0},y_{0}),(x_{1},y_{1}),\dots ,(x_{k},y_{k}),}$  we find ${\displaystyle k}$ polynomials ${\displaystyle P_{i}(x),}$ which satisfy for ${\displaystyle 1\leq i\leq k-1}$:

* ${\displaystyle P_{0}(x_{0})=y_{0}\quad }$ and ${\displaystyle \quad P_{i-1}(x_{i})=y_{i}=P_{i}(x_{i}),}$
* ${\displaystyle P'_{i-1}(x_{i})=P'_{i}(x_{i}),}$
* ${\displaystyle P''_{i-1}(x_{i})=P''_{i}(x_{i}),}$
* ${\displaystyle P''_{0}(x_{0})=P''_{k-1}(x_{k})=0.}$

One such polynomial ${\displaystyle P_{i}}$ is given by a 5-tuple ${\displaystyle (a,b,c,d,x)}$ where ${\displaystyle a,b,c\,}$  and ${\displaystyle d\,}$, correspond to the coefficients as used above and ${\displaystyle x}$ denotes the variable over the appropriate domain ${\displaystyle [x_{i},x_{i+1}]}$

### Computation of Natural Cubic Splines:
Input: a set of ${\displaystyle k+1}$ coordinates
Output: a spline as a set of polynomial pieces, each represented by a 5-tuple.

1. Create a new array a of size k + 1, and for ${\displaystyle i=0,\ldots ,k}$ set ${\displaystyle a_{i}=y_{i}}$
2. Create new arrays b, d and μ each of size k
3. Create a new array h of size k and for ${\displaystyle i=0,\ldots ,k-1}$ set ${\displaystyle h_{i}=x_{i+1}-x_{i}}$
4. Create a new array α of size k-1 and for ${\displaystyle i=1,\ldots ,k-1}$ set ${\displaystyle \alpha _{i}={\tfrac {3}{h_{i}}}(a_{i+1}-a_{i})-{\tfrac {3}{h_{i-1}}}(a_{i}-a_{i-1})}$
5. Create new arrays c, l, and z each of size ${\displaystyle k+1}$.
6. Set ${\displaystyle l_{0}=1,\;\mu _{0}=z_{0}=0}$
7. For ${\displaystyle i=1,\ldots ,k-1\,}$
    1. Set ${\displaystyle l_{i}=2(x_{i+1}-x_{i-1})-h_{i-1}\mu _{i-1}.}$
    2. Set ${\displaystyle \mu _{i}={\tfrac {h_{i}}{l_{i}}}.}$
    3. Set ${\displaystyle z_{i}={\tfrac {\alpha _{i}-h_{i-1}z_{i-1}}{l_{i}}}.}$
8. Set ${\displaystyle l_{k}=1;z_{k}=c_{k}=0.}$
9. For ${\displaystyle j=k-1,k-2,\ldots ,0}$
    1. Set ${\displaystyle c_{j}=z_{j}-\mu _{j}c_{j+1}}$
    2. Set ${\displaystyle b_{j}={\tfrac {a_{j+1}-a_{j}}{h_{j}}}-{\tfrac {h_{j}(c_{j+1}+2c_{j})}{3}}}$
    3. Set ${\displaystyle d_{j}={\tfrac {c_{j+1}-c_{j}}{3h_{j}}}.}$
10. Create the spline as a new set of polynomials and call it output_set. Populate it with k 5-tuples for the polynomials P.
11. For ${\displaystyle i=0,\ldots ,k-1}$
    1. Set Pi,a = ai
    2. Set Pi,b = bi
    3. Set Pi,c = ci
    4. Set Pi,d = di
    5. Set Pi,x = xi
12. Output output_set

```python
import numpy as np
import matplotlib.pyplot as plt
#x = np.array([1,2,3,4,5,6,7], dtype=float)
#y = np.array([2,1,2,3,2,1,0], dtype=float)
x = np.array([3,5,7], dtype=float)
y = np.array([6,1,4], dtype=float)
# 1
k = x.size - 1
a = y
assert a.size == k + 1

# 2
b = np.full(k, np.inf)
d = np.full(k, np.inf)
mu = np.full(k, np.inf)
assert b.size == d.size == mu.size == k

# 3
h = (np.roll(x,-1) - x)[:k]
assert h.size == k

# 4
alpha = np.full(k, np.inf)
for i in range(k):
    if i == 0:
        continue
    alpha[i] = 3 / h[i] * (a[i+1] - a[i]) - 3 / h[i-1] * (a[i] - a[i-1])

assert alpha.size == k
assert alpha[0] == np.inf

# 5
c = np.full(k+1, np.inf)
l = np.full(k+1, np.inf)
z = np.full(k+1, np.inf)

# 6
l[0] = 1
mu[0] = z[0] = 0

# 7
for i in range(1,k):
    l[i] = 2 * (x[i+1] - x[i-1]) - h[i-1]*mu[i-1]
    mu[i] = h[i] / l[i]
    z[i] = (alpha[i] - h[i-1] * z[i-1])/l[i]

# 8
l[k] = 1
z[k] = c[k] = 0

# 9
for j in range(k-1,-1,-1):
    c[j] = z[j] - mu[j] * c[j+1]
    b[j] = (a[j+1] - a[j]) / h[j] - h[j] * (c[j+1] + 2 * c[j]) / 3
    d[j] = (c[j+1] - c[j]) / (3 * h[j])

# 10
output_set = np.full((5,k), np.inf)

# 11
for i in range(k):
    output_set[0,i] = a[i]
    output_set[1,i] = b[i]
    output_set[2,i] = c[i]
    output_set[3,i] = d[i]
    output_set[4,i] = x[i]

print(output_set)
#output_set = np.array([[6,1],[-7/2, -1/2], [0,3/2],[1/4, -1/4],[3,5]])
#print(output_set)

# plot data
X = np.linspace(start=x[0], stop=x[-1], num=100)
Y = np.full(X.size, np.inf)
idx = 0
for k in X:
    for j in range(len(x)-1):
        if x[j] <= k < x[j+1]:
            i = j
            break
    diff = k - output_set[4,i]
    Y[idx] = output_set[0,i] + output_set[1,i] * (diff) + output_set[2,i] * (diff)**2 + output_set[3,i] * (diff)**3
    idx+=1

plt.figure(1)
plt.plot(X, Y)
plt.plot(x, y, 'o')
plt.show()
```

# B-Spline, Aperiodic

The trick was to either intercept the coefficients, i.e. element 1 of the tuple returned by scipy.interpolate.splrep, and to replace them with the control point values before handing them to scipy.interpolate.splev, or, if you are fine with creating the knots yourself, you can also do without splrep and create the entire tuple yourself.

What is strange about this all, though, is that, according to the manual, splrep returns (and splev expects) a tuple containing, among others, a spline coefficients vector with one coefficient per knot. However, according to all sources I found, a spline is defined as the weighted sum of the N_control_points basis splines, so I would expect the coefficients vector to have as many elements as control points, not knot positions.

In fact, when supplying splrep's result tuple with the coefficients vector modified as described above to scipy.interpolate.splev, it turns out that the first N_control_points of that vector actually are the expected coefficients for the N_control_points basis splines. The last degree + 1 elements of that vector seem to have no effect. I'm stumped as to why it's done this way. If anyone can clarify that, that would be great. Here's the source that generates the above plots:

```python
import numpy as np
import matplotlib.pyplot as plt
import scipy.interpolate as si

# Collected Points
points = [[0, 0], [0, 2], [2, 3], [4, 0], [6, 3], [8, 2], [8, 0]];
points = np.array(points)
x = points[:,0]
y = points[:,1]

# Knots
t = range(len(points))
ipl_t = np.linspace(0.0, len(points) - 1, 100)

# Find B-spline representation of 1-D curve
x_tup = si.splrep(t, x, k=3)
y_tup = si.splrep(t, y, k=3)

x_list = list(x_tup)
xl = x.tolist()
x_list[1] = xl + [0.0, 0.0, 0.0, 0.0]

y_list = list(y_tup)
yl = y.tolist()
y_list[1] = yl + [0.0, 0.0, 0.0, 0.0]

# Given knots and coefficents evaluate the value of the smoothing polynomial and its derivatives
x_i = si.splev(ipl_t, x_list)
y_i = si.splev(ipl_t, y_list)

#==============================================================================
# Plot
#==============================================================================

fig = plt.figure()

ax = fig.add_subplot(231)
plt.plot(t, x, '-og')
plt.plot(ipl_t, x_i, 'r')
plt.xlim([0.0, max(t)])
plt.title('Splined x(t)')

ax = fig.add_subplot(232)
plt.plot(t, y, '-og')
plt.plot(ipl_t, y_i, 'r')
plt.xlim([0.0, max(t)])
plt.title('Splined y(t)')

ax = fig.add_subplot(233)
plt.plot(x, y, '-og')
plt.plot(x_i, y_i, 'r')
plt.xlim([min(x) - 0.3, max(x) + 0.3])
plt.ylim([min(y) - 0.3, max(y) + 0.3])
plt.title('Splined f(x(t), y(t))')

ax = fig.add_subplot(234)
for i in range(7):
    vec = np.zeros(11)
    vec[i] = 1.0
    x_list = list(x_tup)
    x_list[1] = vec.tolist()
    x_i = si.splev(ipl_t, x_list)
    plt.plot(ipl_t, x_i)
plt.xlim([0.0, max(t)])
plt.title('Basis splines')
plt.show()
```

# B-Spline, Periodic

Now in order to create a closed curve like the following, which is another Mathematica example that can be found on the web, Closed b-spline curve

it is necessary to set the per parameter in the splrep call, if you use that. After padding the list of control points with degree+1 values at the end, this seems to work well enough, as the images show.

The next peculiarity here, however, is that the first and the last degree elements in the coefficients vector have no effect, meaning that the control points must be put in the vector starting at the second position, i.e. position 1. Only then are the results ok. For degrees k=4 and k=5, that position even changes to position 2.

Here's the source for generating the closed curve:

```python
import numpy as np
import matplotlib.pyplot as plt
import scipy.interpolate as si

points = [[-2, 2], [0, 1], [-2, 0], [0, -1], [-2, -2], [-4, -4], [2, -4], [4, 0], [2, 4], [-4, 4]]

degree = 3

points = points + points[0:degree + 1]
points = np.array(points)
n_points = len(points)
x = points[:,0]
y = points[:,1]

t = range(len(x))
ipl_t = np.linspace(1.0, len(points) - degree, 1000)

x_tup = si.splrep(t, x, k=degree, per=1)
y_tup = si.splrep(t, y, k=degree, per=1)
x_list = list(x_tup)
xl = x.tolist()
x_list[1] = [0.0] + xl + [0.0, 0.0, 0.0, 0.0]

y_list = list(y_tup)
yl = y.tolist()
y_list[1] = [0.0] + yl + [0.0, 0.0, 0.0, 0.0]

x_i = si.splev(ipl_t, x_list)
y_i = si.splev(ipl_t, y_list)

#==============================================================================
# Plot
#==============================================================================

fig = plt.figure()

ax = fig.add_subplot(231)
plt.plot(t, x, '-og')
plt.plot(ipl_t, x_i, 'r')
plt.xlim([0.0, max(t)])
plt.title('Splined x(t)')

ax = fig.add_subplot(232)
plt.plot(t, y, '-og')
plt.plot(ipl_t, y_i, 'r')
plt.xlim([0.0, max(t)])
plt.title('Splined y(t)')

ax = fig.add_subplot(233)
plt.plot(x, y, '-og')
plt.plot(x_i, y_i, 'r')
plt.xlim([min(x) - 0.3, max(x) + 0.3])
plt.ylim([min(y) - 0.3, max(y) + 0.3])
plt.title('Splined f(x(t), y(t))')

ax = fig.add_subplot(234)
for i in range(n_points - degree - 1):
    vec = np.zeros(11)
    vec[i] = 1.0
    x_list = list(x_tup)
    x_list[1] = vec.tolist()
    x_i = si.splev(ipl_t, x_list)
    plt.plot(ipl_t, x_i)
plt.xlim([0.0, 9.0])
plt.title('Periodic basis splines')

plt.show()
```

# B-Spline, Periodic, Higher Degree

Lastly, there is an effect that I can not explain either, and this is when going to degree 5, there is a small discontinuity that appears in the splined curve, see the upper right panel, which is a close-up of that 'half-moon-with-nose-shape'. The source code that produces this is listed below.

```python
import numpy as np
import matplotlib.pyplot as plt
import scipy.interpolate as si

points = [[-2, 2], [0, 1], [-2, 0], [0, -1], [-2, -2], [-4, -4], [2, -4], [4, 0], [2, 4], [-4, 4]]

degree = 5

points = points + points[0:degree + 1]
points = np.array(points)
n_points = len(points)
x = points[:,0]
y = points[:,1]

t = range(len(x))
ipl_t = np.linspace(1.0, len(points) - degree, 1000)

knots = np.linspace(-degree, len(points), len(points) + degree + 1).tolist()

xl = x.tolist()
coeffs_x = [0.0, 0.0] + xl + [0.0, 0.0, 0.0]

yl = y.tolist()
coeffs_y = [0.0, 0.0] + yl + [0.0, 0.0, 0.0]

x_i = si.splev(ipl_t, (knots, coeffs_x, degree))
y_i = si.splev(ipl_t, (knots, coeffs_y, degree))

#==============================================================================
# Plot
#==============================================================================

fig = plt.figure()

ax = fig.add_subplot(231)
plt.plot(t, x, '-og')
plt.plot(ipl_t, x_i, 'r')
plt.xlim([0.0, max(t)])
plt.title('Splined x(t)')

ax = fig.add_subplot(232)
plt.plot(t, y, '-og')
plt.plot(ipl_t, y_i, 'r')
plt.xlim([0.0, max(t)])
plt.title('Splined y(t)')

ax = fig.add_subplot(233)
plt.plot(x, y, '-og')
plt.plot(x_i, y_i, 'r')
plt.xlim([min(x) - 0.3, max(x) + 0.3])
plt.ylim([min(y) - 0.3, max(y) + 0.3])
plt.title('Splined f(x(t), y(t))')

ax = fig.add_subplot(234)
for i in range(n_points - degree - 1):
    vec = np.zeros(11)
    vec[i] = 1.0
    x_i = si.splev(ipl_t, (knots, vec, degree))
    plt.plot(ipl_t, x_i)
plt.xlim([0.0, 9.0])
plt.title('Periodic basis splines')

plt.show()
```

# Algorithm

$N-Metric = \sum_{x,y} (A(x,y) - B(x,y))^2 + c * Distance $

where

$Distance = \sqrt{(x_1 - x_2)^2 + (y_1 - y_2)^2}$

1. Calculate the optimal affine transform
2. Initialize the control points P
3. Repeat
    1. Calculate N-Metric to the non-rigid transform parameter P:

    \begin{equation}
        {\nabla c = \tfrac{\delta {N(\theta, P^l})}{ \delta{P^l}}}$
    \end{equation}
    2. while ${|{\nabla c}| > \epsilon}$ do

      Recalculate the control points $P = P + \mu \tfrac{\nabla c}{||\nabla c||}$

      Recalculate the gradient vector $\nabla c$

4. Increase the control points resulution by calculating new control points $P^{l+1}$ from $P^l$
5. Increase the image resolution
6. Unitil finest level of resolution is reached.

```python
import itk
import sys

fixedImageFile = sys.argv[1]
movingImageFile = sys.argv[2]
outputImageFile = sys.argv[3]
differenceImageAfterFile = sys.argv[4]
differenceImageBeforeFile = sys.argv[5]

PixelType = itk.ctype('float')

fixedImage = itk.imread(fixedImageFile, PixelType)
movingImage = itk.imread(movingImageFile, PixelType)

Dimension = fixedImage.GetImageDimension()
FixedImageType = itk.Image[PixelType, Dimension]
MovingImageType = itk.Image[PixelType, Dimension]

TransformType = itk.TranslationTransform[itk.D, Dimension]
initialTransform = TransformType.New()

# Optimizer
optimizer = itk.RegularStepGradientDescentOptimizerv4.New(
        LearningRate=4,
        MinimumStepLength=0.001,
        RelaxationFactor=0.5,
        NumberOfIterations=200)

# Metric
metric = itk.MeanSquaresImageToImageMetricv4[
    FixedImageType, MovingImageType].New()

registration = itk.ImageRegistrationMethodv4.New(FixedImage=fixedImage,
        MovingImage=movingImage,
        Metric=metric,
        Optimizer=optimizer,
        InitialTransform=initialTransform)

movingInitialTransform = TransformType.New()
initialParameters = movingInitialTransform.GetParameters()
initialParameters[0] = 0
initialParameters[1] = 0
movingInitialTransform.SetParameters(initialParameters)
registration.SetMovingInitialTransform(movingInitialTransform)

identityTransform = TransformType.New()
identityTransform.SetIdentity()
registration.SetFixedInitialTransform(identityTransform)

registration.SetNumberOfLevels(1)
registration.SetSmoothingSigmasPerLevel([0])
registration.SetShrinkFactorsPerLevel([1])

registration.Update()

transform = registration.GetTransform()
finalParameters = transform.GetParameters()
translationAlongX = finalParameters.GetElement(0)
translationAlongY = finalParameters.GetElement(1)

numberOfIterations = optimizer.GetCurrentIteration()

bestValue = optimizer.GetValue()

print("Result = ")
print(" Translation X = " + str(translationAlongX))
print(" Translation Y = " + str(translationAlongY))
print(" Iterations    = " + str(numberOfIterations))
print(" Metric value  = " + str(bestValue))

CompositeTransformType = itk.CompositeTransform[itk.D, Dimension]
outputCompositeTransform = CompositeTransformType.New()
outputCompositeTransform.AddTransform(movingInitialTransform)
outputCompositeTransform.AddTransform(registration.GetModifiableTransform())

resampler = itk.ResampleImageFilter.New(
                Input=movingImage,
                Transform=outputCompositeTransform,
                UseReferenceImage=True,
                ReferenceImage=fixedImage)
resampler.SetDefaultPixelValue(100)

OutputPixelType = itk.ctype('unsigned char')
OutputImageType = itk.Image[OutputPixelType, Dimension]

caster = itk.CastImageFilter[FixedImageType,
        OutputImageType].New(Input=resampler)

writer = itk.ImageFileWriter.New(Input=caster, FileName=outputImageFile)
writer.SetFileName(outputImageFile)
writer.Update()

difference = itk.SubtractImageFilter.New(Input1=fixedImage,
        Input2=resampler)

intensityRescaler = itk.RescaleIntensityImageFilter[FixedImageType,
        OutputImageType].New(
            Input=difference,
            OutputMinimum=itk.NumericTraits[OutputPixelType].min(),
            OutputMaximum=itk.NumericTraits[OutputPixelType].max())

resampler.SetDefaultPixelValue(1)
writer.SetInput(intensityRescaler.GetOutput())
writer.SetFileName(differenceImageAfterFile)
writer.Update()

resampler.SetTransform(identityTransform)
writer.SetFileName(differenceImageBeforeFile)
writer.Update()
```

```python

```


# Input Driven Observations ("GLM-HMM")

Notebook prepared by Zoe Ashwood: feel free to email me with feedback or questions (zashwood at cs dot princeton dot edu).

This notebook demonstrates the "InputDrivenObservations" class, and illustrates its use in the context of modeling decision-making data as in Ashwood et al. (2020) ([Mice alternate between discrete strategies during perceptual
decision-making](https://www.biorxiv.org/content/10.1101/2020.10.19.346353v1.full.pdf)).

Compared to the model considered in the notebook ["2 Input Driven HMM"](https://github.com/lindermanlab/ssm/blob/master/notebooks/2%20Input%20Driven%20HMM.ipynb), Ashwood et al. (2020) assumes a stationary transition matrix where transition probabilities *do not* depend on external inputs. However, observation probabilities now *do* depend on external covariates according to:

$$
\begin{align}
\Pr(y_t = c \mid z_{t} = k, u_t, w_{kc}) =
\frac{\exp\{w_{kc}^\mathsf{T} u_t\}}
{\sum_{c'=1}^C \exp\{w_{kc'}^\mathsf{T} u_t\}}
\end{align}
$$

where $c \in \{1, ..., C\}$ indicates the categorical class for the observation, $u_{t} \in \mathbb{R}^{M}$ is the set of input covariates, and $w_{kc} \in \mathbb{R}^{M}$ is the set of input weights associated with state $k$ and class $c$. These weights, along with the transition matrix and initial state probabilities, will be learned.

In Ashwood et al. (2020), $C = 2$ as $y_{t}$ represents the binary choice made by an animal during a 2AFC (2-Alternative Forced Choice) task. The above equation then reduces to:

$$
\begin{align}
\Pr(y_t = 1 \mid z_{t} = k, u_t, w_{k}) =
\frac{1}
{1 + \exp\{-w_{k}^\mathsf{T} u_t\}}.
\end{align}
$$

and only a single set of weights is associated with each state.

## 1. Setup
The line `import ssm` imports the package for use. Here, we have also imported a few other packages for plotting.

```python
import numpy as np
import numpy.random as npr
import matplotlib.pyplot as plt
import ssm
from ssm.util import one_hot, find_permutation

%matplotlib inline

npr.seed(0)
```

## 2. Input Driven Observations
We create a HMM with input-driven observations and 'standard' (stationary) transitions with the following line:
```python
        ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs", observation_kwargs=dict(C=num_categories), transitions="standard")
```

As in Ashwood et al. (2020), we are going to model an animal's binary choice data during a decision-making task, so we will set `num_categories=2` because the animal only has two options available to it. We will also set `obs_dim = 1` because the dimensionality of the observation data is 1 (if we were also modeling, for example, the binned reaction time of the animal, we could set `obs_dim = 2`).  For the sake of simplicity, we will assume that an animal's choice in a particular state is only affected by the external stimulus associated with that particular trial, and its innate choice bias. Thus, we will set `input_dim = 2` and we will simulate input data that resembles sequences of stimuli in what follows.  In Ashwood et al. (2020), they found that many mice used 3 decision-making states when performing 2AFC tasks. We will, thus, set `num_states = 3`.

### 2a. Initialize GLM-HMM

```python
# Set the parameters of the GLM-HMM
num_states = 3        # number of discrete states
obs_dim = 1           # number of observed dimensions
num_categories = 2    # number of categories for output
input_dim = 2         # input dimensions

# Make a GLM-HMM
true_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
                   observation_kwargs=dict(C=num_categories), transitions="standard")
```

### 2b. Specify parameters of generative GLM-HMM

Let's update the weights and transition matrix for the true GLM-HMM so as to bring the GLM-HMM to the parameter regime that real animals use (according to Ashwood et al. (2020)):

```python
gen_weights = np.array([[[6, 1]], [[2, -3]], [[2, 3]]])
gen_log_trans_mat = np.log(np.array([[[0.98, 0.01, 0.01], [0.05, 0.92, 0.03], [0.02, 0.03, 0.94]]]))
true_glmhmm.observations.params = gen_weights
true_glmhmm.transitions.params = gen_log_trans_mat
```

```python
# Plot generative parameters:
fig = plt.figure(figsize=(8, 3), dpi=80, facecolor='w', edgecolor='k')
plt.subplot(1, 2, 1)
cols = ['#ff7f00', '#4daf4a', '#377eb8']
for k in range(num_states):
    plt.plot(range(input_dim), gen_weights[k][0], marker='o',
             color=cols[k], linestyle='-',
             lw=1.5, label="state " + str(k+1))
plt.yticks(fontsize=10)
plt.ylabel("GLM weight", fontsize=15)
plt.xlabel("covariate", fontsize=15)
plt.xticks([0, 1], ['stimulus', 'bias'], fontsize=12, rotation=45)
plt.axhline(y=0, color="k", alpha=0.5, ls="--")
plt.legend()
plt.title("Generative weights", fontsize = 15)

plt.subplot(1, 2, 2)
gen_trans_mat = np.exp(gen_log_trans_mat)[0]
plt.imshow(gen_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(gen_trans_mat.shape[0]):
    for j in range(gen_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(gen_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.ylabel("state t", fontsize = 15)
plt.xlabel("state t+1", fontsize = 15)
plt.title("Generative transition matrix", fontsize = 15)
```

### 2c. Create external input sequences

Simulate an example set of external inputs for each trial in a session. We will create an array of size `(num_sess x num_trials_per_sess x num_covariates)`. As in Ashwood et al. (2020), for each trial in a session we will include the stimulus presented to the animal at that trial, as well as a '1' as the second covariate (so as to capture the animal's innate bias for one of the two options available to it). We will simulate stimuli sequences so as to resemble the sequences of stimuli in the International Brain Laboratory et al. (2020) task.

```python
num_sess = 20 # number of example sessions
num_trials_per_sess = 100 # number of trials in a session
inpts = np.ones((num_sess, num_trials_per_sess, input_dim)) # initialize inpts array
stim_vals = [-1, -0.5, -0.25, -0.125, -0.0625, 0, 0.0625, 0.125, 0.25, 0.5, 1]
inpts[:,:,0] = np.random.choice(stim_vals, (num_sess, num_trials_per_sess)) # generate random sequence of stimuli
inpts = list(inpts) #convert inpts to correct format
```

### 2d. Simulate states and observations with generative model

```python
# Generate a sequence of latents and choices for each session
true_latents, true_choices = [], []
for sess in range(num_sess):
    true_z, true_y = true_glmhmm.sample(num_trials_per_sess, input=inpts[sess])
    true_latents.append(true_z)
    true_choices.append(true_y)
```

```python
# Calculate true loglikelihood
true_ll = true_glmhmm.log_probability(true_choices, inputs=inpts)
print("true ll = " + str(true_ll))
```

    true ll = -910.4271498215511

## 3. Fit GLM-HMM and perform recovery analysis

### 3a. Maximum Likelihood Estimation

Now we instantiate a new GLM-HMM and check that we can recover the generative parameters in simulated data:

```python
new_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
                   observation_kwargs=dict(C=num_categories), transitions="standard")

N_iters = 200 # maximum number of EM iterations. Fitting with stop earlier if increase in LL is below tolerance specified by tolerance parameter
fit_ll = new_glmhmm.fit(true_choices, inputs=inpts, method="em", num_iters=N_iters, tolerance=10**-4)
```

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=200.0), HTML(value='')))

```python
# Plot the log probabilities of the true and fit models. Fit model final LL should be greater
# than or equal to true LL.
fig = plt.figure(figsize=(4, 3), dpi=80, facecolor='w', edgecolor='k')
plt.plot(fit_ll, label="EM")
plt.plot([0, len(fit_ll)], true_ll * np.ones(2), ':k', label="True")
plt.legend(loc="lower right")
plt.xlabel("EM Iteration")
plt.xlim(0, len(fit_ll))
plt.ylabel("Log Probability")
plt.show()
```

### 3b. Retrieved parameters

Compare retrieved weights and transition matrices to generative parameters. To do this, we may first need to permute the states of the fit GLM-HMM relative to the
generative model.  One way to do this uses the `find_permutation` function from `ssm`:

```python
new_glmhmm.permute(find_permutation(true_latents[0], new_glmhmm.most_likely_states(true_choices[0], input=inpts[0])))
```

Now plot generative and retrieved weights for GLMs (analogous plot to Figure S1c in
Ashwood et al. (2020)):

```python
fig = plt.figure(figsize=(4, 3), dpi=80, facecolor='w', edgecolor='k')
cols = ['#ff7f00', '#4daf4a', '#377eb8']
recovered_weights = new_glmhmm.observations.params
for k in range(num_states):
    if k ==0:
        plt.plot(range(input_dim), gen_weights[k][0], marker='o',
                 color=cols[k], linestyle='-',
                 lw=1.5, label="generative")
        plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                     lw=1.5,  label = "recovered", linestyle = '--')
    else:
        plt.plot(range(input_dim), gen_weights[k][0], marker='o',
                 color=cols[k], linestyle='-',
                 lw=1.5, label="")
        plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                     lw=1.5,  label = '', linestyle = '--')
plt.yticks(fontsize=10)
plt.ylabel("GLM weight", fontsize=15)
plt.xlabel("covariate", fontsize=15)
plt.xticks([0, 1], ['stimulus', 'bias'], fontsize=12, rotation=45)
plt.axhline(y=0, color="k", alpha=0.5, ls="--")
plt.legend()
plt.title("Weight recovery", fontsize=15)
```

Now plot generative and retrieved transition matrices (analogous plot to Figure S1c in
Ashwood et al. (2020)):

```python
fig = plt.figure(figsize=(5, 2.5), dpi=80, facecolor='w', edgecolor='k')
plt.subplot(1, 2, 1)
gen_trans_mat = np.exp(gen_log_trans_mat)[0]
plt.imshow(gen_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(gen_trans_mat.shape[0]):
    for j in range(gen_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(gen_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.ylabel("state t", fontsize = 15)
plt.xlabel("state t+1", fontsize = 15)
plt.title("generative", fontsize = 15)

plt.subplot(1, 2, 2)
recovered_trans_mat = np.exp(new_glmhmm.transitions.log_Ps)
plt.imshow(recovered_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(recovered_trans_mat.shape[0]):
    for j in range(recovered_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(recovered_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.title("recovered", fontsize = 15)
plt.subplots_adjust(0, 0, 1, 1)

```

### 3c. Posterior State Probabilities

Let's now plot $p(z_{t} = k|\mathbf{y}, \{u_{t}\}_{t=1}^{T})$, the posterior state probabilities, which give the probability of the animal being in state k at trial t.

```python
# Get expected states:
posterior_probs = [new_glmhmm.expected_states(data=data, input=inpt)[0]
                for data, inpt
                in zip(true_choices, inpts)]
```

```python
fig = plt.figure(figsize=(5, 2.5), dpi=80, facecolor='w', edgecolor='k')
sess_id = 0 #session id; can choose any index between 0 and num_sess-1
for k in range(num_states):
    plt.plot(posterior_probs[sess_id][:, k], label="State " + str(k + 1), lw=2,
             color=cols[k])
plt.ylim((-0.01, 1.01))
plt.yticks([0, 0.5, 1], fontsize = 10)
plt.xlabel("trial #", fontsize = 15)
plt.ylabel("p(state)", fontsize = 15)
```

With these posterior state probabilities, we can assign trials to states and then plot the fractional occupancy of each state:

```python
# concatenate posterior probabilities across sessions
posterior_probs_concat = np.concatenate(posterior_probs)
# get state with maximum posterior probability at particular trial:
state_max_posterior = np.argmax(posterior_probs_concat, axis = 1)
# now obtain state fractional occupancies:
_, state_occupancies = np.unique(state_max_posterior, return_counts=True)
state_occupancies = state_occupancies/np.sum(state_occupancies)
```

```python
fig = plt.figure(figsize=(2, 2.5), dpi=80, facecolor='w', edgecolor='k')
for z, occ in enumerate(state_occupancies):
    plt.bar(z, occ, width = 0.8, color = cols[z])
plt.ylim((0, 1))
plt.xticks([0, 1, 2], ['1', '2', '3'], fontsize = 10)
plt.yticks([0, 0.5, 1], ['0', '0.5', '1'], fontsize=10)
plt.xlabel('state', fontsize = 15)
plt.ylabel('frac. occupancy', fontsize=15)
```

## 4. Fit GLM-HMM and perform recovery analysis: Maximum A Priori Estimation

Above, we performed Maximum Likelihood Estimation to retrieve the generative parameters of the GLM-HMM in simulated data. In the small data regime, where we do not have many trials available to us, we may instead want to perform Maximum A Priori (MAP) Estimation in order to incorporate a prior term and restrict the range for the best fitting parameters. Unfortunately, what is meant by 'small data regime' is problem dependent and will be affected by the number of states in the generative GLM-HMM, and the specific parameters of the generative model, amongst other things. In practice, we may perform both Maximum Likelihood Estimation and MAP estimation and compare the ability of the fit models to make predictions on held-out data (see Section 5 on Cross-Validation below).

The prior we consider for the GLM-HMM is the product of a Gaussian prior on the GLM weights, $W$, and a Dirichlet prior on the transition matrix, $A$:

$$
\begin{align}
\Pr(W, A) &= \mathcal{N}(W|0, \Sigma) \Pr(A|\alpha) \\&= \mathcal{N}(W|0, diag(\sigma^{2}, \cdots, \sigma^{2})) \prod_{j=1}^{K} \dfrac{1}{B(\alpha)} \prod_{k=1}^{K} A_{jk}^{\alpha -1}
\end{align}
$$

There are two hyperparameters controlling the strength of the prior: $\sigma$ and $\alpha$.  The larger the value of $\sigma$ and if $\alpha = 1$, the more similar MAP estimation will become to Maximum Likelihood Estimation, and the prior term will become an additive offset to the objective function of the GLM-HMM that is independent of the values of $W$ and $A$.  In comparison, setting $\sigma = 2$ and $\alpha = 2$ will result in the prior no longer being independent of $W$ and $\alpha$.

In order to perform MAP estimation for the GLM-HMM with `ssm`, the new syntax is:

```python
ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
             observation_kwargs=dict(C=num_categories,prior_sigma=prior_sigma),
             transitions="sticky", transition_kwargs=dict(alpha=prior_alpha,kappa=0))```

where `prior_sigma` is the $\sigma$ parameter from above, and `prior_alpha` is the $\alpha$ parameter.

```python
# Instantiate GLM-HMM and set prior hyperparameters
prior_sigma = 2
prior_alpha = 2
map_glmhmm = ssm.HMM(num_states, obs_dim, input_dim, observations="input_driven_obs",
             observation_kwargs=dict(C=num_categories,prior_sigma=prior_sigma),
             transitions="sticky", transition_kwargs=dict(alpha=prior_alpha,kappa=0))
```

```python
# Fit GLM-HMM with MAP estimation:
_ = map_glmhmm.fit(true_choices, inputs=inpts, method="em", num_iters=N_iters, tolerance=10**-4)
```

    HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=200.0), HTML(value='')))

Compare final likelihood of data with MAP estimation and MLE to likelihood under generative model (note: we cannot use log_probability that is output of `fit` function as this incorporates prior term, which is not comparable between generative and MAP models). We want to check that MAP and MLE likelihood values are higher than true likelihood; if they are not, this may indicate a poor initialization and that we should refit these models.

```python
true_likelihood = true_glmhmm.log_likelihood(true_choices, inputs=inpts)
mle_final_ll = new_glmhmm.log_likelihood(true_choices, inputs=inpts)
map_final_ll = map_glmhmm.log_likelihood(true_choices, inputs=inpts)
```

```python
# Plot these values
fig = plt.figure(figsize=(2, 2.5), dpi=80, facecolor='w', edgecolor='k')
loglikelihood_vals = [true_likelihood, mle_final_ll, map_final_ll]
colors = ['Red', 'Navy', 'Purple']
for z, occ in enumerate(loglikelihood_vals):
    plt.bar(z, occ, width = 0.8, color = colors[z])
plt.ylim((true_likelihood-5, true_likelihood+15))
plt.xticks([0, 1, 2], ['true', 'mle', 'map'], fontsize = 10)
plt.xlabel('model', fontsize = 15)
plt.ylabel('loglikelihood', fontsize=15)
```

## 5. Cross Validation

To assess which model is better - the model fit via Maximum Likelihood Estimation, or the model fit via MAP estimation - we can investigate the predictive power of these fit models on held-out test data sets.

```python
# Create additional input sequences to be used as held-out test data
num_test_sess = 10
test_inpts = np.ones((num_test_sess, num_trials_per_sess, input_dim))
test_inpts[:,:,0] = np.random.choice(stim_vals, (num_test_sess, num_trials_per_sess))
test_inpts = list(test_inpts)
```

```python
# Create set of test latents and choices to accompany input sequences:
test_latents, test_choices = [], []
for sess in range(num_test_sess):
    test_z, test_y = true_glmhmm.sample(num_trials_per_sess, input=test_inpts[sess])
    test_latents.append(test_z)
    test_choices.append(test_y)
```

```python
# Compare likelihood of test_choices for model fit with MLE and MAP:
mle_test_ll = new_glmhmm.log_likelihood(test_choices, inputs=test_inpts)
map_test_ll = map_glmhmm.log_likelihood(test_choices, inputs=test_inpts)
```

```python
fig = plt.figure(figsize=(2, 2.5), dpi=80, facecolor='w', edgecolor='k')
loglikelihood_vals = [mle_test_ll, map_test_ll]
colors = ['Navy', 'Purple']
for z, occ in enumerate(loglikelihood_vals):
    plt.bar(z, occ, width = 0.8, color = colors[z])
plt.ylim((mle_test_ll-2, mle_test_ll+5))
plt.xticks([0, 1], ['mle', 'map'], fontsize = 10)
plt.xlabel('model', fontsize = 15)
plt.ylabel('loglikelihood', fontsize=15)
```

Here we see that the model fit with MAP estimation achieves higher likelihood on the held-out dataset than the model fit with MLE, so we would choose this model as the best model of animal decision-making behavior (although we'd probably want to perform multiple fold cross-validation to be sure that this is the case in all instantiations of test data).

Let's finish by comparing the retrieved weights and transition matrices from MAP estimation to the generative parameters.

```python
map_glmhmm.permute(find_permutation(true_latents[0], map_glmhmm.most_likely_states(true_choices[0], input=inpts[0])))
```

```python
fig = plt.figure(figsize=(6, 3), dpi=80, facecolor='w', edgecolor='k')
cols = ['#ff7f00', '#4daf4a', '#377eb8']
plt.subplot(1,2,1)
recovered_weights = new_glmhmm.observations.params
for k in range(num_states):
    if k ==0: # show labels only for first state
        plt.plot(range(input_dim), gen_weights[k][0], marker='o',
                 color=cols[k],
                 lw=1.5, label="generative")
        plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                     lw=1.5,  label = 'recovered', linestyle='--')
    else:
        plt.plot(range(input_dim), gen_weights[k][0], marker='o',
                 color=cols[k],
                 lw=1.5, label="")
        plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                     lw=1.5,  label = '', linestyle='--')
plt.yticks(fontsize=10)
plt.ylabel("GLM weight", fontsize=15)
plt.xlabel("covariate", fontsize=15)
plt.xticks([0, 1], ['stimulus', 'bias'], fontsize=12, rotation=45)
plt.axhline(y=0, color="k", alpha=0.5, ls="--")
plt.title("MLE", fontsize = 15)
plt.legend()

plt.subplot(1,2,2)
recovered_weights = map_glmhmm.observations.params
for k in range(num_states):
    plt.plot(range(input_dim), gen_weights[k][0], marker='o',
             color=cols[k],
             lw=1.5, label="", linestyle = '-')
    plt.plot(range(input_dim), recovered_weights[k][0], color=cols[k],
                 lw=1.5,  label = '', linestyle='--')
plt.yticks(fontsize=10)
plt.xticks([0, 1], ['', ''], fontsize=12, rotation=45)
plt.axhline(y=0, color="k", alpha=0.5, ls="--")
plt.title("MAP", fontsize = 15)
```

```python
fig = plt.figure(figsize=(7, 2.5), dpi=80, facecolor='w', edgecolor='k')
plt.subplot(1, 3, 1)
gen_trans_mat = np.exp(gen_log_trans_mat)[0]
plt.imshow(gen_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(gen_trans_mat.shape[0]):
    for j in range(gen_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(gen_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.ylabel("state t", fontsize = 15)
plt.xlabel("state t+1", fontsize = 15)
plt.title("generative", fontsize = 15)

plt.subplot(1, 3, 2)
recovered_trans_mat = np.exp(new_glmhmm.transitions.log_Ps)
plt.imshow(recovered_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(recovered_trans_mat.shape[0]):
    for j in range(recovered_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(recovered_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.title("recovered - MLE", fontsize = 15)
plt.subplots_adjust(0, 0, 1, 1)

plt.subplot(1, 3, 3)
recovered_trans_mat = np.exp(map_glmhmm.transitions.log_Ps)
plt.imshow(recovered_trans_mat, vmin=-0.8, vmax=1, cmap='bone')
for i in range(recovered_trans_mat.shape[0]):
    for j in range(recovered_trans_mat.shape[1]):
        text = plt.text(j, i, str(np.around(recovered_trans_mat[i, j], decimals=2)), ha="center", va="center",
                        color="k", fontsize=12)
plt.xlim(-0.5, num_states - 0.5)
plt.xticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.yticks(range(0, num_states), ('1', '2', '3'), fontsize=10)
plt.ylim(num_states - 0.5, -0.5)
plt.title("recovered - MAP", fontsize = 15)
plt.subplots_adjust(0, 0, 1, 1)
```

notebooks/2b Input Driven Observations (GLM-HMM).ipynb


# Método de la secante
J.J
---
Está dado por las iteraciones
\begin{equation}
x_{i+1} = x_{i} - \frac{f(x_{i})}{d_{i}},
\end{equation}
con
\begin{equation}
d_{i} = \frac{f(x_{i})-f(x_{i-1})}{x_{i} - x_{i-1}}.
\end{equation}

Ejemplo: $f(x) = x^{2} + 3x +1$

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

```python
def f(x):
    f = x**2 + 3*x + 1
    return f
```

```python
X = np.linspace(-3, 1, 100)
plt.plot(X, f(X))
```

```python
e = 0.0001 #error
maxit = 1000 #iteraciones máximas
```

```python
def Secante(x0, x1, func = f, error = e, iterations = maxit):
    it = 0
    while it < maxit:
        it += 1
        x = x1 - ((x1-x0)/(f(x1)-f(x0)))*f(x1)
        if abs(f(x)) > e:
            x0 = x1
            x1 = x
        else:
            break
    return x
```

```python
sol1 = Secante(-3.,-2.)
sol2 = Secante(-1.,0.)
print(sol1)
print(sol2)
```

    -2.618034055727554
    -0.38197424892703863

```python
X = np.linspace(-3, 1, 100)
plt.plot(X, f(X))
plt.plot(sol1,f(sol1),'ro')
plt.plot(sol2,f(sol2),'ro')
plt.show()
```

```python

```


```python
from sympy import init_printing; init_printing(use_latex='mathjax')
import sympy as sym
```

```python
r, theta, k  = sym.symbols('r theta k')
r, theta, k
```

$$\left ( r, \quad \theta, \quad k\right )$$

```python
n = sym.Symbol('n', positive = True, integer=True)
n
```

$$n$$

```python
def lamb(n,k):
    return sym.Symbol('lambda_%s%s'%(n,k), positive = True)
```

```python
lamb(0,k)
```

$$\lambda_{0k}$$

```python
f = 1 - r**2; f
```

$$- r^{2} + 1$$

```python
integrand = f * sym.besselj(n, lamb(n,k) * r) * sym.cos(n *theta) * r
integrand
```

$$r \left(- r^{2} + 1\right) \cos{\left (n \theta \right )} J_{n}\left(\lambda_{nk} r\right)$$

```python
ank = sym.Integral(integrand, (r, 0, 1), (theta, 0, 2*sym.pi))
ank
```

$$\int_{0}^{2 \pi}\int_{0}^{1} r \left(- r^{2} + 1\right) \cos{\left (n \theta \right )} J_{n}\left(\lambda_{nk} r\right)\, dr\, d\theta$$

```python
solution = ank.doit()
solution
```

$$0$$

```python
integ = lambda n: f * sym.besselj(n, lamb(n,k) * r) * sym.cos(n*theta) * r
```

```python
integ(0)
```

$$r \left(- r^{2} + 1\right) J_{0}\left(\lambda_{0k} r\right)$$

```python
a0k = sym.Integral(integ(0), (r, 0, 1), (theta, 0, 2*sym.pi))
a0k
```

$$\int_{0}^{2 \pi}\int_{0}^{1} r \left(- r^{2} + 1\right) J_{0}\left(\lambda_{0k} r\right)\, dr\, d\theta$$

```python
a0k.doit()
```

$$2 \pi \left(- \frac{2 J_{0}\left(\lambda_{0k}\right)}{\lambda_{0k}^{2}} + \frac{4 J_{1}\left(\lambda_{0k}\right)}{\lambda_{0k}^{3}}\right)$$

```python

```


# Curves

__Definition:__ A parametrized differentiable curve is a map $\alpha : I \rightarrow \mathbb{R}^3$ being I an open interval $I \subset \mathbb{R}$ and all its components $\alpha_i \in \mathcal{C}^{\infty}$.

__Example__: The parametrized differentiable curve given by

\begin{equation}
\alpha (t) = ( a \cos(t) + b \cos(3t), a \sin(t) - b \sin(3t), c \sin(2t))
\end{equation}

Can be represented in python as follows:

```python
import math
import tensorflow as tf
import ipywidgets as widgets
import matplotlib.pyplot as plt
```

```python
def curve(a=1, b=1, c=1):
    t = tf.linspace(-2 * math.pi, 2 * math.pi, 100)
    return (
        a * tf.math.cos(t) + b * tf.math.cos(3*t),
        a * tf.math.sin(t) - b * tf.math.sin(3*t),
        c * tf.math.sin(2*t)
    )
```

```python
@widgets.interact
def plot(a=1., b=1., c=1.):
    fig = plt.figure(figsize=(8, 6))
    ax = fig.gca(projection='3d')
    x, y, z = curve(a, b, c)
    ax.plot(x, y, z, label='parametric curve')
    ax.legend()
    plt.show()
```

```python
def plot(a=1., b=1., c=1.):
    fig = plt.figure(figsize=(10, 10))
    ax = fig.gca(projection='3d')

    # make the panes transparent
    ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
    ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
    ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
    # make the grid lines transparent
    ax.xaxis._axinfo["grid"]['color'] =  (1,1,1,0)
    ax.yaxis._axinfo["grid"]['color'] =  (1,1,1,0)
    ax.zaxis._axinfo["grid"]['color'] =  (1,1,1,0)

    steps = [i/10 for i in range(20)]
    for i in steps:
        x, y, z = curve(a + i, b, c)
        ax.plot(x, y, z)
    plt.show()
```

```python
plot()
```


# Portfolio Selection Optimization
This model is an example of the classic [Markowitz portfolio selection optimization model](https://en.wikipedia.org/wiki/Markowitz_model). We want to find the fraction of the portfolio to invest among a set of stocks that balances risk and return. It is a Quadratic Programming (QP) model with vector and matrix data for returns and risk, respectively. This is best suited to a matrix formulation, so we use the Gurobi Python *matrix* interface. The basic model is fairly simple, so we also solve it parametrically to find the efficient frontier.

**Download the Repository** <br />
You can download the repository containing this and other examples by clicking [here](https://github.com/Gurobi/modeling-examples/archive/master.zip).

## Model Formulation
### Parameters

We use the [Greek values](https://en.wikipedia.org/wiki/Greeks_\(finance\)) that are traditional in finance:

- $\delta$: n-element vector measuring the change in price for each stock
- $\sigma$: n x n matrix measuring the covariance among stocks

There is one additional parameter when solving the model parametrically:

- r: target return

### Decision Variables
- $x \ge 0$: n-element vector where each element represents the fraction of the porfolio to invest in each stock

### Objective Function
Minimize the total risk, a convex quadratic function:

\begin{equation}
\min x^t \cdot \sigma \cdot x
\end{equation}

### Constraints

Allocate the entire portfolio: the total investments should be 1.0 (100%), where $e$ is a unit vector (all 1's):

\begin{equation}
e \cdot x = 1
\end{equation}

Return: When we solve the model parametrically for different return values $r$, we add a constraint on the target return:

\begin{equation}
\delta \cdot x = r
\end{equation}

## Python Implementation
### Stock data
Use [yfinance](https://pypi.org/project/yfinance/) library to get the latest 2 years of _actual stock data_ from the 20 most profitable US companies, [according to Wikipedia in April 2021](https://en.wikipedia.org/wiki/List_of_largest_companies_in_the_United_States_by_revenue#List_of_companies_by_profit).

```python
%pip install gurobipy yfinance
```

```python
import yfinance as yf

stocks = ['BRK-A', 'AAPL', 'MSFT', 'JPM', 'GOOG', 'BAC', 'INTC', 'WFC',
          'C', 'VZ', 'FB', 'PFE', 'JNJ', 'WMT', 'XOM',
          'FNMA', 'T', 'UNH', 'CMCSA', 'V' ]

data = yf.download(stocks, period='2y')
```

### Compute Greeks
Using the downloaded stock data, find the delta (return), sigma (covariance) and standard deviation values for stock prices:

```python
import numpy as np

closes = np.transpose(np.array(data.Close)) # matrix of daily closing prices
absdiff = np.diff(closes)                   # change in closing price each day
reldiff = np.divide(absdiff, closes[:,:-1]) # relative change in daily closing price
delta = np.mean(reldiff, axis=1)            # mean price change
sigma = np.cov(reldiff)                     # covariance (standard deviations)
std = np.std(reldiff, axis=1)               # standard deviation
```

## Minimize risk by solving QP model

```python
import gurobipy as gp
from gurobipy import GRB
from math import sqrt

# Create an empty model
m = gp.Model('portfolio')

# Add matrix variable for the stocks
x = m.addMVar(len(stocks))

# Objective is to minimize risk (squared).  This is modeled using the
# covariance matrix, which measures the historical correlation between stocks
portfolio_risk = x @ sigma @ x
m.setObjective(portfolio_risk, GRB.MINIMIZE)

# Fix budget with a constraint
m.addConstr(x.sum() == 1, 'budget')

# Verify model formulation
m.write('portfolio_selection_optimization.lp')

# Optimize model to find the minimum risk portfolio
m.optimize()
```

## Display minimum risk portfolio using Pandas

```python
import pandas as pd
minrisk_volatility = sqrt(m.ObjVal)
minrisk_return = delta @ x.X
pd.DataFrame(data=np.append(x.X, [minrisk_volatility, minrisk_return]),
             index=stocks + ['Volatility', 'Expected Return'],
             columns=['Minimum Risk Portfolio'])
```

## Compute the efficient frontier
Solve the QP parametrically to find the lowest risk portfolio for different expected returns.

```python
# Create an expression representing the expected return for the portfolio
portfolio_return = delta @ x
target = m.addConstr(portfolio_return == minrisk_return, 'target')

# Solve for efficient frontier by varying target return
frontier = np.empty((2,0))
for r in np.linspace(delta.min(), delta.max(), 25):
    target[0].rhs = r
    m.optimize()
    frontier = np.append(frontier, [[sqrt(m.ObjVal)],[r]], axis=1)
```

## Plot results
Use the matplot library to plot the optimized solutions, along with the individual stocks:

```python
import matplotlib.pyplot as plt
#plt.figure(figsize=(10,10))

fig, ax = plt.subplots(figsize=(10,8))

# Plot volatility versus expected return for individual stocks
ax.scatter(x=std, y=delta,
           color='Blue', label='Individual Stocks')
for i, stock in enumerate(stocks):
    ax.annotate(stock, (std[i], delta[i]))

# Plot volatility versus expected return for minimum risk portfolio
ax.scatter(x=minrisk_volatility, y=minrisk_return, color='DarkGreen')
ax.annotate('Minimum\nRisk\nPortfolio', (minrisk_volatility, minrisk_return),
            horizontalalignment='right')

# Plot efficient frontier
ax.plot(frontier[0], frontier[1], label='Efficient Frontier', color='DarkGreen')

# Format and display the final plot
ax.axis([frontier[0].min()*0.7, frontier[0].max()*1.3, delta.min()*1.2, delta.max()*1.2])
ax.set_xlabel('Volatility (standard deviation)')
ax.set_ylabel('Expected Return')
ax.legend()
ax.grid()
plt.show()
```


# Univariate Linear Regression

In this iPython notebook, I'm going to create an Univariate Linear Regression model without using any scikit-learn or other libraries. The objective is to understand better what these libraries actually does under the hood and get a better intuition into how we can actually make these algorithms perform better. Before we begin, I'd like to take a moment and introduce a few notations that we're going to use throughout this notebook.

**Notations:**
- $m$ = Numner of Training Examples
- $X$ = Input Variables
- $y$ = Output Variables
- $(x^{(i)},y^{(i)})$ = $i^{th}$ Training Example

Now let's get started and I'll explain the rest on the way as we proceed with this notebook

```python
# Importing Necessary Libraries
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
```

The above libraries are some bare minimum which I had to import to make this notebook somewhat concise and not ending up covering a lot of stuff which might turn out to be irrelevant for the purpose of this notebook. Also, these libraries are proven to be fast and I would not have to worry too much on things like writing a function for matrix multiplication which might turn out to have a huge time complexity.

```python
# Okay, enough talk nor now, let's start by importing our dataset
# (This dataset is from one of the Octave excercises from Andrew Ng's tutorial on Machine Learning)
data = pd.read_csv("./data.txt", sep=",", names=["Population in 10k", "Profit in $10k"])
```

```python
# Taking a look at the imported data
data.head(5)
```

      Population in 10k
      Profit in $10k

      0
      6.1101
      17.5920

      1
      5.5277
      9.1302

      2
      8.5186
      13.6620

      3
      7.0032
      11.8540

      4
      5.8598
      6.8233

```python
# Splitting columns from our dataset
X = data['Population in 10k']
y = data['Profit in $10k']
```

```python
# Now I wanna plot and see how the dataset looks like
# I'm reducing the marker size a bit, so that the individual points close by don't overlap each other
plt.scatter(X, y, s=10, alpha=0.5)
plt.title('Scatter Plot of the Dataset')
plt.xlabel('Population in 10k')
plt.ylabel('Profit in $10k')
plt.show()
```

```python
# Now I want to convert 'X' into a single column numpy array
X = np.array(X)
X.shape
```

    (97,)

Few more notations I'm using here:
- $R$ = Number of Rows
- $C$ = Number of Columns

As you can see above, the shape turned out to be a vector instead of a RxC matrix or 2D array. If you're not familiar with numpy, or new to numpy array indexing, let me try to elaborate a bit.

##### Shape: (R,)
Here the array is basically single indexed. For example, consider the following example of an array of 8 numbers:

| Data  | 9 | 1 | 5 | 3 | 7 | 3 | 2 | 5 |
| ----- | - | - | - | - | - | - | - | - |
| Index | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |

##### Shape: (R,C)
We want to have a 2D indexing in our array which will basically be a single column matrix (a number of rows with only one column):

| Data      | 9 | 1 | 5 | 3 | 7 | 3 | 2 | 5 |
| --------- | - | - | - | - | - | - | - | - |
| Index (i) | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
| Index (j) | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 |

**TL;DR:** `array.shape` basically returns the number of rows (R) vs the number of columns (R) in the array/matrix

```python
# Reshaping our numpy array
m = y.size
X = X.reshape(m,1)
X.shape
```

    (97, 1)

### Hypothesis
Before we proceed further, it's important to understand what hypothesis is. It is a straight line fit through our dataset, which tries to map the input features (X) to output variables (y).

\begin{align}
h_\theta(x) = \theta _0 + \theta _1 x
\end{align}

**Notations:**
- $h$ = Hypothesis
- $\theta_0$ & $\theta_1$ = Parameters

> A good hypothesis fits the evidence and can be used to make predictions about new observations or new situations.

```python
# Now I want to add a column of 1's to the 0th index of 'X' (Explanation Below)
# Also shift the existing data in the 0th index to the 1st index
X = np.append(np.ones((m, 1), dtype=int), X, axis=1)
X.shape
```

    (97, 2)

I did the above to take into account the intercept term $\theta_0$. I added an extra column of 1's to the matrix 'X', so that I can treat $\theta_0$ as another feature. For this, I needed to place the new column in the 0th column index and shift the existing column of data from 0th to 1st index. Now the shape of the matrix will be (R,2).

```python
# Let's initialize Theta as a 2x1 null array
theta = np.zeros((2,1), dtype=float)
theta.shape
```

    (2, 1)

### Cost Function & Gradient Descent
Gradient Descent attempts to find the local or global optima of a function. We succeed in training a model when the Cost Function reaches the global optimum.

\begin{align}
J(\theta_0, \theta_1) = \frac{1}{2m} \sum^m_{i=1} (h_\theta(x^{(i)})-y^{(i)})^2
\end{align}

> A model learns by minimizing the Cost Function using Gradient Descent

**Algorithm:**
{
$$
\theta_j := \theta_j - \alpha \frac{\partial}{\partial\theta_j} J(\theta_0, \theta_1)
$$
}

Simultaneously update $\theta_1$ and $\theta_2$,
$temp0 := \theta_0 - \alpha \frac{\partial}{\partial \theta_0} J(\theta_0, \theta_1)$
$temp1 := \theta_1 - \alpha \frac{\partial}{\partial \theta_1} J(\theta_0, \theta_1)$
$\theta_0 := temp0$
$\theta_1 := temp1$

**Notation:**
- $\alpha$ = Learning Rate

```python
# Function for Computing Cost Function
def costFunction(theta):
    J_sum = 0
    for i in range(m):
        J_sum += ((theta[0] + theta[1] * X[i][1]) - y[i])**2
    J = (1/(2*m)) * J_sum
    return J
```

```python
# Let's see the cost when theta is 0
costFunction(theta)
```

    array([32.07273388])

```python
# Function for Running Gradient Descent
def gradientDescent(theta, iterations, alpha):
    J_history = np.zeros((iterations,1), dtype=float)
    for i in range(iterations):
        t1 = 0
        t2 = 0
        for k in range(m):
            t1 = t1 + ((theta[0] + theta[1]*X[k][1]) - y[k])
            t2 = t2 + ((theta[0] + theta[1]*X[k][1]) - y[k])*X[k][1]
        theta[0] = theta[0] - (alpha/m)*t1
        theta[1] = theta[1] - (alpha/m)*t2
        J_history[i] = costFunction(theta)
    return theta
```

```python
# Now let's actually run the gradient descent
iterations = 1500
alpha = 0.01
theta = gradientDescent(theta, iterations, alpha)
```

```python
# Let's see what values of theta we got
print("Theta 0: ", theta[0])
print("Theta 1: ", theta[1])
```

    Theta 0:  [-3.63029144]
    Theta 1:  [1.16636235]

```python
# Plotting the linear fit
plt.plot(X[:,1], np.dot(X, theta))
plt.scatter(X[:,1], y, s=10, alpha=0.5)
plt.title('Linear Fit through the Dataset')
plt.xlabel('Population in 10k')
plt.ylabel('Profit in $10k')
plt.show()
```

```python
# Now let's see how our model performs on new data (Let's initialize a few random ones)
new_data_population = np.array([7.2, 18, 24, 13, 26])
new_data = new_data_population.reshape(5,1)
new_data = np.append(np.ones((5, 1), dtype=int), new_data, axis=1)
```

```python
# Predicting the Results
predictions = np.zeros((5,1), dtype=float)
for i in range(5):
    predictions[i][0] = abs(np.dot(new_data[i][:], theta))
    print("Predicted Value of profit for ", new_data_population[i], " population (10k) : $", predictions[i][0], " (10k)")
```

    Predicted Value of profit for  7.2  population (10k) : $ 4.767517483011831  (10k)
    Predicted Value of profit for  18.0  population (10k) : $ 17.364230866636117  (10k)
    Predicted Value of profit for  24.0  population (10k) : $ 24.362404968649606  (10k)
    Predicted Value of profit for  13.0  population (10k) : $ 11.532419114958206  (10k)
    Predicted Value of profit for  26.0  population (10k) : $ 26.69512966932077  (10k)

```python
# Plotting the results
plt.plot(X[:,1], np.dot(X, theta))
plt.scatter(X[:,1], y, s=10, alpha=0.5, color='c')
plt.scatter(new_data_population, predictions, s=150, alpha=0.5, color='r', marker='x')
plt.title('Predicted Results over Existing Dataset')
plt.xlabel('Population in 10k')
plt.ylabel('Profit in $10k')
plt.show()
```

So that's it. Our univariate linear regression model works and it was able to predict the profits for some random populations. Now we can go ahead and create a similar linear regression model for datasets with multiple features, also known as **Multivariate Linear Regression**. But that's going to be another notebook soon.


```python
from sympy import init_session, init_printing, Matrix
init_session()
init_printing()
```

    IPython console for SymPy 0.7.6 (Python 2.7.8-64-bit) (ground types: python)

    These commands were executed:
    >>> from __future__ import division
    >>> from sympy import *
    >>> x, y, z, t = symbols('x y z t')
    >>> k, m, n = symbols('k m n', integer=True)
    >>> f, g, h = symbols('f g h', cls=Function)
    >>> init_printing()

    Documentation can be found at http://www.sympy.org

## Twiss Matrices

```python
B, Bi, R, mu     = symbols('B, B_{inv}, R, mu')
a0, a1, b0, b1   = symbols('alpha_0, alpha_1, beta_0, beta_1', real=True, positive=True)
# V, h, phi, eta, R, C = symbols('V, h, phi, eta, R, C')
# z, dp                = symbols('z, dp', real=True)
# Hc                   = symbols('H_c')

B  = Matrix([[1/sqrt(b0), 0], [a0/sqrt(b0), sqrt(b0)]])
R  = Matrix([[cos(mu), sin(mu)], [-sin(mu), cos(mu)]])
Bi = Matrix([[sqrt(b1), 0], [-a1/sqrt(b1), 1/sqrt(b1)]])
```

```python
expand_trig(simplify(trigsimp(Bi * R * B)))
```

```python
simplify(Bi * Matrix([[1, 0],[0, 1]]) * B)
```

```python
simplify(Bi * Matrix([[0, 1],[-1, 0]]) * B)
```

## Hamiltonian

```python
e, beta, c           = symbols('e, beta, c')
V, h, phi, eta, R, C = symbols('V, h, phi, eta, R, C')
z, dp                = symbols('z, dp', real=True)
Hc                   = symbols('H_c')

zs = C/h/2.

def H(z, dp):
    return -eta*beta*c*dp**2 + e*V/C * cos(h*z/R + phi)
H(z, dp)
```

```python
def ps(z):
    return solve(Eq(H(z,dp), 0), dp)[0]
ps(z)
```

```python
def ps(z):
    return sqrt(cos(z))
#     return -sqrt(e*V/(eta*beta*c*C) * h*z/R)
#     return -sqrt(e*V/(eta*beta*c*C) * cos(h*z/R+phi))
ps(z)
```

```python
integrate(ps(z), (z, -zs, zs))
```


```python
from sympy.physics.mechanics import *
import sympy as sp
mechanics_printing(pretty_print=True)
```

```python
m1, m2, m3, l1, l2, l3 = sp.symbols(r'm_1 m_2 m_3 l_1 l_2 l_3')
t, g, h = sp.symbols('t g h')
v1, v2, v3 = dynamicsymbols(r'\theta_1 \theta_2 \theta_3')
dv1, dv2, dv3 = dynamicsymbols(r'\theta_1 \theta_2 \theta_3', 1)
```

```python
x1 = l1 * sp.sin(v1)
y1 = -l1 * sp.cos(v1)
x2 = x1 + l2 * sp.sin(v2)
y2 = y1 + -l2 * sp.cos(v2)
x3 = x2 + l3 * sp.sin(v3)
y3 = y2 + -l3 * sp.cos(v3)

dx1 = x1.diff(t)
dy1 = y1.diff(t)
dx2 = x2.diff(t)
dy2 = y2.diff(t)
dx3 = x3.diff(t)
dy3 = y3.diff(t)
```

```python
V = (m1 * g * y1) + (m2 * g * y2) + (m3 * g * y3)
T = (sp.Rational(1, 2) * m1 * (dx1**2 + dy1**2)) + (sp.Rational(1, 2) * m2 * (dx2**2+dy2**2)) + (sp.Rational(1, 2) * m3 * (dx3**2 + dy3**2))
L = T - V
```

```python
LM = LagrangesMethod(L, [v1, v2, v3])
```

```python
soln = LM.form_lagranges_equations()
```

```python
solvedsoln = sp.solve((sp.Eq(soln[0]), sp.Eq(soln[1]), sp.Eq(soln[2])), (v1.diff(t, t), v2.diff(t, t), v3.diff(t, t)))
```

```python
solvedsoln
```

```python

```


# Постановка задачи оптимального управления

$\dot{x} = f(x,u,t), \qquad  x \in \mathbb{R}^n, \qquad  u \in \mathbb{U}^m$

$x$ - вектор фазовых переменных, $u$ - вектор управляющих параметров.

Начальные и терминальные условия условия:

$$x(t_0) \in A = \left\{
    \begin{array}\\
       a_1(x,t) = 0 \\
        \dots \\
        a_l(x,t) = 0
    \end{array}
\right. \qquad
x(t_f) \in B = \left\{
    \begin{array}\\
       b_1(x,t) = 0 \\
        \dots \\
        b_r(x,t) = 0
    \end{array}
\right.$$

Функционал
$$\Phi(x(t_f), t_f) \rightarrow \max_u$$

# Вспомогательные конструкции

Вектор сопряжённых переменных:

$$p \in \mathbf{R}^n$$

Функция Гамильтона-Понтрягина (гамильтониан):

$$H(x,p,u,t) = p'\cdot f$$

Дифференциальные уравнения для фазовых и сопряжённых переменных:

$$\begin{equation}\tag{1}
 \left\{
    \begin{array}\\
       \dot{x} = \frac{\partial H}{\partial p} \\
       \dot{p} = -\frac{\partial H}{\partial x}
    \end{array}
\right.
\end{equation}$$

Расширенный вектор фазовых переменных:
$$P = \left(
    \begin{array}\\
       p \\
       -H
    \end{array}
\right)$$

Расширенное терминальное многообразие:
$$B^* = B \cap (\Phi - \Phi^* = 0)$$

Условия трансверсальности:
$$\begin{equation}\tag{2}
P(t_0) \perp A, \qquad P(t_f) \perp B^*
\end{equation}$$

# Принцип маскимума Понтрягина

Для того чтобы некоторое управление $u(t)$ было оптимальным, необходимо, чтобы на соответствующей ему траектории $x(t), p(t)$ для системы (1) с граничными условиями (2) выполнялось условие:
$$H(x,p,u,t) = \max_{v\in\mathbb{U}}H(x,p,v,t), \quad t \in [t_0,t_f]$$

Примечание: в задаче с закреплённым левым концом (старт из фиксированной точки) и свободным левым (терминальное многообразие определяется только заданным моментом времени $t_f$) условия трансверсальности выглядят следующим образом:
$$p(t_f) = \frac{\partial \Phi(x,t)}{\partial x}$$

# Пример (управление одномерным осциллятором)

Уравнения модели:
$$m\cdot \ddot{x} + k\cdot x = F(t) \qquad \rightarrow \qquad
\left\{
    \begin{array}\\
       \dot{x} = v \\
       \dot{v} = -\omega^2 \cdot x + u
    \end{array}
\right.
\qquad \omega^2 = \frac{k}{m}, \qquad \frac{F(t)}{m} = u, \qquad |u| \leq u^{*}
$$

Начальное многообразие:
$$    \begin{array}\\
       x(0) = 0 \\
       v(0) = 0
    \end{array}
\qquad \rightarrow \qquad
x(t_0) \in A = \left\{
    \begin{array}\\
       x = 0 \\
       v = 0 \\
       t = 0
    \end{array}
\right.
$$

Требуется обеспечить максимальное значение координаты $x$ к моменту времени $t_f$, таким образом, функционал задачи:
$$\Phi(x(t_f), t_f) = x(t_f) \rightarrow \max_{u}$$

Функция Гамильтона-Понтрягина:
$$H(x,p,u,t) = p_x \cdot v + p_v \cdot (-\omega^2 \cdot x + u)$$

Управление $u$ входит в $H$ линейно, а значит, чтобы удовлетворить условиям принципа максимума
$$u(t) = u^{*} \cdot sign(p_v)$$

Дифференциальные уравнения на сопряженные переменные:
$$
\left\{
    \begin{array}\\
       \dot{p_x} = -\frac{\partial H}{\partial x} = \omega^2 p_v \\
       \dot{p_v} = -\frac{\partial H}{\partial v} = - p_x
    \end{array}
\right.
$$

Условия трансверсальности на левом конце:
$$P(t_0) = \lambda_x \cdot\frac{\partial x}{\partial \left(
    \begin{array}\\
       x \\
       v \\
       t
    \end{array}
\right)} +
\lambda_v \cdot\frac{\partial v}{\partial \left(
    \begin{array}\\
       x \\
       v \\
       t
    \end{array}
\right)} +
\lambda_t \cdot\frac{\partial t}{\partial \left(
    \begin{array}\\
       x \\
       v \\
       t
    \end{array}
\right)}=
\left(
    \begin{array}\\
       \lambda_x \\
       \lambda_v \\
       \lambda_t
    \end{array}
\right),
$$

то есть точке ортогонален любой вектор и никаких дополнительных ограничений на векотр сопряжённых переменных на левом конце ($t = 0$) не возникает.

Многообразие $B^{*}$ определяется как
$$B^{*} = \left\{
    \begin{array}\\
       t = t_f \\
       x - x^{*} = 0
    \end{array}
\right. ,$$
где $x^{*}$ неизвестное оптимальное значение функционала.

Условие трансверсальности на правом конце (заметьте, что множитель Лагранжа при функционале равен 1):
$$P(t_f) = \frac{\partial \Phi}{\partial
\left(
    \begin{array}\\
       x \\
       v \\
       t
    \end{array}
\right)} +
\lambda_{tf} \cdot \frac{\partial t}{\partial
\left(
    \begin{array}\\
       x \\
       v \\
       t
    \end{array}
\right)} =
\left(
    \begin{array}\\
       1 \\
       0 \\
       \lambda_{tf}
    \end{array}
\right).$$

Таким образом, уравнения на сопряжённые переменные имеют следующие решения:
$$
    \begin{array}\\
       p_x(t) = \cos{(\omega \cdot (t - t_f))} \\
       p_v(t) = - \frac{1}{\omega} \sin{(\omega \cdot (t - t_f))}
    \end{array}
$$

Синтезированное оптимальное управление:
$$ u(t) = u^{*} \cdot sign(-\sin{(\omega \cdot (t - t_f))})$$.

```python
import numpy as np
from numpy import cos, sin, tan
import scipy.integrate as sci
import matplotlib.pyplot as plt
```

```python
class Parameters(object):
    pass

# Physical parameter
param = Parameters()  # instance of class Parameters
param.m = 1      # mass
param.k = 2      # spring constant
param.f_max = 1  # max force
param.u_max = param.f_max / param.m
param.omega = np.sqrt(param.k / param.m)
param.omega2 = param.k / param.m

# Simulation parameter
sim_param = Parameters()  # instance of class Parameters
sim_param.t0 = 0          # start time
sim_param.tf = 10         # final time
sim_param.dt = 0.01       # step-size
```

```python
def rhs(t, x, p):
    dxdt = np.zeros(2)

    r = x[0]   # position vector
    v = x[1]   # velocity vector

    u = control(t, p)  # control vector

    # dxdt = f(x, u, t):
    dxdt[0] = v
    dxdt[1] = - p.omega2 * r + u

    # return state derivative
    return dxdt
```

```python
def control(t, p):

    u = p.u_max * np.sign(-np.sin(p.omega * (t - sim_param.tf)))

    return u
```

```python
def plot_data(x, u, t, fig_width, fig_height):
    """Plotting function of simulated state and actions
    Args:
        x(ndarray) : state-vector trajectory
        u(ndarray) : control vector trajectory
        t(ndarray) : time vector
        fig_width  : figure width in cm
        fig_height : figure height in cm
        save (bool): save figure (default: False)
    Returns: None
    """
    # creating a figure with 2 subplots, that share the x-axis
    fig1, (ax1, ax2) = plt.subplots(2)

    # plot coords in subplot 1
    ax1.plot(t, x[:, 0], color='r')

    # plot velocities in subplot 2
    ax2.plot(t, np.rad2deg(u[:]), label=r'$u(t)$', lw=1, color='b')

    # Grids
    ax1.grid(True)
    ax2.grid(True)

    # set the labels on the x and y axis and the titles
    ax1.set_title('Position')
    ax1.set_ylabel(r'm')
    ax1.set_xlabel(r't in s')
    ax2.set_title('Control')
    ax2.set_ylabel(r'm/s^2')
    ax2.set_xlabel(r't in s')

    # automatically adjusts subplot to fit in figure window
    plt.tight_layout()

    return None
```

```python
# time vector
tt = np.arange(sim_param.t0, sim_param.tf, sim_param.dt)

# initial state
x0 = [0, 0]

# simulation
sol = sci.solve_ivp(lambda t, x: rhs(t, x, param), (sim_param.t0, sim_param.tf), x0, t_eval=tt)
y_traj = sol.y.T
u_traj = control(tt, param)

# plot
plot_data(y_traj, u_traj, tt, 12, 16)
plt.show()
```

# Численный метод решения задач оптимального управления

Формула для малых приращений функционала задачи оптимального управления:
$$
\begin{equation}\tag{3}
\Delta \Phi(x(t_f), t_f) = p'(t_0)\cdot \Delta x(t_0) + \int_{t_0}^{t_f} (H(x,p,u+\Delta u,t) - H(x,p,u,t)) dt + O(\varepsilon^2),
\end{equation}$$
где $\Delta x(t_0) = O(\varepsilon)$ и $\Delta u(t) = O(\varepsilon)$.

На основе этой формулы представляется возможным реализовать следующую итеративную процедуру (на примере задачи с фиксированным левым и свободным правым концом):
1. зафиксировать некоторое управление $u(t)$
2. решить задачу Коши с уравнениями модели и получить $x(t)$
3. решить задачу Коши с уравнениями для сопряжённых переменных и условиями трансверсальности на правом конце и получить
$p(t)$
4. разбить интервал $[t_0, t_f]$ на равные достаточно малые подинтервалы и выбрать тот из них, для которого изменение $u(t)$, оставляющее $u$ в области допустимых управлений, обеспечивает максимальное приращение функционала в силу формулы (3)
5. изменить управление $u(t)$ в соответствии с п.4 и повторить пп. 2-4

Процедура повторяется, пока удаётся обеспечивать приращение функционала, то есть улучшать управление в смысле поставленной задачи.

Примечание: уравнения для $p$ придётся интегрировать в обращенном времени, т.к. условия для $p$ определены на правом конце. Пример интегрирования дифференциальных уравнений в обращенном времени приведён в следующей ячейке.

```python
# функция, возвращающая непрерывное по времени управление, на основе определенного массива u(t)
def u(t):
    if t > sim_param.tf:
        raise ValueError
    i = int((t - sim_param.t0)//sim_param.dt)
    return u0[i]

#функция, вычисляющая правую часть уравнений движения осциллятора
def rhs_x(t, x, params):
    dxdt = np.zeros(2)
    dxdt[0] = x[1]
    dxdt[1] = u(t) - params.omega2 * x[0]
    return dxdt

#функция, вычисляющая правую часть уравнений на сопряжённые переменные
def rhs_p(t, p, params):
    dpdt = np.zeros(2)
    dpdt[0] = params.omega2 * p[1]
    dpdt[1] = -p[0]
    return dpdt

#функция Гамильтона-Понтрягина
def H(t, x, p, u, params):
    return p[:, 0] * x[:, 1] + p[:, 1]*(u - params.omega2 * x[:, 0])
```

```python
time = np.arange(sim_param.t0, sim_param.tf + sim_param.dt, sim_param.dt)
u0 = np.zeros(time.shape[0])
x0 = [0, 0]
pf = [1, 0]

#интегрируем уравнения движения слева направо
sx = sci.solve_ivp(lambda t, x: rhs_x(t, x, param), (sim_param.t0, sim_param.tf), x0, t_eval = time)
x = sx.y.T

#интегрируем уравнения движения в обращенном времени (от правого конца к левому,
#т.к. условия трансверсальности даны на правом конце)
sp = sci.solve_ivp(lambda t, p: rhs_p(t, p , param), (sim_param.tf, sim_param.t0), pf, t_eval = sim_param.tf - time)
p = sp.y.T[::-1] # получен в обратном порядке
```

```python
u_min = -np.ones(time.shape[0]) * param.u_max
u_max = np.ones(time.shape[0]) * param.u_max
```

```python
#Вычисляем гамильтониан на полученной для фиксированного управления u0(t) траектории x(t), p(t)
H0 = H(time, x, p, u0, param)

#а также гамильтонианы на той же траектории, но для управлений принадлежащих верхней (u_max) и нижней (u_min) границам
H_min = H(time, x, p, u_min, param)
H_max = H(time, x, p, u_max, param)
```

```python
# вычисляем интервал на котором изменение управления с фиксированного u0 на максимальное
# даст наибольшее локальное приращение гамильтониану, а значит и функционалу задачи
i = np.argmax(H_max - H0)

if H_max[i] - H0[i] > 0:
    u0[i] = param.u_max

# то же для изменения с u0 на u_min
j = np.argmax(H_min - H0)

if H_min[i] - H0[i] > 0:
    u0[j] = -param.u_max

i, j
```

    (445, 223)

Так выглядела одна итерация вычислительной процедуры. В следующей ячейке все описанные выше операции повторяются в цикле. Цикл может содержать некоторое наперед заданное число итераций или строиться по условию while (пока удаётся увеличивать функционал).

```python
F = []

for i in range(u0.shape[0]):
    sx = sci.solve_ivp(lambda t, x: rhs_x(t, x, param), (sim_param.t0, sim_param.tf), x0, t_eval = time)
    x = sx.y.T

    F.append(x[-1, 0])

    sp = sci.solve_ivp(lambda t, p: rhs_p(t, p , param), (sim_param.tf, sim_param.t0), pf, t_eval = sim_param.tf - time)
    p = sp.y.T[::-1] # получен в обратном порядке

    H0 = H(time, x, p, u0, param)
    H_min = H(time, x, p, u_min, param)
    H_max = H(time, x, p, u_max, param)

    i = np.argmax(H_max - H0)
    u0[i] = param.u_max
    if H_max[i] - H0[i] > 0:
        u0[i] = param.u_max

    j = np.argmax(H_min - H0)
    if H_min[j] - H0[j] > 0:
        u0[j] = -param.u_max
```

Строим графики переменной $x$, траектории системы на фазовой плоскости ${x,v}$, полученного закона управления $u(t)$ и график изменения функционала $\Phi$ с числом итераций

```python
plt.plot(time, x[:, 0])
```

```python
plt.plot(x[:, 0], x[:, 1])
```

```python
plt.plot(time, u0)
```

```python
plt.plot(F)
```

```python

```

Lecture Jupyter Notebooks/Lecture 11. Pontryagin's Maximum Principle.ipynb


```python

```

```javascript
%%javascript
MathJax.Hub.Config({
    TeX: { equationNumbers: { autoNumber: "AMS" } }
});
```

    <IPython.core.display.Javascript object>

# Exercise 1

Let us consider the sequence $U_n$ given by
\begin{equation}\label{fib}
\left\lbrace
\begin{array}{ll}
U_0 &= 1,\\
U_1 &= 2,\\
U_{n} &=-3U_{n-1} +U_{n-2}, \;\; \forall\; n=2,3,4\cdots
\end{array}\right.
\end{equation}

Write a python function named  SeqTerms that takes as input an integer $n,\;\;n\geq 0$ and return an array of the first $n$ terms (i.e. $U_0, \cdots, U_{n-1}$) of the sequence \eqref{fib}.

```python
import numpy as np
def Seq(n):
    a=1
    b=2
    if n==0:
        return 1
    if n==1:
        return 2
    for i in range(2,n+1):
        c=-3*b+a
        a=b
        b=c
    return c
Seq(2)
def SeqTerms(n):
    l=[]
    g=np.vectorize(Seq)
    for i in range(n):
        l+=[Seq(i)]
    return l
SeqTerms(4)
```

    [1, 2, -5, 17]

# Exercise 2

Let $\{ x_k\}$ be a partition of $[a,b]$ such that $a=x_0<x_1<\cdots<x_{N-1}<x_N=b$ and $H$ be the length of the $k$-th subinterval ($H = x_k - x_{k-1}$),
then we have
$$\int_a^bf(x)dx \approx \sum_{k=1}^N \frac{f(x_{k-1})+f(x_k)}{2}H = A$$

1. Write a function named Trap that takes $a,b,N, f$ as inputs and return A

```python
def trap(a,b,N,f):
    C=np.linspace(a,b,N+1)
    g=np.vectorize(f)
    A=g(C)
    S=0
    for i in range(1,len(A)):
        S+=A[i]+A[i-1]
    K=1/2*S*((b-a)/N)
    return K
f= lambda x: x**3+7
trap(0,1,10**6,f)
```

    7.250000000000129

2. Write a Python code to compute and display an approximation $Aquad$ of  the integral bellow using the Python function $quad$
$$A = \int_{0}^{2} \dfrac{x^3+5x-20}{x^2+3}dx$$

```python
from scipy.integrate import quad
a = 0
b = 2
f = lambda x: (x**3+5*x-20)/(x**2+3)
Aquad= quad(f, a, b)[0]
print(Aquad)
```

    -7.049316535735796

3. write a Python function ErrorTrap that takes $M$ as input and return an arrays $ErrorInt$ and $ListN$. Here, $ErrorInt$ contains the absolute errors between $Aquad$ and the approximation of the integral $A$ obtained using the function Trap for all positve intergers $N$ in $ListN$ the set of all multiples of 10 less or equal to $M$.

```python
def ErrorTrap(M):
    u= lambda x: abs(quad(f,0,2)[0]-trap(0,2,x,f))
    ListN=[]
    #ErrorInt=np.zeros(M)
    for i in range(1,M+1):
        if i%10==0:
            ListN+=[i]
    g=np.vectorize(u)
    ErrorInt=g(ListN)
    return ErrorInt, ListN
ErrorTrap(30)
```

    (array([0.0030795, 0.0007707, 0.0003426]), [10, 20, 30])

4. Plot the output $ErrorInt$ against $ListN$ for $M=200$

```python
𝐸𝑟𝑟𝑜𝑟𝐼𝑛𝑡 , 𝐿𝑖𝑠𝑡𝑁 = ErrorTrap(200)
print(𝐿𝑖𝑠𝑡𝑁)
print(ErrorInt)
```

    [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140, 150, 160, 170, 180, 190, 200]
    [3.07950054e-03 7.70701307e-04 3.42601551e-04 1.92726674e-04
     1.23349010e-04 8.56605200e-05 6.29349176e-05 4.81848732e-05
     3.80721757e-05 3.08385649e-05 2.54864800e-05 2.14157629e-05
     1.82477772e-05 1.57340710e-05 1.37061368e-05 1.20464185e-05
     1.06708827e-05 9.51816892e-06 8.54262634e-06 7.70972323e-06]

# Exercise 3

1. Write code to solve the following system of ordinary differential equations using the Python function odeint.

$$
\begin{cases}
\dfrac{dx_1}{dt}& = & -\dfrac{1}{2}x_1\\\\
\dfrac{dx_2}{dt}& = & \dfrac{1}{2}x_1-\dfrac{1}{4}x_2\\\\
\dfrac{dx_3}{dt}& = & \dfrac{1}{4}x_2-\dfrac{1}{6}x_3
\end{cases}, \text{ on } [0,4]
$$

Subject to the initial conditions $x_1(0) = 1, x_2(0) = 1, x_3(0) = 1$.

```python
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

# function that returns dz/dt
def model(z,t):
    x_1,x_2,x_3 = z
    dx_1dt = -1/2*x_1
    dx_2dt = 1/2*x_1 -1/4*x_2
    dx_3dt = 1/4*x_2-1/6*x_3
    return dx_1dt,dx_2dt,dx_3dt

# initial condition
z0 = [1,1,1]

# time points
a = 0
b = 4
N = 100
t = np.linspace(a,b,N+1)

# solve ODE
z = odeint(model,z0,t)

x_1 = z[:,0]
x_2 = z[:,1]
x_3=z[:,2]

plt.plot(t,x_1,'b-')
plt.plot(t,x_3,'r--')
plt.plot(t,x_2,'green');

```

```python
def f(z,t):
    x1,x2,x3=z
    dx1dt=-1/2*z[0]
    dx2dt=1/2*z[0]-1/4*z[1]
    dx3dt=1/4*z[1]-1/6*z[2]
    return dx1dt, dx2dt,dx3dt
f(6,7)
```

2. The exact solution of the above system of ODEs is given by

$$
\begin{cases}
x_1(t)& = & e^{-t/2}\\
x_2(t)& = & -2e^{-t/2}+3e^{-t/4}\\
x_3(t)& = & \dfrac{3}{2}e^{-t/2} - 9e^{-t/4} + \dfrac{17}{2}e^{-t/6}
\end{cases}
$$

Use $Subplot$ to plot side by side

- each exact and approximate solution in the same window
- and their absolute error vs the time

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

# x_1(t)=np.exp(-t/2)
# x_2(t)=-2*np.exp(-t/2)+3*np.exp(-t/4)
# x_3(t)=3/2*np.exp(-t/2)-9*np.exp(-t/4)+17/2*np.exp(-t/6)
# #plot results
plt.subplot(3,1,1)
plt.plot(t,np.exp(-t/2),'b')
plt.plot(t,x_1,'y--')
plt.xlabel('time')
plt.ylabel('x_1(t)')
plt.show()

#plot results
plt.subplot(3,1,2)
plt.plot(t,-2*np.exp(-t/2)+3*np.exp(-t/4),'y-')
plt.plot(t,x_2,'g--')
plt.xlabel('time')
plt.ylabel('x_2(t)')
plt.show()

plt.subplot(3,1,3)
plt.plot(t,3/2*np.exp(-t/2)-9*np.exp(-t/4)+17/2*np.exp(-t/6),'r-')
plt.plot(t,x_3,'b--')
plt.xlabel('time')
plt.ylabel('x_3(t)')
plt.show()

#plot results
# plt.subplot(3,1,3)
# plt.plot(x,y)
# plt.xlabel('x')
# plt.ylabel('y')
# plt.show()

```

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

# x_1(t)=np.exp(-t/2)
# x_2(t)=-2*np.exp(-t/2)+3*np.exp(-t/4)
# x_3(t)=3/2*np.exp(-t/2)-9*np.exp(-t/4)+17/2*np.exp(-t/6)
# #plot results

plt.subplot(3,1,1)
plt.title("Absolute Error vs Times")
#plt.plot(t,np.exp(-t/2),'b')
plt.plot(t,abs(x_1-np.exp(-t/2)),'b-')
plt.xlabel('time')
plt.ylabel('absolute error of x_1')
plt.show()

#plot results
plt.subplot(3,1,2)
#plt.plot(t,-2*np.exp(-t/2)+3*np.exp(-t/4),'g-')
plt.plot(t,abs(x_2+2*np.exp(-t/2)-3*np.exp(-t/4)),'g-')
plt.xlabel('time')
plt.ylabel('absolute error of x_2')
plt.show()

plt.subplot(3,1,3)
#plt.plot(t,3/2*np.exp(-t/2)-9*np.exp(-t/4)+17/2*np.exp(-t/6),'r-')
plt.plot(t,abs(x_3-3/2*np.exp(-t/2)+9*np.exp(-t/4)-17/2*np.exp(-t/6)),'r-')
plt.xlabel('time')
plt.ylabel('absolute error of x_3')
plt.show()

#plot results
# plt.subplot(3,1,3)
# plt.plot(x,y)
# plt.xlabel('x')
# plt.ylabel('y')
# plt.show()

```

# Exercise 4

Let $\{ t_k\}$ be a partition of $[a,b]$ such that $a=t_1<t_2<\cdots<t_{N}=b$ and $H$ be the constant length of the $k$-th subinterval ($H = t_k - t_{k-1}$). Let us consider initial value problem

\begin{equation}\label{eul2}
  \begin{cases}
    \dfrac{dz}{dt} = f(z,t),      & \quad \text{on } [a, b]\\\\
    z(a) = c,
  \end{cases}
\end{equation}
where $z,f,c\in R^M$ i.e. $z = [x_1, x_2,\cdots, x_{M}]$, $c = [x_1(a), x_2(a),\cdots, x_{M}(a)]$ and $f = [f_1, f_2,\cdots, f_{M}]$. Note that \eqref{eul2} is a the general form of system of ODEs.

Let $t, z_k,Z$ defined as follows $$t=[t_1,t_2,\cdots,t_{N-1},t_{N}],\quad z_k = [x_1(t_k), x_2(t_k),\cdots, x_{M}(t_k)], \quad
Z =\begin{pmatrix}
x_1(t_1)& x_2(t_1)&\cdots& x_{M}(t_1)\\
x_1(t_2)& x_2(t_2)&\cdots& x_{M}(t_2)\\
\vdots& \vdots&\ddots& \vdots\\
x_1(t_{N})& x_2(t_{N})&\cdots& x_{M}(t_{N})
\end{pmatrix}
$$

1. Write a python function  EulerOdeSys  that takes $f,c,t$ and return the solution $Z$ of the initial value problem \eqref{eul2} using Euler method i.e.
$$ z_{k+1} = z_k + Hf(z_k,t_k) $$

```python
def Euler(f,c,t):
    N=len(t)
    M=len(c)
    Z=np.zeros((N,M))
    Z[0]=c
    for i in range(1,N+1):
        z2=z1+(t[i+1]-t[i])*f(z2,t[i])
        z1=z2

```

```python
np.shape?
```

```python
np.zeros?
```

```python

```

```python

```

```python
def f(z,t):
    x_1,x_2,x_3= z
    dx_1dt = -1/2*x_1
    dx_2dt = 1/2*x_1 -1/4*x_2
    dx_3dt = 1/4*x_2-1/6*x_3
    return np.array([dx_1dt,dx_2dt,dx_3dt])

def EulerOdeSys(f,c,t):
    n=len(t)
    Z = np.zeros((len(t),)+ np.shape(c))
    Z[0]= c
    for i in range(n-1):
        h =(t[i+1] - t[i])
        Z[i+1]= Z[i]+ h*f(Z[i],t[i])
    return Z
# def f(x,y):
#     return x+y
c=[1,1,1]
t=np.linspace(0,4,100)
EulerOdeSys(f,c,t)
plt.plot(t,EulerOdeSys(f,c,t));
```

2. Write a python function  RK4OdeSys  that takes $f,c,t$ and return the solution $Z$ of the initial value problem (1) using the fourth order Runge-Kutta method i.e.

\begin{equation}
\begin{cases}
k_1 = f(z_k,t_k),\\\\
k_2 = f(z_k+H\dfrac{k_1}{2}, t_k + \dfrac{H}{2}),\\\\
k_3 = f(z_k+H\dfrac{k_2}{2}, t_k + \dfrac{H}{2}),\\\\
k_4 = f(z_k+Hk_3, t_k + H),\\\\
z_{k+1} = z_k + \dfrac{H}{6}(k_1+2k_2+2k_3+k_4)
\end{cases}
\end{equation}

```python
def RK4OdeSys(f,c,t):
    n = len (t)
    Z = np.zeros((len(t),)+ np.shape(c))
    Z[0]= c
    for i in range (n-1):
        k1 = f(Z[i] ,t[i])
        h =(t[i+1] - t[i])/2
        k2 = f(Z[i]+ h*k1 , t[i]+h)
        k3 = f(Z[i]+ h*k2 , t[i]+h)
        k4 = f(Z[i]+2*h*k3 ,t[i]+2*h )
        Z[i+1]= Z[i]+ h/3*(k1 +2*k2 +2*k3+k4 )
    return Z

# def f(x,y):

#     return
c=[1,1,1]
t=np.linspace(0,4,100)
RK4OdeSys(f,c,t)
plt.plot(RK4OdeSys(f,c,t));
```

3. Solve the system of ODEs in $Exercise2$ using your function  EulerOdeSys  and  RK4OdeSys

```python

```

4. By plotting the absolute error in the approximate and exact solutions, tell us which function gives a more accurate solution of a system of ODEs.

```python

```

    1.791759469228055

# Exercise 5

Let consider us consider the function  primes  that takes $n$ as input and return a list of primes less than $n$

```python
# This cell is only to import the labraries

import numpy as np
import time

def primes(n):
    """ Returns  a list of primes < n """
    sieve = [True] * (n//2)
    for i in range(3,int(n**0.5)+1,2):
        if sieve[i//2]:
            sieve[i*i//2::i] = [False] * ((n-i*i-1)//(2*i)+1)
    return [2] + [2*i+1 for i in range(1,n//2) if sieve[i]]

```

For any integer $n>0$ and a prime number $p$, define $\nu_p(n)$ as the greatest integer $r$ such that $p^r$ divides $n$.
Define $$ D(n,m) = \sum_{p\; prime} \Bigl| \nu_p(n) - \nu_p(m)\Bigr| $$

For example $D(14,24)=4$.

Furthermore, define

$$S(N)  = \sum_{n=1}^{N}\sum_{m=1}^{N}D(n,m).$$

You are given $S(10)=210$.

1. Write an efficient python function, Func_S , that takes $N$ as input and return the value $S(N)$.

```python
from math import floor
from math import log as ln
def nu(n,p):
    L=[]
    for i in range(floor(ln(n)//ln(p))+2):
        if n%(p**i)==0:
            L+=[i]
    return L[-1]

def D(n,m):
    list_prime=primes(max(m,n)+1)
    SumD=0
    for i in list_prime:
        SumD+=abs(nu(n,i)-nu(m,i))

    return SumD

print(D(14,15))

def Func_S(N):
    s=0
    for i in range(1,N+1):
        for j in range(1,N+1):
            #if j!=i:
            s=s+D(i,j)
    return s

Func_S(10)
nu(7,23)

```

    4

    0

2. Compute $S(10)$ and display its computational time

```python
N = 10
time_start = time.perf_counter()
S = Func_S(N)
time_elapsed = (time.perf_counter() - time_start)
print('S({}) =  {}'.format(N,S))
print('computational Time = ', time_elapsed)
```

    S(10) =  210
    computational Time =  0.0012173870054539293

3. Compute $S(100)$ and display its computational time

```python
N = 100
time_start = time.perf_counter()
S = Func_S(N)
time_elapsed = (time.perf_counter() - time_start)
print('S({}) =  {}'.format(N,S))
print('computational Time = ', time_elapsed)
```

    S(100) =  37018
    computational Time =  0.4434022469940828

4. Compute $S(1000)$ and display its computational time

```python
N = 1000
time_start = time.perf_counter()
S = Func_S(N)
time_elapsed = (time.perf_counter() - time_start)
print('S({}) =  {}'.format(N,S))
print('computational Time = ', time_elapsed)
```

    S(1000) =  4654406
    computational Time =  256.6060328760068

5. Compute $S(10000)$ and display its computational time

```python
N = 10000
time_start = time.perf_counter()
S = Func_S(N)
time_elapsed = (time.perf_counter() - time_start)
print('S({}) =  {}'.format(N,S))
print('computational Time = ', time_elapsed)
```

6. Compute $S(100000)$ and display its computational time

```python
N = 100000
time_start = time.perf_counter()
S = Func_S(N)
time_elapsed = (time.perf_counter() - time_start)
print('S({}) =  {}'.format(N,S))
print('computational Time = ', time_elapsed)
```

7. Compute $S(1000000)$ and display its computational time

```python
N = 1000000
time_start = time.perf_counter()
S = Func_S(N)
time_elapsed = (time.perf_counter() - time_start)
print('S({}) =  {}'.format(N,S))
print('computational Time = ', time_elapsed)
```

# Exercise 6
1. Read the Covid-19 dataset

```python
import pandas as pd
import numpy as np
a=pd.read_csv('Covid-19.csv')
a
```

      Date_reported
      Country_code
      Country
      WHO_region
      New_cases
      Cumulative_cases
      New_deaths
      Cumulative_deaths

      0
      2020-01-03
      AF
      Afghanistan
      EMRO
      0
      0
      0
      0

      1
      2020-01-04
      AF
      Afghanistan
      EMRO
      0
      0
      0
      0

      2
      2020-01-05
      AF
      Afghanistan
      EMRO
      0
      0
      0
      0

      3
      2020-01-06
      AF
      Afghanistan
      EMRO
      0
      0
      0
      0

      4
      2020-01-07
      AF
      Afghanistan
      EMRO
      0
      0
      0
      0

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      149542
      2021-09-20
      ZW
      Zimbabwe
      AFRO
      199
      127938
      4
      4567

      149543
      2021-09-21
      ZW
      Zimbabwe
      AFRO
      248
      128186
      2
      4569

      149544
      2021-09-22
      ZW
      Zimbabwe
      AFRO
      0
      128186
      0
      4569

      149545
      2021-09-23
      ZW
      Zimbabwe
      AFRO
      618
      128804
      23
      4592

      149546
      2021-09-24
      ZW
      Zimbabwe
      AFRO
      330
      129134
      8
      4600

149547 rows × 8 columns

2. Drop the Country code column

```python
del a['Country_code']
a
```

      Date_reported
      Country
      WHO_region
      New_cases
      Cumulative_cases
      New_deaths
      Cumulative_deaths

      0
      2020-01-03
      Afghanistan
      EMRO
      0
      0
      0
      0

      1
      2020-01-04
      Afghanistan
      EMRO
      0
      0
      0
      0

      2
      2020-01-05
      Afghanistan
      EMRO
      0
      0
      0
      0

      3
      2020-01-06
      Afghanistan
      EMRO
      0
      0
      0
      0

      4
      2020-01-07
      Afghanistan
      EMRO
      0
      0
      0
      0

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      149542
      2021-09-20
      Zimbabwe
      AFRO
      199
      127938
      4
      4567

      149543
      2021-09-21
      Zimbabwe
      AFRO
      248
      128186
      2
      4569

      149544
      2021-09-22
      Zimbabwe
      AFRO
      0
      128186
      0
      4569

      149545
      2021-09-23
      Zimbabwe
      AFRO
      618
      128804
      23
      4592

      149546
      2021-09-24
      Zimbabwe
      AFRO
      330
      129134
      8
      4600

149547 rows × 7 columns

3. Randomely choose three different countries

```python
a.sample(n=3)
b=a['Country']
rand=b.sample(n=3)
rand
# b=a.sample(n=3)
# b
```

    65561                Italy
    102871            Paraguay
    102512    Papua New Guinea
    Name: Country, dtype: object

4. Select and display the records for those three countries

```python
q=a[a['Country'].isin(rand)]
q
```

      Date_reported
      Country
      WHO_region
      New_cases
      Cumulative_cases
      New_deaths
      Cumulative_deaths

      64993
      2020-01-03
      Italy
      EURO
      0
      0
      0
      0

      64994
      2020-01-04
      Italy
      EURO
      0
      0
      0
      0

      64995
      2020-01-05
      Italy
      EURO
      0
      0
      0
      0

      64996
      2020-01-06
      Italy
      EURO
      0
      0
      0
      0

      64997
      2020-01-07
      Italy
      EURO
      0
      0
      0
      0

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      103479
      2021-09-20
      Paraguay
      AMRO
      42
      459622
      3
      16126

      103480
      2021-09-21
      Paraguay
      AMRO
      28
      459650
      2
      16128

      103481
      2021-09-22
      Paraguay
      AMRO
      15
      459665
      4
      16132

      103482
      2021-09-23
      Paraguay
      AMRO
      30
      459695
      4
      16136

      103483
      2021-09-24
      Paraguay
      AMRO
      25
      459720
      2
      16138

1893 rows × 7 columns

5. Calculate and display the sum, the average of the cumulative cases of each WHO region.

```python
M=a.groupby('WHO_region').mean()
print("the average of cumulative case of each WHO region is\n",M['Cumulative_cases'])
S=a.groupby('WHO_region').sum()
print("the Sum of cumulative case of each WHO region is\n",S['Cumulative_cases'])
```

    the average of cumulative case of each WHO region is
     WHO_region
    AFRO     3.853838e+04
    AMRO     5.795662e+05
    EMRO     2.228024e+05
    EURO     3.936270e+05
    Other    6.911569e+02
    SEARO    1.181455e+06
    WPRO     4.465386e+04
    Name: Cumulative_cases, dtype: float64
    the Sum of cumulative case of each WHO region is
     WHO_region
    AFRO      1215886016
    AMRO     20479550865
    EMRO      3092942346
    EURO     15399475886
    Other         436120
    SEARO     8200477089
    WPRO       986180579
    Name: Cumulative_cases, dtype: int64

6. Calculate and display sum, the average of the cumulative deaths of each WHO region.

```python
M=a.groupby('WHO_region').mean()
print("the average of cumulative deaths of each WHO region is\n",M['Cumulative_deaths'])
S=a.groupby('WHO_region').sum()
print("the Sum of cumulative case of each WHO region is\n",S['Cumulative_deaths'])
```

    the average of cumulative deaths of each WHO region is
     WHO_region
    AFRO       916.600856
    AMRO     15814.192891
    EMRO      4672.818470
    EURO      8890.024871
    Other       11.554675
    SEARO    17467.325457
    WPRO       723.480281
    Name: Cumulative_deaths, dtype: float64
    the Sum of cumulative case of each WHO region is
     WHO_region
    AFRO      28918757
    AMRO     558810320
    EMRO      64868066
    EURO     347795553
    Other         7291
    SEARO    121240706
    WPRO      15978062
    Name: Cumulative_deaths, dtype: int64

7. Produce plots that look like the following three figures. Pay attention to the annotations.

7.a.

```python
import seaborn as sns
sns.boxplot(x="Country", y="New_cases", data=q)
sns.stripplot(x="Country", y="New_cases", data=q);#, jitter=True, edgecolor="gray")
plt.legend()
```

```python
a.groupby('WHO_region')['Cumulative_cases',"Cumulative_deaths"].sum().plot.bar(grid=False);
```

7.b.

```python
import matplotlib.pyplot as plt
sns.lineplot(x="Date_reported", y="Cumulative_cases", hue="Country",linewidth=5,data=q)

```

7.c.

```python

```


# Maximizing Concentration of an Intermediate in a Batch Reactor

A desired product $B$ is produced as intermediate in a series reaction

\begin{align}
A \overset{k_A}{\longrightarrow} B \overset{k_B}{\longrightarrow} C
\end{align}

where $A$ is a raw material and $C$ is a undesired b
y-product. The reaction operates at temperature where the rate constants are $k_A = 0.5\ \mbox{min}^{-1}$ and $k_A = 0.1\ \mbox{min}^{-1}$. The raw material is available as a solution with concenration $C_{A,f} = 2.0\ \mbox{moles/liter}$.

A 100 liter tank is available which can be operated as a batch reactor. At what point should the operation be terminated in order to maximize the concentration of $B$?

```
!pip install -q pyomo
!wget -N -q "https://ampl.com/dl/open/ipopt/ipopt-linux64.zip"
!unzip -o -q ipopt-linux64

ipopt_executable = '/content/ipopt'
```

## Mathematical Model

A material balance for an isothermal stirred batch reactor with a volume $V = 40$ liters and an initial concentration $C_{A,f}$ is given by

\begin{align}
V\frac{dC_A}{dt} & = - V k_A C_A \\
V\frac{dC_B}{dt} & = V k_A C_A - V k_B C_B
\end{align}

Eliminating the common factor $V$

\begin{align}
\frac{dC_A}{dt} & = - k_A C_A \\
\frac{dC_B}{dt} & = k_A C_A - k_B C_B
\end{align}

With an initial concentration $C_{A,f}$. A numerical solution to these equations is shown in the following cell.

```
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import odeint

V = 40     # liters
kA = 0.5   # 1/min
kB = 0.1   # l/min
CAf = 2.0  # moles/liter

def batch(X, t):
    CA, CB = X
    dCA_dt = -kA*CA
    dCB_dt = kA*CA - kB*CB
    return [dCA_dt, dCB_dt]

t = np.linspace(0,30,200)
soln = odeint(batch, [CAf,0], t)
plt.plot(t, soln)
plt.xlabel('time / minutes')
plt.ylabel('concentration / moles per liter')
plt.title('Batch Reactor')
plt.legend(['$C_A$','$C_B$'])
```

## Optimization with `scipy.minimize_scalar`

To find the maximum value, we first write a function to compute $C_B$ for any value of time $t$.

```
def CB(tf):
    soln = odeint(batch, [CAf, 0], [0, tf])
    return soln[-1][1]
```

We gain use `minimize_scalar` to find the value of $t$ that minimizes the negative value of $C_B(t)$.|

```
from scipy.optimize import minimize_scalar
minimize_scalar(lambda t: -CB(t), bracket=[0,50])
```

         fun: -1.3374806339222158
        nfev: 23
         nit: 19
     success: True
           x: 4.023594924340666

```
tmax = minimize_scalar(lambda t: -CB(t), bracket=[0,50]).x

print('Concentration c_B has maximum', CB(tmax), 'moles/liter at time', tmax, 'minutes.')
```

    Concentration c_B has maximum 1.3374806339222158 moles/liter at time 4.023594924340666 minutes.

## Solution Using Pyomo

The variable to be found is the time $t_f$ corresponding to the maximum concentration of $B$. For this purpose we introduce a scaled time

$$\tau = \frac{t}{t_f}$$

so that $\tau=1$ as the desired solution. The problem then reads

\begin{align}
\max_{t_f} C_B(\tau=1)
\end{align}

subject to

\begin{align}
\frac{dC_A}{d\tau} & = - t_f k_A C_A \\
\frac{dC_B}{d\tau} & = t_f(k_A C_A - k_B C_B)
\end{align}

The solution to this problem is implemented as a solution to the following Pyomo model.

```
from pyomo.environ import *
from pyomo.dae import *

V   = 40    # liters
kA  = 0.5   # 1/min
kB  = 0.1   # l/min
cAf = 2.0   # moles/liter

m = ConcreteModel()

m.tau = ContinuousSet(bounds=(0, 1))

m.tf = Var(domain=NonNegativeReals)
m.cA = Var(m.tau, domain=NonNegativeReals)
m.cB = Var(m.tau, domain=NonNegativeReals)

m.dcA = DerivativeVar(m.cA)
m.dcB = DerivativeVar(m.cB)

m.odeA = Constraint(m.tau,
    rule=lambda m, tau: m.dcA[tau] == m.tf*(-kA*m.cA[tau]) if tau > 0 else Constraint.Skip)
m.odeB = Constraint(m.tau,
    rule=lambda m, tau: m.dcB[tau] == m.tf*(kA*m.cA[tau] - kB*m.cB[tau]) if tau > 0 else Constraint.Skip)

m.ic = ConstraintList()
m.ic.add(m.cA[0]  == cAf)
m.ic.add(m.cB[0]  == 0)

m.obj = Objective(expr=m.cB[1], sense=maximize)

TransformationFactory('dae.collocation').apply_to(m)
SolverFactory('ipopt', executable=ipopt_executable).solve(m)
print('Concentration c_B has maximum', m.cB[1](), 'moles/liter at time', m.tf(), 'minutes.')
```

    Concentration c_B has maximum 1.3374805810221073 moles/liter at time 4.023594178375687 minutes.

```

```