# Sistema de Navegación Inteligente de un Hexápodo para la Evasión de Obstaculos (SINHEO)

## Cinemática Inversa para el Hexapodo
Este documento contiene el progreso de la Interfaz Gráfica de Usuario (GUI) construida en Python 3 para generar y visualizar los movimientos de un robot Hexapodo a partir de sus dimensiones y especificaciones del usuario.

Inicialmente utilizaremos el modulo **Tkinter**, el cual es una biblioteca gráfica **Tcl/Tk** de Python para realizar nuestra GUI. Para ello, debemos conocer ciertos conceptos básicos los cuales se ilustraran con el siguiente ejemplo:

```python
import tkinter as tk

# Testing window that display the current version of Tkinter
# tk._test()

# Read the whole documentation at http://effbot.org/tkinterbook/tkinter-index.htm

# To initialize Tkinter, we create a Tk root widget (an ordinary window).
root = tk.Tk()
# We create a label widget as child to the root window.
theLabel = tk.Label(root, text="This is too easy")
# We call the pack method on this widget. This tells it to size itself to fit the given text, and make itself visible.
theLabel.pack()
# We display the window using the Tkinter event loop. The program will stay in the event loop until we close the window
root.mainloop()
```

Inicialmente, se importa el módulo Tkinter con el alias **tk** a tráves de la siguiente línea de código:

```python
import tkinter as tk
```

Toda GUI en Tkinter necesita de una ventana raíz o *root* donde visualizar los otros componentes hijos o _widgets_ para nuestra aplicación. En este ejemplo, creamos un elemento de tipo *label*, el cual puede mostrar tanto textos como iconos o alguna otra imagen. Para este caso, especificamos que es un texto y a su vez al componente madre al que pertenece, _root_.

Luego llamamos el método `.pack()` en nuestro elemento para autoajustar su tamaño en función del texto y hacerlo visible. Aún así, no se visualizará nada debido hasta no ejecutar el **ciclo de eventos Tkinter**, el cual se encarga de manejar cualquier acción del usuario con la interfaz (clic, teclear, mensajes, etc).

## Definición de los marcos de referencia y modelos cinemáticos

### Marco de referencia global
El marco de referencia global para está simulación estará ubicado en una de las esquinas de la simulación (o terreno), el cual denotaremos por la letra **G** en todas nuestras notaciones. Además, los ejes de este marco de referencia global seguirán la regla de la mano derecha.

### Marco de referencia para el cuerpo (_Trunk_)
En el caso de la araña, se sitúa el origen de los ejes coordenados a la misma altura del panel inferior del robot, ubicandose en su centro geométrico. Se toma la denotación de los giros en los ejes **x - y - z** como **roll - pitch - yaw**, siendo estos relativos al marco de referencia global. En base a ello, las respectivas matrices de rotación se muestrán en las siguientes ecuaciones con sus respectivos ángulos:

\begin{align}
Roll: R_x(\gamma) =
\begin{bmatrix}
    1 & 0 & 0 \\
    0 & cos(\gamma) & -sin(\gamma) \\
    0 & sin(\gamma) & cos(\gamma)
\end{bmatrix}
,\
Pitch: R_y(\beta) =
\begin{bmatrix}
    cos(\beta) & 0 & sin(\beta) \\
    0 & 1 & 0 \\
    -sin(\beta) & 0 & cos(\beta)
\end{bmatrix}
,\
Yaw: R_z(\alpha) =
\begin{bmatrix}
    cos(\alpha) & -sin(\alpha) & 0 \\
    sin(\alpha) & cos(\alpha) & 0 \\
    0 & 0 & 1
\end{bmatrix}
\end{align}

### Definición de los paramétros Denavit - Hartenberg para la cinemática de una pata.
Basandonos en la convención Denavit - Hartenberg para describir la geometria de una cadena cinemática / manipulador robótico; los cuales se definen en la tabla de abajo, podemos obtener el modelo cinemático para una pata. Las notaciones para cada eslabón (o servo) están basadas en los insectos, siendo el primer servo (Une el cuerpo con la pata) denominado ***Coxa***, el segundo servo ***Femur*** y el último servo ***Tibia***.

| Parametro | Definición |
| :---: | :---: |
| $\theta_{\mathrm{i}}$ | Es el ángulo comprendido entre los ejes $x_{\mathrm{i - 1}}$ y $x_{\mathrm{i}}$, medido en torno al eje $z_{\mathrm{i - 1}}$, en el sentido de la regla de la mano derecha. |
| $\mathrm{d_i}$ | Es la desviación del eslabón, es decir, la distancia comprendida desde el eje $x_{\mathrm{i - 1}}$ al eje $x_{\mathrm{i}}$ a lo largo del eje $z_{\mathrm{i - 1}}$. |
| $\mathrm{a_i}$	| Es la longitud cinemática del eslabón, es decir, la distancia entre el eje $z_{\mathrm{i - 1}}$ y el $z_{\mathrm{i}}$ a lo largo del eje $x_{\mathrm{i}}$ |
| $\alpha_{\mathrm{i}}$ | Es el ángulo comprendido entre los ejes $z_{\mathrm{i - 1}}$ y $z_{\mathrm{i}}$, medido en torno al eje $x_{\mathrm{i}}$, en sentido de la regla de la mano derecha. |

La obtención de los parametros se realiza siguiendo los pasos de la metodologia en la figura siguiente, con la visualización de cada eje coordenado, rotaciones y la tabla resumiendo los valores / variables del modelo cinemático.

| Eslabón | $\theta_{\mathrm{i}}$ | $\mathrm{a_i}$ | $\mathrm{d_i}$ | $\alpha_{\mathrm{i}}$ |
| :---: | :---: | :---: | :---: | :---: |
| 1 | $\theta_1$ | $\mathrm{L_1}$ | $\mathrm{d_L}$ | $\pi/2$ |
| 2 | $\theta_2$ | $\mathrm{L_2}$ | 0 | $\pi$ |
| 3 | $\theta_3$ | $\mathrm{L_3}$ | 0 | 0 |

#### 1. Importación de modulos / librerias de Python

```python
# Importamos la librerias de Numpy y SymPy
from sympy import *
import numpy as np
import matplotlib.pyplot as plt
import math
```

#### 2. Definición de simbolos

```python
# Impresión en formato latex para las ecuaciones de la libreria SymPy.
init_printing(use_latex='mathjax')

# Simbolos generales para las matrices, ángulos y variables.
alpha, beta, gamma = symbols('alpha beta gamma')
g_k, r_c, theta_i, alpha_i, a_i, d_i = symbols('gamma_k r_c theta_i alpha_i a_i d_i')
# Simbolos especificos.
# Ángulos de los servos (variables)
t_1, t_2, t_3 = symbols('theta_1 theta_2 theta_3')
# Longitud de los eslabones de las patas.
L_1, L_2, L_3, d_L = symbols('L_1 L_2 L_3 d_L')
# Esta es la traslación desde el origen coordenadas globales hasta el centro del robot.
X_t, Y_t, Z_t = symbols('X_t Y_t Z_t')
# Coordenadas de origen de una pata
x_l, y_l, z_l = symbols('x_l y_l z_l')
```

#### 3. Matrices de rotación / traslación globales y locales

```python
# Matrices de rotación alrededor de cada eje con sus variables
RotMatrX = rot_axis1(gamma).T
RotMatrY = rot_axis2(beta).T
RotMatrZ = rot_axis3(alpha).T
# Matriz de rotación completa (Roll -> Pitch -> Yaw) o "balance, cabeceo, guiñada"
RotMatrRPY = RotMatrZ*RotMatrY*RotMatrX
# Matriz homogénea de transformación de cuerpo a coordenadas globales (B -> G)
TransfMatrRPY = RotMatrRPY.col_insert(3, Matrix([X_t, Y_t, Z_t])).row_insert(3, Matrix([[0, 0, 0, 1]]))

# Matriz general homogénea de traslación / rotación del eje central de la araña a cada pata.
# Llamado tambien ángulo local de guiñada o "Yaw"
TransfMatrBP = Matrix([[cos(g_k), -sin(g_k), 0, r_c*cos(g_k)],
                       [sin(g_k), cos(g_k), 0, r_c*sin(g_k)],
                       [0, 0, 1, 0],
                       [0, 0, 0, 1]])
# Matriz general homogénea de transformación del eslabón i al eslabón i-1.
TransfMatrLink = Matrix([[cos(theta_i), -sin(theta_i)*cos(alpha_i), sin(theta_i)*sin(alpha_i), a_i*cos(theta_i)],
                         [sin(theta_i), cos(theta_i)*cos(alpha_i), -cos(theta_i)*sin(alpha_i), a_i*sin(theta_i)],
                         [0, sin(alpha_i), cos(alpha_i), d_i],
                         [0, 0, 0, 1]])
# Matrices de transformación para cada eslabón en una pata.
# En orden: Matrix de eslabón 1 -> 0, 2 -> 1, 3 -> 2
TransfMatrL1 = TransfMatrLink.subs([(theta_i, t_1), (alpha_i, pi/2), (d_i, d_L), (a_i, L_1)])
TransfMatrL2 = TransfMatrLink.subs([(theta_i, t_2), (alpha_i, pi), (d_i, 0), (a_i, L_2)])
TransfMatrL3 = TransfMatrLink.subs([(theta_i, t_3), (alpha_i, 0), (d_i, 0), (a_i, L_3)])

# Matriz de transformación total desde elemento final de la pata hasta el cuerpo (Eslabón 0).
TransfMatrPaw = simplify(TransfMatrL1*TransfMatrL2*TransfMatrL3)
```

#### 4. Cinemática directa

```python
# Matrices de transformación de pata al centro del cuerpo.
# Basicamente es multiplicar la matriz de transformación total (anterior) con la de transformación desde el eje central.
# Notación general para las patas: R - Derecha, L - Izquierda // F - Frontal, M - Media, R - Trasera.
# En dirección contrareloj, empezando por la izquierda: 1 - LF, 2 - LM, 3 - LR, 4 - RR, 5 - RM, 6 - RF
TransfMatrLF = TransfMatrBP.subs(g_k, 2*pi/3)*TransfMatrPaw
TransfMatrLM = TransfMatrBP.subs(g_k, pi)*TransfMatrPaw
TransfMatrLR = TransfMatrBP.subs(g_k, 4*pi/3)*TransfMatrPaw
TransfMatrRR = TransfMatrBP.subs(g_k, 5*pi/3)*TransfMatrPaw
TransfMatrRM = TransfMatrBP.subs(g_k, 2*pi)*TransfMatrPaw
TransfMatrRF = TransfMatrBP.subs(g_k, pi/3)*TransfMatrPaw

# Matrices de transformación global para cada pata para obtener las coordenadas finales de cada una.
TransfMatrLF_G = TransfMatrRPY*TransfMatrLF*Matrix([0, 0, 0, 1])
TransfMatrLM_G = TransfMatrRPY*TransfMatrLM*Matrix([0, 0, 0, 1])
TransfMatrLR_G = TransfMatrRPY*TransfMatrLR*Matrix([0, 0, 0, 1])
TransfMatrRR_G = TransfMatrRPY*TransfMatrRR*Matrix([0, 0, 0, 1])
TransfMatrRM_G = TransfMatrRPY*TransfMatrRM*Matrix([0, 0, 0, 1])
TransfMatrRF_G = TransfMatrRPY*TransfMatrRF*Matrix([0, 0, 0, 1])
```

#### 5. Cinemática Inversa global a pata

```python
# Algoritmo para la Cinemática Inversa del robot
# Transformación Inversa - Global hasta cuerpo (G -> B)
TransfMatrRPY_IK = RotMatrRPY.T.col_insert(3, RotMatrRPY.T*Matrix([X_t, Y_t, Z_t])).row_insert(3, Matrix([[0, 0, 0, 1]]))
# Esto es equivalente a la siguiente operación.
# TransfMatrRPY_IK = simplify(TransfMatrRPY**-1)

# Toca simplificar este pedazo, ya que utilizo mucho esta variable como tal.
TransfMatrGBP = TransfMatrRPY_IK*TransfMatrBP

# Transformación Inversa - Cuerpo hasta la pata (Eslabón 0)
TransfMatrLF_IK = TransfMatrBP.subs(g_k, 2*pi/3)**-1
TransfMatrLM_IK = TransfMatrBP.subs(g_k, pi)**-1
TransfMatrLR_IK = TransfMatrBP.subs(g_k, 4*pi/3)**-1
TransfMatrRR_IK = TransfMatrBP.subs(g_k, 5*pi/3)**-1
TransfMatrRM_IK = TransfMatrBP.subs(g_k, 2*pi)**-1
TransfMatrRF_IK = TransfMatrBP.subs(g_k, pi/3)**-1

# Transformación Inversa - Global hasta la pata (Eslabón 0).
# Se multiplica cada transformación por una matriz cuyas variables seran las coordenadas del elemento final por pata.
TransfMatrLF_Coxa = simplify(TransfMatrRPY_IK*TransfMatrLF_IK*Matrix([x_l, y_l, z_l, 1]))
TransfMatrLM_Coxa = simplify(TransfMatrRPY_IK*TransfMatrLM_IK*Matrix([x_l, y_l, z_l, 1]))
TransfMatrLR_Coxa = simplify(TransfMatrRPY_IK*TransfMatrLR_IK*Matrix([x_l, y_l, z_l, 1]))
TransfMatrRR_Coxa = simplify(TransfMatrRPY_IK*TransfMatrRR_IK*Matrix([x_l, y_l, z_l, 1]))
TransfMatrRM_Coxa = simplify(TransfMatrRPY_IK*TransfMatrRM_IK*Matrix([x_l, y_l, z_l, 1]))
TransfMatrRF_Coxa = simplify(TransfMatrRPY_IK*TransfMatrRF_IK*Matrix([x_l, y_l, z_l, 1]))
```

