# Symbolic Mathematics in Python

There are times when you need to solve a difficult problem symbollically or analytically. If you have ever used Wolfram Alpha, then you have already done this. Sympy is a python library that allows you to do symbolic mathematics in python.

```python
import sympy as sym
```

## 1. Introduction

### Example 1.1

If you try to write the follwing in python by itself, you will get an error telling you x is undefined:
$$x-x$$

```python
x-x
```

(The error above is on purpose). Variables in python need to be defined before you can say something specific about them

```python
x=102
x-x
```

    0

If you are trying to show that $x-x=0$ is true for any $x$, the above answer would not be valid. Instead you can use a symbolic expression to show that it is true

**First we define the variable as a symmbolic expression**

```python
x = sym.symbols('x')
```

**Now we can use the variable in a symbolic expression**

```python
x-x
```

    0

### Example 1.2

Sympy can be used to perform algebreic operations (among other things). Consider the following expression: $$(3a-4b)^3$$

We can use symppy to expand the expression algebraically.

**First we need to define the variables as symbolic expressions**

```python
a,b = sym.symbols('a,b')
```

**Side note** Notice that the left hand side of the epression has two variables being defined. Python can define more than one variable at a time:

```python
x1,y1 =10,20
print(x1)
print(y1)
```

    10
    20

**Back to the expression** We can define an expression using the variables $a$ and $b$.

```python
expr = (3*a-4*b)**3
print(expr)
```

    (3*a - 4*b)**3

We can also make it look nicer in our notebook. This doesn't affect the math, but it makes our notebook more readable.

```python
sym.init_printing()
```

```python
expr
```

**Now we expand the function algebreically**

```python
expr.expand()
```

Sympy can also factor the equation

```python
sym.factor(26*a**3-108*a**2*b+144*a*b**2-64*b**3)
```

If you want to copy and paste a result, you print the result.

```python
print(sym.factor(26*a**3-108*a**2*b+144*a*b**2-64*b**3))
```

    2*(a - 2*b)*(13*a**2 - 28*a*b + 16*b**2)

You can also chain together functions

```python
expr.expand().factor()
```

### Exercise 1.1

Show that the following two expressions are true.
$$(2w-3z)(2w+3z)=4w^2-9z^2$$
$$(2w-3z)^2\ne4w^2-9z^2$$

```python
w,z = sym.symbols('w,z')
expr = (2*w-3*z)*(2*w+3*z)
expr.expand()
```

$\displaystyle 4 w^{2} - 9 z^{2}$

```python
#(2𝑤−3𝑧)(2𝑤+3𝑧)=4𝑤2−9𝑧2
```

```python
w,z = sym.symbols('w,z')
expr = (2*w-3*z)**2
expr.expand()
```

$\displaystyle 4 w^{2} - 12 w z + 9 z^{2}$

```python
#(2𝑤−3𝑧)2≠4𝑤2−9𝑧2
```

## 2. Solving Equations

Sympy can be used to symbolilically solve equations. As before, you need to define which variables are symbols

### Example 2.1

Use sympy to solve the following equation
$$ax^3+bx^2+cx+d=0$$

```python
# Define the variables
a,b,c,d,x = sym.symbols('a,b,c,d,x')
```

```python
# Define the expression
expr=a*x**3+b*x**2+c*x+d
expr
```

We can use the `solvset` function to solve this equation

```python
solutions=sym.solveset(expr,x)
```

```python
print(solutions)
```

    {-(-3*c/a + b**2/a**2)/(3*(sqrt(-4*(-3*c/a + b**2/a**2)**3 + (27*d/a - 9*b*c/a**2 + 2*b**3/a**3)**2)/2 + 27*d/(2*a) - 9*b*c/(2*a**2) + b**3/a**3)**(1/3)) - (sqrt(-4*(-3*c/a + b**2/a**2)**3 + (27*d/a - 9*b*c/a**2 + 2*b**3/a**3)**2)/2 + 27*d/(2*a) - 9*b*c/(2*a**2) + b**3/a**3)**(1/3)/3 - b/(3*a), -(-3*c/a + b**2/a**2)/(3*(-1/2 - sqrt(3)*I/2)*(sqrt(-4*(-3*c/a + b**2/a**2)**3 + (27*d/a - 9*b*c/a**2 + 2*b**3/a**3)**2)/2 + 27*d/(2*a) - 9*b*c/(2*a**2) + b**3/a**3)**(1/3)) - (-1/2 - sqrt(3)*I/2)*(sqrt(-4*(-3*c/a + b**2/a**2)**3 + (27*d/a - 9*b*c/a**2 + 2*b**3/a**3)**2)/2 + 27*d/(2*a) - 9*b*c/(2*a**2) + b**3/a**3)**(1/3)/3 - b/(3*a), -(-3*c/a + b**2/a**2)/(3*(-1/2 + sqrt(3)*I/2)*(sqrt(-4*(-3*c/a + b**2/a**2)**3 + (27*d/a - 9*b*c/a**2 + 2*b**3/a**3)**2)/2 + 27*d/(2*a) - 9*b*c/(2*a**2) + b**3/a**3)**(1/3)) - (-1/2 + sqrt(3)*I/2)*(sqrt(-4*(-3*c/a + b**2/a**2)**3 + (27*d/a - 9*b*c/a**2 + 2*b**3/a**3)**2)/2 + 27*d/(2*a) - 9*b*c/(2*a**2) + b**3/a**3)**(1/3)/3 - b/(3*a)}

```python
solutions
```

What if I need help. You can do this with any python function. `function?`

```python
# Run this command to see a help box
sym.solveset?
```

### Exercise 2.1

Use the `solveset` function to solve the following chemical problem.

Phosgene gas, $\text{COCl}_2$, dissociates at high temperatures according to the following equilibrium:

$$ \text{COCl}_2 \rightleftharpoons  \text{CO} + \text{Cl}_2 $$

At $\text{400 C}$, the equilibrium constant $K_c=8.05$.

If you start with a  $\text{0.250 M}$ phosgene sample at $\text{400 C}$, determine the concentrations of all species at equilibrium.

```python
x = sym.symbols('x')
k = 8.05
expr=(((x**2)/(0.25-x))-k)
solutions=sym.solveset(expr,x)
print(solutions)
```

    {-8.29268379803378, 0.242683798033777}

Why did you pick your answer?

```python
#answer: 0.243
    #concentrations can't be negative
```

## 3. Calculus

We can use also Sympy to differentiate and integrate. Let us experiment with differentiating the following expression:

$$x ^ 2 - \cos(x)$$

```python
sym.diff(x ** 2 - sym.cos(x), x)
```

Similarly we can integrate:

```python
sym.integrate(x ** 2 - sym.cos(x), x)
```

We can also carry out definite integrals:

```python
sym.integrate(x ** 2 - sym.cos(x), (x, 0, 5))
```

### Exercise 3.1

Use Sympy to calculate the following:

1. $\frac{d(x ^2 + xy - \ln(y))}{dy}$
1. $\int_0^5 e^{2x}\;dx$

```python
x,y = sym.symbols('x,y')
sym.diff((((x ** 2 )+ (x*y)) - sym.ln(y))/(y),x,y)
```

$\displaystyle \frac{1 - \frac{2 x + y}{y}}{y}$

```python
x,y = sym.symbols('x,y')
sym.integrate((sym.E**(2*x)), (x, 0, 5))
```

$\displaystyle - \frac{1}{2} + \frac{e^{10}}{2}$

### Exercise 3.2
Solve the following definate integral
$$\int\limits_{ - \infty }^\infty  {\frac{1}{{\sigma \sqrt {2\pi } }}{e^{ - \frac{1}{2}{{\left( {\frac{{x - \mu }}{\sigma }} \right)}^2}}}}$$

Hint, the sympy symbol for infinity is `oo`

```python
x,m,o,oo = sym.symbols('x,m,o,oo')
sym.integrate(((1/(o*(sym.sqrt(2*sym.pi)))*(sym.E**((-1/2)*((x-m)/o)**2)))), (x, -oo, oo))
```

$\displaystyle 0.353553390593274 \sqrt{2} \operatorname{erf}{\left(- \frac{0.707106781186548 \left(- m - oo\right)}{o} \right)} + 0.353553390593274 \sqrt{2} \operatorname{erf}{\left(\frac{0.707106781186548 \left(- m + oo\right)}{o} \right)}$

Lookup Gaussian functions: https://en.wikipedia.org/wiki/Gaussian_function
Does your answer maake sense?

## 4. Plotting with Sympy

Finally Sympy can be used to plot functions. Note that this makes use of [matplotlib](http://matplotlib.org/).

Let us plot $x^2$:

```python
expr = x ** 2
p = sym.plot(expr)
```

### Exercise 4.1 Plot the following function:

1. $y=x + cos(x)$
1. ${\frac{1}{{ \sqrt {2\pi } }}{e^{ - \frac{x^2}{2}}}}$

```python
expr = y = x + sym.cos(x)
p = sym.plot(expr)
```

```python
expr = ((1/(sym.sqrt(2*sym.pi)))*(sym.E**(-(x**2)/2)))
p = sym.plot(expr)
```

# Lecture

```python

```

## L1. Hydrogen Atom

Sympy has built in modules for the eigenfunctions of the hydrogen atom.

```python
import sympy.physics.hydrogen
import numpy as np
```

You can caluclate the eigenvalues ($E$) in Hartrees

`sym.physics.hydrogen.E_nl(n,Z)`

```python
sym.physics.hydrogen.E_nl(1,1)
```

    -1/2

We can use a loop to print out many energies

```python
for n in range(1,5):
    print(sym.physics.hydrogen.E_nl(n,1))
```

    -1/2
    -1/8
    -1/18
    -1/32

We can plot the hydrogen radial wavefunction (1s orbital)

```python
sympy.symbols('r')
sympy.physics.hydrogen.R_nl(1, 0, r, 1)
```

    2*exp(-r)

```python
sym.plot(sympy.physics.hydrogen.R_nl(1, 0, r, 1),(r,0,10.50))
```

And the probablity distribution function

```python
sympy.symbols('r')
prob_1s=sympy.physics.hydrogen.R_nl(1, 0, r, 1)*sympy.physics.hydrogen.R_nl(1, 0, r, 1)
prob_1s
```

    4*exp(-2*r)

```python
sym.plot(prob_1s,(r,0,10))
```

Plot a 2s orbital

```python
sympy.symbols('r')
prob_2s=sympy.physics.hydrogen.R_nl(2, 0, r, 1)*sympy.physics.hydrogen.R_nl(2, 0, r, 1)
prob_2s
```

```python
sym.plot(prob_2s,(r,0,10))
```

We can change the range to see the node better.

```python
sym.plot(prob_2s,(r,1,8))
```

Notice the node!

### Exercise L1.1

Plot the radial distriubution function for a 2p, 3s, 4s, and 3d orbital.

```python
import sympy.physics.hydrogen
import numpy as np
r = sym.symbols('r')
prob_2p=sympy.physics.hydrogen.R_nl(2, 1, r, 1)*sympy.physics.hydrogen.R_nl(2, 1, r, 1)
prob_2p
sym.plot(prob_2p,(r,0,10))
```

```python
import sympy.physics.hydrogen
import numpy as np
r = sym.symbols('r')
prob_3s=sympy.physics.hydrogen.R_nl(3, 0, r, 1)*sympy.physics.hydrogen.R_nl(3, 0, r, 1)
prob_3s
sym.plot(prob_3s,(r,1,6))
```

```python
import sympy.physics.hydrogen
import numpy as np
r = sym.symbols('r')
prob_4s=sympy.physics.hydrogen.R_nl(4, 0, r, 1)*sympy.physics.hydrogen.R_nl(4, 0, r, 1)
prob_4s
sym.plot(prob_4s,(r,1,6))
```

```python
import sympy.physics.hydrogen
import numpy as np
r = sym.symbols('r')
prob_3d=sympy.physics.hydrogen.R_nl(3, 2, r, 1)*sympy.physics.hydrogen.R_nl(3, 2, r, 1)
prob_3d
sym.plot(prob_3d,(r,0,20))
```

```python

```


### ¿Cómo vibra un tambor cuando lo golpeas?
> Analizar el problema de la membrana vibrante permite entender el funcionamiento de instrumentos de percusión tales como los tambores, timbales e incluso sistemas biológicos como el tímpano.

Referencias:
 - https://en.wikipedia.org/wiki/Bessel_function
 - https://es.wikipedia.org/wiki/Vibraciones_de_una_membrana_circular
 - https://docs.scipy.org/doc/scipy-0.18.1/reference/tutorial/special.html

___

Considere un tambor _(membrana)_ de radio $a$, entonces la función de onda en $\mathbb{R}^2$ para este sistema se puede escribir como,
$$ \frac{1}{v}\frac{\partial^2 u}{\partial t^2} = \frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2} $$

donde $u\equiv u(x,y,t)$ es el desplazamiento transversal_(elevación)_ y $v$ es la rapidez de propagación de la onda.

La forma habitual de encontrar la solución a esta ecuación es primero hacer un cambio de coordenadas, de cartesianas a polares y posteriormente considerar el método de separación de variable, tal que

$$ u(r, \theta, t) = R(r) \Theta(\theta) T (t)$$

Esta sustitución da como resultado tres ecuaciones diferenciales, una para cada variable de separación. Y cuyas soluciones nos ayudan a escribir los **modos normales**.

$$u_{nk}(r,\theta, t) = J_{n}(\lambda_{nk} r)(a_{nk}\cos{n\theta} + b_{nk}\sin{n\theta})\cos{(v\lambda_{nk}t)}$$
$$u^{*}_{nk}(r,\theta, t) = J_{n}(\lambda_{nk}r)(a^{*}_{nk}\cos{n\theta} + b^{*}_{nk}\sin{n\theta})\sin{(v\lambda_{nk}t)})$$

para $n = 0,1,2,\dots$, $k = 1,2,3,\dots$, donde $J_{n}$ es la **función de Bessel** de orden $n$ de primera clase.  Además,
$$\lambda_{nk} = \frac{\alpha_{nk}}{a}$$

donde $\alpha_{nk}$ es el k-ésimo cero de $J_{n}(\lambda a)=0$. Esto es consecuencia de que $u$ sea cero en la frontera de la membrana, $r = a$.

Los coeficientes $a_{nk} , b_{nk}, a^{*}_{nk}$ y $b^{*}_{nk}$ se determinan de tal forma que cumplan con las condiciones iniciales:

$$u(r,\theta, 0) = f(r,\theta)$$
$$u_{t}(r,\theta, 0) = g(r,\theta)$$