#### 6. Cinemática inversa para cada pata de manera geométrica
Ya con la transformación de la coordenada global final para cada pata hasta el eslabón 0 (Unión de la pata con el cuerpo), procedemos a resolver de manera geometrica los ángulos de las patas necesarios para lograr dichas coordenadas.

El ángulo $\theta_{\mathrm{1}}$ que debe formar el servo 1 (*Coxa*) es muy fácil de sacar con solo las coordenadas $x_k$ y $y_k$ del elemento final de la pata $k$. Aún así, hay que tener cuidado en el ángulo a calcular debido a los siguientes casos:

- La posición final se encuentra justo debajo de la articulación 1: Se generan múltiples soluciones debido a que la rotación de este servo no afectará la solución.
- La posición final se encuentra cerca del centro del robot: Puede generarse inversión del eje y por ende, se debe aplicar correciones para evitar dicha situación.

\begin{align}
\Theta_1(k) =  \begin{Bmatrix}
 \Theta_1(k-1) & ,x_L = 0 \\
 atan2(y_L,x_L) + 180^{\circ} & ,x_L < 0\\
 atan2(y_L,x_L) & ,otros
\end{Bmatrix}
\end{align}

Siendo $\Theta_1(k)$ el ángulo en el tiempo $k$, $x_L$ la componente en el eje x del efector final de la pata y $y_L$ su componente en el eje y.

Los ángulos $\Theta_2(k)$ y $\Theta_3(k)$ se puede encontrar a partir de la transformaciòn de las coordenadas finales de la pata desde el cuerpo hasta el sistema de coordenadas del *Coxa*, mediante la matriz de transformaciòn inversa $T_0^1 = (T_1^0)^{-1}$. Con ello, podemos obtener geometricamente las ecuaciones de los àngulos como se muestra en la siguiente figura:

\begin{align}
\Theta_2(k) = acos \left ( \frac{(L_2)^2+b^2-(L_3)^2}{2*L_2*b} \right) + atan2\left ( z_c(k),x_c(k) \right )
\newline
\Theta_3(k) = 180^{\circ} - acos\left ( \frac{x_c(k)^{2} + z_c(k)^{2} - \left ( L_2\right )^{2} - \left ( L_3 \right )^{2}}{-2*L_3*L_2} \right )
\end{align}

#### 7. Inicialización datos de entrada

```python
# Este bloque de código realiza las computaciones de los 3 ángulos theta necesarios para la cinematica
# de la pata en cuestión en base a las ecuaciones descritas anteriormente.

# Variables geometricas del robot.
lin_1 = 38.5; lin_2 = 56.2; lin_3 = 143.5;
dlin = 0; r_bot = 137.0; x_init = 0; y_init = 0; z_init = 143.5
alf = 0; bet = 0; gam = 0;

# Expresión general utilizada para la substitución.
inVal = [(alpha, alf), (beta, bet), (gamma, gam),
         (L_1, lin_1), (L_2, lin_2), (L_3, lin_3),
         (d_L, dlin), (r_c, r_bot),
         (X_t, x_init), (Y_t, y_init), (Z_t, z_init)]

# Creación del array con las coordenadas de cada pata.
CoordXYZ_Paw = np.array([[-116.0, 240.92, -20],
                [-232.0, 0, 0],
                [-116.0, -160.92, -20],
                [116.0, -200.92, 0],
                [232.0, 40.0, -20],
                [116.0, 200.92, 0]])

# Creación de lista con todas las substituciones de las coordenadas por pata.
inVal_G = [inVal] * 6
for index, value in enumerate(CoordXYZ_Paw) :
    inVal_G[index] = inVal_G[index] + [(x_l, value[0]), (y_l, value[1]), (z_l, value[2])]

# Impresión de los elementos de dicha lista para verificar su contenido.
# for index, value in enumerate(inVal_G) :
#     print("Valores de substitución para la pata {}: {} ".format(index+1, value))
```

#### 8. Cálculo de Theta 1 _(Coxa)_

```python
# Los siguientes bloques de código nos permitiran realizar el calculo del ángulo theta_1
# en base a la transformación de las coordenadas finales desde el marco global hasta el eslabón 0,
# siendo estás denotadas como coordenadas x_L, y_L, z_L con los respectivos indices de la pata.

# Creamos un array utilizando la libreria NumPy para guardar las coordenadas x, y de cada pata.
# Las columnas serán el número de la pata y las filas serán las coordenadas.
CoordXY_Coxa = np.zeros((3,6))

for i in range(CoordXY_Coxa.shape[0]) :
    CoordXY_Coxa[i][0] = TransfMatrLF_Coxa[i].subs(inVal_G[0])
    CoordXY_Coxa[i][1] = TransfMatrLM_Coxa[i].subs(inVal_G[1])
    CoordXY_Coxa[i][2] = TransfMatrLR_Coxa[i].subs(inVal_G[2])
    CoordXY_Coxa[i][3] = TransfMatrRR_Coxa[i].subs(inVal_G[3])
    CoordXY_Coxa[i][4] = TransfMatrRM_Coxa[i].subs(inVal_G[4])
    CoordXY_Coxa[i][5] = TransfMatrRF_Coxa[i].subs(inVal_G[5])

CoordXY_Coxa = np.trunc(CoordXY_Coxa)

print("\033[1m" + "COORDENADAS GLOBALES X, Y, Z POR PATA" + "\033[0m")
for index, value in enumerate(CoordXYZ_Paw) :
    print("Pata N°{} - X: {: ^10.2f} | Y: {: ^10.2f} | Z: {: ^10.2f}"
          .format(index+1, value[0], value[1], value[2]))

print("\n\033[1m" + "ÁNGULOS THETA 1 CALCULADOS / COORDENADAS X, Y, Z RELATIVAS AL CUERPO" + "\033[0m")
Theta_1 = []
for index, value in enumerate(CoordXY_Coxa.T) :
    Theta_1 += [math.ceil(math.atan2(value[1], value[0]) * 180 / math.pi)]
    print("Pata N°{} = {: ^6.2f} | X: {: ^10.2f} | Y: {: ^10.2f} | Z: {: ^10.2f}"
          .format(index+1, Theta_1[index], value[0], value[1], value[2]))
```

    COORDENADAS GLOBALES X, Y, Z POR PATA
    Pata N°1 - X:  -116.00   | Y:   240.92   | Z:   -20.00
    Pata N°2 - X:  -232.00   | Y:    0.00    | Z:    0.00
    Pata N°3 - X:  -116.00   | Y:  -160.92   | Z:   -20.00
    Pata N°4 - X:   116.00   | Y:  -200.92   | Z:    0.00
    Pata N°5 - X:   232.00   | Y:   40.00    | Z:   -20.00
    Pata N°6 - X:   116.00   | Y:   200.92   | Z:    0.00

    ÁNGULOS THETA 1 CALCULADOS / COORDENADAS X, Y, Z RELATIVAS AL CUERPO
    Pata N°1 = -8.00  | X:   129.00   | Y:   -20.00   | Z:   123.00
    Pata N°2 =  0.00  | X:   95.00    | Y:    0.00    | Z:   143.00
    Pata N°3 = -17.00 | X:   60.00    | Y:   -19.00   | Z:   123.00
    Pata N°4 =  0.00  | X:   95.00    | Y:   -0.00    | Z:   143.00
    Pata N°5 = 23.00  | X:   95.00    | Y:   40.00    | Z:   123.00
    Pata N°6 =  0.00  | X:   95.00    | Y:    0.00    | Z:   143.00

#### 9. Cálculo de Theta 2 _(Femur)_

```python
# Transformación Inversa - Eslabón 0 hasta Eslabón 1.
TransfMatrLF_Femur = (TransfMatrL1**-1)*TransfMatrLF_Coxa
TransfMatrLM_Femur = (TransfMatrL1**-1)*TransfMatrLM_Coxa
TransfMatrLR_Femur = (TransfMatrL1**-1)*TransfMatrLR_Coxa
TransfMatrRR_Femur = (TransfMatrL1**-1)*TransfMatrRR_Coxa
TransfMatrRM_Femur = (TransfMatrL1**-1)*TransfMatrRM_Coxa
TransfMatrRF_Femur = (TransfMatrL1**-1)*TransfMatrRF_Coxa
```

```python
CoordXZ_Femur = np.zeros((3,6))

# Realizamos un ciclo para obtener los valores de las coordenadas x, z referidas al eslabón 0.
# En este caso guardamos la coordenada y para simplificar la obtención de z.
for j in range(CoordXZ_Femur.shape[0]) :
    CoordXZ_Femur[j][0] = TransfMatrLF_Femur[j].subs(inVal_G[0] + [(t_1, Theta_1[0])])
    CoordXZ_Femur[j][1] = TransfMatrLM_Femur[j].subs(inVal_G[1] + [(t_1, Theta_1[1])])
    CoordXZ_Femur[j][2] = TransfMatrLR_Femur[j].subs(inVal_G[2] + [(t_1, Theta_1[2])])
    CoordXZ_Femur[j][3] = TransfMatrRR_Femur[j].subs(inVal_G[3] + [(t_1, Theta_1[3])])
    CoordXZ_Femur[j][4] = TransfMatrRM_Femur[j].subs(inVal_G[4] + [(t_1, Theta_1[4])])
    CoordXZ_Femur[j][5] = TransfMatrRF_Femur[j].subs(inVal_G[5] + [(t_1, Theta_1[5])])

CoordXZ_Femur = np.around(CoordXZ_Femur, decimals=2)
print("\n\033[1m" + "ÁNGULOS THETA 2 CALCULADOS / COORDENADAS X, Y, Z RELATIVAS AL CUERPO" + "\033[0m")

Theta_2 = []
for index, value in enumerate(CoordXZ_Femur.T) :
    b_len = math.sqrt(math.pow(value[0], 2) + math.pow(value[1], 2))
    Theta_2 += [( math.acos( (math.pow(lin_2, 2) - math.pow(lin_3, 2) + math.pow(b_len,2)) / (2*lin_2*b_len) )
                - math.atan2(value[1], value[0]) ) * 180 / math.pi]
    print("Pata N°{} = {: ^6.2f} | X: {: ^10.2f} | Y: {: ^10.2f} | Z: {: ^10.2f}"
          .format(index+1, Theta_2[index], value[0], value[1], value[2]))
```

    ÁNGULOS THETA 2 CALCULADOS / COORDENADAS X, Y, Z RELATIVAS AL CUERPO
    Pata N°1 = -13.88 | X:   -37.57   | Y:   123.50   | Z:  -131.17
    Pata N°2 = -0.00  | X:   56.50    | Y:   143.50   | Z:    0.00
    Pata N°3 = -42.96 | X:   -74.34   | Y:   123.50   | Z:   52.53
    Pata N°4 = -0.00  | X:   56.50    | Y:   143.50   | Z:    0.00
    Pata N°5 = -86.22 | X:  -122.97   | Y:   123.50   | Z:   -59.08
    Pata N°6 = -0.00  | X:   56.50    | Y:   143.50   | Z:   -0.00

#### 10. Cálculo de Theta 3 _(Tibia)_

```python
print("\n\033[1m" + "ÁNGULOS THETA 3 CALCULADOS / COORDENADAS X, Y, Z RELATIVAS AL CUERPO" + "\033[0m")
Theta_3 = []
for index, value in enumerate(CoordXZ_Femur.T) :
    b_len = math.sqrt(math.pow(value[0], 2) + math.pow(value[1], 2))
    Theta_3 += [180 - math.acos((math.pow(lin_2, 2) + math.pow(lin_3, 2) - math.pow(b_len,2)) / (2*lin_2*lin_3)) * 180 / math.pi]
    print("Pata N°{} = {: ^6.2f} | X: {: ^10.2f} | Y: {: ^10.2f} | Z: {: ^10.2f}"
          .format(index+1, Theta_3[index], value[0], value[1], value[2]))
```

    ÁNGULOS THETA 3 CALCULADOS / COORDENADAS X, Y, Z RELATIVAS AL CUERPO
    Pata N°1 = 116.06 | X:   -37.57   | Y:   123.50   | Z:  -131.17
    Pata N°2 = 89.88  | X:   56.50    | Y:   143.50   | Z:    0.00
    Pata N°3 = 100.62 | X:   -74.34   | Y:   123.50   | Z:   52.53
    Pata N°4 = 89.88  | X:   56.50    | Y:   143.50   | Z:    0.00
    Pata N°5 = 65.76  | X:  -122.97   | Y:   123.50   | Z:   -59.08
    Pata N°6 = 89.88  | X:   56.50    | Y:   143.50   | Z:   -0.00

#### 11. Modelo en 3 dimensiones de nuestro robot.

##### 11.1 Librerías y utilidades personalizadas.

```python
'''
========================
Simulación 3D (Hexapodo)
========================

Demostración en 3D de la cinemática inversa del Hexapodo.
'''
# Importamos la libreria para hacer plot en 3D y la libreria para colores
from mpl_toolkits.mplot3d import Axes3D

# Creación de matrix para la base del hexapodo utilizando matriz de transformación pata a cuerpo.
HexapodFrame = np.zeros((3, int(360/10)))
for i in range(HexapodFrame.shape[1]):
    for j in range(3) :
        HexapodFrame[j,i] = TransfMatrGBP.subs([(g_k, np.radians(360/(HexapodFrame.shape[1] - 1) * i))] + inVal).col(-1)[j]

# Etiquetas
xlab = 'Coordenadas X'
ylab = 'Coordenadas Y'
zlab = 'Coordenadas Z'

def multiplots(ax, data1, data2, **kwargs) :
    if (data1.shape[0] == 1) :
        out = ax.plot([data1[0], data2[0,0]],
            [data1[1], data2[1,1]],
            zs=[data1[2], data2[2,2]], **kwargs)
    else :
        for i in range(6) :
            out = ax.plot([data1[0,i], data2[0,i]],
                [data1[1,i], data2[1,i]],
                zs=[data1[2,i], data2[2,i]], **kwargs)
    return out
```

##### 11.2 Computo de las coordenadas finales por eslabón.

```python
inVal3 = [(alpha, alf), (beta, bet), (gamma, gam),
          (L_1, lin_1), (L_2, lin_2), (L_3, lin_3),
          (d_L, dlin), (r_c, r_bot), (X_t, x_init), (Y_t, y_init), (Z_t, z_init)]

# Coordenadas iniciales: centro
CoordStart_c = np.ones((3,6))*np.array([x_init, y_init, z_init])[:, np.newaxis]
# Coordenadas finales: centro
CoordEnd_c = np.zeros((3,6))
# Coordenadas finales: coxa
CoordEnd_1 = np.zeros((3,6))
# Coordenadas finales: femur
CoordEnd_2 = np.zeros((3,6))
# Coordenadas finales: ribia
CoordEnd_3 = np.zeros((3,6))

for i in range(3) :
    for j in range(6) :
        TransfTemp = (TransfMatrRPY*TransfMatrBP).subs([(g_k, (j+2)*pi/3)] + inVal3)
        SubAngles = [(t_1, Theta_1[j] * math.pi / 180), (t_2, Theta_2[j] * math.pi / 180), (t_3, Theta_3[j] * math.pi / 180)]
        CoordEnd_c[i][j] = (TransfTemp*Matrix([0, 0, 0, 1]))[i]
        CoordEnd_1[i][j] = (TransfTemp*TransfMatrL1*Matrix([0, 0, 0, 1])).subs(inVal3 + SubAngles)[i]
        CoordEnd_2[i][j] = (TransfTemp*TransfMatrL1*TransfMatrL2*Matrix([0, 0, 0, 1])).subs(inVal3 + SubAngles)[i]
        CoordEnd_3[i][j] = (TransfTemp*TransfMatrL1*TransfMatrL2*TransfMatrL3*Matrix([0, 0, 0, 1])).subs(inVal3 + SubAngles)[i]

# Generación de línea extra para diferenciar el frente.
CoordFront = (TransfMatrRPY*TransfMatrBP).subs([(g_k, pi/2)] + inVal3)*Matrix([0, 0, 0, 1])
```

##### 11.2A paso extra para pruebas rapidas

```python
alf = -pi/4; bet = 0; gam = 0;
inVal3 = [(alpha, alf), (beta, bet), (gamma, gam),
          (L_1, lin_1), (L_2, lin_2), (L_3, lin_3),
          (d_L, dlin), (r_c, r_bot), (X_t, x_init), (Y_t, y_init), (Z_t, z_init)]

Theta_1 = [-45, 0, -45, 0, -45, 0]
Theta_2 = [30, 0, 30, 0, 30, 0]
Theta_3 = [110, 90, 110, 90, 110, 90]

# Coordenadas iniciales: centro
CoordStart_c = np.ones((3,6))*np.array([x_init, y_init, z_init])[:, np.newaxis]
# Coordenadas finales: centro
CoordEnd_c = np.zeros((3,6))
# Coordenadas finales: coxa
CoordEnd_1 = np.zeros((3,6))
# Coordenadas finales: femur
CoordEnd_2 = np.zeros((3,6))
# Coordenadas finales: ribia
CoordEnd_3 = np.zeros((3,6))

for i in range(3) :
    for j in range(6) :
        TransfTemp = (TransfMatrRPY*TransfMatrBP).subs([(g_k, (j+2)*pi/3)] + inVal3)
        SubAngles = [(t_1, Theta_1[j] * math.pi / 180), (t_2, Theta_2[j] * math.pi / 180), (t_3, Theta_3[j] * math.pi / 180)]
        CoordEnd_c[i][j] = (TransfTemp*Matrix([0, 0, 0, 1]))[i]
        CoordEnd_1[i][j] = (TransfTemp*TransfMatrL1*Matrix([0, 0, 0, 1])).subs(inVal3 + SubAngles)[i]
        CoordEnd_2[i][j] = (TransfTemp*TransfMatrL1*TransfMatrL2*Matrix([0, 0, 0, 1])).subs(inVal3 + SubAngles)[i]
        CoordEnd_3[i][j] = (TransfTemp*TransfMatrL1*TransfMatrL2*TransfMatrL3*Matrix([0, 0, 0, 1])).subs(inVal3 + SubAngles)[i]

# Generación de línea extra para diferenciar el frente.
CoordFront = (TransfMatrRPY*TransfMatrBP).subs([(g_k, pi/2)] + inVal3)*Matrix([0, 0, 0, 1])
```

##### 11.3 Modelo 3D final.

```python
# Creación de las figuras.
fig, axes = plt.subplots(2, 2, subplot_kw=dict(projection='3d'), figsize=(14, 10), dpi=100, facecolor='w')

for axe in axes.flat :
    multiplots(axe, CoordFront.T, CoordStart_c, c='k', marker='o', markevery=2)
    multiplots(axe, CoordStart_c, CoordEnd_c, c='k')
    multiplots(axe, CoordEnd_c, CoordEnd_1, c='r')
    multiplots(axe, CoordEnd_1, CoordEnd_2, c='b')
    multiplots(axe, CoordEnd_2, CoordEnd_3, c='g')
    axe.set_xlabel(xlab); axe.set_ylabel(ylab); axe.set_zlabel(zlab)

# Configuración de las vistas
axes.flat[0].view_init(30, 110)
axes.flat[1].view_init(90, 90)
axes.flat[1].set_zlabel('')
axes.flat[1].set_zticklabels([])

axes.flat[2].view_init(0, 180)
axes.flat[2].set_xlabel('')
axes.flat[2].set_xticklabels([])

axes.flat[3].view_init(0, 90)
axes.flat[3].set_ylabel('')
axes.flat[3].set_yticklabels([])

plt.show()

print("\n\033[1m" + "COORDENADAS X, Y, Z - COXA" + "\033[0m"); print(np.round(CoordEnd_1))
print("\n\033[1m" + "COORDENADAS X, Y, Z - FEMUR" + "\033[0m"); print(np.round(CoordEnd_2))
print("\n\033[1m" + "COORDENADAS X, Y, Z - TIBIA" + "\033[0m"); print(np.round(CoordEnd_3))
```


# Tutorial 4: Multiple linear regression and polynomial regression
**Week 1, Day 3: Model Fitting**

**By Neuromatch Academy**

**Content creators**: Pierre-Étienne Fiquet, Anqi Wu, Alex Hyafil with help from Byron Galbraith, Ella Batty

**Content reviewers**: Lina Teichmann, Saeed Salehi, Patrick Mineault, Michael Waskom

**Our 2021 Sponsors, including Presenting Sponsor Facebook Reality Labs**

---
# Tutorial Objectives

*Estimated timing of tutorial: 35 minutes*

This is Tutorial 4 of a series on fitting models to data. We start with simple linear regression, using least squares optimization (Tutorial 1) and Maximum Likelihood Estimation (Tutorial 2). We will use bootstrapping to build confidence intervals around the inferred linear model parameters (Tutorial 3). We'll finish our exploration of regression models by generalizing to multiple linear regression and polynomial regression (Tutorial 4). We end by learning how to choose between these various models. We discuss the bias-variance trade-off (Tutorial 5) and Cross Validation for model selection (Tutorial 6).

In this tutorial, we will generalize the regression model to incorporate multiple features.
- Learn how to structure inputs for regression using the 'Design Matrix'
- Generalize the MSE for multiple features using the ordinary least squares estimator
- Visualize data and model fit in multiple dimensions
- Fit polynomial regression models of different complexity
- Plot and evaluate the polynomial regression fits

```python
# @title Tutorial slides

# @markdown These are the slides for the videos in all tutorials today
from IPython.display import IFrame
IFrame(src=f"https://mfr.ca-1.osf.io/render?url=https://osf.io/2mkq4/?direct%26mode=render%26action=download%26mode=render", width=854, height=480)
```

```python
# @title Video 1: Multiple Linear Regression and Polynomial Regression
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV11Z4y1u7cf", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="d4nfTki6Ejc", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)
```

    Tab(children=(Output(), Output()), _titles={'0': 'Youtube', '1': 'Bilibili'})

---
# Setup

```python
# Imports

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

```python
#@title Figure Settings
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Plotting Functions

def evaluate_fits(order_list, mse_list):
  """ Compare the quality of multiple polynomial fits
  by plotting their MSE values.

  Args:
    order_list (list): list of the order of polynomials to be compared
    mse_list (list): list of the MSE values for the corresponding polynomial fit
  """
  fig, ax = plt.subplots()
  ax.bar(order_list, mse_list)
  ax.set(title='Comparing Polynomial Fits', xlabel='Polynomial order', ylabel='MSE')

def plot_fitted_polynomials(x, y, theta_hat):
  """ Plot polynomials of different orders

  Args:
    x (ndarray): input vector of shape (n_samples)
    y (ndarray): vector of measurements of shape (n_samples)
    theta_hat (dict): polynomial regression weights for different orders
  """

  x_grid = np.linspace(x.min() - .5, x.max() + .5)

  plt.figure()

  for order in range(0, max_order + 1):
    X_design = make_design_matrix(x_grid, order)
    plt.plot(x_grid, X_design @ theta_hat[order]);

  plt.ylabel('y')
  plt.xlabel('x')
  plt.plot(x, y, 'C0.');
  plt.legend([f'order {o}' for o in range(max_order + 1)], loc=1)
  plt.title('polynomial fits')
  plt.show()
```

---
# Section 1: Multiple Linear Regression

*Estimated timing to here from start of tutorial: 8 min*

This video covers linear regression with multiple inputs (more than 1D) and polynomial regression.

 Click here for text recap of video

Now that we have considered the univariate case and how to produce confidence intervals for our estimator, we turn to the general linear regression case, where we can have more than one regressor, or feature, in our input.

Recall that our original univariate linear model was given as

\begin{align}
y = \theta x + \epsilon
\end{align}

where $\theta$ is the slope and $\epsilon$ some noise. We can easily extend this to the multivariate scenario by adding another parameter for each additional feature

\begin{align}
y = \theta_0 + \theta_1 x_1 + \theta_2 x_2 + ... +\theta_d x_d + \epsilon
\end{align}

where $\theta_0$ is the intercept and $d$ is the number of features (it is also the dimensionality of our input).

We can condense this succinctly using vector notation for a single data point

\begin{align}
y_i = \boldsymbol{\theta}^{\top}\mathbf{x}_i + \epsilon
\end{align}

and fully in matrix form

\begin{align}
\mathbf{y} = \mathbf{X}\boldsymbol{\theta} + \mathbf{\epsilon}
\end{align}

where $\mathbf{y}$ is a vector of measurements, $\mathbf{X}$ is a matrix containing the feature values (columns) for each input sample (rows), and $\boldsymbol{\theta}$ is our parameter vector.

This matrix $\mathbf{X}$ is often referred to as the "[design matrix](https://en.wikipedia.org/wiki/Design_matrix)".

We want to find an optimal vector of paramters $\boldsymbol{\hat\theta}$. Recall our analytic solution to minimizing MSE for a single regressor:

\begin{align}
\hat\theta = \frac{\sum_{i=1}^N x_i y_i}{\sum_{i=1}^N x_i^2}.
\end{align}

The same holds true for the multiple regressor case, only now expressed in matrix form

\begin{align}
\boldsymbol{\hat\theta} = (\mathbf{X}^\top\mathbf{X})^{-1}\mathbf{X}^\top\mathbf{y}.
\end{align}

This is called the [ordinary least squares](https://en.wikipedia.org/wiki/Ordinary_least_squares) (OLS) estimator.

For this tutorial we will focus on the two-dimensional case ($d=2$), which allows us to easily visualize our results. As an example, think of a situation where a scientist records the spiking response of a retinal ganglion cell to patterns of light signals that vary in contrast and in orientation. Then contrast and orientation values can be used as features / regressors to predict the cells response.

In this case our model can be writen for a single data point as:

\begin{align}
y = \theta_0 + \theta_1 x_1 + \theta_2 x_2 + \epsilon
\end{align}

or for multiple data points in matrix form where

\begin{align}
\mathbf{X} =
\begin{bmatrix}
1 & x_{1,1} & x_{1,2} \\
1 & x_{2,1} & x_{2,2} \\
\vdots & \vdots & \vdots \\
1 & x_{n,1} & x_{n,2}
\end{bmatrix},
\boldsymbol{\theta} =
\begin{bmatrix}
\theta_0 \\
\theta_1 \\
\theta_2 \\
\end{bmatrix}
\end{align}

When we refer to $x_{i, j}$, we mean that it is the i-th data point and the j-th feature of that data point.

For our actual exploration dataset we shall set $\boldsymbol{\theta}=[0, -2, -3]$ and draw $N=40$ noisy samples from $x \in [-2,2)$. Note that setting the value of $\theta_0 = 0$ effectively ignores the offset term.

```python
# @markdown Execute this cell to simulate some data