donde el primer termino es la geometría inicial y el segundo la rapidez inicial. Se puede demostrar que las expresiones para estos coeficientes se pueden escribir como:

\begin{align}
a_{0k} &= \frac{1}{\pi a^2 J_{1}^{2}(\alpha_{0k})}\int_{0}^{2\pi}\int_{0}^{a}\; f(r,\theta)\, J_{0}(\lambda_{0k}r)\, r \, dr \, d\theta\\
a_{nk} &= \frac{2}{\pi a^2 J_{n+1}^{2}(\alpha_{nk})}\int_{0}^{2\pi}\int_{0}^{a}\; f(r,\theta)\, J_{n}(\lambda_{nk}r)\cos(n\theta)\, r \, dr \, d\theta\\
b_{nk} &= \frac{2}{\pi a^2 J_{n+1}^{2}(\alpha_{nk})}\int_{0}^{2\pi}\int_{0}^{a}\; f(r,\theta)\, J_{n}(\lambda_{nk}r)\sin(n\theta)\, r \, dr \, d\theta
\end{align}

Y similarmente,

\begin{align}
a^{*}_{0k} &= \frac{1}{\pi \,v\, \alpha_{0k}\,a J_{1}^{2}(\alpha_{0k})}\int_{0}^{2\pi}\int_{0}^{a}\; g(r,\theta)\, J_{0}(\lambda_{0k}r)\, r \, dr \, d\theta\\
a^{*}_{nk} &= \frac{2}{\pi\, v\,\alpha_{0k}\, a J_{n+1}^{2}(\alpha_{nk})}\int_{0}^{2\pi}\int_{0}^{a}\; g(r,\theta)\, J_{n}(\lambda_{nk}r)\cos(n\theta)\, r \, dr \, d\theta\\
b^{*}_{nk} &= \frac{2}{\pi\, v\,\alpha_{0k}\, a J_{n+1}^{2}(\alpha_{nk})}\int_{0}^{2\pi}\int_{0}^{a}\; g(r,\theta)\, J_{n}(\lambda_{nk}r)\sin(n\theta)\, r \, dr \, d\theta
\end{align}

Estamos familiarizados con la función `coseno`, pero no tanto con la función de Bessel. Entonces, nuestra primera actividad será conocer su comportamiento.

\begin{align}
u(r,\theta, t) &= \sum_{n=0}^{\infty}\sum_{k = 1}^{\infty}J_{n}(\lambda_{nk} r)(a_{nk}\cos{n\theta} + b_{nk}\sin{n\theta})\cos{(v\lambda_{nk}t)}\\
&+  \sum_{n=0}^{\infty}\sum_{k = 1}^{\infty}J_{n}(\lambda_{nk}r)(a^{*}_{nk}\cos{n\theta} + b^{*}_{nk}\sin{n\theta})\sin{(v\lambda_{nk}t)})
\end{align}

```python
%matplotlib inline
import matplotlib.pyplot as plt
from scipy import special
import numpy as np
from ipywidgets import *
```

```python
r = np.linspace(0, 10,100)
for n in range(5):
    plt.plot(r, special.jn(n, r), label = '$J_{%s}(r)$'%n)
plt.xlabel('$r$', fontsize = 18)
plt.ylabel('$J_{n}(r)$', fontsize = 18)
plt.axhline(y = 0, color = 'k')
plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5), prop={'size': 14})
plt.show()
```

 Por simplicidad vamos a suponer  que $r = 1$ y  determinar los ceros, significa encontrar todas las intersecciones de las curvas anteriores con el eje horizontal.

### Ejemplo: Caso radialmente simétrico

Suponga que $a = 1$, $v = 1$ y que las condiciones iniciales son:

$$ f(r,\theta) = 1- r^4\quad\quad g(r,\theta) = 0$$

Dado que la rapidez inicial es cero, entonces $a^{*}_{nk} = b^{*}_{nk} = 0$

```python
def f_shape(r):
    return 1 - r**4
```

```python
r = np.r_[0:1:100j]
angle = np.r_[0:2*np.pi:200j]
r_shape = f_shape(r)
```

```python
u = np.array([np.full(len(angle), radi) for radi in r_shape])
x = np.array([var_r * np.cos(angle) for var_r in r])
y = np.array([var_r * np.sin(angle) for var_r in r])
```

```python
plt.figure(figsize = (6, 5))
plt.pcolor(x, y, u_good, cmap = 'viridis')
plt.axis('off')
plt.colorbar()
plt.show()
```

Y la solución para el desplazamiento en el tiempo es simplemente,
\begin{equation}
u(r,\theta, t) = \sum_{n=0}^{\infty}\sum_{k = 1}^{\infty}J_{n}(\lambda_{nk} r)(a_{nk}\cos{n\theta} + b_{nk}\sin{n\theta})\cos{(v\lambda_{nk}t)}
\end{equation}

Entonces, solo será necesario encontrar $a_{nk}$ y $b_{nk}$.

\begin{align}
a_{0k} &= \frac{1}{\pi a^2 J_{1}^{2}(\alpha_{0k})}\int_{0}^{2\pi}\int_{0}^{a}\; f(r,\theta)\, J_{0}(\lambda_{0k}r)\, r \, dr \, d\theta\\
a_{nk} &= \frac{2}{\pi a^2 J_{n+1}^{2}(\alpha_{nk})}\int_{0}^{2\pi}\int_{0}^{a}\; f(r,\theta)\, J_{n}(\lambda_{nk}r)\cos(n\theta)\, r \, dr \, d\theta\\
b_{nk} &= \frac{2}{\pi a^2 J_{n+1}^{2}(\alpha_{nk})}\int_{0}^{2\pi}\int_{0}^{a}\; f(r,\theta)\, J_{n}(\lambda_{nk}r)\sin(n\theta)\, r \, dr \, d\theta
\end{align}

Para resolver estas integrales haremos uso de `sympy`. Iniciemos con $a_{nk}$.

### Primer caso $n>0$

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

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

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

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

$$n$$

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

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

$$\lambda_{0k}$$

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

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

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

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

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

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

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

$$0$$

Entonces para cualquier $n>0$ no se tiene contribución.

### Segundo caso $n=0$

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

```python
integ(0)
```

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

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

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

```python
a0k_solution = a0k.doit()
a0k_solution
```

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

```python
a0k_sol = 1/(sym.pi*sym.besselj(1, lamb(0,k))**2)*a0k_solution
a0k_sol
```

$$\frac{1}{J^{2}_{1}\left(\lambda_{0k}\right)} \left(- \frac{8 J_{0}\left(\lambda_{0k}\right)}{\lambda_{0k}^{2}} + \frac{32 J_{1}\left(\lambda_{0k}\right)}{\lambda_{0k}^{3}} + \frac{64 J_{0}\left(\lambda_{0k}\right)}{\lambda_{0k}^{4}} - \frac{128 J_{1}\left(\lambda_{0k}\right)}{\lambda_{0k}^{5}}\right)$$

```python
sym.simplify(a0k_sol)
```

$$\frac{1}{\lambda_{0k}^{5} J^{2}_{1}\left(\lambda_{0k}\right)} \left(- 8 \lambda_{0k}^{3} J_{0}\left(\lambda_{0k}\right) + 32 \lambda_{0k}^{2} J_{1}\left(\lambda_{0k}\right) + 64 \lambda_{0k} J_{0}\left(\lambda_{0k}\right) - 128 J_{1}\left(\lambda_{0k}\right)\right)$$

### Y para $b_{nk}$

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

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

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

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

```python
solution_b = bnk.doit()
solution_b
```

$$0$$

### ¿Qué sucede para $n = 0$?

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

```python
integ_b(0)
```

$$0$$

Ahhh! el integrando es cero, entonces este termino también es cero.

\begin{equation}
u(r,\theta, t) = \sum_{k = 1}^{\infty} a_{0k}J_{0}(\lambda_{0k} r)\cos{(v\lambda_{0k}t)}
\end{equation}

```python
a0k_sol
```

$$\frac{1}{J^{2}_{1}\left(\lambda_{0k}\right)} \left(- \frac{8 J_{0}\left(\lambda_{0k}\right)}{\lambda_{0k}^{2}} + \frac{32 J_{1}\left(\lambda_{0k}\right)}{\lambda_{0k}^{3}} + \frac{64 J_{0}\left(\lambda_{0k}\right)}{\lambda_{0k}^{4}} - \frac{128 J_{1}\left(\lambda_{0k}\right)}{\lambda_{0k}^{5}}\right)$$

**Primero vamos a programar para algún modo $k$. **

```python
def a0k_sym(lambd):
    solucion = (-8*special.jn(0, lambd)/lambd**2
    +32*special.jn(1, lambd)/lambd**3 +
     64*special.jn(0, lambd)/lambd**4 -
        128*special.jn(1, lambd)/lambd**5)/special.jn(1, lambd)**2
    return solucion
```

```python
def tambor(v, kth_zero, nt, t):
    r = np.r_[0:1:100j]
    angle = np.r_[0:2*np.pi:200j]
    ceros = special.jn_zeros(0, nt)
    lambd = ceros[kth_zero]
    u_r = a0k_sym(lambd)*special.jn(0, lambd * r) * np.cos(lambd * v * t)
    u = np.array([np.full(len(angle), u_rs) for u_rs in u_r])
    x = np.array([var_r * np.cos(angle) for var_r in r])
    y = np.array([var_r * np.sin(angle) for var_r in r])
    return x, y, u
```

```python
x1, y1, u1 = tambor(1, 0, 15, 7)
```

```python
plt.figure(figsize = (6, 5))
plt.pcolor(x1 , y1 , u1, cmap = 'viridis')
plt.axis('off')
plt.colorbar()
plt.show()
```

```python
def tambor_nk(t = 0, kth=0):
    fig = plt.figure(figsize = (6,5))
    ax = fig.add_subplot(1, 1, 1)
    x, y, u = tambor(1, kth, 50, t)
    im = ax.pcolor(x, y, u, cmap = 'viridis')
    ax.set_xlim(xmin = -1.1, xmax = 1.1)
    ax.set_ylim(ymin = -1.1, ymax = 1.1)
    plt.colorbar(im)
    plt.axis('off')
    fig.canvas.draw()
interact_manual(tambor_nk, t = (0, 15,.01), n = (0, 10, 1), kth = (0, 10, 1));
```

```python
def tambor_nk(t = 0, kth=0):
    fig = plt.figure(figsize = (6,5))
    ax = fig.add_subplot(1, 1, 1)
    x, y, u = tambor(1, kth, 50, t)
    im = ax.pcolor(x, y, u, cmap = 'viridis', vmax = 1.2, vmin = -1.2)
    ax.set_xlim(xmin = -1.1, xmax = 1.1)
    ax.set_ylim(ymin = -1.1, ymax = 1.1)
    plt.colorbar(im)
    plt.axis('off')
    fig.canvas.draw()
interact(tambor_nk, t = (0, 15,.01), n = (0, 10, 1), kth = (0, 10, 1));
```

**Y ahoara, la solución completa. **

```python
def tambor_n_allk(v, nk_zeros, t):
    r = np.r_[0:1:100j]
    angle = np.r_[0:2*np.pi:200j]
    ceros = special.jn_zeros(0, nk_zeros)
    lambd = ceros[0]
    u_r = a0k_sym(lambd)*special.jn(0, lambd * r) * np.cos(lambd * v * t)
    u0 = np.array([np.full(len(angle), u_rs) for u_rs in u_r])
    for cero in range(1, nk_zeros):
        lambd = ceros[cero]
        u_r = a0k_sym(lambd)*special.jn(0, lambd * r) * np.cos(lambd * v * t)
        u = np.array([np.full(len(angle), u_rs) for u_rs in u_r])
        u0 += u
    x = np.array([var_r * np.cos(angle) for var_r in r])
    y = np.array([var_r * np.sin(angle) for var_r in r])
    return x, y, u0
```

```python
def tambor_0(t = 0):
    fig = plt.figure(figsize = (6,5))
    ax = fig.add_subplot(1, 1, 1)
    x, y, u = tambor_n_allk(1, 15, t)
    im = ax.pcolor(x, y, u, cmap = 'viridis')
    ax.set_xlim(xmin = -1.1, xmax = 1.1)
    ax.set_ylim(ymin = -1.1, ymax = 1.1)
    plt.colorbar(im)
    plt.axis('off')
    fig.canvas.draw()
interact_manual(tambor_0, t = (0, 15,.01));
```

Fíjise bien, la condición inicial en $t = 0$, se cumple para la  solución encontrada.

### Tarea

- **Problema 1.** Suponga que $a = 1$, $v = 1$ y que las condiciones iniciales son:

$$ f(r,\theta) = (1- r^4)\cos(\theta)\quad\quad g(r,\theta) = 0$$

- **Problema 2.** uponga que $a = 2$, $v = 1$ y que las condiciones iniciales son:
$$ f(r,\theta) = 0 \quad\quad g(r,\theta) = r^2(2- r)\sin^8\left(\frac{\theta}{2}\right)$$

## Fin Modulo 1

```python

```

```python

```

```python

```

### Versión antigua (Lo único que tal vez sirva de algo, son los programas)
Tal vez sirvan de algo para sus tareas.

Entonces, primero veamos algunos modos normales del sistema. Por ejemplo(demasiado simplificado),

$$u(r,\theta, t)_{nk} = J_{n}(\lambda_{nk} r)\,\cos(n\theta)\,\cos(\lambda_{nk} v t)$$

La siguiente función se aplica caso simplificado.

```python
def tambor(n, r_max, v, kth_zero, nt, t):
    r = np.r_[0:r_max:100j]
    angle = np.r_[0:2*np.pi:200j]
    ceros = special.jn_zeros(0, nt)
    lamb = ceros[kth_zero]
    u = np.array([special.jn(n, lamb* var_r) * np.cos(n * angle)
                 * np.cos(lamb * v * t) for var_r in r])
    x = np.array([var_r * np.cos(angle) for var_r in r])
    y = np.array([var_r * np.sin(angle) for var_r in r])
    return x, y, u
```

Entonces, por ejemplo si $n = 1$, $a = 1$, $v = 1$,  $k = 1$ y $t= 0$. Este sería el modo de vibración $(n,k)\rightarrow (1,1)$.

```python
x, y, u = tambor(1, 1, 1, 0, 15, 0)
```

```python
plt.figure(figsize = (6, 5))
plt.pcolor(x, y, u, cmap = 'viridis')
plt.axis('off')
plt.colorbar()
plt.show()
```

Ahora, veamos como lucen todos demás modos de vibración $(n,k)$.