# Set random seed for reproducibility
np.random.seed(1234)

# Set parameters
theta = [0, -2, -3]
n_samples = 40

# Draw x and calculate y
n_regressors = len(theta)
x0 = np.ones((n_samples, 1))
x1 = np.random.uniform(-2, 2, (n_samples, 1))
x2 = np.random.uniform(-2, 2, (n_samples, 1))
X = np.hstack((x0, x1, x2))
noise = np.random.randn(n_samples)
y = X @ theta + noise

ax = plt.subplot(projection='3d')
ax.plot(X[:,1], X[:,2], y, '.')

ax.set(
    xlabel='$x_1$',
    ylabel='$x_2$',
    zlabel='y'
)
plt.tight_layout()
```

## Coding Exercise 1: Ordinary Least Squares Estimator

In this exercise you will implement the OLS approach to estimating $\boldsymbol{\hat\theta}$ from the design matrix $\mathbf{X}$ and measurement vector $\mathbf{y}$. You can use the `@` symbol for matrix multiplication, `.T` for transpose, and `np.linalg.inv` for matrix inversion.

```python
def ordinary_least_squares(X, y):
  """Ordinary least squares estimator for linear regression.

  Args:
    x (ndarray): design matrix of shape (n_samples, n_regressors)
    y (ndarray): vector of measurements of shape (n_samples)

  Returns:
    ndarray: estimated parameter values of shape (n_regressors)
  """
  ######################################################################
  ## TODO for students: solve for the optimal parameter vector using OLS
  # Fill out function and remove
  #raise NotImplementedError("Student exercise: solve for theta_hat vector using OLS")
  ######################################################################

  # Compute theta_hat using OLS
  theta_hat = np.linalg.inv(X.T @ X)@ X.T @y

  return theta_hat

theta_hat = ordinary_least_squares(X, y)
print(theta_hat)
```

    [ 0.13861386 -2.09395731 -3.16370742]

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial4_Solution_25849be9.py)

After filling in this function, you should see that $\boldsymbol{\hat\theta}$ = [ 0.13861386, -2.09395731, -3.16370742]

Now that we have our $\boldsymbol{\hat\theta}$, we can obtain $\hat{\mathbf{y}}$ and thus our mean squared error.

```python
# Compute predicted data
theta_hat = ordinary_least_squares(X, y)
y_hat = X @ theta_hat

# Compute MSE
print(f"MSE = {np.mean((y - y_hat)**2):.2f}")
```

    MSE = 0.91

Finally, the following code will plot a geometric visualization of the data points (blue) and fitted plane.

```python
# @markdown Execute this cell to visualize data and predicted plane

theta_hat = ordinary_least_squares(X, y)
xx, yy = np.mgrid[-2:2:50j, -2:2:50j]
y_hat_grid = np.array([xx.flatten(), yy.flatten()]).T @ theta_hat[1:]
y_hat_grid = y_hat_grid.reshape((50, 50))

ax = plt.subplot(projection='3d')
ax.plot(X[:, 1], X[:, 2], y, '.')
ax.plot_surface(xx, yy, y_hat_grid, linewidth=0, alpha=0.5, color='C1',
                cmap=plt.get_cmap('coolwarm'))

for i in range(len(X)):
  ax.plot((X[i, 1], X[i, 1]),
          (X[i, 2], X[i, 2]),
          (y[i], y_hat[i]),
          'g-', alpha=.5)

ax.set(
    xlabel='$x_1$',
    ylabel='$x_2$',
    zlabel='y'
)
plt.tight_layout()
```

---
# Section 2:  Polynomial Regression

So far today, you learned how to predict outputs from inputs by fitting a linear regression model. We can now model all sort of relationships, including in neuroscience!

One potential problem with this approach is the simplicity of the model. Linear regression, as the name implies, can only capture a linear relationship between the inputs and outputs. Put another way, the predicted outputs are only a weighted sum of the inputs. What if there are more complicated computations happening? Luckily, many more complex models exist (and you will encounter many more over the next 3 weeks). One model that is still very simple to fit and understand, but captures more complex relationships, is **polynomial regression**, an extension of linear regression.

 Click here for text recap of relevant part of video

Since polynomial regression is an extension of linear regression, everything you learned so far will come in handy now! The goal is the same: we want to predict the dependent variable $y$ given the input values $x$. The key change is the type of relationship between inputs and outputs that the model can capture.

Linear regression models predict the outputs as a weighted sum of the inputs:

\begin{align}
y = \theta_0 + \theta x + \epsilon
\end{align}

With polynomial regression, we model the outputs as a polynomial equation based on the inputs. For example, we can model the outputs as:

\begin{align}
y & = \theta_0 + \theta_1 x + \theta_2 x^2 + \theta_3 x^3 + \epsilon
\end{align}

We can change how complex a polynomial is fit by changing the order of the polynomial. The order of a polynomial refers to the highest power in the polynomial. The equation above is a third order polynomial because the highest value x is raised to is 3. We could add another term ($+ \theta_4 x^4$) to model an order 4 polynomial and so on.

First, we will simulate some data to practice fitting polynomial regression models. We will generate random inputs $x$ and then compute y according to $y = x^2 - x - 2 $, with some extra noise both in the input and the output to make the model fitting exercise closer to a real life situation.

```python
# @markdown Execute this cell to simulate some data

# setting a fixed seed to our random number generator ensures we will always
# get the same psuedorandom number sequence
np.random.seed(121)
n_samples = 30
x = np.random.uniform(-2, 2.5, n_samples)  # inputs uniformly sampled from [-2, 2.5)
y =  x**2 - x - 2   # computing the outputs

output_noise = 1/8 * np.random.randn(n_samples)
y += output_noise  # adding some output noise

input_noise = 1/2 * np.random.randn(n_samples)
x += input_noise  # adding some input noise

fig, ax = plt.subplots()
ax.scatter(x, y)  # produces a scatter plot
ax.set(xlabel='x', ylabel='y');
```

## Section 2.1: Design matrix for polynomial regression

*Estimated timing to here from start of tutorial: 16 min*

Now we have the basic idea of polynomial regression and some noisy data, let's begin! The key difference between fitting a linear regression model and a polynomial regression model lies in how we structure the input variables.

Let's go back to one feature for each data point. For linear regression, we used $\mathbf{X} = \mathbf{x}$ as the input data, where $\mathbf{x}$ is a vector where each element is the input for a single data point. To add a constant bias (a y-intercept in a 2-D plot), we use $\mathbf{X} = \big[ \boldsymbol 1, \mathbf{x} \big]$, where $\boldsymbol 1$ is a column of ones.  When fitting, we learn a weight for each column of this matrix. So we learn a weight that multiples with column 1 - in this case that column is all ones so we gain the bias parameter ($+ \theta_0$).

This matrix $\mathbf{X}$ that we use for our inputs is known as a **design matrix**. We want to create our design matrix so we learn weights for $\mathbf{x}^2, \mathbf{x}^3,$ etc. Thus, we want to build our design matrix $X$ for polynomial regression of order $k$ as:

\begin{align}
\mathbf{X} = \big[ \boldsymbol 1 , \mathbf{x}^1, \mathbf{x}^2 , \ldots , \mathbf{x}^k \big],
\end{align}

where $\boldsymbol{1}$ is the vector the same length as $\mathbf{x}$ consisting of of all ones, and $\mathbf{x}^p$ is the vector $\mathbf{x}$ with all elements raised to the power $p$. Note that $\boldsymbol{1} = \mathbf{x}^0$ and $\mathbf{x}^1 = \mathbf{x}$.

If we have inputs with more than one feature, we can use a similar design matrix but include all features raised to each power. Imagine that we have two features per data point: $\mathbf{x}_m$ is a vector of one feature per data point and  $\mathbf{x}_n$ is another.  Our design matrix for a polynomial regression would be:
\begin{align}
\mathbf{X} = \big[ \boldsymbol 1 , \mathbf{x}_m^1, \mathbf{x}_n^1, \mathbf{x}_m^2 , \mathbf{x}_n^2\ldots , \mathbf{x}_m^k , \mathbf{x}_n^k \big],
\end{align}

### Coding Exercise 2.1: Structure design matrix

Create a function (`make_design_matrix`) that structures the design matrix given the input data and the order of the polynomial you wish to fit. We will print part of this design matrix for our data and order 5.

```python
def make_design_matrix(x, order):
  """Create the design matrix of inputs for use in polynomial regression

  Args:
    x (ndarray): input vector of shape (n_samples)
    order (scalar): polynomial regression order

  Returns:
    ndarray: design matrix for polynomial regression of shape (samples, order+1)
  """
  ########################################################################
  ## TODO for students: create the design matrix ##
  # Fill out function and remove
  # raise NotImplementedError("Student exercise: create the design matrix")
  ########################################################################

  # Broadcast to shape (n x 1) so dimensions work
  if x.ndim == 1:
    x = x[:, None]

  #if x has more than one feature, we don't want multiple columns of ones so we assign
  # x^0 here
  design_matrix = np.ones((x.shape[0], 1))

  # Loop through rest of degrees and stack columns (hint: np.hstack)
  for degree in range(1, order + 1):
      design_matrix = np.hstack((design_matrix, x**degree))

  return design_matrix

order = 5
X_design = make_design_matrix(x, order)

print(X_design[0:2, 0:2])
```

    [[ 1.         -1.51194917]
     [ 1.         -0.35259945]]

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial4_Solution_5e30078a.py)

You should see that the printed section of this design matrix is `[[ 1.         -1.51194917]
 [ 1.         -0.35259945]]`

## Section 2.2: Fitting polynomial regression models

*Estimated timing to here from start of tutorial: 24 min*

Now that we have the inputs structured correctly in our design matrix, fitting a polynomial regression is the same as fitting a linear regression model! All of the polynomial structure we need to learn is contained in how the inputs are structured in the design matrix. We can use the same least squares solution we computed in previous exercises.

### Coding Exercise 2.2: Fitting polynomial regression models with different orders

Here, we will fit polynomial regression models to find the regression coefficients ($\theta_0, \theta_1, \theta_2,$ ...) by solving the least squares problem. Create a function `solve_poly_reg` that loops over different order polynomials (up to `max_order`), fits that model, and saves out the weights for each. You may invoke the `ordinary_least_squares` function.

We will then qualitatively inspect the quality of our fits for each order by plotting the fitted polynomials on top of the data. In order to see smooth curves, we evaluate the fitted polynomials on a grid of $x$ values (ranging between the largest and smallest of the inputs present in the dataset).

```python
def solve_poly_reg(x, y, max_order):
  """Fit a polynomial regression model for each order 0 through max_order.

  Args:
    x (ndarray): input vector of shape (n_samples)
    y (ndarray): vector of measurements of shape (n_samples)
    max_order (scalar): max order for polynomial fits

  Returns:
    dict: fitted weights for each polynomial model (dict key is order)
  """

  # Create a dictionary with polynomial order as keys,
  # and np array of theta_hat (weights) as the values
  theta_hats = {}

  # Loop over polynomial orders from 0 through max_order
  for order in range(max_order + 1):

    ##################################################################################
    ## TODO for students: Create design matrix and fit polynomial model for this order
    # Fill out function and remove
    #raise NotImplementedError("Student exercise: fit a polynomial model")
    ##################################################################################

    # Create design matrix
    X_design =  make_design_matrix(x, order)

    # Fit polynomial model
    this_theta = ordinary_least_squares(X_design, y)

    theta_hats[order] = this_theta

  return theta_hats