```python
def tambor_nk(t = 0, n = 0, kth=0):
    fig = plt.figure(figsize = (6,5))
    ax = fig.add_subplot(1, 1, 1)
    x, y, u = tambor(n, 1, 1, kth, 15, t)
    im = ax.pcolor(x, y, u, cmap = 'viridis')
    ax.set_xlim(xmin = -1.1, xmax = 1.1)
    ax.set_ylim(ymin = -1.1, ymax = 1.1)
    plt.colorbar(im)
    plt.axis('off')
    fig.canvas.draw()
interact_manual(tambor_nk, t = (0, 15,.01), n = (0, 10, 1), kth = (0, 10, 1));
```

Ahora, tal vez nos interesaría conocer el comportamiento de la membrana cuando sumamos sobre un conjunto de modos $k$. Es decir,

$$u(r,\theta, t)_{n} =\sum_{k = 1}u(r,\theta, t)_{nk} =  \sum_{k = 1}J_{n}(\lambda_{nk} r)\,\cos(n\theta)\,\cos(\lambda_{nk} v t) $$
La manera usual de hacer esto es considerar la suma en series de Fourier, es decir a esta suma le falta un coeficiente $A_{nk}$, pero por simplicidad aquí no vamos a considerar este término.

Una posible función para realizar esto sería,

```python
def tambor_n_allk(n, r_max, v, nk_zeros, t):
    r = np.r_[0:r_max:100j]
    angle = np.r_[0:2*np.pi:200j]
    ceros = special.jn_zeros(0, nk_zeros)
    lamb = ceros[0]
    u0 = np.array([special.jn(n, lamb* var_r) * np.cos(n * angle)
                  * np.cos(lamb * v * t) for var_r in r])
    for cero in range(1, nk_zeros):
        lamb = ceros[cero]
        u = np.array([special.jn(n, lamb* var_r) * np.cos(n * angle)
                 * np.cos(lamb * v * t) for var_r in r])
        u0 += u
    x = np.array([var_r * np.cos(angle) for var_r in r])
    y = np.array([var_r * np.sin(angle) for var_r in r])
    return x, y, u0
```

```python
def tambor_n(t = 0, n = 0):
    fig = plt.figure(figsize = (6,5))
    ax = fig.add_subplot(1, 1, 1)
    x, y, u = tambor_n_allk(n, 1, 1, 15, t)
    im = ax.pcolor(x, y, u, cmap = 'viridis')
    ax.set_xlim(xmin = -1.1, xmax = 1.1)
    ax.set_ylim(ymin = -1.1, ymax = 1.1)
    plt.colorbar(im)
    plt.axis('off')
    fig.canvas.draw()
interact_manual(tambor_n, t = (0, 15,.01), n = (0, 10, 1));
```

Por último, nos queda el caso cuando sumamos sobre todos los modos $n$. Es decir,

```python

```

```python
def order_n(n, ceros, nk_zeros, angle, v, r, t):
    lamb = ceros[0]
    u0 = np.array([special.jn(n, lamb* var_r) * np.cos(n * angle)
                      * np.cos(lamb * v * t) for var_r in r])
    for cero in range(1, nk_zeros):
        lamb = ceros[cero]
        u = np.array([special.jn(n, lamb* var_r) * np.cos(n * angle)
                     * np.cos(lamb * v * t) for var_r in r])
        u0 += u
    return u0
```

```python
def tambor(orden_n, r_max, v, nk_zeros, t):
    r = np.r_[0:r_max:100j]
    angle = np.r_[0:2*np.pi:100j]
    ceros = special.jn_zeros(0, nk_zeros)
    u0 = order_n(0, ceros, nk_zeros,  angle, v, r, t)
    for n in range(1, orden_n):
        u = order_n(n, ceros, nk_zeros, angle, v, r, t)
        u0 += u
    x = np.array([var_r * np.cos(angle) for var_r in r])
    y = np.array([var_r * np.sin(angle) for var_r in r])
    return x, y, u0
```

```python
x, y, u = tambor(10, 1, 1, 5, 5)
```

```python
plt.figure(figsize = (5, 5))
plt.pcolor(x, y, u, cmap = 'inferno')
plt.axis('on')
plt.show()
```


# Input Driven HMM

This notebook is a simple example of an HMM with exogenous inputs.  The inputs modulate the probability of discrete state transitions via a multiclass logistic regression. Let $z_t \in \{1, \ldots, K\}$ denote the discrete latent state at time $t$ and $u_t \in \mathbb{R}^U$ be the exogenous input at time~$t$.  The transition probability is given by,

$$
\begin{align}
\Pr(z_t = k \mid z_{t-1} = j, u_t) =
\frac{\exp\{\log P_{j,k} + w_k^\mathsf{T} u_t\}}
{\sum_{k'=1}^K \exp\{\log P_{j,k'} + w_{k'}^\mathsf{T} u_t\}}.
\end{align}
$$

The parameters of the transition model are $P \in \mathbb{R}_+^{K \times K}$, a baseline set of (unnormalized) transition weights, and $W \in \mathbb{R}^{K \times U}$, a set of input weights.

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

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

%matplotlib inline

npr.seed(0)
sns.set(palette="colorblind")
```

## 2. Create an Input Driven HMM
SSM is designed to be modular, so that the user can easily mix and match different types of transitions and observations.
We create an input-driven HMM with the following line:
```python
true_hmm = ssm.HMM(num_states, obs_dim, input_dim,
               observations="categorical", observation_kwargs=dict(C=num_categories),
               transitions="inputdriven")
```

Let's look at what each of the arguments do. The first three arguments specify the number of states, and the dimensionality of the observations and inputs.

**Setting the observation model**
For this example, we have set `observations="categorical"`, which means each observation will take on one of a discrete set of values, i.e $y_t \in \{1, \ldots, C \}$.
For categorical observations, the observations are drawn from a multinomial distribution, with parameters depending on the current state. Assuming $z_t = k$,the observations are a vector $y \in \mathbb{R}^D$, where $y_i \sim \text{mult} (\lambda_{k,i})$, where $\lambda_{k,i}$ is the multinomal parameter associated with coordinate $i$ of the observations in state $k$. Note that each observation variable is independent from the others.

For categorical observations, we also specify the number of discrete observations possible (in this case 3). We do this by creating a dictionary where the keys are the keyword arguments which we want to pass to the observation model. For categorical observations, there is just one keyword argument, `C`, which specifies the number of categories. This is set using `observation_kwargs=dict(C=num_categories)`.

The observations keyword argument should be one of : `"gaussian", "poisson" "studentst", "exponential", "bernoulli", "autoregressive", "robust_autoregressive"`.
**NOTE:**
Setting the observations as "autoregressive" means that each observation will be dependent on the prior observation, as well as on the input (if the input is nonzero). By constrast, the standard "inputdriven" transitions are not affected by previous observations or directly by the inputs.

**Setting the transition model**
In order to create an HMM with exogenous inputs, we set ```transitions="inputdriven"```. This means that the baseline transition matrix $P$ is modified according to a Generalized Linear Model, as described at the top of the page.

SSM support many transition models, set by keyword argument to the constructor of the class. The keyword argument should be one of: `"standard", "sticky", "inputdriven", "recurrent", "recurrent_only", "rbf_recurrent", "nn_recurrent".` We're working on creating standalone documentation to describe these in more detail. For most users, the stationary and input driven transition classes should suffice.

**Creating inputs and sampling**
After creating our HMM object, we create an input array called `inpt` which is simply a jittered sine wave. We also increase the transition weights so that it will be clear (for demonstration purposes) that the input is changing the transition probabilities. In this case, we will actually increase the weights such that the transitions appear almost deterministic.

```python
# Set the parameters of the HMM
time_bins = 1000 # number of time bins
num_states = 2    # number of discrete states
obs_dim = 1    # data dimension
input_dim = 1    # input dimension
num_categories = 3    # number of output types/categories

# Make an HMM
true_hmm = ssm.HMM(num_states, obs_dim, input_dim,
               observations="categorical", observation_kwargs=dict(C=num_categories),
               transitions="inputdriven")

# Optionally, turn up the input weights to exaggerate the effect
true_hmm.transitions.Ws *= 3

# Create an exogenous input
inpt = np.sin(2 * np.pi * np.arange(time_bins) / 50)[:, None] + 1e-1 * npr.randn(time_bins, input_dim)

# Sample some data from the HMM
true_states, obs = true_hmm.sample(time_bins, input=inpt)

# Compute the true log probability of the data, summing out the discrete states
true_lp = true_hmm.log_probability(obs, inputs=inpt)

# By default, SSM returns categorical observations as a list of lists.
# We convert to a 1D array for plotting.
obs_flat = np.array([x[0] for x in obs])
```

```python
# Plot the data
plt.figure(figsize=(8, 5))

plt.subplot(311)
plt.plot(inpt)
plt.xticks([])
plt.xlim(0, time_bins)
plt.ylabel("input")

plt.subplot(312)
plt.imshow(true_states[None, :], aspect="auto")
plt.xticks([])
plt.xlim(0, time_bins)
plt.ylabel("discrete\nstate")
plt.yticks([])

# Create Cmap for visualizing categorical observations
plt.subplot(313)
plt.imshow(obs_flat[None,:], aspect="auto", )
plt.xlim(0, time_bins)
plt.ylabel("observation")
plt.grid(b=None)
plt.show()
```

### 2.1 Exercise: EM for the input-driven HMM
There are a few good references that derive the EM algorithm for the case of a vanilla HMM (e.g Machine Learning and Pattern Recognition, by Chris Bishop, and [this tutorial](https://www.ece.ucsb.edu/Faculty/Rabiner/ece259/Reprints/tutorial%20on%20hmm%20and%20applications.pdf) by Lawrence Rabiner. How should the EM updates change for the case of input-driven HMMs?

## 3. Fit an input-driven HMM to data
Below, we'll show to fit an input-driven HMM from data. We'll treat the samples generated above as a dataset, and try to learn the appropriate HMM parameters from this dataset.

We create a new HMM object here, with the same parameters as the HMM in Section 1:
```python
hmm = ssm.HMM(num_states, obs_dim, input_dim,
          observations="categorical", observation_kwargs=dict(C=num_categories),
          transitions="inputdriven")
```

We fit the dataset simply by calling the `fit` method:
```python
hmm_lps = hmm.fit(obs, inputs=inpt, method="em", num_em_iters=N_iters)
```
Here, the variable `hmm_lps` will be set to a list of log-probabilities at each step of the EM-algorithm, which we'll use to check convergence.

```python
# Now create a new HMM and fit it to the data with EM
N_iters = 100
hmm = ssm.HMM(num_states, obs_dim, input_dim,
          observations="categorical", observation_kwargs=dict(C=num_categories),
          transitions="inputdriven")

# Fit
hmm_lps = hmm.fit(obs, inputs=inpt, method="em", num_iters=N_iters)
```

    HBox(children=(IntProgress(value=0), HTML(value='')))

### 3.1 Permute the latent states, check convergence
As in the vanilla-HMM notebook, we need to find a permutation of the latent states from our new hmm such that they match the states from the true HMM above. SSM accomplishes this with two function calls: first, we call `find_permutation(true_states, inferred_states)` which returns a list of indexes into states.

Then, we call `hmm.permute(permuation)` with the results of our first function call. Finally, we set `inferred_states` to be the underlying states we predict given the data.

Below, we plot the results of the `fit` function in order to check convergence of the EM algorithm. We see that the log-probability from the EM algorithm approaches the true log-probability of the data (which we have stored as `lp_true`).

```python
# Find a permutation of the states that best matches the true and inferred states
hmm.permute(find_permutation(true_states, hmm.most_likely_states(obs, input=inpt)))
inferred_states = hmm.most_likely_states(obs, input=inpt)
```

```python
# Plot the log probabilities of the true and fit models
plt.plot(hmm_lps, label="EM")
plt.plot([0, N_iters], true_lp * np.ones(2), ':k', label="True")
plt.legend(loc="lower right")
plt.xlabel("EM Iteration")
plt.xlim(0, N_iters)
plt.ylabel("Log Probability")
plt.show()
```

### 3.3 Exercise: Change the Fitting Method
As an experiment, try fitting the same dataset using another fitting method. The two other fitting methods supported for HMMs are "sgd" and "adam", which you can set by passing `method="sgd"` and `method="adam"` respectively. For these methods, you'll probably need to increase the number of iteratations to around 1000 or so.

After fitting with a different method, re-run the two cells above to generate a plot. How does the convergence of these other methods converge to EM?

```python
# Plot the true and inferred states
plt.figure(figsize=(8, 3.5))

plt.subplot(211)
plt.imshow(true_states[None, :], aspect="auto")
plt.xticks([])
plt.xlim(0, time_bins)
plt.ylabel("true\nstate")
plt.yticks([])

plt.subplot(212)
plt.imshow(inferred_states[None, :], aspect="auto")
plt.xlim(0, time_bins)
plt.ylabel("inferred\nstate")
plt.yticks([])
plt.show()
```

## 4. Visualize the Learned Parameters
After calling `fit`, our new HMM object will have parameters updated according to the dataset. We can get a sense of whether we successfully learned these parameters by comparing them to the _true_ parameters which generated the data.

Below, we plot the baseline log transition probabilities (the log of the state-transition matrix) as well as the input weights $w$.

```python
# Plot the true and inferred input effects
plt.figure(figsize=(8, 4))

vlim = max(abs(true_hmm.transitions.log_Ps).max(),
           abs(true_hmm.transitions.Ws).max(),
           abs(hmm.transitions.log_Ps).max(),
           abs(hmm.transitions.Ws).max())

plt.subplot(141)
plt.imshow(true_hmm.transitions.log_Ps, vmin=-vlim, vmax=vlim, cmap="RdBu", aspect=1)
plt.xticks(np.arange(num_states))
plt.yticks(np.arange(num_states))
plt.title("True\nBaseline Weights")
plt.grid(b=None)

plt.subplot(142)
plt.imshow(true_hmm.transitions.Ws, vmin=-vlim, vmax=vlim, cmap="RdBu", aspect=num_states/input_dim)
plt.xticks(np.arange(input_dim))
plt.yticks(np.arange(num_states))
plt.title("True\nInput Weights")
plt.grid(b=None)

plt.subplot(143)
plt.imshow(hmm.transitions.log_Ps, vmin=-vlim, vmax=vlim, cmap="RdBu", aspect=1)
plt.xticks(np.arange(num_states))
plt.yticks(np.arange(num_states))
plt.title("Inferred\nBaseline Weights")
plt.grid(b=None)

plt.subplot(144)
plt.imshow(hmm.transitions.Ws, vmin=-vlim, vmax=vlim, cmap="RdBu", aspect=num_states/input_dim)
plt.xticks(np.arange(input_dim))
plt.yticks(np.arange(num_states))
plt.title("Inferred\nInput Weights")
plt.grid(b=None)
plt.colorbar()
plt.show()
```


# Let's draw samples from a nested Gaussian Process

We will draw samples from a $(m, n) = (2, 1)$ nested Gaussian process following [Zhu and Dunson (2013)](http://amstat.tandfonline.com/doi/abs/10.1080/01621459.2013.838568).

For our first trick, we assume values for the random walk and observation noise variabilities and sample from the state space model

$$\begin{align}
Y_j &\sim Z_j \theta_j + \epsilon_j \\
\theta_{j+1} &= G_j \theta_j + \omega_j \\
\theta &= \begin{bmatrix} U \\ DU \\ A \end{bmatrix} \\
\epsilon_j &\sim \mathcal{N}(0, \sigma^2_\epsilon) \\
\omega_j &\sim \mathcal{N}(0, W_j) \\
Z_j &=
    \begin{pmatrix}
    1 & 0 & 0
    \end{pmatrix} \\
G_j &=
    \begin{pmatrix}
    1 & \delta_j & \frac{\delta_j^2}{2} \\
    0 & 1 & \delta_j \\
    0 & 0 & 1
    \end{pmatrix} \\
W_j & =
    \begin{pmatrix}
    \frac{\delta_j^3}{3} \sigma_U^2 + \frac{\delta_j^5}{20} \sigma_A^2 &
    \frac{\delta_j^2}{2} \sigma_U^2 + \frac{\delta_j^4}{8} \sigma_A^2 &
    \frac{\delta_j^3}{6} \sigma_A^2 \\
    \frac{\delta_j^2}{2} \sigma_U^2 + \frac{\delta_j^4}{8} \sigma_A^2 &
    \delta_j \sigma_U^2 + \frac{\delta_j^3}{3} \sigma_A^2 &
    \frac{\delta_j^2}{2} \sigma_A^2 \\
    \frac{\delta_j^3}{6} \sigma_A^2 &
    \frac{\delta_j^2}{2} \sigma_A^2 &
    \delta_j \sigma_A^2
    \end{pmatrix} \\
\theta_0 &\sim \mathcal{N}(0, \mathrm{diag}(\sigma^2_\mu, \sigma^2_\mu, \sigma^2_\alpha))
\end{align}$$

In the above, we take $\delta_j \equiv t_{j + 1} - t_j$ and assume $\sigma^2_\mu$ and $\sigma^2_\alpha$ are specified as part of the prior.

# Implementation

```julia
using NGPtools
using PyPlot
plt.style[:use]("ggplot")
srand(12347);
```

    WARNING: could not import Base.help into PyCall

## Set up system matrices and dimensions

```julia
Nt = 100
tobs = linspace(0, 1, Nt);
```

```julia
Np = 3
Nst = 3  #number of states

σϵ = 0.7
σA = 50
σU = .25
σμ = 3
σα = 1.

δ = diff(tobs)[1];
```

## Create fake data

```julia
y, α = generate(Np, Nst, length(tobs), δ, σϵ, σU, σA, σμ, σα);
```

```julia
figure(figsize=(10, 4))

for idx in 1:Np
    scatter(tobs, y[idx,:] , label="data");
    plot(tobs, α[idx,:]')
end
xlim(tobs[1], tobs[end]);
```

```julia
Ns = 500  # number of samples

α_samples = sample(y, Ns, Nst, δ, σϵ, σU, σA, σμ, σα);
```

```julia
trans = 5 / Ns
figure(figsize=(10, 25))
th_names = [L"U", L"DU", L"A"]

for obs in 1:Np, idx in 1:3
    subplot(Nst, 1, idx)
    plot(α[Np * (idx - 1) + obs, :]', linestyle="--", label="Actual")
    plot(squeeze(α_samples[Np * (idx - 1) + obs, :, :], 1), color="k", alpha=trans)
    title(th_names[idx])
    legend()
end
```

```julia

```


# Conceptual and Numerical Approaches in Climate Science

## Example lab: Ice-albedo feedback and Snowball Earth

Welcome!  In this lab we will apply a simple climate model, called an "Energy Balance Model" or EBM, to evaluate the climate impacts of changing ice cover.

We will use the Jupyter environment we've used for previous labs, and we'll integrate knowledge from lecture on albedo and previous lectures on Earth's atmosphere and land cover.

We will study one of the most extreme cases of the ice-albedo feedback hinted at in the geologic record: Snowball Earth.

***what is snowball earth, nice images, more explanation***

```python
#  Ensure compatibility with Python 2 and 3
from __future__ import print_function, division
```

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import climlab
import climlab_snowball as clsn
```

## Contents

1. [Setting up an energy balance model](#section1)
2. [Ice advance and retreat in the EBM](#section2)
3. [Snowball Earth: onset](#section3)
4. [Escape from the Snowball](#section4)

***description of simple energy balance model, what is going on in our 1D Earth, what we will see***

____________

## 1. Setting up an energy balance model
____________

As we saw in lecture, Earth's energy budget is an important control on the climate we experience.  Radiation energy coming in from the sun can be reflected (in the atmosphere or from the surface) or absorbed.  The portion of incoming solar radiation $Q$ that is reflected, $Q_{reflected}$, is set by the **planetary albedo**, $\alpha$.

As we saw in lecture,

```python
Q = 341.3 # W/m2, the incoming solar radiation
Q_reflected = 101.9 # W/m2, the reflected shortwave radiation
alpha = Q_reflected/Q

print('The planetary albedo is {}'.format(alpha))
```

    The planetary albedo is 0.29856431292118374

The portion of incoming solar radiation that is not reflected is the **Absorbed Shortwave Radiation**,
\begin{equation}
ASR=Q-Q_{reflected}=Q (1-\alpha),
\end{equation}
controlled by the albedo $\alpha$.

The heat energy emitted to space at the top of the atmosphere is the **Outgoing Longwave Radiation**, $OLR$.  The total **energy budget** of the Earth system is the balance between energy going out ($OLR$) and coming in ($ASR$):

\begin{align}
\frac{dE}{dt} &= ASR - OLR \\
    &= Q (1-\alpha) - OLR,
\end{align}
where we see $\alpha$ is a key parameter.

The simplest climate models use this equation as their starting point.  Today we will set up such an **energy balance model** in Climlab and use it to explore the ice-albedo feedback.

### Our EBM: latitude-dependent, annual mean

Here, we'll use an energy balance model that accounts for radiation fluxes that vary with latitude $\phi$, but averages out seasonal changes in an annual mean.

The equation the model will solve for us is below.
\begin{align*}
C(\phi) \frac{\partial T_s}{\partial t} = & ~(1-\alpha) ~ Q - \left( A + B~T_s \right) +  \\
& \frac{D}{\cos\phi } \frac{\partial }{\partial \phi} \left(\cos\phi  ~ \frac{\partial T_s}{\partial \phi} \right)
\end{align*}

*1.1. With your lab partner, diagram the physical meaning of each term in the equation.*

*1.2. How do you expect the heat capacity, $C(\phi)$, to vary with latitude?  Why?*

____________

## 2. Interactive snow and ice line in the EBM
____________

### Temperature-dependent ice line

Let the surface albedo be larger wherever the temperature is below some threshold $T_f$:

$$ \alpha\left(\phi, T(\phi) \right) = \left\{\begin{array}{ccc}
\alpha_0 + \alpha_2 P_2(\sin\phi) & ~ & T(\phi) > T_f \\
a_i & ~ & T(\phi) \le T_f \\
\end{array} \right. $$

```python
#  for convenience, set up a dictionary with our reference parameters
param = {'A':210, 'B':2, 'a0':0.3, 'a2':0.078, 'ai':0.62, 'Tf':-10.}
model1 = climlab.EBM_annual(name='Annual EBM with ice line',
                            num_lat=180, D=0.55, **param )
print( model1)
```

    climlab Process of type <class 'climlab.model.ebm.EBM_annual'>.
    State variables and domain shapes:
      Ts: (180, 1)
    The subprocess tree:
    Annual EBM with ice line: <class 'climlab.model.ebm.EBM_annual'>
       LW: <class 'climlab.radiation.aplusbt.AplusBT'>
       insolation: <class 'climlab.radiation.insolation.AnnualMeanInsolation'>
       albedo: <class 'climlab.surface.albedo.StepFunctionAlbedo'>
          iceline: <class 'climlab.surface.albedo.Iceline'>
          warm_albedo: <class 'climlab.surface.albedo.P2Albedo'>
          cold_albedo: <class 'climlab.surface.albedo.ConstantAlbedo'>
       SW: <class 'climlab.radiation.absorbed_shorwave.SimpleAbsorbedShortwave'>
       diffusion: <class 'climlab.dynamics.meridional_heat_diffusion.MeridionalHeatDiffusion'>

Because we provided a parameter `ai` for the icy albedo, our model now contains several sub-processes contained within the process called `albedo`. Together these implement the step-function formula above.

The process called `iceline` simply looks for grid cells with temperature below $T_f$ and adjusts their albedo.

```python
print( model1.param)
```

    {'timestep': 350632.51200000005, 'S0': 1365.2, 's2': -0.48, 'A': 210, 'B': 2, 'D': 0.55, 'Tf': -10.0, 'water_depth': 10.0, 'a0': 0.3, 'a2': 0.078, 'ai': 0.62}

```python
model1.integrate_years(5)
```

    Integrating for 450 steps, 1826.2110000000002 days, or 5 years.
    Total elapsed time is 5.000000000000044 years.

```python
f = clsn.state_plot(model1)
```

Grey shading on the above plots indicates where there is ice in the model.  We can find the same information by querying the model `icelat` attribute:

```python
model1.icelat
```

    array([-70.,  70.])

### Sudden cooling

What happens if we force the model to be colder? Let's store our current model state, then make a new model like it that will decrease the temperature by 20 $^{\circ}C$ everywhere.

```python
## Assign current model diagnostics to separate variables
Tequil = np.array(model1.Ts)
ALBequil = np.array(model1.albedo)
OLRequil = np.array(model1.OLR)
ASRequil = np.array(model1.ASR)
```

```python
m2 = climlab.process_like(model1)
m2.Ts -= 20.
m2.compute_diagnostics()
f3 = clsn.state_plot(m2)
```

```python
f2 = clsn.state_plot(model1)
```

*2.1. Compare and contrast the global patterns of temperature and albedo in this new, colder climate versus the previous simulation.*

*2.2. Repeat the comparison with a warmer climate.  What is responsible for the differences?*

Let's look at the radiative effect - how does the absorbed shortwave radiation change with this colder climate?

```python
my_ticks = [-90,-60,-30,0,30,60,90]
lat = model1.lat

fig = plt.figure( figsize=(12,5) )

ax1 = fig.add_subplot(1,2,1)
ax1.plot(lat, Tequil, label='equil')
ax1.plot(lat, m2.state['Ts'], label='pert' )
ax1.grid()
ax1.legend()
ax1.set_xlim(-90,90)
ax1.set_xticks(my_ticks)
ax1.set_xlabel('Latitude')
ax1.set_ylabel('Temperature (degC)')

ax2 = fig.add_subplot(1,2,2)
ax2.plot( lat, ASRequil, label='equil')
ax2.plot( lat, m2.diagnostics['ASR'], label='pert' )
ax2.grid()
ax2.legend()
ax2.set_xlim(-90,90)
ax2.set_xticks(my_ticks)
ax2.set_xlabel('Latitude')
ax2.set_ylabel('ASR (W m$^{-2}$)')
```

This tells us that making the climate colder, and allowing the ice edge to advance, tends to decrease the absorbed shortwave radiation (ASR).  That is, ice advance due to cooling is a ***positive feedback*** that will tend to lead to more cooling and more ice advance.

____________

## 3. Snowball Earth onset
____________

In section 2 we forced the model to be cooler everywhere.  Now, we'll examine some real conditions of the past that could have forced a cooling.

The radiation coming in from the sun has not always been what it was today.  Millions of years ago, the Sun was less bright, and as a result there was less solar energy entering the Earth system.  Luckily, it is easy to use our ClimLab energy balance model to investigate a past climate with a weaker sun.

The model parameter `S0` is the solar constant, describing the flux of solar radiation.

```python
model1.subprocess.insolation.S0
```

    1365.2

```python
model1.icelat
```

    array([-70.,  70.])

Let's decrease the solar constant to examine past conditions.

```python
model1.subprocess.insolation.S0 = 1300.
model1.integrate_years(10.)
model1.icelat
```

    Integrating for 900 steps, 3652.4220000000005 days, or 10.0 years.
    Total elapsed time is 14.999999999999647 years.

    array([-54.,  54.])

```python
f3 = clsn.state_plot(model1)
```

*3.1. Find a value or values of S0 that results in 100% ice cover (entire plot shaded).*

*3.2. Compare the value of S0 that you found with the inferred history of S0 (from the textbook or another source).  At what times in Earth's history was total ice cover most likely?*

*3.3. Based on our discussion of the atmosphere, describe an alternative mechanism that could induce cooling sufficient for total ice cover.*

____________

## 4. Escape from the Snowball
____________

With ice covering nearly the entire surface of the Earth, the planetary albedo would be very high.  So, how come we do not still live on Snowball Earth today?

*4.1. Plot the temperature, albedo, and ice edge latitude for a range of solar constants $S_0$.  Use information from Hoffman & Li to indicate present-day and Marinoan/Sturtian conditions on the plot.*

```python

```

*4.2. Find a value of the solar constant $S_0$ for which there is no ice cover.  Show model input and output that support your inference that there is no ice cover for this value.*

```python

```

*4.3. Suppose we were in a Snowball Earth event with present-day solar constant.  By how much would the solar constant have to increase to break out of this Snowball Earth?  How does that increase compare with the known variability of the solar constant?*

*4.4. Based on our class discussions and your reading, describe two Earth system processes (other than solar variability) that could spur a transition from Snowball Earth to current conditions.*

*4.5. A second hypothesis to explain geologic evidence for low-paleolatitude glaciation is called "Slushball Earth", in which thick ice cover was mostly confined to land, with slushy or partially open oceans.  How would you expand on the simulations in this lab to test the plausibility of the Slushball Earth hypothesis?*

____________

## Credits

This notebook was developed by [Lizz Ultee](http://ehultee.github.io), based in large part on the `ClimateModeling_courseware` resources of author of [Brian E. J. Rose](http://www.atmos.albany.edu/facstaff/brose/index.html).
____________

```python

```


# Calculating the time dependent solution for $\eta$ and $\theta$ domain
\begin{equation}\label{p1introhpa} \Psi _{\eta \eta \tau }  = \frac{1}{{R_s }}\Psi _{\eta \eta \eta \eta }  + \Psi _{\eta \eta \eta } \Psi _\theta   - \Psi _{\eta \eta \theta } \Psi _\eta, \\
\quad \Psi(\eta=\mp 1) = 0, \qquad \Psi_\eta(\eta=\mp 1) = - sin\theta, \qquad
-1 \le \eta \le1, \quad 0 \le \theta \le \pi
\end{equation}

The numerical solution found in this notebook is only stable for $0 < R_s \lesssim 200$ with a shock like discontinuity in the $\theta$-direction developing at $R_s \approx 50$

One other aspect that could be investigated is "A_general_method_to_remove_the_stiffness_of_PDEs" using Jens Eggers method.

```julia
using DifferentialEquations
using Plots
using BenchmarkTools
using DelimitedFiles
using LaTeXStrings
using SparseArrays
using LinearAlgebra
```

In this code the vector $\Psi_j$ is used to store the solution at the current time step and $\Psi_{jnext}$ is psi value at the next time step

```julia
pyplot()
function Ψ_smallr(θ::Float64, η::Float64 ,r::Float64)
    return(-η*(η^2-1)*sin(θ)/2.0 + r*η*(η^2-1)^2.0*(2.0+η^2)*sin(2*θ)/560.0 +
        r^2.0*η*(η^2-1)^2*((-591 -2294*η^2 + 161*η^4 + 1428*η^6)*sin(θ) +
            3.0*(423 + 166*η^2 - 553*η^4 - 84*η^6)sin(3*θ))/62092800.0)
end

function Ψ_smallr_d1(θ::Float64, η::Float64 ,r::Float64)
    return(-(3.0*η^2-1)*sin(θ)/2.0 + r*(7*η^6-9*η^2+2)*sin(2*θ)/560.0 +
        r^2.0*(3*(-197 - 1112*η^2 + 6930*η^4 - 2772*η^6 - 8085*η^8 + 5236*η^10)*sin(θ) + (1269 - 6120*η^2 - 6930*η^4 + 24948*η^6 - 10395*η^8 - 2772*η^10)*sin(3*θ))/62092800.0)
end

function  Ψ_smallr_d2(θ::Float64, η::Float64 ,r::Float64)
    return(-3.0*η*sin(θ) + r*(42*η^5-18*η)*sin(2*θ)/560.0 +
        r^2.0*(3*( - 2224*η + 27720*η^3 - 16632*η^5 - 64680*η^7 + 52360*η^9)*sin(θ) +
            (- 12240*η - 27720*η^3 + 149688*η^5 - 83160*η^7 - 27720*η^9)*sin(3*θ))/62092800.0)
end

function  Ψ_smallr_d3(θ::Float64, η::Float64 ,r::Float64)
    return(-3.0*sin(θ) + r*(210*η^4-18)*sin(2*θ)/560.0 +
        r^2.0*(3*( - 2224 + 83160*η^2 - 83160*η^4 - 452760*η^6 + 471240*η^8)*sin(θ) + (- 12240 - 83160*η^2 + 748440*η^4 - 582120*η^6 - 249480*η^8)*sin(3*θ))/62092800.0)
end

function Ψ_bigr(θ::Float64, η::Float64)
    return(sin(pi*η)*sin(θ)/pi)
end

function Ψ_bigr_d1(θ::Float64, η::Float64)
    return(cos(pi*η)*sin(θ))
end

function  Ψ_bigr_d2(θ::Float64, η::Float64)
    return(-sin(pi*η)*sin(θ)*pi)
end

function  Ψ_bigr_d3(θ::Float64, η::Float64)
     return(-cos(pi*η)*sin(θ)*pi^2)
end

function ψ_η_norm(ψ, θ_step_max::Int64)
    norm_sum = 0.0
    for counter = 0:θ_step_max
        norm_sum += norm(ψ_j[:,counter+1])
    end
    return norm_sum
end

function ψ_θ_norm(ψ, η_step_max::Int64)
    norm_sum = 0.0
    norm_sum += norm(ψ_j[η_step_max-1,:])
    norm_sum += norm(ψ_j[η_step_max,:])
    norm_sum += norm(ψ_j[η_step_max+1,:])
    return norm_sum
end

function ψ_ηηθ_core_wind(ψ, n::Int64, m::Int64)
    #If \psi_eta < 0 then upwind with forward difference
    #This function should only be called for 2<n<N-1
    derivative = 0.0
    if ψ[n+1,m] < ψ[n-1,m]    #forward difference
        derivative = -( -ψ[n-2,m  ]/12 +  4*ψ[n-1,m  ]/3  -5*ψ[n,m  ]/2 + 4*ψ[n+1,m  ]/3 -ψ[n+2,m  ]/12)+
                      ( -ψ[n-2,m+1]/12 +  4*ψ[n-1,m+1]/3  -5*ψ[n,m+1]/2 + 4*ψ[n+1,m+1]/3 -ψ[n+2,m+1]/12)

    else  #backward difference is the default
        derivative = -(  -ψ[n-2,m-1]/12 +  4*ψ[n-1,m-1]/3 -5*ψ[n,m-1]/2 + 4*ψ[n+1,m-1]/3 -ψ[n+2,m-1]/12)+
                       ( -ψ[n-2,m]/12 +    4*ψ[n-1,m]/3   -5*ψ[n,m]/2 + 4*ψ[n+1,m]/3     -ψ[n+2,m]/12)
    end
    return derivative
end

function ψ_θ_core_wind(ψ, n::Int64, m::Int64)
    #If \psi_eta_eta_eta > 0 then upwind with forward difference
    #This function should only be called for 2<n<N-1
    ψ_ηηη = -ψ[n-2,m]/2.0 +     ψ[n-1,m]      -ψ[n+1,m]   + ψ[n+2,m]/2.0
    derivative = 0.0
    if (ψ_ηηη < 0.0 ) #forward difference
        derivative = -ψ[n,m] + ψ[n,m+1]
    else  #backward difference is the default
        derivative = -ψ[n,m-1] + ψ[n,m]
    end
    return derivative
end

function ψ_θ_boundary_wind(ψ, n::Int64, m::Int64, ψ_ηηη::Float64)
    #If \psi_eta_eta_eta > 0 then upwind with forward difference
    #This function should only be called for 2<n<N-1
    derivative = 0.0
    if (ψ_ηηη < 0.0 ) #forward difference
        derivative = -ψ[n,m] + ψ[n,m+1]
    else   #backward difference is the default
        derivative = -ψ[n,m-1] + ψ[n,m]
    end
    return derivative
end

function ψ_ηηθ_boundary_wind(ψ, n::Int64, m::Int64, Δθ::Float64, Δη::Float64, N::Int64, ψ_η::Float64 )
    #If \psi_eta_eta_eta > 0 then upwind with forward difference
    #This function should only be called for n= 2 and n = N
    derivative = 0.0
    bc=2*Δη*sin(Δθ*(m-1))
    bcp1=2*Δη*sin(Δθ*m)
    bcm1=2*Δη*sin(Δθ*(m-2))

    if n==2
        if ψ_η < 0.0    #forward difference

        derivative=-(-(ψ_j[2,m ]+bc  )/12 + 4*ψ_j[n-1,m  ]/3 -5*ψ_j[n,m  ]/2 + 4*ψ_j[n+1,m  ]/3 -ψ_j[n+2,m  ]/12)+
                    (-(ψ_j[2,m+1]+bcp1)/12 + 4*ψ_j[n-1,m+1]/3 -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)

        else   #backward difference is the default
        derivative=-(-(ψ_j[2,m-1]+bcm1)/12 + 4*ψ_j[n-1,m-1]/3 -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)+
                    (-(ψ_j[2,m  ]+bc  )/12 + 4*ψ_j[n-1,m  ]/3 -5*ψ_j[n,m  ]/2 + 4*ψ_j[n+1,m  ]/3 -ψ_j[n+2,m  ]/12)
        end

    elseif n == N

        if ψ_η < 0.0   #forward difference
        derivative=-(-ψ_j[n-2,m  ]/12 + 4*ψ_j[n-1,m  ]/3 -5*ψ_j[n,m  ]/2 + 4*ψ_j[n+1,m  ]/3 -(ψ_j[N,m  ]-bc)/12)+
                    (-ψ_j[n-2,m+1]/12 + 4*ψ_j[n-1,m+1]/3 -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -(ψ_j[N,m+1]-bcp1)/12)
        else #backward difference
        derivative=-(-ψ_j[n-2,m-1]/12 + 4*ψ_j[n-1,m-1]/3 -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -(ψ_j[N,m-1]-bcm1)/12)+
                    (-ψ_j[n-2,m  ]/12 + 4*ψ_j[n-1,m  ]/3 -5*ψ_j[n,m  ]/2 + 4*ψ_j[n+1,m  ]/3 -(ψ_j[N,m  ]-bc  )/12)
        end
    else
        println("Error in psietaetaboundary wind")
    end
    return derivative
end

function make_ψ_jnext_on_boundaries(ψ_j, ψ_jnext, LUA, parameters)
    (Δθ, Δη, R, N, M, rc, rd1, rd2, rd3, d1, d2, d3) = parameters
    for m = 2:M-2:M #For the first theta profile and last profile use second order difference formula for theta near boundaries

    #at n=1 psi[1,:] = 0 and we are only calculating the interior points from n=2 to n=N
    n=2
    bc=2*Δη*sin(Δθ*(m-1))
    bcp1=2*Δη*sin(Δθ*m)
    bcm1=2*Δη*sin(Δθ*(m-2))

    rhs = @view ψ_jnext[2:N,m]
    ψ_ηηη = (-(ψ_j[2,m]+bc)/2 +     ψ_j[n-1,m]      -ψ_j[n+1,m]   + ψ_j[n+2,m]/2)
    ψ_η   = ((ψ_j[2,m]+bc)/12 -  2*ψ_j[n-1,m]/3  +2*ψ_j[n+1,m]/3 - ψ_j[n+2,m]/12)
    rhs[1] = rd3*(ψ_j[2,m]+bc)    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*ψ_j[n+2,m]+
        rc*ψ_ηηη*ψ_θ_boundary_wind(ψ_j, n, m, ψ_ηηη)-
        rc*ψ_ηηθ_boundary_wind(ψ_j, n, m, Δθ , Δη, N, ψ_η)*ψ_η - d3*bc

        for n = 3:N-1
        rhs[n-1] = rd3*ψ_j[n-2,m]    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*ψ_j[n+2,m]+
                rc*(-ψ_j[n-2,m]/2 +     ψ_j[n-1,m]      -ψ_j[n+1,m]   + ψ_j[n+2,m]/2)*ψ_θ_core_wind(ψ_j, n, m)-
                rc*ψ_ηηθ_core_wind(ψ_j,n, m)*(ψ_j[n-2,m]/12 -  2*ψ_j[n-1,m]/3  +2*ψ_j[n+1,m]/3 - ψ_j[n+2,m]/12)

        end#End of the n loop
    #at n=N psi[N+1,:] = 0

    n=N
    ψ_ηηη = (-ψ_j[n-2,m]/2 +     ψ_j[n-1,m]      -ψ_j[n+1,m]   + (ψ_j[N,m]-bc)/2)
    ψ_η   =  (ψ_j[n-2,m]/12 -  2*ψ_j[n-1,m]/3  +2*ψ_j[n+1,m]/3 - (ψ_j[N,m]-bc)/12)
    rhs[N-1] = rd3*ψ_j[n-2,m]    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*(ψ_j[N,m]-bc)+
            rc*ψ_ηηη*ψ_θ_boundary_wind(ψ_j, n, m, ψ_ηηη)-
            rc*ψ_ηηθ_boundary_wind(ψ_j, n, m, Δθ , Δη, N, ψ_η)*ψ_η  +d3*bc

    #ψ_jnext[2:N,m] = LUA\rhs
    ldiv!(LUA,rhs)
    end#End of the m loop
end

function make_ψ_jnext_in_core_region(ψ_j, ψ_jnext, LUA, parameters)
    (Δθ, Δη, R, N, M, rc, rd1, rd2, rd3, d1, d2, d3) = parameters
    for m = 3:M-1

    #at n=1 psi[1,:] = 0 and we are only calculating the interior points from n=2 to n=N
    n=2
    bc=2*Δη*sin(Δθ*(m-1))
    bcp1=2*Δη*sin(Δθ*m)
    bcp2=2*Δη*sin(Δθ*(m+1))
    bcm1=2*Δη*sin(Δθ*(m-2))
    bcm2=2*Δη*sin(Δθ*(m-3))

    rhs = @view ψ_jnext[2:N,m]
    ψ_ηηη =    (-(ψ_j[2,m]+bc)/2 +     ψ_j[n-1,m]      -ψ_j[n+1,m]   + ψ_j[n+2,m]/2)
    ψ_η   =    ( (ψ_j[2,m]+bc)/12 -  2*ψ_j[n-1,m]/3  +2*ψ_j[n+1,m]/3 - ψ_j[n+2,m]/12)
    rhs[1] = rd3*(ψ_j[2,m]+bc)    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*ψ_j[n+2,m]+
            rc*ψ_ηηη*ψ_θ_boundary_wind(ψ_j, n, m, ψ_ηηη)-
            rc*ψ_ηηθ_boundary_wind(ψ_j, n, m, Δθ , Δη, N, ψ_η)*ψ_η -d3*bc

        for n = 3:N-1
            rhs[n-1] = rd3*ψ_j[n-2,m]    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*ψ_j[n+2,m]+
            rc*(-ψ_j[n-2,m]/2 +     ψ_j[n-1,m]      -ψ_j[n+1,m]   + ψ_j[n+2,m]/2)*ψ_θ_core_wind(ψ_j,n, m)-
            rc*ψ_ηηθ_core_wind(ψ_j,n, m)*(ψ_j[n-2,m]/12 -  2*ψ_j[n-1,m]/3  +2*ψ_j[n+1,m]/3 - ψ_j[n+2,m]/12)
        end#End of the n loop

    #at n=N psi[N+1,:] = 0
    n=N
    ψ_ηηη = (-ψ_j[n-2,m]/2 +     ψ_j[n-1,m]      -ψ_j[n+1,m]   + (ψ_j[N,m]-bc)/2)
    ψ_η   =  (ψ_j[n-2,m]/12 -  2*ψ_j[n-1,m]/3  +2*ψ_j[n+1,m]/3 - (ψ_j[N,m]-bc)/12)
    rhs[N-1] = rd3*ψ_j[n-2,m]    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*(ψ_j[N,m]-bc)+
            rc*ψ_ηηη*ψ_θ_boundary_wind(ψ_j, n, m, ψ_ηηη)-
            rc*ψ_ηηθ_boundary_wind(ψ_j, n, m, Δθ , Δη, N, ψ_η)*ψ_η +d3*bc

    #ψ_jnext[2:N,m] = LUA\rhs
    ldiv!(LUA,rhs)
    end#End of the m loop
end

```

    make_ψ_jnext_in_core_region (generic function with 1 method)

```julia
#Change these to constants after

const N = 200  # Total number of intervals (steps-1) in the eta direction
const M = 256  # Total number of intervals (steps-1) theta direction
const J = 50000    # Total number of time steps
ψ_current = zeros(N+1,M+1)
ψ_next =zeros(N+1,M+1)
ψ_j = ψ_current
ψ_jnext = ψ_next
const Δη = 2.0/N
const Δθ = pi/M
const Δt = 0.001
const α = 0.5
const c = Δt/Δη^2
const rc    =  Δt/(Δη*Δθ)

```

```julia
#wireframe(range(0,stop = pi,length = M+1),range(-1,stop = 1,length = N+1),ψ_j)

```

```julia
R = Float64(150.0) # Reynolds number initialised, but will be re-initialised see below
ψ_current[:,1] .=0.0
ψ_current[:,M+1] .=0.0
ψ_current[1,:] .= 0.0
ψ_current[N+1,:] .= 0.0
for n in 2:N
    for m in 2:M
        #ψ_current[n,m] = Ψ_smallr(Δθ*(m-1),-1.0+Δη*(n-1),R)
        ψ_current[n,m] = 0
    end
end
combined_plot =plot()
plot_array = Array{Plots.Plot{Plots.PyPlotBackend},1}()
```

    Plots.Plot{Plots.PyPlotBackend}[]

```julia
reynolds_values = [20.0]
for rey_number in reynolds_values

R = rey_number
d3 =     -1.0/12.0 -   α*c/R
d2 =      4.0/3.0 + 4*α*c/R
d1 =     -5.0/2.0  -6*α*c/R

A =spdiagm(-2=>ones(N-3)*d3 ,-1 => ones(N-2)*d2, 0 => ones(N-1)*d1, 1 => ones(N-2)*d2,2 => ones(N-3)*d3)

A[1,1] = d1 + d3 #Boundary condition psi@i=2 = 0, psi_eta@i=2 = -sin(theta) add -2d3\delta\eta\sine(theta) to RHS
A[end,end] = d1 + d3 #Boundary condition psi@i=N+1 = 0, psi_eta@i=N+1 = -sin(theta) add 2d3\delta\eta\sine(theta) to RHS
LUA = lu(A)
varinfo(r"A")

#Forming the righthand side vector
rd3 =      -1.0/12.0 + (1-α)*c/R
rd2 =      4.0/3.0 - 4*(1-α)*c/R
rd1 =     -5.0/2.0 + 6*(1-α)*c/R

parameters =  (Δθ, Δη, R, N, M, rc, rd1, rd2, rd3, d1, d2, d3)
percent = 5.0
θstep_of_norm = 5
ηstep_of_norm = round(Int, (1+0.6)*N/2) + 1
old_norm = expected_η_norm = 0.0
for counter =1:θstep_of_norm
        expected_η_norm += norm([Ψ_smallr(Δθ*counter,-1.0+Δη*i,R) for i = 1:N])
end
expected_η_norm *= Δη*Δθ*10^6

#Only theta profiles from m =2 to m = M need to be calculated with one sided theta derivatives second order.
#11111111111111111111111111111111111111111111111111111111

@time for j =1:J

    make_ψ_jnext_on_boundaries(ψ_j, ψ_jnext, LUA, parameters)
    make_ψ_jnext_in_core_region(ψ_j,ψ_jnext, LUA, parameters)

    if( mod(j,2) == 1)
        ψ_j    = ψ_next
        ψ_jnext=ψ_current
    else
        ψ_j    = ψ_current
        ψ_jnext=ψ_next
    end #End of if statement

    if( j > round(percent*J/100)-1)
        #Calculate a simple norm of psi
        psi_norm1 = ψ_η_norm(ψ_j, θstep_of_norm)*Δη*Δθ*10^6
        psi_norm2 = ψ_θ_norm(ψ_j, round(Int, 1.6*N/2))*Δη*Δθ*10^4
        print(string(percent)*"% t="*string(round(j*Δt; digits=3))*" norm1 near theta0 = "*string(round(psi_norm1; digits=6))*" Expected norm1= "*string(round(expected_η_norm; digits=2)))
        print(" norm2 near eta0.6 = "*string(round(psi_norm2; digits=6))*"\n")
        percent+=5
        if isnan(psi_norm1 )
           println("\n Psi is no longer a number, exiting... ")
           exit()
        elseif(abs(psi_norm1-old_norm)<0.00001)
            println("\n Psi has converged, exiting... ")
            break
       end
        old_norm = psi_norm1
    end

end #End of time loop j
println("Reynolds number= "*string(R))
#At this point inverse of A should be applied to the rhs vector to produce the next iterate of psi at the j+1 time step
eta_profile =0.6
theta=range(0, stop=pi,length= M+1)
push!(plot_array,plot(theta,ψ_j[round(Int,(eta_profile+1)/Δη)+1,:], label="Julia "*L"R_s ="*string(R), linestyle=:solid, xlabel = L"\theta", ylabel = L"\Psi",size=(1000,500),yguidefontrotation=-90))
plot!(combined_plot, theta,ψ_j[round(Int,(eta_profile+1)/Δη)+1,:], label="Julia "*L"R_s ="*string(R), linestyle=:solid, xlabel = L"\theta", ylabel = L"\Psi",size=(1000,500),yguidefontrotation=-90)
end
```

    5.0% t=2.5 norm1 near theta0 = 53.872564 Expected norm1= 54.62 norm2 near eta0.6 = 7.933603
    10.0% t=5.0 norm1 near theta0 = 53.793182 Expected norm1= 54.62 norm2 near eta0.6 = 7.951104
    15.0% t=7.5 norm1 near theta0 = 53.792602 Expected norm1= 54.62 norm2 near eta0.6 = 7.9531
    20.0% t=10.0 norm1 near theta0 = 53.792597 Expected norm1= 54.62 norm2 near eta0.6 = 7.953311

     Psi has converged, exiting...
     35.985808 seconds (267.81 M allocations: 30.931 GiB, 5.05% gc time)
    Reynolds number= 20.0

```julia
plot!(combined_plot,xtickfont = font(12))
xticks!(([0:π/6:π;], ["0", "\\pi/6", "\\pi/3", "\\pi/2", "2\\pi/3", "5\\pi/6","\\pi"]))
title!(L"\Psi _{\eta \eta t }  = \frac{1}{{R_s }}\Psi _{\eta \eta \eta \eta }  + \Psi _{\eta \eta \eta } \Psi _\theta   - \Psi _{\eta \eta \theta } \Psi _\eta,\ \ \   \Psi(\eta=\pm 1)=0,\ \ \Psi_{\eta}(\eta=\pm 1)=-sin\theta  ")
display(combined_plot)
```

```julia
Plots.default(overwrite_figure=false)
#display(plot_array[end])
#display(plot_array[end-1])
plot!(plot_array[end-1],plot_array[end])
xticks!(([0:π/6:π;], ["0", "\\pi/6", "\\pi/3", "\\pi/2", "2\\pi/3", "5\\pi/6","\\pi"]))
```

```julia
open("psi_solution.txt", "w") do io
           writedlm(io, ψ_current)
       end

```

```julia
m=50
MatrixPsi=readdlm("psifin.m",Float64,);
eta_steps = size(MatrixPsi)[1]
theta_steps = size(MatrixPsi)[2]
xinterval=range(-1,stop=1,length=eta_steps-2)
θ_scalefactor = round(Int,N/theta_steps)
eta=range(-1, stop=1,length= N+1)
plot(eta,ψ_j[:,m],linecolor=:orange, label="Julia", linestyle=:solid)
plot!(title= L"\theta="*string(m)*L"\pi/"*string(theta_steps)*L"\quad R_s= "*string(R),size=(1100,500), legend=:outertopright, xlabel =L"\eta", ylabel=L"\Psi", yguidefontrotation=-90)
plot!(x->Ψ_smallr(Δθ*(m-1),x,R),-1,1, label="Small R analytical", linestyle=:dot, linecolor=:blue)
plot!(xinterval,MatrixPsi[2:end-1,θ_scalefactor*(m-1)+1], label="Fortran FD", linestyle=:dash, linecolor=:black)
plot!(x->Ψ_bigr((m-1)*Δθ,x),-1,1, label="Large R analytical",  linestyle=:dash, linecolor=:red)
```

```julia
eta_profile =0.02
delta_eta_steps = 2/(eta_steps-3)
theta=range(0, stop=pi,length= M+1)
theta2=range(0, stop=pi,length= theta_steps)
plot_julia = plot(theta,ψ_j[round(Int,(eta_profile+1)/Δη)+1,:],linecolor=:orange, label="Julia R ="*string(R), linestyle=:solid,xlabel = L"\theta", ylabel = L"\Psi",size=(1000,400),yguidefontrotation=-90, legend=:outertopright)
xticks!(plot_julia,([0:π/6:π;], ["0", "\\pi/6", "\\pi/3", "\\pi/2", "2\\pi/3", "5\\pi/6","\\pi"]))
title!(plot_julia,L"\theta"*" profile at "*L"\eta= "*string(eta_profile))
plot!(x->Ψ_smallr(x,eta_profile,R),0,pi, label="Small R analytical", linestyle=:dot, linecolor=:blue,xlabel = L"\theta", ylabel = L"\Psi",size=(1000,400))
#plot!(x->Ψ_bigr(x,eta_profile),0,pi, label="Large R analytical",  linestyle=:dash, linecolor=:red,)
plot!(theta2,MatrixPsi[round(Int,(eta_profile+1)/delta_eta_steps)+2,:], linestyle=:dash, label="Fortran winding",linecolor=:black)
```

```julia
eta_profile =0.01
theta=range(0, stop=pi,length= M+1)
julia_theta_profile = plot(theta,ψ_j[round(Int,(eta_profile+1)/Δη)+1,:],linecolor=:blue , label="Julia R ="*string(R), linestyle=:solid)
scatter!(julia_theta_profile,theta,ψ_j[round(Int,(eta_profile+1)/Δη)+1,:],linecolor=:blue , label="Julia R ="*string(R), linestyle=:solid,  markersize=4)
plot!(julia_theta_profile,xticks=([0:π/6:π;], ["0", "\\pi/6", "\\pi/3", "\\pi/2", "2\\pi/3", "5\\pi/6","\\pi"]), title=L"\theta"*" profile at "*L"\eta= "*string(eta_profile),size=(1100,400),legend=:outertopright, xlabel = L"\theta", ylabel = L"\Psi",yguidefontrotation=-90)
```

```julia

```


```python
import networkx as nx
import sympy as sym
import numpy as np

class ElectricalCircuit():
    """ Electrical circuit modelling using an undirect weighted MultiGraph.

    Features:
    ----------
    - Captures and analyzes an electrical circuit given by it's branches,
    in aim to get the current value which flows for each branch.

    - Applies Kirchhoff Laws over the nodes and circuit loops,
    and resolves the linear resulting system, if possible.

    - Uses an internal graph representation of the circuit, shaping the
    nodes, branches and meshes of the circuit as the nodes, edges and
    cycles of the graph.

    Edge weights will contain the branches information, only supporting
    (for the moment):

    a) the resistance value 'R':

        The total resistance value inside the branch. It's zero by default,
        when there are no components with an appreciable resistance value.

    b) the voltage value 'V':

        The total pottential difference given by the voltage sources inside
        the branch. Because the sources have two positions (direct or inverse),
        we consider a positive value when the source is direct in respect the
        ascending order '<' between edge connected nodes, and negative if not.

    c) the intensity value 'I':

        The intensity value of the current that pass throught the branch.
        Represents one unknown of the Kirchhoff linear system, and its value
        is updated after every change. If the system isn't stable, just
        outputs a '???' value.

    Attributes:
    ------------
    nodes    : list of the circuit nodes.

    branches : dictionary of the circuit branches. Contains the above
               'R', 'V' and 'I' values.

    meshes   : list of the circuit meshes or loops, as the base cicles
               of the internal graph.

    With these three attributes (and the subjacent connection), one can solve
    the circuit by just extracting manually the necessary equations.

    Instead, this class only needs the branches dictionary passed as an initial
    parameter, and it will calculate the intensities automatically. Of course,
    it still needs the work of translating correctly all the circuit components.

    """

    # Constants
    V_units = 'V'
    R_units = 'kO'
    I_units = 'mA'

    # Constructor
    def __init__(self, edgelist = []):
        """ Builds the circuit by passing a list of branches with their attributes.
            Empty circuit by default. """

        self._circuit = nx.MultiGraph()
        self._meshes = []

        # Add all branches
        for edge in edgelist:
            self.add_branch((edge[0], edge[1]), **(edge[2].copy()))

    # Branch adder
    def add_branch(self, edge, **edgedict):
        """ Adds a new branch to the circuit. Also sets its 'R' and 'V' values. """
        # Insertion preconditions + preprocessing
        self.__order_between_nodes(edge)
        self.__supported_dict_keys(edgedict)
        self.__supported_value_types(edgedict)

        # Insertion
        self._circuit.add_edge(edge[0], edge[1], **edgedict)

        # Updates
        self.__mesh_update()
        self.__kirchhoff_update()

    def __order_between_nodes(self, edge):
        """ Edge has to respect node's ascending order.
        Raises an error if edge[0] ≥ edge[1]. """
        if edge[0] >= edge[1]: raise AssertionError(
        "Branch does not satisfy the order '<' between nodes.")

    def __supported_dict_keys(self, edgedict):
        """ Edge dictionary must have two of the supported variables:
        'V', 'R', 'I', to determine the unknown variable. """
        if ('V' in edgedict and 'R' in edgedict and 'I' not in edgedict):
            edgedict['unknown'] = 'I' # Intensity as the unknown variable

        elif ('V' not in edgedict and 'R' in edgedict and 'I' in edgedict):
            edgedict['unknown'] = 'V' # Voltage as the unknown variable

        elif ('V' in edgedict and 'R' not in edgedict and 'I' in edgedict):
            edgedict['unknown'] = 'R' # Resistance as the unknown variable

        else: raise AttributeError("Branch dictionary must have one and only one" +
                    " undefined key variable (the unknown).")

    def __supported_value_types(self, edgedict):
        """ Edge dictionary values must be converted to symbolic expression. """
        if 'V' in edgedict: edgedict['V'] = sym.simplify(edgedict['V'])
        if 'R' in edgedict: edgedict['R'] = sym.simplify(edgedict['R'])
        if 'I' in edgedict: edgedict['I'] = sym.simplify(edgedict['I'])

    # Branch deleter
    def del_branch(self, edge):
        """ Removes the branch specified from the circuit. Has to give a key. """
        self._circuit.remove_edge(edge[0], edge[1], key = edge[2])

        # Updates
        self.__mesh_update()
        self.__kirchhoff_update()

    # Getter for nodes
    @property
    def nodes(self):
        """ Getter for the nodes list of the circuit. """
        return list(self._circuit.nodes)

    # Getter for branches
    @property
    def branches(self):
        """ Getter for the edges dictionary of the circuit. """
        return dict(self._circuit.edges)

    # Getter for meshes
    @property
    def meshes(self):
        """ Getter for the base cycles list of the circuit. """
        return self._meshes

    # Voltage difference
    def potential_diff(self, nodeA, nodeB):
        """ Returns the pottential difference between nodeA and nodeB, as the
        pottential of nodeA with respect nodeB: V_AB = (V_A - V_B). """

        diff = 0

        # For all edges in a biconnected path from B to A,
        nodes = nx.shortest_path(self._circuit, nodeB, nodeA)
        for branch in [(nodes[i], nodes[i+1], 0) for i in range(len(nodes) - 1)]:

            # intensity direction
            if branch in self.branches:
                diff += self._circuit.edges[branch]['V']
                diff -= self._circuit.edges[branch]['I'] * self._circuit.edges[branch]['R']

            # intensity inverse direction
            if (branch[1], branch[0], branch[2]) in self.branches:
                diff -= self._circuit.edges[branch]['V']
                diff += self._circuit.edges[branch]['I'] * self._circuit.edges[branch]['R']

        return diff

    # Str representation
    def __str__(self):
        """ Shows the circuit as a pictographic drawing, by calling the
        circuit_view or branches_view depending on its size and type. """

        # Full view representation
        if (len(self._circuit.nodes) <= 4):
            return self.circuit_view()

        # Short view representation
        else: return self.branches_view()

    # Circuit view (not sure if well labeled)
    def circuit_view(self):
        """ Full circuit view, prints nodes as column separators with
        homogeneous branches. Thinked for short circuits (4 or less nodes). """

        node_idx = {node:i for i, node in enumerate(self.nodes)} # Index remembering
        sep = "\t\t\t\t\t" # Separator between nodes

        # Str representation. Default branch as default lines, for branch separations
        circuit_pattern = sep.join([f"{node}" for node in self.nodes]) + "\n"

        # For each branch between nodes
        for j, branch in enumerate(self._circuit.edges):
            # get the branch nodes index.
            min_idx = min(node_idx[branch[0]], node_idx[branch[1]])
            max_idx = max(node_idx[branch[0]], node_idx[branch[1]])

            # Join and append branch components
            branch_components = []
            for component, value in self._circuit.edges[branch].items():
                # If V is not the unknown
                if component == 'V' and value:
                    branch_components.append("(" + str(value) + " " + self.V_units + ")")
                if component == 'R' and value:
                    branch_components.append("[" + str(value) + " " + self.R_units + "]")
                if component == 'I' and value:
                    branch_components.append("\\" + str(value) + " " + self.I_units + "\\")
            branch_components = "---".join(branch_components)
            branch_components = branch_components.center(40 * (max_idx - min_idx) - 1, "-")

            # Add two default separator lines
            branch_pattern = [""] + [sep for node in self.nodes][:-1] + [""]
            circuit_pattern += "|".join(branch_pattern) + "\n"
            circuit_pattern += "|".join(branch_pattern) + "\n"

            # Add the branch line
            circuit_pattern += "|".join(branch_pattern[:min_idx + 1])
            circuit_pattern += "+" + branch_components + "+"
            circuit_pattern += "|".join(branch_pattern[max_idx + 1:]) + "\n"

        return circuit_pattern

    # Branches view
    def branches_view(self):
        """ Simplified circuit view, as a row list of branches. Suitable for
        simbolic type circuits, or dense circuits (more than 4 nodes). """

        circuit_pattern = "" # Str representation of the branches

        # For each branch between nodes
        for j, branch in enumerate(self._circuit.edges):

            # Join and append branch components
            branch_components = []
            for component, value in self._circuit.edges[branch].items():
                # If V is not the unknown
                if component == 'V' and value:
                    branch_components.append("(" + str(value) + " " + self.V_units + ")")
                if component == 'R' and value:
                    branch_components.append("[" + str(value) + " " + self.R_units + "]")
                if component == 'I' and value:
                    branch_components.append("\\" + str(value) + " " + self.I_units + "\\")
            branch_components = "--------".join(branch_components)
            circuit_pattern += str(branch) + ": (" + str(branch[0]) + ")-------"
            circuit_pattern += branch_components + "-------(" + str(branch[1]) + ")\n\n"
        return circuit_pattern

    # Cycle base calculator
    def __mesh_update(self):
        """ Mesh update, every time that the circuit has changed.
        Gets all the graph base cycles, combined with the multiple edge cycles.
        """

        # Get all base cycles for the undirect multiple graph as:
        self._base_circuit = nx.Graph(self._circuit)
        self._meshes = []

        # the sum of all the base cycles for the undirect simple graph
        for cycle in nx.cycle_basis(self._base_circuit):
            self._meshes += [[(cycle[i-1], cycle[i], 0) for i in range(len(cycle))]]

        # with all the inner cycles for the undirect multiple graph
        for edge in self._base_circuit.edges:
            n = len(self._circuit.adj[edge[0]][edge[1]])
            self._meshes += [[(edge[0], edge[1], i), (edge[1], edge[0], i + 1)]
                             for i in range(n - 1)]

    # Kirchhoff solver
    def __kirchhoff_update(self):
        """ Applies Kirchhoff Laws for soving the circuit, by updating
        the unknowed values for each branch. """

        # Number of nodes and edges.
        N = len(self._circuit.nodes)
        E = len(self._circuit.edges)

        if not E: return # For empty circuits!!

        # Sympy matrixes for the lineal resulting system AX = B.
        A = sym.Matrix.zeros(E)
        B = sym.Matrix.zeros(E, 1)
        X = sym.symbols(f'x:{E}')

        # Apply Kirchhoff laws to get the linear system
        for i in range(E):
            for j, branch in enumerate(self._circuit.edges):

                # 1st Kirchhoff Law (charge conservation):
                # the sum of currents meeting at a node is zero.
                if i < (N - 1):
                    if (branch[1] == list(self._circuit.nodes)[i]): # intensity in

                        # Be sure of treating the correct unknown ('I', 'V' or 'R')
                        if (self._circuit.edges[branch]['unknown'] == 'I'): A[i, j] = 1
                        else: B[i] -= self._circuit.edges[branch]['I']

                    if (branch[0] == list(self._circuit.nodes)[i]): # intensity out

                        # Be sure of treating the correct unknown ('I', 'V' or 'R')
                        if (self._circuit.edges[branch]['unknown'] == 'I'): A[i, j] = -1
                        else: B[i] += self._circuit.edges[branch]['I']

                # 2nd Kirchhoff Law (energy conservation):
                # the directed sum of the voltages around any mesh is zero.
                else:
                    # direct branch
                    if branch in self.meshes[i - (N - 1)]:

                        # Be sure of treating the correct unknown ('I', 'V' or 'R')
                        if (self._circuit.edges[branch]['unknown'] == 'I'):
                            A[i, j] = - self._circuit.edges[branch]['R']
                            B[i] -= self._circuit.edges[branch]['V']

                        if (self._circuit.edges[branch]['unknown'] == 'V'):
                            A[i, j] = 1
                            B[i] += self._circuit.edges[branch]['R'] * self._circuit.edges[branch]['I']

                        if (self._circuit.edges[branch]['unknown'] == 'R'):
                            A[i, j] = - self._circuit.edges[branch]['I']
                            B[i] -= self._circuit.edges[branch]['V']

                    # inverse branch
                    if (branch[1], branch[0], branch[2]) in self.meshes[i - (N - 1)]:

                        # Be sure of treating the correct unknown ('I', 'V' or 'R')
                        if (self._circuit.edges[branch]['unknown'] == 'I'):
                            A[i, j] = self._circuit.edges[branch]['R']
                            B[i] += self._circuit.edges[branch]['V']

                        if (self._circuit.edges[branch]['unknown'] == 'V'):
                            A[i, j] = -1
                            B[i] -= self._circuit.edges[branch]['R'] * self._circuit.edges[branch]['I']

                        if (self._circuit.edges[branch]['unknown'] == 'R'):
                            A[i, j] = self._circuit.edges[branch]['I']
                            B[i] += self._circuit.edges[branch]['V']

        # Solves the linear system for symbolic coefficients
        # or numeric values and gets the solution
        solution_space = sym.linsolve((A, B), X)
        self.solved = bool(solution_space)
        for solution in solution_space: break

        # Adds the solutions back to branches
        for j, branch in enumerate(self._circuit.edges):
            variable = self._circuit.edges[branch]['unknown']
            if self.solved:  self._circuit.edges[branch][variable] = solution[j]
            else:  self._circuit.edges[branch][variable] = sym.S.NaN
```

#### Tests:

```python
# Empty circuit test
empty_circuit = ElectricalCircuit()
print('Nodes:', empty_circuit.nodes)
print('Branches:', empty_circuit.branches)
print('Meshes:', empty_circuit.meshes)
print(empty_circuit)
```

    Nodes: []
    Branches: {}
    Meshes: []

```python
# Base circuit test
base_circuit = ElectricalCircuit([(1, 2, {'R': 'R1', 'V': 3}), (1, 2, {'R':0, 'V':'V2'}), (2, 3, {'R': 34, 'V': -3})])
print('Nodes:', base_circuit.nodes)
print('Branches:', base_circuit.branches)
print('Meshes:', base_circuit.meshes)
print(base_circuit)
```

    Nodes: [1, 2, 3]
    Branches: {(1, 2, 0): {'R': R1, 'V': 3, 'unknown': 'I', 'I': (-V2 + 3)/R1}, (1, 2, 1): {'R': 0, 'V': V2, 'unknown': 'I', 'I': (V2 - 3)/R1}, (2, 3, 0): {'R': 34, 'V': -3, 'unknown': 'I', 'I': 0}}
    Meshes: [[(1, 2, 0), (2, 1, 1)]]
    1					2					3
    |					|					|
    |					|					|
    +--[R1 kO]---(3 V)---\(-V2 + 3)/R1 mA\--+					|
    |					|					|
    |					|					|
    +-------(V2 V)---\(V2 - 3)/R1 mA\-------+					|
    |					|					|
    |					|					|
    |					+------------[34 kO]---(-3 V)-----------+

```python
# __order_between_nodes test
try:
    base_circuit.add_branch((2, 1), R = 20, V = 30)
except AssertionError as e:
    print('Test error message: ' + str(e))
```

    Test error message: Branch does not satisfy the order '<' between nodes.

```python
# __supported_dict_keys test
try:
    base_circuit.add_branch((1, 2), V = 30)
except AttributeError as e:
    print('Test error message: ' + str(e))
```

    Test error message: Branch dictionary must have one and only one undefined key variable (the unknown).

```python
# __supported_value_types test
try:
    base_circuit.add_branch((1, 2), V = (), I = 38)
except TypeError as e:
    print('Test error message: ' + str(e))
```

    Test error message: bad operand type for unary -: 'Tuple'

```python
# add_branch, del_branch test

empty_circuit.add_branch((1, 2), R = 0.0, V = 0.0)
print(empty_circuit)

try:
    empty_circuit.del_branch((1, 2))
except IndexError as e:
    print('Test error message: ' + str(e))

empty_circuit.del_branch((1, 2, 0))
print(empty_circuit)
```

    1					2
    |					|
    |					|
    +---------------------------------------+

    Test error message: tuple index out of range
    1					2

```python
# Numeric correctness test

# Intensity unknown
numeric = ElectricalCircuit()
numeric.add_branch((1, 2), R = 0.0, V = -3)
numeric.add_branch((1, 2), R = 3, V = 10)
numeric.add_branch((1, 3), R = 10, V = -20)
numeric.add_branch((2, 3), R = 0.0, V = 100)
numeric.add_branch((2, 3), R = 200, V = 200)
print(numeric)

# Voltage unknown
numeric2 = ElectricalCircuit()
numeric2.add_branch((1, 2), R = 0.0, I = -3)
numeric2.add_branch((1, 2), R = 3, I = 10)
numeric2.add_branch((1, 3), R = 10, V = -20)
numeric2.add_branch((2, 3), R = 0.0, V = 100)
numeric2.add_branch((2, 3), R = 200, V = 200)
print(numeric2)

# Resistance unknown
numeric3 = ElectricalCircuit()
numeric3.add_branch((1, 2), V = 0.0, I = -3)
numeric3.add_branch((1, 2), V = 3, I = 10)
numeric3.add_branch((1, 3), V = 10, R = -20)
numeric3.add_branch((2, 3), V = 0.0, R = 100)
numeric3.add_branch((2, 3), V = 200, R = 200)
print(numeric3)
```

    1					2					3
    |					|					|
    |					|					|
    +----------(-3 V)---\221/30 mA\---------+					|
    |					|					|
    |					|					|
    +------[3 kO]---(10 V)---\13/3 mA\------+					|
    |					|					|
    |					|					|
    +------------------------[10 kO]---(-20 V)---\-117/10 mA\-----------------------+
    |					|					|
    |					|					|
    |					+----------(100 V)---\56/5 mA\----------+
    |					|					|
    |					|					|
    |					+-----[200 kO]---(200 V)---\1/2 mA\-----+

    1					2					3
    |					|					|
    |					|					|
    +-----------\-3 mA\---(-50 V)-----------+					|
    |					|					|
    |					|					|
    +-------[3 kO]---\10 mA\---(-20 V)------+					|
    |					|					|
    |					|					|
    +--------------------------[10 kO]---(-20 V)---\-7 mA\--------------------------+
    |					|					|
    |					|					|
    |					+----------(100 V)---\13/2 mA\----------+
    |					|					|
    |					|					|
    |					+-----[200 kO]---(200 V)---\1/2 mA\-----+

    1					2					3
    |					|					|
    |					|					|
    +-----------\-3 mA\---[90 kO]-----------+					|
    |					|					|
    |					|					|
    +-----(3 V)---\10 mA\---[-267/10 kO]----+					|
    |					|					|
    |					|					|
    +--------------------------(10 V)---[-20 kO]---\-7 mA\--------------------------+
    |					|					|
    |					|					|
    |					+-----------[100 kO]---\4 mA\-----------+
    |					|					|
    |					|					|
    |					+------(200 V)---[200 kO]---\3 mA\------+

```python
# Symbolic support test

#  Intensity unknown
symbolic = ElectricalCircuit()
symbolic.add_branch(('A', 'B'), R = '2 * R', V = 'V2')
symbolic.add_branch(('A', 'C'), R = '2 * R', V = 0)
symbolic.add_branch(('A', 'D'), R = '2 * R', V = 'V4')
symbolic.add_branch(('B', 'C'), R = 'R', V = '- V3')
symbolic.add_branch(('B', 'C'), R = '2 * R', V = 'V1')
symbolic.add_branch(('B', 'D'), R = '2 * R', V = 0)
symbolic.add_branch(('B', 'D'), R = '2 * R', V = 0)
print(symbolic.branches_view())

# Voltage unknown
symbolic2 = ElectricalCircuit()
symbolic2.add_branch((1, 2), R = 'R', V = 'V1')
symbolic2.add_branch((1, 2), R = 'R', I = 'I2')
symbolic2.add_branch((1, 2), R = 'R', V = 0)
print(symbolic2)

# Resistance unknown
symbolic3 = ElectricalCircuit()
symbolic3.add_branch((1, 2), V = 'V1', I = 'I1')
symbolic3.add_branch((1, 2), V = '- V2', R = 'R2')
symbolic3.add_branch((1, 2), I = 'I3', R = 'R3')
print(symbolic3)
```

    ('A', 'B', 0): (A)-------[2*R kO]--------(V2 V)--------\(3*V1 + 17*V2 - 6*V3 - 8*V4)/(58*R) mA\-------(B)

    ('A', 'C', 0): (A)-------[2*R kO]--------\(-5*V1 - 9*V2 + 10*V3 - 6*V4)/(58*R) mA\-------(C)

    ('A', 'D', 0): (A)-------[2*R kO]--------(V4 V)--------\(V1 - 4*V2 - 2*V3 + 7*V4)/(29*R) mA\-------(D)

    ('B', 'C', 0): (B)-------[R kO]--------(-V3 V)--------\(-8*V1 + 3*V2 - 13*V3 + 2*V4)/(29*R) mA\-------(C)

    ('B', 'C', 1): (B)-------[2*R kO]--------(V1 V)--------\(21*V1 + 3*V2 + 16*V3 + 2*V4)/(58*R) mA\-------(C)

    ('B', 'D', 0): (B)-------[2*R kO]--------\(-V1 + 4*V2 + 2*V3 - 7*V4)/(58*R) mA\-------(D)

    ('B', 'D', 1): (B)-------[2*R kO]--------\(-V1 + 4*V2 + 2*V3 - 7*V4)/(58*R) mA\-------(D)

    1					2
    |					|
    |					|
    +[R kO]---(V1 V)---\(-I2*R + V1)/(2*R) mA\+
    |					|
    |					|
    +-[R kO]---\I2 mA\---(3*I2*R/2 + V1/2 V)+
    |					|
    |					|
    +----[R kO]---\-(I2*R + V1)/(2*R) mA\---+

    1					2
    |					|
    |					|
    +(V1 V)---\I1 mA\---[(-R2*(I1 + I3) + V1 + V2)/I1 kO]+
    |					|
    |					|
    +---(-V2 V)---[R2 kO]---\-I1 - I3 mA\---+
    |					|
    |					|
    +\I3 mA\---[R3 kO]---(I3*R3 + R2*(I1 + I3) - V2 V)+

```python
# Potential difference test

# Numeric support
print('Potential diff V12 for numeric2 circuit:', numeric3.potential_diff(1, 2))

# Symbolic support
print('Potential diff VAB for symbolic circuit:', symbolic.potential_diff('A', 'B'))
```

    Potential diff V12 for numeric2 circuit: -270
    Potential diff VAB for symbolic circuit: 3*V1/29 - 12*V2/29 - 6*V3/29 - 8*V4/29

```python
# views test:

print("Empty circuit:\n")
print(empty_circuit.branches_view())
print(empty_circuit.circuit_view())

print("\nBase circuit:\n")
print(base_circuit.branches_view())
print(base_circuit.circuit_view())

print("\nNumeric circuit:\n")
print(numeric.branches_view())
print(numeric.circuit_view())

print("\nSymbolic circuit:\n")
print(symbolic.branches_view())
print(symbolic.circuit_view())
```

    Empty circuit:

    1					2

    Base circuit:

    (1, 2, 0): (1)-------[R1 kO]--------(3 V)--------\(-V2 + 3)/R1 mA\-------(2)

    (1, 2, 1): (1)-------(V2 V)--------\(V2 - 3)/R1 mA\-------(2)

    (1, 2, 2): (1)-------\38 mA\-------(2)

    (2, 3, 0): (2)-------[34 kO]--------(-3 V)-------(3)

    1					2					3
    |					|					|
    |					|					|
    +--[R1 kO]---(3 V)---\(-V2 + 3)/R1 mA\--+					|
    |					|					|
    |					|					|
    +-------(V2 V)---\(V2 - 3)/R1 mA\-------+					|
    |					|					|
    |					|					|
    +----------------\38 mA\----------------+					|
    |					|					|
    |					|					|
    |					+------------[34 kO]---(-3 V)-----------+

    Numeric circuit:

    (1, 2, 0): (1)-------(-3 V)--------\221/30 mA\-------(2)

    (1, 2, 1): (1)-------[3 kO]--------(10 V)--------\13/3 mA\-------(2)

    (1, 3, 0): (1)-------[10 kO]--------(-20 V)--------\-117/10 mA\-------(3)

    (2, 3, 0): (2)-------(100 V)--------\56/5 mA\-------(3)

    (2, 3, 1): (2)-------[200 kO]--------(200 V)--------\1/2 mA\-------(3)

    1					2					3
    |					|					|
    |					|					|
    +----------(-3 V)---\221/30 mA\---------+					|
    |					|					|
    |					|					|
    +------[3 kO]---(10 V)---\13/3 mA\------+					|
    |					|					|
    |					|					|
    +------------------------[10 kO]---(-20 V)---\-117/10 mA\-----------------------+
    |					|					|
    |					|					|
    |					+----------(100 V)---\56/5 mA\----------+
    |					|					|
    |					|					|
    |					+-----[200 kO]---(200 V)---\1/2 mA\-----+

    Symbolic circuit:

    ('A', 'B', 0): (A)-------[2*R kO]--------(V2 V)--------\(3*V1 + 17*V2 - 6*V3 - 8*V4)/(58*R) mA\-------(B)

    ('A', 'C', 0): (A)-------[2*R kO]--------\(-5*V1 - 9*V2 + 10*V3 - 6*V4)/(58*R) mA\-------(C)

    ('A', 'D', 0): (A)-------[2*R kO]--------(V4 V)--------\(V1 - 4*V2 - 2*V3 + 7*V4)/(29*R) mA\-------(D)

    ('B', 'C', 0): (B)-------[R kO]--------(-V3 V)--------\(-8*V1 + 3*V2 - 13*V3 + 2*V4)/(29*R) mA\-------(C)

    ('B', 'C', 1): (B)-------[2*R kO]--------(V1 V)--------\(21*V1 + 3*V2 + 16*V3 + 2*V4)/(58*R) mA\-------(C)

    ('B', 'D', 0): (B)-------[2*R kO]--------\(-V1 + 4*V2 + 2*V3 - 7*V4)/(58*R) mA\-------(D)

    ('B', 'D', 1): (B)-------[2*R kO]--------\(-V1 + 4*V2 + 2*V3 - 7*V4)/(58*R) mA\-------(D)

    A					B					C					D
    |					|					|					|
    |					|					|					|
    +[2*R kO]---(V2 V)---\(3*V1 + 17*V2 - 6*V3 - 8*V4)/(58*R) mA\+					|					|
    |					|					|					|
    |					|					|					|
    +--------------[2*R kO]---\(-5*V1 - 9*V2 + 10*V3 - 6*V4)/(58*R) mA\-------------+					|
    |					|					|					|
    |					|					|					|
    +-------------------------------[2*R kO]---(V4 V)---\(V1 - 4*V2 - 2*V3 + 7*V4)/(29*R) mA\-------------------------------+
    |					|					|					|
    |					|					|					|
    |					+[R kO]---(-V3 V)---\(-8*V1 + 3*V2 - 13*V3 + 2*V4)/(29*R) mA\+					|
    |					|					|					|
    |					|					|					|
    |					+[2*R kO]---(V1 V)---\(21*V1 + 3*V2 + 16*V3 + 2*V4)/(58*R) mA\+					|
    |					|					|					|
    |					|					|					|
    |					+---------------[2*R kO]---\(-V1 + 4*V2 + 2*V3 - 7*V4)/(58*R) mA\---------------+
    |					|					|					|
    |					|					|					|
    |					+---------------[2*R kO]---\(-V1 + 4*V2 + 2*V3 - 7*V4)/(58*R) mA\---------------+

```python
# thevenin test

thevenin = ElectricalCircuit()
thevenin.add_branch(('A', 'B'), R = '2 * R', V = 'V2')
thevenin.add_branch(('A', 'C'), R = '2 * R', V = 0)
thevenin.add_branch(('A', 'D'), R = '2 * R', V = 'V4')
thevenin.add_branch(('B', 'C'), R = 'R', V = '- V3')
thevenin.add_branch(('B', 'C'), R = '2 * R', V = 'V1')
thevenin.add_branch(('B', 'D'), R = '2 * R', V = 0)
thevenin.add_branch(('B', 'D'), R = '2 * R', V = 0)
thevenin.add_branch(('A', 'B'), R = '7 * R', V = 'V6 - V5')
print(thevenin)
```

    A					B					C					D
    |					|					|					|
    |					|					|					|
    +[2*R kO]---(V2 V)---\(21*V1 + 143*V2 - 42*V3 - 56*V4 + 24*V5 - 24*V6)/(454*R) mA\+					|					|
    |					|					|					|
    |					|					|					|
    +[7*R kO]---(-V5 + V6 V)---\(3*V1 - 12*V2 - 6*V3 - 8*V4 - 29*V5 + 29*V6)/(227*R) mA\+					|					|
    |					|					|					|
    |					|					|					|
    +----[2*R kO]---\(-41*V1 - 63*V2 + 82*V3 - 42*V4 + 18*V5 - 18*V6)/(454*R) mA\---+					|
    |					|					|					|
    |					|					|					|
    +---------------------[2*R kO]---(V4 V)---\(7*V1 - 28*V2 - 14*V3 + 57*V4 + 8*V5 - 8*V6)/(227*R) mA\---------------------+
    |					|					|					|
    |					|					|					|
    |					+[R kO]---(-V3 V)---\(-62*V1 + 21*V2 - 103*V3 + 14*V4 - 6*V5 + 6*V6)/(227*R) mA\+					|
    |					|					|					|
    |					|					|					|
    |					+[2*R kO]---(V1 V)---\(165*V1 + 21*V2 + 124*V3 + 14*V4 - 6*V5 + 6*V6)/(454*R) mA\+					|
    |					|					|					|
    |					|					|					|
    |					+-----[2*R kO]---\(-7*V1 + 28*V2 + 14*V3 - 57*V4 - 8*V5 + 8*V6)/(454*R) mA\-----+
    |					|					|					|
    |					|					|					|
    |					+-----[2*R kO]---\(-7*V1 + 28*V2 + 14*V3 - 57*V4 - 8*V5 + 8*V6)/(454*R) mA\-----+

-------


### Worksheet for generating gradients, velocities, sources, and boundary conditions

```
import sympy
```

This is for coupled equations of the form:

    - div( k0(u,v) grad(u) ) = f0
    - div( k1(u,v) grad(v) ) = f1

```
x,y = sympy.var('x,y')
```

```
# user input!
u = 2*x**2 + y**2
v = (sympy.sin(sympy.pi/4 * (x+y)))**2

# note exp() here is very useful, as we need kr > 0
ku = sympy.exp(u*v)
kv = sympy.exp((u+v)/2)

print "u = ", sympy.ccode(u)
print "v = ", sympy.ccode(v)
print ""
print "ku = ", sympy.ccode(ku)
print "kv = ", sympy.ccode(kv)
```

    u =  2*pow(x, 2) + pow(y, 2)
    v =  pow(sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y)), 2)

    ku =  exp((2*pow(x, 2) + pow(y, 2))*pow(sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y)), 2))
    kv =  exp(pow(x, 2) + (1.0L/2.0L)*pow(y, 2) + (1.0L/2.0L)*pow(sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y)), 2))

```
# print the gradients, in c-code

grad_u = sympy.diff(u, x), sympy.diff(u,y)
grad_v = sympy.diff(v, x), sympy.diff(v,y)

print "dudx = ", sympy.ccode(grad_u[0])
print "dudy = ", sympy.ccode(grad_u[1])
print "dvdx = ", sympy.ccode(grad_v[0])
print "dvdy = ", sympy.ccode(grad_v[1])

```

    dudx =  4*x
    dudy =  2*y
    dvdx =  (1.0L/2.0L)*M_PI*sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y))*cos(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y))
    dvdy =  (1.0L/2.0L)*M_PI*sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y))*cos(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y))