max_order = 5
theta_hats = solve_poly_reg(x, y, max_order)

# Visualize
plot_fitted_polynomials(x, y, theta_hats)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial4_Solution_f5217dbd.py)

*Example output:*

## Section 2.3: Evaluating fit quality

*Estimated timing to here from start of tutorial: 29 min*

As with linear regression, we can compute mean squared error (MSE) to get a sense of how well the model fits the data.

We compute MSE as:

\begin{align}
\mathrm{MSE} = \frac 1 N ||\mathbf{y} - \hat{\mathbf{y}}||^2 = \frac 1 N \sum_{i=1}^N (y_i - \hat y_i)^2
\end{align}

where the predicted values for each model are given by $ \hat{\mathbf{y}} = \mathbf{X}\boldsymbol{\hat\theta}$.

*Which model (i.e. which polynomial order) do you think will have the best MSE?*

### Coding Exercise 2.3: Compute MSE and compare models

We will compare the MSE for different polynomial orders with a bar plot.

```python
mse_list = []
order_list = list(range(max_order + 1))

for order in order_list:

  X_design = make_design_matrix(x, order)

  ########################################################################
  ## TODO for students
  # Fill out function and remove
  #raise NotImplementedError("Student exercise: compute MSE")
  ########################################################################

  # Get prediction for the polynomial regression model of this order
  y_hat = X_design @ theta_hats[order]

  # Compute the residuals
  residuals = y - y_hat

  # Compute the MSE
  mse = np.mean(residuals ** 2)

  mse_list.append(mse)

# Visualize MSE of fits
evaluate_fits(order_list, mse_list)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial4_Solution_89324713.py)

*Example output:*

---
# Summary

*Estimated timing of tutorial: 35 minutes*

* Linear regression generalizes naturally to multiple dimensions
* Linear algebra affords us the mathematical tools to reason and solve such problems beyond the two dimensional case

* To change from a linear regression model to a polynomial regression model, we only have to change how the input data is structured

* We can choose the complexity of the model by changing the order of the polynomial model fit

* Higher order polynomial models tend to have lower MSE on the data they're fit with

**Note**: In practice, multidimensional least squares problems can be solved very efficiently (thanks to numerical routines such as LAPACK).

---
# Notation

\begin{align}
x &\quad \text{input, independent variable}\\
y &\quad \text{response measurement, dependent variable}\\
\epsilon &\quad \text{measurement error, noise contribution}\\
\theta &\quad \text{slope parameter}\\
\hat{\theta} &\quad \text{estimated slope parameter}\\
\mathbf{x} &\quad \text{vector of inputs where each element is a different data point}\\
\mathbf{X} &\quad \text{design matrix}\\
\mathbf{y} &\quad \text{vector of measurements}\\
\mathbf{\hat y} &\quad \text{vector of estimated measurements}\\
\boldsymbol{\theta} &\quad \text{vector of parameters}\\
\boldsymbol{\hat\theta} &\quad \text{vector of estimated parameters}\\
d &\quad \text{dimensionality of input}\\
N &\quad \text{number of samples}\\
\end{align}

**Suggested readings**

[Introduction to Applied Linear Algebra – Vectors, Matrices, and Least Squares](http://vmls-book.stanford.edu/)
Stephen Boyd and Lieven Vandenberghe


# Joint TV for mutiple PET images
This demonstration shows how to do a synergistic reconstruction of two PET images. Both PET images show the same underlying anatomy but of course with different tracer distribution. In order to make use of this similarity a joint total variation (TV) operator is used as a regularisation in an iterative image reconstruction approach.

This demo is a jupyter notebook, i.e. intended to be run step by step.
You could export it as a Python file and run it one go, but that might
make little sense as the figures are not labelled.

First Authors: Christoph Kolbitsch, Evangelos Papoutsellis, Edoardo Pasca
Updated for PET: Sam Porter \
First version: 16th of June 2021
Updated: 26nd of June 2021
Updated for PET: 31st of March 2022

CCP SyneRBI Synergistic Image Reconstruction Framework (SIRF).
Copyright 2021 Rutherford Appleton Laboratory STFC.
Copyright 2021 Physikalisch-Technische Bundesanstalt.

This is software developed for the Collaborative Computational
Project in Positron Emission Tomography and Magnetic Resonance imaging
(http://www.ccppetmr.ac.uk/).

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

# Initial set-up

```python
# Make sure figures appears inline
%matplotlib notebook
%matplotlib inline
```

```python
# Make sure everything is installed that we need
!pip install brainweb nibabel --user
!pip install ipython --user
!pip install plotly --user
```

```python
# Initial imports etc
import numpy
from numpy.linalg import norm
import matplotlib.pyplot as plt
from matplotlib import animation

import os
import brainweb
from tqdm.auto import tqdm
import glob

# Import SIRF functionality
import notebook_setup
import sirf.Gadgetron as mr
from sirf.Utilities import examples_data_path
import sirf.STIR as pet
import plotly.io as pio

# Import CIL functionality
from cil.framework import  BlockDataContainer, BlockGeometry
from cil.optimisation.functions import Function, OperatorCompositionFunction, KullbackLeibler
from cil.optimisation.operators import GradientOperator, CompositionOperator, LinearOperator, ProjectionMap
from cil.optimisation.algorithms import GD
from IPython.display import HTML
from sirf_exercises import cd_to_working_dir

msg = pet.MessageRedirector('info.txt', 'warnings.txt', 'errors.txt')

data_path = examples_data_path('PET')

cd_to_working_dir("Synergistic", "JTV_PET")
```

# Utilities

```python
# First define some handy function definitions
# To make subsequent code cleaner, we have a few functions here. You can ignore
# ignore them when you first see this demo.

def plot_2d_image(idx,vol,title,clims=None,cmap="viridis"):
    """Customized version of subplot to plot 2D image"""
    plt.subplot(*idx)
    plt.imshow(vol,cmap=cmap)
    if not clims is None:
        plt.clim(clims)
    plt.colorbar(shrink = 0.3)
    plt.title(title)
    plt.axis("off")

def crop_and_fill(templ_im, vol):
    """Crop volumetric image data and replace image content in template image object"""
    # Get size of template image and crop
    idim_orig = templ_im.as_array().shape
    idim = (1,)*(3-len(idim_orig)) + idim_orig
    offset = (numpy.array(vol.shape) - numpy.array(idim)) // 2
    vol = vol[offset[0]:offset[0]+idim[0], offset[1]:offset[1]+idim[1], offset[2]:offset[2]+idim[2]]

    # Make a copy of the template to ensure we do not overwrite it
    templ_im_out = templ_im.copy()

    # Fill image content
    templ_im_out.fill(numpy.reshape(vol, idim_orig))
    return(templ_im_out)

def make_cylindrical_FOV(image):
    """truncate to cylindrical FOV"""
    filter = pet.TruncateToCylinderProcessor()
    filter.apply(image)
    return image
```

### Joint TV reconstruction of two PET/SPECT images

Assume we want to reconstruct two PET\SPECT images $u$ and $v$ and utilse the similarity between both images using a joint TV ($JTV$) operator we can formulate the reconstruction problem as:

$$
\begin{equation}
(u^{*}, v^{*}) \in \underset{u,v}{\operatorname{argmin}}\biggl \{ \mathcal{D}\bigl(A_1u,g_1) +  \mathcal{D}\bigl(A_2v,g_2) + \alpha\,\mathrm{JTV}_{\eta, \lambda}(u, v) \biggl \}
\end{equation}
$$

* $\mathcal{D}\bigl(Au,g)$ `Data Fidelity Term` - we'll use the Kullbach Liebler divergence:
$$ \mathcal{D}(g|Au) = \sum_i \log g_i\frac{g_i}{(Au)_i} -(Au+\eta)_i + g_i$$ $$
* $JTV_{\eta, \lambda}(u, v) = \sum \sqrt{ \lambda|\nabla u|^{2} + (1-\lambda)|\nabla v|^{2} + \eta^{2}}$
* $A_{1}$, $A_{2}$: __PET/SPECT__ `AcquisitionModel`
* $g_{1}$, $g_{2}$: __PET/SPECT__ `AcquisitionData`

### Solving this problem

In order to solve the above minimization problem, we will use an alternating minimisation approach, where one variable is fixed and we solve wrt to the other variable:

$$
\begin{align*}
u^{k+1} & = \underset{u}{\operatorname{argmin}} \biggl \{ \mathcal{D}_{KL}\bigl(A_1u,g_1) + \alpha_{1}\,\mathrm{JTV}_{\eta, \lambda}(u, v^{k}) \biggl \} \quad \text{subproblem 1}\\
v^{k+1} & = \underset{v}{\operatorname{argmin}} \biggl \{ \mathcal{D}_{KL}\bigl(A_2v,g_2) + \alpha_{2}\,\mathrm{JTV}_{\eta, 1-\lambda}(u^{k+1}, v) \biggl \} \quad \text{subproblem 2}\\
\end{align*}$$

We are going to use a gradient descent approach to solve each of these subproblems alternatingly.

The *regularisation parameter* `alpha` should be different for each subproblem. But not to worry at this stage. Maybe we should use $\alpha_{1}, \alpha_{2}$ in front of the two JTVs and a $\lambda$, $1-\lambda$ for the first JTV and $1-\lambda$, $\lambda$, for the second JTV with $0<\lambda<1$.

This notebook builds on several other notebooks and hence certain steps will be carried out with minimal documentation. If you want more explainations, then we would like to ask you to refer to the corresponding notebooks which are mentioned in the following list. The steps we are going to carry out are

  - (A) Get an FDG and amyloid PET image from brainweb which we are going to use as ground truth $u_{gt}$ and $v_{gt}$ for our reconstruction (further information: `introduction` notebook)

  - (B) Create __PET/SPECT__ `AcquisitionModel` $A_{1}$ and $A_{2}$ (further information: `acquisition_model_mr_pet_ct` notebook)

  - (C) Set up the joint TV reconstruction problem

  - (D) Solve the joint TV reconstruction problem (further information on gradient descent: `gradient_descent_mr_pet_ct` notebook)

# (A) Get brainweb data

We will download and use data from the brainweb.

```python
fname, url= sorted(brainweb.utils.LINKS.items())[0]
files = brainweb.get_file(fname, url, ".")
data = brainweb.load_file(fname)

brainweb.seed(1337)
```

```python
for f in tqdm([fname], desc="mMR ground truths", unit="subject"):
    vol = brainweb.get_mmr_fromfile(f, petNoise=1, t1Noise=0.75, t2Noise=0.75, petSigma=1, t1Sigma=1, t2Sigma=1)
    vol_amyl = brainweb.get_mmr_fromfile(f, petNoise=1, t1Noise=0.75, t2Noise=0.75, petSigma=1, t1Sigma=1, t2Sigma=1, PetClass=brainweb.Amyloid)
```

```python
uMap_arr = vol['uMap']
amyl_arr = vol_amyl['PET']
fdg_arr = vol['PET']

# Reduce values in uMamp to make more realistic
uMap_arr /= numpy.max(uMap_arr)*10

# Normalise fdg and amyloid images
fdg_arr /= numpy.max(fdg_arr)
amyl_arr /= numpy.max(amyl_arr)
```

```python
# Display it
plt.figure();
slice_show = fdg_arr.shape[0]//2
plot_2d_image([1,3,1], fdg_arr[slice_show, :, :], 'FDG')
plot_2d_image([1,3,2], amyl_arr[slice_show, :, :], 'Amyloid')
plot_2d_image([1,3,3], uMap_arr[slice_show, :, :], 'uMap', cmap="Greys_r")
```

Ok, we got to two images with FDG and Amyloid tracer distribution BUT they brain looks a bit small. Spoiler alert: We are going to reconstruct images with a FOV 120x120 voxels

In order to do this we are going to use an image `rescale` from the skimage package and simply rescale the image by a factor 2 and then crop it. To speed things up, we are going to already select a single slice because also our MR scan is going to be 2D.

```python
# Select central slice
central_slice = fdg_arr.shape[0]//2
fdg_arr = fdg_arr[central_slice, :, :]
amyl_arr = amyl_arr[central_slice, :, :]
uMap_arr = uMap_arr[central_slice, :, :]
# Select a central ROI with 120x120
# We could also skip this because it is automaticall done by crop_and_fill()
# but we would like to check if we did the right thing
idim = [120, 120]
offset = (numpy.array(fdg_arr.shape) - numpy.array(idim)) // 2
fdg_arr = fdg_arr[offset[0]:offset[0]+idim[0], offset[1]:offset[1]+idim[1]]
amyl_arr = amyl_arr[offset[0]:offset[0]+idim[0], offset[1]:offset[1]+idim[1]]
uMap_arr = uMap_arr[offset[0]:offset[0]+idim[0], offset[1]:offset[1]+idim[1]]