```
qu = (ku*grad_u[0], ku*grad_u[1])
qv = (kv*grad_v[0], kv*grad_v[1])
print "qu_x = ", sympy.ccode(qu[0])
print "qu_y = ", sympy.ccode(qu[1])
print "qv_x = ", sympy.ccode(qv[0])
print "qv_y = ", sympy.ccode(qv[1])

```

    qu_x =  4*x*exp((2*pow(x, 2) + pow(y, 2))*pow(sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y)), 2))
    qu_y =  2*y*exp((2*pow(x, 2) + pow(y, 2))*pow(sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y)), 2))
    qv_x =  (1.0L/2.0L)*M_PI*exp(pow(x, 2) + (1.0L/2.0L)*pow(y, 2) + (1.0L/2.0L)*pow(sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y)), 2))*sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y))*cos(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y))
    qv_y =  (1.0L/2.0L)*M_PI*exp(pow(x, 2) + (1.0L/2.0L)*pow(y, 2) + (1.0L/2.0L)*pow(sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y)), 2))*sin(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y))*cos(M_PI*((1.0L/4.0L)*x + (1.0L/4.0L)*y))

```

```

```

```

fu = - (sympy.diff(qu[0],x) + sympy.diff(qu[1],y))
fv = - (sympy.diff(qv[0],x) + sympy.diff(qv[1],y))

print "fu = ", sympy.ccode(fu)
print "fv = ", sympy.ccode(fv)

```

```