# Now we make sure our image is of shape (1, 120, 120) again because in __SIRF__ even 2D images
# are expected to have 3 dimensions.
fdg_arr = fdg_arr[numpy.newaxis,...]
amyl_arr = amyl_arr[numpy.newaxis,...]
uMap_arr = uMap_arr[numpy.newaxis,...]
```

```python
templ_sino = pet.AcquisitionData(os.path.join(data_path, 'template_sinogram.hs'))
dim = fdg_arr.shape
im = pet.ImageData(templ_sino)
voxel_size=im.voxel_sizes()
im.initialise(dim,voxel_size)
fdg = crop_and_fill(im, fdg_arr)
amyl = crop_and_fill(im, amyl_arr)
uMap = crop_and_fill(im, uMap_arr)
```

```python
# Display it
plt.figure();
slice_show = fdg_arr.shape[0]//2
plot_2d_image([1,3,1], numpy.squeeze(fdg.as_array()), 'FDG')
plot_2d_image([1,3,2], numpy.squeeze(amyl.as_array()), 'Amyloid',)
plot_2d_image([1,3,3], numpy.squeeze(uMap.as_array()), 'uMap', cmap="Greys_r")
```

Now, that looks better. Now we have got images we can use for our simulation.

# (B) Simulate noisy AcquisitionData

```python
def get_acquisition_model(uMap, templ_sino, detector_efficiency = 1, attenuation = True):
    """ Function to create PET acquisition model """
    am = pet.AcquisitionModelUsingRayTracingMatrix()
    # Set up sensitivity due to attenuation
    am.set_num_tangential_LORs(5)
    if attenuation is True:
        asm_attn = pet.AcquisitionSensitivityModel(uMap, am)
        asm_attn.set_up(templ_sino)
        bin_eff = templ_sino.get_uniform_copy(detector_efficiency)
        print('applying attenuation (please wait, may take a while)...')
        asm_attn.unnormalise(bin_eff)
        asm = pet.AcquisitionSensitivityModel(bin_eff)
        am.set_acquisition_sensitivity(asm)
    am.set_up(templ_sino, uMap)
    return am

```

```python
def add_noise(proj_data,noise_factor = 0.05):
    """ Function to add noise to PET/SPECT Acquistion Data """
    proj_data_arr = proj_data.as_array() / noise_factor
    # Data should be >=0 anyway, but add abs just to be safe
    proj_data_arr = numpy.abs(proj_data_arr)
    noisy_proj_data_arr = numpy.random.poisson(proj_data_arr).astype('float32');
    noisy_proj_data = proj_data.copy()
    noisy_proj_data.fill(noisy_proj_data_arr);
    return noisy_proj_data*noise_factor
```

Next we are going to create the two __PET__ `AcquisitionModel` $A_{1}$ and $A_{2}$

```python
acq_model = get_acquisition_model(uMap, templ_sino)
anorm = acq_model.norm() #norm of the operator to avoid any scaling
```

and simulate `AcquisitionData` $g_{1}$ and $g_{2}$

```python
# Forward project our images into data space
data_fdg = acq_model.forward(fdg)
data_amyl = acq_model.forward(amyl)
c_slice = data_fdg.shape[1]//2 # central slice
pet.show_2D_array("sino",data_fdg.as_array()[0,c_slice,:,:])
```

Lastly we are going to add some noise

```python
# Add noise
noisy_data_fdg = add_noise(data_fdg)
noisy_data_amyl = add_noise(data_amyl)
pet.show_2D_array("noisy sino",noisy_data_fdg.as_array()[0,c_slice,:,:])
```

Just to check we are going to apply the backward/adjoint operation to do a simply image reconstruction.

```python
# Simple reconstruction
u_simple = acq_model.backward(noisy_data_fdg)/anorm
v_simple = acq_model.backward(noisy_data_amyl)/anorm
```

```python
# Display it
plt.figure();
plot_2d_image([1,2,1], numpy.abs(u_simple.as_array())[0, :, :], '$u_{simple}$')
plot_2d_image([1,2,2], numpy.abs(v_simple.as_array())[0, :, :], '$v_{simple}$')
```

These images look quite poor compared to the ground truth input images, because they are reconstructed by a simple back-projection.

# (C) Set up the joint TV reconstruction problem

So far we have used mainly __SIRF__ functionality, now we are going to use __CIL__ in order to set up the reconstruction problem and then solve it. In order to be able to reconstruct both $u$ and $v$ at the same time, we will make use of `BlockDataContainer`.

In the following we define the `SmoothJointTV` class. Our plan is to use a modified Gradient descent (`GD`) algorithm to solve the above problems. This implements the `__call__` method required to monitor the objective value and the `gradient` method that evaluates the gradient of `JTV`.

For the two subproblems, the first variations with respect to $u$ and $v$ variables are:

$$
\begin{equation}
\begin{aligned}
& \biggl(1-\frac{g_{1}}{A_{1}u+\eta}\biggl) - \alpha_{1} \mathrm{div}\bigg( \frac{\nabla u}{|\nabla(u, v)|_{2,\eta,\lambda}}\bigg)\\
& \biggl(1-\frac{g_{2}}{A_{2}v+\eta}\biggl) - \alpha_{2} \mathrm{div}\bigg( \frac{\nabla v}{|\nabla(u, v)|_{2,\eta,1-\lambda}}\bigg)
\end{aligned}
\end{equation}
$$

where $$|\nabla(u, v)|_{2,\eta,\lambda} = \sqrt{ \lambda|\nabla u|^{2} + (1-\lambda)|\nabla v|^{2} + \eta^{2}}.$$

and $ \eta $ means we can avoid dividing by zero

```python
class SmoothJointTV(Function):

    def __init__(self, eta, axis, lambda_par, domain_geometry):

        r'''
        :param eta: smoothing parameter making SmoothJointTV differentiable
        '''

        super(SmoothJointTV, self).__init__(L=8)

        # smoothing parameter
        self.eta = eta

        # GradientOperator
        self.grad = GradientOperator(domain_geometry)

        # Which variable to differentiate
        self.axis = axis

        if self.eta<=0:
            raise ValueError('Need positive value for eta')

        self.lambda_par=lambda_par

    def __call__(self, x):

        r""" x is BlockDataContainer that contains (u,v). Actually x is a BlockDataContainer that contains 2 BDC.
        """
        if not isinstance(x, BlockDataContainer):
            raise ValueError('__call__ expected BlockDataContainer, got {}'.format(type(x)))

        tmp = (self.lambda_par*self.grad.direct(x[0]).pnorm(2).power(2) + (1-self.lambda_par)*self.grad.direct(x[1]).pnorm(2).power(2)+\
              self.eta**2).sqrt().sum()

        return tmp

    def gradient(self, x, out=None):

        denom = (self.lambda_par*self.grad.direct(x[0]).pnorm(2).power(2) + (1-self.lambda_par)*self.grad.direct(x[1]).pnorm(2).power(2)+\
              self.eta**2).sqrt()

        if self.axis==0:
            num = self.lambda_par*self.grad.direct(x[0])
        else:
            num = (1-self.lambda_par)*self.grad.direct(x[1])

        if out is None:
            tmp = self.grad.range.allocate()
            tmp[self.axis].fill(self.grad.adjoint(num.divide(denom)))
            return tmp
        else:
            self.grad.adjoint(num.divide(denom), out=out[self.axis])
```

We now slightly alter the gradient descent `update` method in order to enforce positivity. This update projects our image onto the set of positive numbers and modifies our algorithm slightly to something known as the Gradient Projection algorithm.

```python
def updateGD(self):
        '''Single iteration of Gradient Descent'''

        self.objective_function.gradient(self.x, out=self.x_update)
        self.x_update.multiply(self.step_size, out = self.x_update)
        self.x.subtract(self.x_update, out = self.x)

        # remove any negative values
        self.x.add(self.x.abs(), out = self.x) ## new line
        self.x.divide(2, out = self.x)            ## new line

GD.update = updateGD
```

Now we are going to put everything together and define our two objective functions which solve the two subproblems which we defined at the beginning

```python
# Parameters
alpha1 = 0.5
alpha2 = 0.5
lambda_par = 0.5
eta = 1e-5
eta_sino = templ_sino.get_uniform_copy(eta) # uniform sino for KL diveregence smoothing

# BlockGeometry for the two modalities
bg = BlockGeometry(u_simple, v_simple)

# Projection map to image u or v
L1 = ProjectionMap(bg, index=0)
L2 = ProjectionMap(bg, index=1)

# Fidelity terms based on the acqusition data
f1 = KullbackLeibler(b=noisy_data_fdg, eta=eta_sino)
f2 = KullbackLeibler(b=noisy_data_amyl, eta=eta_sino)

# JTV for each of the subproblems
JTV1 = alpha1*SmoothJointTV(eta=eta, axis=0, lambda_par = lambda_par , domain_geometry = uMap)
JTV2 = alpha2*SmoothJointTV(eta=eta, axis=1, lambda_par = 1-lambda_par, domain_geometry = uMap)

# Compose the two objective functions
op1 = CompositionOperator(acq_model, L1)
op2 = CompositionOperator(acq_model, L2)

objective1 = OperatorCompositionFunction(f1 , op1) + JTV1
objective2 = OperatorCompositionFunction(f2 , op2) + JTV2
```

# (D) Solve the joint TV reconstruction problem

```python
# We start with a cylindrical uniform image
init = make_cylindrical_FOV(fdg.get_uniform_copy(1))
x0 = BlockDataContainer(init,init)

step_size = 0.25

# Here we add the functionality for a preconditioned gradient step to improve the convergence rate
sens_im_arr = acq_model.backward(templ_sino.get_uniform_copy(1)).as_array()
sens_im_arr[sens_im_arr == 0] = 1
sens_im = fdg.clone().fill(sens_im_arr)

precond_im = fdg.get_uniform_copy(1)/sens_im

outeriter = 4
inneriter = 4

# We are also going to log the value of the objective functions
obj1_val_it = []
obj2_val_it = []
# and the images
images = []

for i in range(outeriter):

    precond = (x0.copy()/sens_im).abs()

    step = precond * step_size
    images.append(x0.copy())

    gd1 = GD(x0, objective1, step_size=step, \
          max_iteration = inneriter, update_objective_interval = 1)
    gd1.run(verbose=1)

    # We skip the first one because it gets repeated
    obj1_val_it.extend(gd1.objective[1:])

    # Here we are going to do a little "trick" in order to better see, when each subproblem is optimised, we
    # are going to append NaNs to the objective function which is currently not optimised. The NaNs will not
    # show up in the final plot and hence we can nicely see each subproblem.
    obj2_val_it.extend(numpy.ones_like(gd1.objective[1:])*numpy.nan)

    gd2 = GD(gd1.solution, objective2, step_size=step, \
          max_iteration = inneriter, update_objective_interval = 1)
    gd2.run(verbose=1)

    obj2_val_it.extend(gd2.objective[1:])
    obj1_val_it.extend(numpy.ones_like(gd2.objective[1:])*numpy.nan)

    x0 = gd2.solution.clone()

    print('* * * * * * Outer Iteration ', i+1, ' * * * * * *\n')
images.append(x0.copy())
print('* * * * * * Reconstruction Complete * * * * * *\n')
```

Finally we can look at the images $u_{jtv}$ and $v_{jtv}$ and compare them to the simple reconstruction $u_{simple}$ and $v_{simple}$ and the original ground truth images.

```python
u_jtv = numpy.squeeze(numpy.abs(x0[0].as_array()))
v_jtv = numpy.squeeze(numpy.abs(x0[1].as_array()))

plt.figure()
plot_2d_image([2,3,1], numpy.squeeze(numpy.abs(u_simple.as_array())), '$u_{simple}$')
plot_2d_image([2,3,2], u_jtv, '$u_{JTV}$')
plot_2d_image([2,3,3], numpy.squeeze(numpy.abs(fdg_arr)), '$u_{gt}$')

plot_2d_image([2,3,4], numpy.squeeze(numpy.abs(v_simple.as_array())), '$v_{simple}$')
plot_2d_image([2,3,5], v_jtv, '$v_{JTV}$')
plot_2d_image([2,3,6], numpy.squeeze(numpy.abs(amyl_arr)), '$v_{gt}$')
```

And let's look at the objective functions

```python
plt.figure()
plt.plot(obj1_val_it, 'o-', label='subproblem 1')
plt.plot(obj2_val_it, '+-', label='subproblem 2')
plt.xlabel('Number of iterations')
plt.ylabel('Value of objective function')
plt.title('Objective functions')
plt.legend()

# Logarithmic y-axis
plt.yscale('log')
```

Finally, let's have a look at how the images changed throughout the reconstruction

```python
pio.renderers.default = 'iframe_connected'
plt.ioff()
# choose which subproblem to display
subprob = 0

fig, ax = plt.subplots()
arrays = []
ims = []
for i, image in enumerate(images):
    array = numpy.squeeze(image[subprob].as_array())
    arrays.append(array)
    im = ax.imshow(array, animated = True, vmin=0, vmax=fdg_arr.max())
    im.set_clim(0,array.max())
    title = plt.text(0.5,1.01,str(i), ha="center",va="bottom",color=numpy.random.rand(3),
                     transform=ax.transAxes, fontsize="large")
    if i == 0:
        ax.imshow(array)
        fig.colorbar(im, ax=ax)
    ims.append([im,title,])

ani = animation.ArtistAnimation(fig, ims, interval = 750, blit = True, repeat_delay = 1000)
HTML(ani.to_jshtml())
```

```python
#%% delete temporary files
wdpath = os.getcwd()
for filename in glob.glob(os.path.join(wdpath, "tmp*")):
    os.remove(filename)
```

# Next steps

The above is a good demonstration for a synergistic image reconstruction of two different images. The following gives a few suggestions of what to do next and also how to extend this notebook to other applications.

## Number of iterations
In our problem we have several regularisation parameters such as $\alpha_{1}$, $\alpha_{2}$ and $\lambda$. In addition, the number of inner iterations for each subproblem (currently set to 3) and the number of outer iterations (currently set to 10) also determine the final solution. Of course, for infinite number of total iterations it shouldn't matter but usually we don't have that much time.

__TODO__: Change the number of iterations and see what happens to the objective functions. For a given number of total iterations, do you think it is better to have a high number of inner or high number of outer iterations? Why? Does this also depend on the undersampling factor?

## Spatial misalignment
In the above example we simulated our data such that there is a perfect spatial match between $u$ and $v$. For real world applications this usually cannot be assumed.

__TODO__: Add spatial misalignment between $u$ and $v$. This can be achieved e.g. by calling `numpy.roll` on `amyl_arr` before calling `v_gt = crop_and_fill(im_mr, amyl_arr)`. What is the effect on the reconstructed images? For a more "advanced" misalignment, have a look at notebook `BrainWeb`.

__TODO__: One way to minimize spatial misalignment is to use image registration to ensure both $u$ and $v$ are well aligned. In the notebook `sirf_registration` you find information about how to register two images and also how to resample one image based on the spatial transformation estimated from the registration. Try to use this to correct for the misalignment you introduced above. For a real world example, at which point in the code would you have to carry out the registration+resampling? (some more information can also be found at the end of [MAPEM_Bowsher notebook](MAPEM_Bowsher.ipynb))

## Pathologies
The images $u$ and $v$ show the same anatomy, just with a different distributions. Clinically more useful are of course images which show complementary image information.

__TODO__: Add a pathology to either $u$ and $v$ and see how this effects the reconstruction. For something more advanced, have a look at the notebook `BrainWeb`.

## Single anatomical prior
So far we have alternated between two reconstruction problems. Another option is to do a single regularised reconstruction and simply use a previously reconstructed image for regularisation.

__TODO__: Adapt the above code such that $u$ is reconstructed first without regularisation and is then used for a regularised reconstruction of $v$ without any further updates of $u$.

## Other regularisation options
In this example we used a TV-based regularisation, but of course other regularisers could also be used, such as directional TV.

__TODO__: Have a look at the __CIL__ notebook `02_Dynamic_CT` and adapt the `SmoothJointTV` class above to use directional TV.

## Other optimisation algorithms
In this example we used a Gradient Descent for ease of understanding, but other optimisation algorithms will be more appropriate and more efficient.

__TODO__: Have a look at the __CIL__ Documentaion and see if you can implement other algorithms such as PDHG or ADMM.

## Resolution Modelling
In this example we do not use any resolution modelling for PET

__TODO__: Implement resolution modelling and see what difference it makes to the reconstructed images


```python
from sympy.abc import *
import sympy
from sympy import Eq, Function, Symbol
sympy.init_printing()

Mr = Function('M_r', positive=True)(r)
P = Function('P')(r)
rho = Function('rho')(r)
phi = Function('phi')(r)
mH = Symbol('m_H', positive=True)
T_avg = Symbol('T_{avg}')
rho_C = Symbol('rho_C', positive=True)
xi = Function('xi')
```

```python
n = 0
1 / xi**2
```


# Scenario D - Peakshape Variation

In this scenario the peakshape in a spectrum with a fixed number of peaks is varied from Gaussian (n = 0.0) to Lorentzian (n = 1.0). All datasets contain 3 peaks and the noise level is kept constant at 1%.

The model used in the inference of the parameters is formulated as follows:

\begin{equation}
\large y = f(x) = \sum\limits_{m=1}^M \big[A_m \cdot f_{pseudo-Voigt}(x)\big] + \epsilon
\end{equation}

where:

\begin{equation}
\large f_{pseudo-Voigt}(x) = \eta \cdot \frac{\sigma_m^2}{(x-\mu_m)^2 + \sigma_m^2} + (1 - \eta) \cdot e^{-\frac{(x-\mu_m)^2}{2\cdot\sigma_m^2}}
\end{equation}

```python
%matplotlib inline
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pymc3 as pm
import arviz as az

#az.style.use('arviz-darkgrid')

print('Running on PyMC3 v{}'.format(pm.__version__))
```

    WARNING (theano.tensor.blas): Using NumPy C-API based implementation for BLAS functions.

    Running on PyMC3 v3.8

## Import local modules

```python
import os
import sys
sys.path.append('../../modules')
import datagen as dg
import models as mdl
import results as res
import figures as fig
import settings as cnf
```

## Local configuration

```python
# output for results and images
out_path      = './output_peakshape_plot'
file_basename = out_path + '/scenario_peakshape'

conf = {}

# scenario name
conf['scenario'] = 'peakshape variation'

# initialization method for sampler
conf['init_mode'] = 'adapt_diag'

# probabilistic model (priors)
conf['prior_model'] = 'lognormal'

# provide peak positions to the model as testvalues ('yes'/'no')
conf['peak_info'] = 'yes'

# model mode ('train'/eval')
conf['model_mode'] = 'eval'

# data mode ('generate'/'preload')
conf['data_mode'] = 'generate'

# dataset directory (needed for 'preload' data mode)
#conf['dataset_dir'] = './input_datasets'

# number of cores to run sampling chains on
conf['ncores'] = 2

# number of samples per chain
conf['nsamples'] = 2000
```

```python
# if the output dir does not exist, create it
if not os.path.exists(out_path):
    os.makedirs(out_path)

conf
```

    {'scenario': 'peakshape variation',
     'init_mode': 'adapt_diag',
     'prior_model': 'lognormal',
     'peak_info': 'yes',
     'model_mode': 'eval',
     'data_mode': 'generate',
     'ncores': 2,
     'nsamples': 2000}

## Save configuration

```python
cnf.save(out_path, conf)
```

# Generate data and plot

```python
# list of wavelengths (x-values)
xval = [i for i in range(200, 400, 2)]

ldata  = []
lpeaks = []
lpeakshape = []

# number of spectra per baseline variation
nsets  = 1

# peakshape weight factors (0 = Gauss, 1 = Lorentz)
#peakshapes = [0.0, 0.25, 0.5, 0.75, 1.0]
peakshapes = [0.0, 1.0]
lpeakshape = [ps for ps in peakshapes for i in range(nsets)]

# total number of datasets
tsets = nsets * len(peakshapes)

if conf['model_mode'] == 'train' and conf['data_mode'] == 'generate':
    # generate the datasets
    for ps in peakshapes:
        for i in range(nsets):
            df, peaks, _ = dg.data_generator(xvalues=xval, nsamples=15, npeaks=3, peakshape=ps)
            ldata.append(df)
            lpeaks.append(peaks)

    # save data and peak information to disk
    for i in range(len(ldata)):
        ldata[i].to_csv(out_path + '/dataset_%02d.csv' % (i+1), index=False)
    dg.data_save(out_path + '/peakinfo.csv', lpeaks)

elif conf['model_mode'] == 'train' and conf['data_mode'] == 'preload':
    # load pre-generated datasets from disk
    ldata, lpeaks, _ = dg.data_load(tsets, conf['dataset_dir'])

else:
    # load data from disk
    if conf['data_mode'] == 'preload':
        ldata, lpeaks, _ = dg.data_load(tsets, conf['dataset_dir'])
    else:
        ldata, lpeaks, _ = dg.data_load(tsets, out_path)
```

```python
print("total number of peakshape variations    : {0}".format(len(peakshapes)))
print("total number of datasets per peakshape  : {0}".format(nsets))
print("total number of datasets per model      : {0}".format(tsets))
print("total number of inference runs          : {0}".format(nsets*len(peakshapes)**2))
```

    total number of peakshape variations    : 2
    total number of datasets per peakshape  : 1
    total number of datasets per model      : 2
    total number of inference runs          : 4

```python
# plot datasets
fig.plot_datasets(ldata, lpeaks, dims=(int(tsets/2),2), figure_size=(12,int(tsets*(1.8))),
                            savefig='yes', fname=file_basename, scenario='peakshape', labels=lpeakshape)
```

# Initialize models and run inference

```python
# convert pandas data to numpy arrays
x_val = np.array(xval, dtype='float32')

# store dataset y-values in list
cols = ldata[0].columns
y_val = [ldata[i][cols].values for i in range(len(ldata))]
```

```python
# initialize models and run inference
models = []
traces = []

# list of tuples with (model peakshape, data peakshape) combination
lmodbase = []

# actual run number
run = 1

# total number of inference runs
truns = nsets * len(peakshapes)**2

for ps in peakshapes:
    if conf['model_mode'] == 'train':
        # for each baseline model run inference on all spectra
        print("running: n = {0} peakshape model".format(ps))

    for i in range(len(ldata)):
        if conf['peak_info'] == 'yes':
            #plist = lpeaks[i].flatten()
            plist = np.array(lpeaks[i], dtype='float32').flatten()
            plist.sort()
            model_g = mdl.model_pvoigt(xvalues=x_val, observations=y_val[i], npeaks=3, peakshape=ps,
                                  mu_peaks=plist, pmodel=conf['prior_model'])
        else:
            model_g = mdl.model_pvoigt(xvalues=x_val, observations=y_val[i], npeaks=3, peakshape=ps,
                                                  pmodel=conf['prior_model'])

        models.append(model_g)

        with model_g:
            if conf['model_mode'] == 'train':
                print("({2}/{3}) running inference on dataset #{0}/{1} [{4} model: {5} data]"
                      .format(i+1,len(ldata),run,truns,ps,lpeakshape[i]))
                trace_g = pm.sample(conf['nsamples'], init=conf['init_mode'], cores=conf['ncores'])
                lmodbase += [(ps,lpeakshape[i])]
                traces.append(trace_g)
                # save inference results
                pm.backends.text.dump(out_path + '/traces_%02d' % (run), trace_g)
            else:
                # load traces from disk
                print("loading dataset #{0}/{1}".format(run,truns))
                trace_g = pm.backends.text.load(out_path + '/traces_%02d' % (run))
                traces.append(trace_g)
            run += 1
```

    loading dataset #1/4
    loading dataset #2/4
    loading dataset #3/4
    loading dataset #4/4

# Model visualization

```python
# display first model
pm.model_to_graphviz(models[0])
```

```python
# save model figure as image
img = pm.model_to_graphviz(models[0])
img.render(filename=file_basename + '_model', format='png');
```

# Collect results and save

```python
# posterior predictive traces
ppc = [pm.sample_posterior_predictive(traces[i], samples=500, model=models[i]) for i in range(len(traces))]
```

    /home/johan/VirtualEnv/ppsda/lib/python3.6/site-packages/pymc3/sampling.py:1247: UserWarning: samples parameter is smaller than nchains times ndraws, some draws and/or chains may not be represented in the returned posterior predictive sample
      "samples parameter is smaller than nchains times ndraws, some draws "
    100%|██████████| 500/500 [00:07<00:00, 65.82it/s]
    /home/johan/VirtualEnv/ppsda/lib/python3.6/site-packages/pymc3/sampling.py:1247: UserWarning: samples parameter is smaller than nchains times ndraws, some draws and/or chains may not be represented in the returned posterior predictive sample
      "samples parameter is smaller than nchains times ndraws, some draws "
    100%|██████████| 500/500 [00:07<00:00, 66.62it/s]
    100%|██████████| 500/500 [00:07<00:00, 68.72it/s]
    /home/johan/VirtualEnv/ppsda/lib/python3.6/site-packages/pymc3/sampling.py:1247: UserWarning: samples parameter is smaller than nchains times ndraws, some draws and/or chains may not be represented in the returned posterior predictive sample
      "samples parameter is smaller than nchains times ndraws, some draws "
    100%|██████████| 500/500 [00:06<00:00, 72.26it/s]

```python
# various plots to inspect the inference results
varnames = mdl.get_varnames(traces[0])

#az.plot_trace(traces[2], varnames, compact=True);
#az.plot_trace(traces[2], varnames, divergences='top');
#az.plot_autocorr(traces[0], varnames);
#az.plot_posterior(traces[8], varnames);

#for idx, trace in enumerate(traces):
#    az.plot_forest(trace, var_names = varnames, r_hat=True, ess=True);

#az.summary(traces[20], varnames)
```

```python
if conf['model_mode'] == 'train':
    # collect the results and display
    df = res.get_results_summary(traces, ppc, y_val, varnames, epsilon_real=0.05, sets=tsets, labels=lmodbase)
else:
    # load results from disk
    df = pd.read_csv(file_basename + '.csv')
    df.index += 1

    # create list of tuples with model/data combinations
    lm = df['model'].to_list()
    ld = df['data'].to_list()
    lmodbase = list(zip(lm,ld))
df
```

      r_hat
      mcse
      ess
      bfmi
      r2
      waic
      epsilon
      epsilon_real
      model
      data

      1
      1.0
      0.0000
      4279.1
      1.069652
      0.999777
      -4813.777589
      0.048500
      0.05
      0.0
      0.0

      2
      1.0
      0.0015
      5495.1
      1.088961
      0.961493
      4566.554773
      1.103881
      0.05
      0.0
      1.0

      3
      1.0
      0.0008
      5053.3
      1.002693
      0.968342
      2498.818329
      0.552781
      0.05
      1.0
      0.0

      4
      1.0
      0.0000
      3846.0
      1.073311
      0.999914
      -4798.013943
      0.048727
      0.05
      1.0
      1.0

```python
if conf['model_mode'] == 'train':
    # save results to .csv
    df.to_csv(file_basename + '.csv', index=False)
```

# Plot posterior

```python
fig.plot_posterior(x_val, ldata, traces, ppc, dims=(int(truns/2),2), figure_size=(12,int(truns*(1.8))),
            savefig='yes', fname=file_basename, showpeaks='no', sets=tsets, labels=lmodbase, scenario='peakshape')
```

```python
fig.plot_posterior(x_val, ldata, traces, ppc, dims=(int(truns/2),2), figure_size=(12,int(truns*(1.8))),
            savefig='yes', fname=file_basename, showpeaks='yes', sets=tsets, labels=lmodbase, scenario='peakshape')
```

```python
cnf.close(out_path)
```

# Plot posterior (single dataset)

```python
# dataset to plot
n_dataset = 1

# trace number in matrix
n_trace = 1

fig.plot_posterior_single(x_val, ldata[n_dataset], ppc[n_trace], 10, file_basename, hpd_color='lightgray')
```

```python

```


```python
import numpy as np
from numpy import linalg as LA
import math
import random

#任意次元のベクトルを生成する関数
#nは生成したい次元
def qubit(n):

    e=np.zeros(n,dtype='float64')

    e[0]=random.random()
    unf=1-e[0]

    for i in range(1,n-1):
        e[i]=random.uniform(0,unf)
        unf=unf-e[i]
    e[n-1]=unf

    return e

#ランダムなベクトルを規格化し、状態ベクトルになる．
def normal(e):
    n=e.size      #配列の次元を確認する関数

    E=np.zeros(n,dtype='float64')

    for i in range(0,n):
        E[i]=np.sqrt(e[i])
    return E

#ベクトルの間のテンソル積
def tensor(a,b):
    n1=a.size
    n2=b.size

    P=np.zeros(n1*n2,dtype='float64')
    for i in range(0,n1):
        for j in range(0,n2):
            P[i*n2+j]=a[i]*b[j]
    return P

#xはn1次元のベクトル、yはn2次元のベクトルとおく、zの次元はn1*n2である
x=normal(qubit(2))
y=normal(qubit(2))
z=tensor(x,y)

print('状態１は')
print(x)
print('状態2は')
print(y)
print('合成系の状態は')
print(z)

#ノルムの確認

print(LA.norm(x,2))

print(LA.norm(y,2))

print(LA.norm(z,2))
```

    状態１は
    [0.9428247  0.33328905]
    状態2は
    [0.54873968 0.83599328]
    合成系の状態は
    [0.51736532 0.78819511 0.18288893 0.27862741]
    1.0
    1.0
    0.9999999999999999

# 純粋状態の密度演算子

\begin{equation}
\rho_{1}=|n1\rangle\langle n1|,\rho_{2}=|n2\rangle\langle n2|,
\end{equation}

合成系の密度演算子は

\begin{equation}
\rho=|n1\otimes n2\rangle\langle n1\otimes n2|
\end{equation}

```python
rho=np.dot(z,z.T)
print(rho)
print(z.T)
```

    0.9999999999999998
    [0.51736532 0.78819511 0.18288893 0.27862741]

```python
print(z.T)
```

    [0.51736532 0.78819511 0.18288893 0.27862741]

```python

#密度演算子を生成する関数
def density(a):
    dimension=a.size

    b=np.zeros([dimension,dimension],dtype='float64')
    for i in range(0,dimension):
        for j in range(0,dimension):
            b[i][j]=a[i]*a[j]

    return b

#トレースを計算する関数
def trace(a):
    sum=0;
    n=a[0].size
    for i in range(0,n):
        sum+=a[i][i]

    return sum

rho=density(z)
print(rho)
print(trace(rho))
```

    [[0.26766688 0.40778482 0.09462039 0.14415216]
     [0.40778482 0.62125153 0.14415216 0.21961276]
     [0.09462039 0.14415216 0.03344836 0.05095787]
     [0.14415216 0.21961276 0.05095787 0.07763323]]
    0.9999999999999998

```python
m1=np.array([1,0])
m2=np.array([0,1])
print(density(m1))
```

    [[1. 0.]
     [0. 0.]]

# 計算基底における同時測定

計算基底について、

\begin{equation}
|00\rangle,  |01\rangle,  |10\rangle,  |11\rangle
\end{equation}

```python
print(tensor(m1,m2))
```

    [0. 1. 0. 0.]

```python
#00 ketの確率を計算する

m=tensor(m1,m2)

def prob(base,state):
    z=np.dot(base,state)
    return math.pow(z,2)

#zは合成系の状態ベクトルである
prob1=prob(tensor(m1,m1),z)
prob2=prob(tensor(m1,m2),z)
prob3=prob(tensor(m2,m1),z)
prob4=prob(tensor(m2,m2),z)

print(prob1,prob2,prob3,prob4)
print(prob1+prob2+prob3+prob4)

#print(prob(m,z))
```

    0.5444983619682794 0.16483890869705195 0.22311724838525626 0.06754548094941248
    1.0

# 計算基底における部分測定

```python
print(np.dot(m1,x))
```

    0.8422216279966522

```python
parcial1=np.dot(m1,x)*y

print(parcial1)
print(math.pow(LA.norm(parcial1,2),2))
```

    [0.73790132 0.40600358]
    0.7093372706653315

# 密度演算子による部分測定

```python
rho_m=trace(density(y))
print(trace(rho_m*density(x)*density(m1)))
```

    0.7093372706653313

# 密度演算子による同時測定

状態$\rho$の下での同時測定は

\begin{equation}
    Pr(E=m,F=n | \rho)=Tr(\rho E_{m}\otimes F_{n})
\end{equation}

```python
#行列式同士のテンソルを作る関数
#P135

#n1=a.size
 #   n2=math.pow(n2,2)

 #   e=np.zeros([n2,n2],dtype='float64')

  #  for i in range(0,n1):
   #     for j in range(0,n1):

```

```python
print(x)
print(LA.norm(x,2))
rho1=density(x)
print(rho1)

```

    [0.9428247  0.33328905]
    1.0
    [[0.88891841 0.31423315]
     [0.31423315 0.11108159]]

```python
E1=0.4*density(m1)
print(E1)
```

    [[0.4 0. ]
     [0.  0. ]]

```python
E2=0.4*density(np.array([1/np.sqrt(2),1/np.sqrt(2)]))
print(E2)
```

    [[0.2 0.2]
     [0.2 0.2]]

```python
E3=np.array([[1,0],[0,1]])-E2-E1
print(E3)
```

    [[ 0.4 -0.2]
     [-0.2  0.8]]

```python
probi1=trace(np.dot(rho1,E1))
print(probi1)
probi2=trace(np.dot(rho1,E2))
probi3=trace(np.dot(rho1,E3))
print(probi2,probi3)
print(probi1+probi2+probi3)
```

    0.35556736235065745
    0.32569326064375625 0.31873937700558636
    1.0


```
from sympy import I, Matrix, symbols, Symbol, eye
from datetime import datetime
import numpy as np
import pandas as pd

```

```
'''
Rotas
'''
C = Matrix([[0,"Cp",0],[0,0,"Cv"],[0,0,0]])
D = Matrix([[0,"Dp",0],[0,0,"Dv"],[0,0,0]])
```

```
'''
Semáforo
'''
S = Matrix([[0, "Sp"],["Sv", 0]])
```

```
def kronSum(A,B):
  m = np.size(A,1)
  n = np.size(B,1)
  A = np.kron(A,np.eye(n))
  B = np.kron(np.eye(m),B)
  return A + B
```

```
momento_inicio = datetime.now()
'''
Algebra de rotas
'''
rotas = kronSum(C,D)

'''
Algebra de seções
'''
# NÃO NECESSÁRIA: UMA SÓ SEÇÃO

'''
Algebra de sistema
'''
sistema = np.kron(rotas, S)
# calcula tempo de processamento
tempo_processamento = datetime.now() - momento_inicio
```

```
sistema = pd.DataFrame(data=sistema,index=list(range(1,np.size(sistema,0)+1)), columns=list(range(1,np.size(sistema,1)+1)))
```

```
sistema.shape
```

    (18, 18)

```
print(tempo_processamento)
```

    0:00:00.010831

```
sistema
```

      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      13
      14
      15
      16
      17
      18

      1
      0
      0
      0
      1.0*Dp*Sp
      0
      0
      0
      1.0*Cp*Sp
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      2
      0
      0
      1.0*Dp*Sv
      0
      0
      0
      1.0*Cp*Sv
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      3
      0
      0
      0
      0
      0
      1.0*Dv*Sp
      0
      0
      0
      1.0*Cp*Sp
      0
      0
      0
      0
      0
      0
      0
      0

      4
      0
      0
      0
      0
      1.0*Dv*Sv
      0
      0
      0
      1.0*Cp*Sv
      0
      0
      0
      0
      0
      0
      0
      0
      0

      5
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Cp*Sp
      0
      0
      0
      0
      0
      0

      6
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Cp*Sv
      0
      0
      0
      0
      0
      0
      0

      7
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Dp*Sp
      0
      0
      0
      1.0*Cv*Sp
      0
      0
      0
      0

      8
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Dp*Sv
      0
      0
      0
      1.0*Cv*Sv
      0
      0
      0
      0
      0

      9
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Dv*Sp
      0
      0
      0
      1.0*Cv*Sp
      0
      0

      10
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Dv*Sv
      0
      0
      0
      1.0*Cv*Sv
      0
      0
      0

      11
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Cv*Sp

      12
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Cv*Sv
      0

      13
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Dp*Sp
      0
      0

      14
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Dp*Sv
      0
      0
      0

      15
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Dv*Sp

      16
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      1.0*Dv*Sv
      0

      17
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      18
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

```
momento_inicio = datetime.now()
colunas = ['denode', 'paranode', 'aresta']
grafo = pd.DataFrame(columns=colunas)
r = 1
c = 1
for j in range(np.size(sistema,0)):
  for i in range(np.size(sistema,0)):
    if sistema.loc[r,c]==0 and c < np.size(sistema,0):
      c += 1
    elif c < np.size(sistema,0):
      grafo.loc[len(grafo)+1] = (r, c, sistema.loc[r,c])
      c += 1
    else:
      c = 1
      r += 1
tempo_processamento = datetime.now() - momento_inicio
print(tempo_processamento)
```

    0:00:00.107340

```
grafo
```

      denode
      paranode
      aresta

      1
      1
      4
      1.0*Dp*Sp

      2
      1
      8
      1.0*Cp*Sp

      3
      2
      3
      1.0*Dp*Sv

      4
      2
      7
      1.0*Cp*Sv

      5
      3
      6
      1.0*Dv*Sp

      6
      3
      10
      1.0*Cp*Sp

      7
      4
      5
      1.0*Dv*Sv

      8
      4
      9
      1.0*Cp*Sv

      9
      5
      12
      1.0*Cp*Sp

      10
      6
      11
      1.0*Cp*Sv

      11
      7
      10
      1.0*Dp*Sp

      12
      7
      14
      1.0*Cv*Sp

      13
      8
      9
      1.0*Dp*Sv

      14
      8
      13
      1.0*Cv*Sv

      15
      9
      12
      1.0*Dv*Sp

      16
      9
      16
      1.0*Cv*Sp

      17
      10
      11
      1.0*Dv*Sv

      18
      10
      15
      1.0*Cv*Sv

      19
      12
      17
      1.0*Cv*Sv

      20
      13
      16
      1.0*Dp*Sp

      21
      14
      15
      1.0*Dp*Sv

      22
      16
      17
      1.0*Dv*Sv

```
grafo.to_csv('grafo.csv', sep=";")
```

```
from google.colab import files
files.download('grafo.csv')
```

    <IPython.core.display.Javascript object>

    <IPython.core.display.Javascript object>