```python
from sympy import init_printing; init_printing();
from sympy import symbols, sin, cos, tan, acos, sqrt, Eq, Function, Array
from einsteinpy.symbolic import MetricTensor, ChristoffelSymbols, RiemannCurvatureTensor

from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all" # display all expression in one cell instead of the last one
```

## Cylindrical Coordinates

Transform Relation from $(r, \theta, \phi)$ to $(x, y, z)$:

\begin{equation}
\left\{\begin{matrix}
x =& R &\sin\phi \\
y =& R &\sin\phi \\
z =& Z &
\end{matrix}\right.
\end{equation}

Transform Relation from $(x, y, z)$ to $(r, \theta, \phi)$:
\begin{equation}
\left\{\begin{matrix}
 & R^2 =& x^2+y^2 & \\
 & Z =& z & \\
\tan & \phi  =& y/x &
\end{matrix}\right.
\end{equation}

Many programming languages have `atan2(y, x)` to avoid the function `atan(y/x)` which is not smart enough.

```python
R = symbols('R', positive=True); Z, phi = symbols('Z phi', real=True)
x, y, z = [f(R, Z, phi) for f in symbols('x, y, z', real=True, cls=Function)]
```

```python
eq_list = [
    Eq(x, R * cos(phi)),
    Eq(y, R * sin(phi)),
    Eq(z, Z)]
for eq in eq_list: eq
diffed_eq_list = []
for eq in eq_list:
    for var in [R, Z, phi]:
        diffed_eq_list.append(
            Eq(eq.lhs.diff(var),
               eq.rhs.diff(var)))# .subs(x**2+y**2+z**2, r**2).simplify() )
```

```python
[eq for eq in diffed_eq_list[0:3]]
[eq for eq in diffed_eq_list[3:6]]
[eq for eq in diffed_eq_list[6:9]]
```

\begin{equation}
\vec{r}_{1} = \frac{\partial \vec{r}}{\partial R} = \frac{\partial (x\hat{x}+y\hat{y}+z\hat{z})}{\partial R},\quad
\vec{r}_{2} = \frac{\partial \vec{r}}{\partial Z} = \frac{\partial (x\hat{x}+y\hat{y}+z\hat{z})}{\partial Z},\quad
\vec{r}_{3} = \frac{\partial \vec{r}}{\partial \phi} = \frac{\partial (x\hat{x}+y\hat{y}+z\hat{z})}{\partial \phi}
\end{equation}

```python
r_i = [None, None, None] # The three components would be
for i, var in enumerate([R, Z, phi]):
    r_i_temp = Array([x.diff(var), y.diff(var), z.diff(var)])
    for eq in diffed_eq_list:
        r_i_temp = r_i_temp.subs(eq.lhs, eq.rhs)
    r_i[i] = r_i_temp; r_i[i]

```

```python
from sympy import trigsimp
# define the metric for 3d cylindrical coordinates
metric = [[0 for i in range(3)] for i in range(3)]
for i in range(3):
    for j in range(3):
        metric[i][j] = (lambda vec1, vec2:
                        vec1[0]*vec2[0] + vec1[1]*vec2[1] + vec1[2]*vec2[2])(r_i[i], r_i[j])
        metric[i][j] = trigsimp(metric[i][j])
        # metric[i][j]
# creating metric object
m_obj = MetricTensor(metric, (R, Z, phi), config='ll'); m_obj.tensor()
g = m_obj.tensor().tomatrix().det(); g
```

```python
ch = ChristoffelSymbols.from_metric(m_obj)
ch.config
```

```python
# Calculating Riemann Tensor from Christoffel Symbols
rm1 = RiemannCurvatureTensor.from_christoffels(ch)
rm1.tensor()
```

```python
# Calculating Riemann Tensor from Metric Tensor
rm2 = RiemannCurvatureTensor.from_metric(m_obj)
rm2.tensor()
```

## Construct $\rho$ and $\vec{J}$ Field

```python
from sympy import exp, Matrix, sqrt, trigsimp
from einsteinpy.symbolic.tensor import BaseRelativityTensor
from einsteinpy.symbolic.vector import GenericVector

```

```python
# define the Levi-Civita tensor for 3d cylindrical coordinates
levi_civita_arr_lll = [[[0 for k in range(3)] for j in range(3)] for i in range(3)]
for i in range(3):
    for j in range(3):
        for k in range(3):
            levi_civita_arr_lll[i][j][k] = Matrix([r_i[i], r_i[j], r_i[k]]).det() / sqrt(g)
            levi_civita_arr_lll[i][j][k] = trigsimp(levi_civita_arr_lll[i][j][k])
```

```python
levi_civita = BaseRelativityTensor(
    levi_civita_arr_lll,
    (R,Z,phi), config='lll', parent_metric=m_obj)
levi_civita.tensor()
```

```python
density_df = BaseRelativityTensor(
    arr=exp(-R**2),
    syms=(R,Z,phi), config='', parent_metric=m_obj); density_df.arr

sqrt_g_df = BaseRelativityTensor(
    arr=sqrt(g),
    syms=(R,Z,phi), config='', parent_metric=m_obj); sqrt_g_df.arr

nabla_density_df = GenericVector(
    arr=[exp(-R**2).diff(var) for var in [R,Z,phi]],
    syms=(R,Z,phi), config='l', parent_metric=m_obj)
nabla_density_df

B_comp_list = [f(R,Z,phi) for f in symbols("B^1:4", real=True, cls=Function)]
B_df = GenericVector(
    arr=B_comp_list,
    syms=(R,Z,phi), config='u', parent_metric=m_obj)
B_df.tensor()
```

```python
J_df = GenericVector([0, exp(-R**2), exp(-R**2)], (R,Z,phi), config='u', parent_metric=m_obj)
J_df
# j_df = j_df.change_config('l')
J2_df = GenericVector([0,exp(-R**2),0], (R,Z,phi), config='u', parent_metric=m_obj)
```

```python
J_df == J2_df
J_df.tensor() == j2_df.tensor()
```

```python
from einsteinpy.symbolic.tensor import tensor_product
tensor_product(density_df, J_df)
```

```python
def tensor_cross(tensor1, tensor2, i=None, j=None):
    from einsteinpy.symbolic.tensor import tensor_product as prod
    middle_tensor = \
        prod(
            prod(
                prod(levi_civita, sqrt_g_df),
                tensor1, i=0,j=0
                ),
            tensor2, i=0,j=0
            )
#     middle_tensor.change_config('llluu')
    return middle_tensor
```

```python

```

```python
# nabla_density == j x B
nabla_density_df.tensor(), nabla_density_df.config
tensor_cross(J_df, B_df).tensor() # j x B
tensor_cross(J_df, B_df).config
```

```python
ideal_MHD_eq = Eq(
    nabla_density_df.tensor(),
    tensor_cross(J_df, B_df).tensor())
```

```python
def divide_Array_Eq(eq):
    from sympy import Eq
    assert(eq.lhs.shape == eq.rhs.shape)
    eq_shape = eq.lhs.shape
    arr_order = len(eq_shape)
    if arr_order > 1: # Not yet tested for high order tensor
        eq_list = []
        for i in range(eq_shape[0]):
            eq_list.append(
                divide_Array_Eq(Eq(eq.lhs[i], eq.rhs[i]))
            )
        return eq_list
    elif arr_order == 1:
        return [Eq(eq.lhs[i], eq.rhs[i]) for i in range(eq_shape[0])]
    else:# arr_order == 0
        return eq
```

```python
from sympy import solve, nonlinsolve
B_sol = solve(
    divide_Array_Eq(ideal_MHD_eq),
    B_df.tensor().tolist())
B_sol
```

```python
B_comp_list
B_comp_N_list = B_comp_list
for i in range(len(B_comp_N_list)):
    B_comp_N_list[i] = B_comp_N_list[i].subs(B_sol)
    B_comp_N_list[i] = B_comp_N_list[i].subs({B_comp_list[1]: 4, B_comp_list[2]: 6})
B_comp_N_list

B_N_df = GenericVector(
    arr=B_comp_N_list,
    syms=(R,Z,phi), config='u', parent_metric=m_obj)
B_N_df.tensor()
```

```python
from scipy.integrate import ode as scipy_ode
import ray; ray.init()
```

```python

```

```python
from scipy.integrate import ode as _scipy_ode

# @ray.remote
class FieldTracer(_scipy_ode): # FieldTracer is a kind of ode integrator.
    def __init__(self, vec_field):
        from sympy.utilities.lambdify import lambdify
        assert(vec_field.config == 'u')
        A__1, A__2, A__3 = vec_field.tensor().tolist()
        A__1 = lambdify(vec_field.syms, A__1)
        A__2 = lambdify(vec_field.syms, A__2)
        A__3 = lambdify(vec_field.syms, A__3)
        def dfdt(t, x):
            du1 = A__1(*x)
            du2 = A__2(*x)
            du3 = A__3(*x)
            return [du1, du2, du3]
        _scipy_ode.__init__(self, dfdt)

    def integrate_until(self, t1, dt):
        import numpy as np
        num_of_t = int(t1 / dt); num_of_t
        u_arr = np.empty((num_of_t, 3)); u_arr[0, :] = self._y
        t_arr = np.arange(num_of_t) * dt
        i = 0
        while self.successful() and i < num_of_t-1:
            i += 1
            u_arr[i, :] = self.integrate(self.t+dt)
        return u_arr, t_arr

    def set_f_params(self, *args, **kwargs):
        raise DeprecationWarning("The set_f_params method has been removed from the FieldTracer class.")

    def set_jac_params(self, *args, **kwargs):
        raise DeprecationWarning("The set_jac_params method has been removed from the FieldTracer class.")

# tracers = [FieldTracer.remote(B_N_df)for i in range(16)]
t = FieldTracer(B_N_df).set_integrator('vode', method='bdf')
# [t.set_integrator.remote('vode', method='bdf') for t in tracers]

u0, t0 = (10.0, 0.0, 0.0), 0.0
t.set_initial_value(u0, t0)
# [t.set_initial_value.remote(u0, t0) for t in tracers]
RZphi_arr, t_arr = t.integrate_until(10, 0.05)
# ray.get([t.integrate_until.remote(10, 0.05) for t in tracers])

```

```python
from sympy.utilities.lambdify import lambdify
import numpy as np
B__1, B__2, B__3 = B_N_df.tensor().tolist()
# B__1_arr = lambdify(B_N_df.syms, B__1, modules='numpy')(RZphi_arr[:, 0], RZphi_arr[:, 1], RZphi_arr[:, 2])
# B__2_arr = lambdify(B_N_df.syms, B__2, modules='numpy')(RZphi_arr[:, 0], RZphi_arr[:, 1], RZphi_arr[:, 2])
# B__3_arr = lambdify(B_N_df.syms, B__3, modules='numpy')(RZphi_arr[:, 0], RZphi_arr[:, 1], RZphi_arr[:, 2])
B__1_arr = np.zeros(len(RZphi_arr[:,0]))
B__2_arr = np.ones(len(RZphi_arr[:,0])) * 4
B__3_arr = np.ones(len(RZphi_arr[:,0])) * 6

```

```python
r_i_arr = [0 for i in range(3)]
for i in range(3):
    r_i_arr[i] = lambdify(B_N_df.syms, r_i[i], modules='numpy')(
        RZphi_arr[:, 0], RZphi_arr[:, 1], RZphi_arr[:, 2])
    r_i_arr[i][2] = np.zeros_like(r_i_arr[i][1])

B_arr = B__1_arr[:, None] * np.stack(r_i_arr[0], axis=-1) + \
        B__2_arr[:, None] * np.stack(r_i_arr[1], axis=-1) + \
        B__3_arr[:, None] * np.stack(r_i_arr[2], axis=-1)

```

```python
import plotly.graph_objects as go
import numpy as np

fig = go.Figure(data=go.Cone(
    x=RZphi_arr[:, 0] * np.cos(RZphi_arr[:, 2]),
    y=RZphi_arr[:, 0] * np.sin(RZphi_arr[:, 2]),
    z=RZphi_arr[:, 1],
    u=B_arr[:, 0],
    v=B_arr[:, 1],
    w=B_arr[:, 2],
    colorscale='Blues',
    sizemode="absolute"
#     marker=dict(
#         size=4,
#         color=p_arr,
#         colorscale='Viridis',
#     ),
#     line=dict(
#         color=p_arr,
#         width=2
#     )
))

fig.update_layout(
    scene=dict(
        aspectratio = dict( x=1, y=1, z=1.0 ),
        aspectmode = 'data'))
# fig.show()

```

```python

```

```python

```


# Frequency domain Acoustic FWI

## Acoustic wave equation
We start with the first order acoustic wave-equations expressed in particle velocity $v_z$, $v_x$ and pressure $p$ as
$$
\begin{align}
\frac{\partial v_z}{\partial t} &= b \frac{\partial p}{\partial z} \\ \frac{\partial v_x}{\partial t} &= b \frac{\partial p}{\partial x} \\
\frac{\partial p}{\partial t} &= k ( \frac{\partial v_z}{\partial z} + \frac{\partial v_x}{\partial x}) + s(t)
\end{align}    (1)
$$

## Split PML absorbing boundary condition
To accomdate the split PML boundary conditions, we split the pressure term $p$ into two unphysical terms $p_z$ and $p_x$ with $p=p_z + p_x$. Equation (1) can be reformulated as
$$
\begin{align}
\frac{\partial v_z}{\partial t} + \gamma_z v_z &= b \frac{\partial (p_z+P_x)}{\partial z} \\
\frac{\partial v_x}{\partial t} + \gamma_x v_x &= b \frac{\partial (p_z+p_x)}{\partial x} \\
\frac{\partial p_z}{\partial t} + \gamma_z p_z &= k  \frac{\partial v_z}{\partial z} \\
\frac{\partial p_x}{\partial t} + \gamma_x p_x &= k  \frac{\partial v_x}{\partial x} + s(t) \\
\end{align}    (2)
$$

where $b$ is the bouyancy which equal to $b = \frac{1}{\rho}$, $k$ is the bulk modulus, $\gamma_z$ and $\gamma_x$ are the PML coefficients. The source term is additive, so it can be add to either one of the last two equations of Equation (2).

$\gamma(x) = c_{pml} cos(\frac{\pi}{2} \frac{x}{L})$ where $x$ is the local coordinate in the PML layer, the origin is located at the outer edges of the model, $L$ is the width of PML layer. the scalar $c_{pml}$ is a scalar defined by trial and error. After disrectization of the PML coefficients, both $\gamma_z$ and $\gamma_x$ are a 1D vector.

## Transform to frequency domain
Perform Fourier transform to both side of Equation (2), we can get the wave-equations in frequency domain
$$
\begin{align}
iw v_z + \gamma_z v_z &= b \frac{\partial (p_z+P_x)}{\partial z}  \\
iw v_x + \gamma_x v_x &= b \frac{\partial (p_z+p_x)}{\partial x}  \\
iw p_z + \gamma_z p_z &= k  \frac{\partial v_z}{\partial z}       \\
iw p_x + \gamma_x p_x &= k  \frac{\partial v_x}{\partial x} + s(w)\\
\end{align}    (3)
$$
where $i$ is the imagary unit and $w$ is the radian frequency $w = 2 \pi f$, where $f$ is frequency. Combine the two terms on the left side, we can get
$$
\begin{align}
(iw + \gamma_z) v_z &= b \frac{\partial (p_z+P_x)}{\partial z}  \\
(iw + \gamma_x) v_x &= b \frac{\partial (p_z+p_x)}{\partial x}  \\
(iw + \gamma_z) p_z &= k  \frac{\partial v_z}{\partial z}       \\
(iw + \gamma_x) p_x &= k  \frac{\partial v_x}{\partial x} + s(w)\\
\end{align}    (4)
$$
We define new variable $\epsilon_z = 1+ \frac{\gamma_z}{iw} = 1 - \frac{i \gamma_z}{w}$ and $\epsilon_x = 1+ \frac{\gamma_x}{iw} = 1 - \frac{i \gamma_x}{w}$. As we have discussed for $\gamma_z$ and $\gamma_x$, $\epsilon_z$ and $\epsilon_x$ are also two 1D vectors.

So Equation (4) can be transformed to
$$
\begin{align}
iw \epsilon_z v_z &= b \frac{\partial (p_z+P_x)}{\partial z}  \\
iw \epsilon_x v_x &= b \frac{\partial (p_z+p_x)}{\partial x}  \\
iw \epsilon_z p_z &= k  \frac{\partial v_z}{\partial z}       \\
iw \epsilon_x p_x &= k  \frac{\partial v_x}{\partial x} + s(w)\\
\end{align} (5)
$$

We can furture transform the Equations (5) into follow form which facilitate to derive the finite-difference stencil.
$$
\begin{align}
iw v_z &= \frac{b}{\epsilon_z} \frac{\partial p}{\partial z}  \\
iw v_x &= \frac{b}{\epsilon_x} \frac{\partial p}{\partial x}  \\
\frac{iw}{k} p_z &= \frac{1}{\epsilon_z} \frac{\partial v_z}{\partial z}       \\
\frac{iw}{k} p_x &= \frac{1}{\epsilon_x} \frac{\partial v_x}{\partial x} + \frac{s(w)}{k \epsilon_x}\\
\end{align} (6)
$$
Note: the main purpose of moving $\epsilon_z$ and $\epsilon_x$ to the right side is that we can summ the last two equations to recover the pressure componnent $p$.

## Discretization
We use stagger grid to discretize the wave-equation. we place the pressure component $p_z$ and $p_x$, bulk modulus $k$, buoyancy $b$ on integer grid, partical velocity $v_z$ and $v_x$ on half grid along vertical and horizontal direction, respectively. So the discretized field is given as
$$
v_z[i+\frac{1}{2},j] \\
v_x[i,j+\frac{1}{2}] \\
p_z[i,j],\, p_x[i,j],\, b[i,j],\, k[i,j]
$$

With this discretizatino, assume the grid size is $h$ in both vertical and horizontal directions. The partial derivative is approximated by second-order finite-difference method. so we can get
$$
\begin{align}
\frac{\partial v_z}{\partial z}[i,j] &= \frac{v_z[i+\frac{1}{2},j] - v_z[i-\frac{1}{2},j]}{h} \\
\frac{\partial v_x}{\partial x}[i,j] &= \frac{v_x[i,j+\frac{1}{2}] - v_z[i,j-\frac{1}{2}]}{h} \\
\end{align}  (7)\\
$$

From the first two equations of Equation (6) and Moving the $iw$ term from left side of Equation (6) to the right hand side, we can get
$$
\begin{align}
v_z[i+\frac{1}{2},j] &= \frac{b[i+\frac{1}{2},j]}{iw\, \epsilon_z[i+\frac{1}{2}]} \frac{1}{h}(p[i+1,j] - p[i,j]) \\
v_z[i-\frac{1}{2},j] &= \frac{b[i-\frac{1}{2},j]}{iw\, \epsilon_z[i-\frac{1}{2}]} \frac{1}{h}(p[i,j] - p[i-1,j]) \\
v_x[i,j+\frac{1}{2}] &= \frac{b[i,j+\frac{1}{2}]}{iw\, \epsilon_x[j+\frac{1}{2}]} \frac{1}{h}(p[i,j+1] - p[i,j]) \\
v_x[i,j-\frac{1}{2}] &= \frac{b[i,j-\frac{1}{2}]}{iw\, \epsilon_x[j-\frac{1}{2}]} \frac{1}{h}(p[i,j] - p[i,j-1])
\end{align}  (8)
$$

Using the expressions in Equation (8) to replace the terms in equation (7), we can get a new system as
$$
\begin{align}
\frac{\partial v_z}{\partial z}[i,j] &=   \frac{b[i+\frac{1}{2},j]}{iw\, \epsilon_z[i+\frac{1}{2}]} \frac{1}{h^2}(p[i+1,j] - p[i,j]) - \frac{b[i-\frac{1}{2},j]}{iw\, \epsilon_z[i-\frac{1}{2}]} \frac{1}{h^2}(p[i,j] - p[i-1,j])\\
\frac{\partial v_x}{\partial x}[i,j] &= \frac{b[i,j+\frac{1}{2}]}{iw\, \epsilon_x[j+\frac{1}{2}]} \frac{1}{h^2}(p[i,j+1] - p[i,j])   -   \frac{b[i,j-\frac{1}{2}]}{iw\, \epsilon_x[j-\frac{1}{2}]} \frac{1}{h^2}(p[i,j] - p[i,j-1]) \\
\end{align}  (9)\\
$$

The last two equations of Equation (6) can be discretized as
$$
\begin{align}
\frac{iw}{k[i,j]} p_z[i,j] &= \frac{1}{\epsilon_z[i]} \frac{\partial v_z}{\partial z}[i,j]       \\
\frac{iw}{k[i,j]} p_x[i,j] &= \frac{1}{\epsilon_x[j]} \frac{\partial v_x}{\partial x}[i,j] + \frac{s(w)}{k[i,j] \epsilon_x[j]}\\
\end{align}  (10)
$$

We know the source is always inject inside the computational area, so $\gamma_x = 0$ and $\epsilon_x = 1$, so equations (10) can be simplified as
$$
\begin{align}
\frac{iw}{k[i,j]} p_z[i,j] &= \frac{1}{\epsilon_z[i]} \frac{\partial v_z}{\partial z}[i,j]       \\
\frac{iw}{k[i,j]} p_x[i,j] &= \frac{1}{\epsilon_x[j]} \frac{\partial v_x}{\partial x}[i,j] + \frac{s(w)}{k[i,j]}\\
\end{align}  (11)
$$

Using the expressions in equations (9), equations (11) can be represented as
$$
\begin{align}
\frac{iw}{k[i,j]} p_z[i,j] &= \frac{1}{\epsilon_z[i] h^2} \left(  \frac{b[i+\frac{1}{2},j]}{iw\, \epsilon_z[i+\frac{1}{2}]} (p[i+1,j] - p[i,j]) - \frac{b[i-\frac{1}{2},j]}{iw\, \epsilon_z[i-\frac{1}{2}]} (p[i,j] - p[i-1,j])  \right)       \\
\frac{iw}{k[i,j]} p_x[i,j] &= \frac{1}{\epsilon_x[j] h^2} \left( \frac{b[i,j+\frac{1}{2}]}{iw\, \epsilon_x[j+\frac{1}{2}]} (p[i,j+1] - p[i,j])   -   \frac{b[i,j-\frac{1}{2}]}{iw\, \epsilon_x[j-\frac{1}{2}]} (p[i,j] - p[i,j-1]) \right) + \frac{s(w)}{k[i,j]}\\
\end{align}  (11)
$$

Summing these two equations and multiply $iw$ on both side of equations (11), we get
$$
\begin{align}
\frac{-w^2}{k[i,j]} p[i,j] &= \frac{1}{\epsilon_z[i] h^2} \left(  \frac{b[i+\frac{1}{2},j]}{\epsilon_z[i+\frac{1}{2}]} (p[i+1,j] - p[i,j]) - \frac{b[i-\frac{1}{2},j]}{\epsilon_z[i-\frac{1}{2}]} (p[i,j] - p[i-1,j])  \right) \\
&+ \frac{1}{\epsilon_x[j] h^2} \left( \frac{b[i,j+\frac{1}{2}]}{\epsilon_x[j+\frac{1}{2}]} (p[i,j+1] - p[i,j])   -   \frac{b[i,j-\frac{1}{2}]}{\epsilon_x[j-\frac{1}{2}]} (p[i,j] - p[i,j-1]) \right) \\
&+ \frac{iw \, s(w)}{k[i,j]}
\end{align}  (11)
$$

The bouyancy at intermidiate grid point is computed via arithmatic average, for example
$$b[i+\frac{1}{2},j] = \frac{b[i,j] + b[i+1,j]}{2}$$

To facilitate building the Helmhotz operator as a sparse matrix, we can move the pressure component to the right size of the equation and the source term to the left side of the equation, we get
$$
\begin{align}
&\frac{1}{\epsilon_z[i] h^2} \left(  \frac{b[i+\frac{1}{2},j]}{\epsilon_z[i+\frac{1}{2}]} (p[i+1,j] - p[i,j]) - \frac{b[i-\frac{1}{2},j]}{\epsilon_z[i-\frac{1}{2}]} (p[i,j] - p[i-1,j])  \right) \\
&+ \frac{1}{\epsilon_x[j] h^2} \left( \frac{b[i,j+\frac{1}{2}]}{\epsilon_x[j+\frac{1}{2}]} (p[i,j+1] - p[i,j])   -   \frac{b[i,j-\frac{1}{2}]}{\epsilon_x[j-\frac{1}{2}]} (p[i,j] - p[i,j-1]) \right) \\
&+ \frac{w^2}{k[i,j]} p[i,j] = -\frac{iw \, s(w)}{k[i,j]}
\end{align}  (12)
$$

So the coefficients in front of
$$
\begin{align}
p[i,j] & \rightarrow \frac{w^2}{k[i,j]} - \frac{1}{\epsilon_z[i] h^2} \left( \frac{b[i+\frac{1}{2},j]}{\epsilon_z[i+\frac{1}{2}]} +  \frac{b[i-\frac{1}{2},j]}{\epsilon_z[i-\frac{1}{2}]} \right) - \frac{1}{\epsilon_x[j] h^2} \left( \frac{b[i,j+\frac{1}{2}]}{\epsilon_x[j+\frac{1}{2}]} + \frac{b[i,j-\frac{1}{2}]}{\epsilon_x[j-\frac{1}{2}]} \right) \\
p[i+1,j] & \rightarrow \frac{1}{\epsilon_z[i] h^2} \frac{b[i+\frac{1}{2},j]}{\epsilon_z[i+\frac{1}{2}]} \\
p[i-1,j] & \rightarrow \frac{1}{\epsilon_z[i] h^2} \frac{b[i-\frac{1}{2},j]}{\epsilon_z[i-\frac{1}{2}]} \\
p[i,j+1] & \rightarrow \frac{1}{\epsilon_x[j] h^2} \frac{b[i,j+\frac{1}{2}]}{\epsilon_x[j+\frac{1}{2}]} \\
p[i,j-1] & \rightarrow \frac{1}{\epsilon_x[j] h^2} \frac{b[i,j-\frac{1}{2}]}{\epsilon_x[j-\frac{1}{2}]}
\end{align}
$$

## Rotated stagger grid

This method still put the pressure component, the new coordinate system is rotated $45$ degrees.
$$
\begin{align}
iw v_z &= \frac{\sqrt{2}}{2} \frac{b}{\epsilon_z} \left(\frac{\partial p}{\partial z^{\prime}} -\frac{\partial p}{\partial x^{\prime}} \right)  \\
iw v_x &= \frac{\sqrt{2}}{2} \frac{b}{\epsilon_x} \left( \frac{\partial p}{\partial z^{\prime}} + \frac{\partial p}{\partial x^{\prime}}  \right)   \\
\frac{iw}{k} p_z &= \frac{\sqrt{2}}{2\epsilon_z} \left( \frac{\partial v_z}{\partial z^{\prime}} - \frac{\partial v_z}{\partial x^{\prime}} \right)     \\
\frac{iw}{k} p_x &= \frac{\sqrt{2}}{2\epsilon_x} \left( \frac{\partial v_x}{\partial z^{\prime}} + \frac{\partial v_x}{\partial x^{\prime}} \right) + \frac{s(w)}{k}
\end{align}    (13)
$$

We first approximate the spatial partial derivative of partical velocity by finite difference, so we have
$$
\begin{align}
\frac{\partial v_z}{\partial z^{\prime}}[i,j] &= \frac{v_z[i+\frac{1}{2},j+\frac{1}{2}] - v_z[i-\frac{1}{2},j-\frac{1}{2}] }{\sqrt{2}h} \\
\frac{\partial v_z}{\partial x^{\prime}}[i,j] &= \frac{v_z[i-\frac{1}{2},j+\frac{1}{2}] - v_z[i+\frac{1}{2},j-\frac{1}{2}] }{\sqrt{2}h} \\
\frac{\partial v_x}{\partial z^{\prime}}[i,j] &= \frac{v_x[i+\frac{1}{2},j+\frac{1}{2}] - v_x[i-\frac{1}{2},j-\frac{1}{2}] }{\sqrt{2}h} \\
\frac{\partial v_x}{\partial x^{\prime}}[i,j] &= \frac{v_x[i-\frac{1}{2},j+\frac{1}{2}] - v_x[i+\frac{1}{2},j-\frac{1}{2}] }{\sqrt{2}h} \\
\end{align} (14)
$$

From the first equation of the system (13):
$$
iw v_z = \frac{\sqrt{2}}{2} \frac{b}{\epsilon_z} \left(\frac{\partial p}{\partial z^{\prime}} -\frac{\partial p}{\partial x^{\prime}} \right)
$$

we represent particle velocity component $v_z$ by pressure component as
$$
\begin{align}
v_z[i-\frac{1}{2}, j-\frac{1}{2}] &= \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{2h\, iw \epsilon_z[i-\frac{1}{2}]} \left(p[i,j] - p[i-1,j-1] + p[i,j-1] - p[i-1,j] \right) \\
v_z[i+\frac{1}{2}, j-\frac{1}{2}] &= \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{2h\, iw \epsilon_z[i+\frac{1}{2}]} \left(p[i+1,j] - p[i,j-1] + p[i+1,j-1] - p[i,j] \right) \\
v_z[i-\frac{1}{2}, j+\frac{1}{2}] &= \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{2h\, iw \epsilon_z[i-\frac{1}{2}]} \left(p[i,j+1] - p[i-1,j] + p[i,j] - p[i-1,j+1] \right) \\
v_z[i+\frac{1}{2}, j+\frac{1}{2}] &= \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{2h\, iw \epsilon_z[i+\frac{1}{2}]} \left(p[i+1,j+1] - p[i,j] + p[i+1,j] - p[i,j+1] \right) \\
\end{align} (15)
$$

From the second equation of system (13):
$$
iw v_x = \frac{\sqrt{2}}{2} \frac{b}{\epsilon_x} \left( \frac{\partial p}{\partial z^{\prime}} + \frac{\partial p}{\partial x^{\prime}}  \right)
$$

By taking advantage of this equation, we can also represent $v_x$ via pressure component as:
$$
\begin{align}
v_x[i-\frac{1}{2}, j-\frac{1}{2}] &= \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{2h\, iw \epsilon_x[j-\frac{1}{2}]} \left(p[i,j] - p[i-1,j-1] + p[i-1,j] - p[i,j-1] \right) \\
v_x[i+\frac{1}{2}, j-\frac{1}{2}] &= \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{2h\, iw \epsilon_x[j-\frac{1}{2}]} \left(p[i+1,j] - p[i,j-1] + p[i,j] - p[i+1,j-1] \right) \\
v_x[i-\frac{1}{2}, j+\frac{1}{2}] &= \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{2h\, iw \epsilon_x[j+\frac{1}{2}]} \left(p[i,j+1] - p[i-1,j] + p[i-1,j+1] - p[i,j] \right) \\
v_x[i+\frac{1}{2}, j+\frac{1}{2}] &= \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{2h\, iw \epsilon_x[j+\frac{1}{2}]} \left(p[i+1,j+1] - p[i,j] + p[i,j+1] - p[i+1,j] \right) \\
\end{align} (16)
$$

Summing the last two equations of the system (13), we can get
$$
\frac{iw}{k} p = \frac{\sqrt{2}}{2\epsilon_z} \left( \frac{\partial v_z}{\partial z^{\prime}} - \frac{\partial v_z}{\partial x^{\prime}} \right) + \frac{\sqrt{2}}{2\epsilon_x} \left( \frac{\partial v_x}{\partial z^{\prime}} + \frac{\partial v_x}{\partial x^{\prime}} \right) + \frac{s(w)}{k} \,\,\, (17)
$$

Using the expressions in equation (14) to replace the terms in equation (16), we can get
$$
\begin{align}
\frac{iw}{k[i,j]}p[i,j] & = \frac{1}{2h \epsilon_z[i]} \left( v_z[i+\frac{1}{2}, j+\frac{1}{2}] - v_z[i-\frac{1}{2},j-\frac{1}{2}] + v_z[i+\frac{1}{2}, j-\frac{1}{2}] - v_z[i-\frac{1}{2}, j+\frac{1}{2}] \right) \\
&+ \frac{1}{2h \epsilon_x[j]} \left(v_x[i+\frac{1}{2},j+\frac{1}{2}]- v_x[i-\frac{1}{2},j-\frac{1}{2}] + v_x[i-\frac{1}{2},j+\frac{1}{2}] - v_x[i+\frac{1}{2},j-\frac{1}{2}]\right) + \frac{s}{k[i,j]} \\
\end{align} (18)
$$

The next step is to use the expression in equations (15, 16) to replace the particle velocity terms in above equation, take out the terms $2h$ and $iw$,  then we can get a equation only with pressure component.

$$
\begin{align}
\frac{-w^2}{k[i,j]}p[i,j] &= \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} \left(p[i+1,j+1] - p[i,j] + p[i+1,j] - p[i,j+1] \right) \\
 &- \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} \left(p[i,j] - p[i-1,j-1] + p[i,j-1] - p[i-1,j] \right) \\
 &+ \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} \left(p[i+1,j] - p[i,j-1] + p[i+1,j-1] - p[i,j] \right) \\
 &- \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} \left(p[i,j+1] - p[i-1,j] + p[i,j] - p[i-1,j+1] \right) \\
 &+ \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} \left(p[i+1,j+1] - p[i,j] + p[i,j+1] - p[i+1,j] \right) \\
 &- \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} \left(p[i,j] - p[i-1,j-1] + p[i-1,j] - p[i,j-1] \right) \\
 &+ \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} \left(p[i,j+1] - p[i-1,j] + p[i-1,j+1] - p[i,j] \right) \\
 &- \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} \left(p[i+1,j] - p[i,j-1] + p[i,j] - p[i+1,j-1] \right) \\
 &+ \frac{iw \,S}{K[i,j]}
\end{align} (19)
$$

To facilitate coding, we group the common terms together, we get a new equations
$$
\begin{align}
(\frac{-w^2}{k[i,j]} &+ \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} \\
                          &+ \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} ) p[i,j] \\
                          &- (\frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]}) p[i+1,j+1] \\
                          &- (\frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]})p[i-1,j-1]  \\
                          &- (\frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]})p[i+1,j-1] \\
                          &- (\frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]})p[i-1,j+1] \\
                          &-(\frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} +\frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} -\frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} -\frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} )p[i+1,j]  \\
                          &-(\frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} - \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} - \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} )p[i,j+1] \\
                          &-(\frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} - \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} - \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} )p[i,j-1] \\
                          &-(\frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} - \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} - \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} )p[i-1,j] \\
                          &= \frac{iw\,S}{k[i,j]}
\end{align} (20)
$$

## Mixed grid
By combine the finte-difference grid from cartesian and rotated system, we can get the mixed-grid finite-difference stencial. This process includes two steps:

### Mass acceleration averaging
$$
\begin{align}
\frac{w^2}{k[i,j]}p[i,j] = w^2( &c \cdot \frac{p[i,j]}{k[i,j]} \\
+ &d (\frac{p[i+1,j]}{k[i+1,j]}+\frac{p[i-1,j]}{k[i-1,j]}+\frac{p[i,j+1]}{k[i,j+1]}+\frac{p[i,j-1]}{k[i,j-1]}) \\
+ & \frac{1-c-4d}{4} (\frac{p[i+1,j+1]}{k[i+1,j+1]}+\frac{p[i-1,j+1]}{k[i-1,j+1]}+\frac{p[i+1,j-1]}{k[i+1,j-1]}+ \frac{p[i-1,j-1]}{k[i-1,j-1]}) ) \\
\end{align}
$$

### Weighted averaging of the two stencil
The weight between the two grid is determined by parameter $\alpha$
$$
\frac{w^2}{k}p + a \cdot \left( \frac{\partial}{\partial z}(b \frac{\partial }{\partial z}) + \frac{\partial}{\partial x}(b \frac{\partial }{\partial x}) \right)p + (1-a) \cdot \left( \frac{\partial}{\partial z^{\prime}}(b\frac{\partial }{\partial z^{\prime}}) + \frac{\partial}{\partial x^{\prime}}(b \frac{\partial }{\partial x^{\prime}}) \right)p  = -\frac{iw\, S}{k}
$$

where $c=0.6287326$, $d=0.092816675$, $\alpha=0.5617365$. so the coefficient $\frac{1-c-4d}{4} = 1.75e-7$.

At here, we give the coefficients for the 9 grid points
$$
\begin{align}
p[i,j] \rightarrow &c \frac{w^2}{k[i,j]} \\
                   -\alpha \left\{  \frac{1}{\epsilon_z[i] h^2}\left(\frac{b[i+\frac{1}{2},j]}{\epsilon_z[i+\frac{1}{2}]} +  \frac{b[i-\frac{1}{2},j]}{\epsilon_z[i-\frac{1}{2}]} \right) + \frac{1}{\epsilon_x[j] h^2} \left( \frac{b[i,j+\frac{1}{2}]}{\epsilon_x[j+\frac{1}{2}]} + \frac{b[i,j-\frac{1}{2}]}{\epsilon_x[j-\frac{1}{2}]}\right) \right\} \\
                   &- (1-\alpha) \left\{ \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} ) \right\} \\
&p[i-1,j-1] \rightarrow \frac{1-c-4d}{4} \frac{w^2}{k[i-1,j-1]} + (1-\alpha) \left\{ \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]}  \right\} \\
&p[i,j-1] \rightarrow d \frac{w^2}{k[i,j-1]} + \alpha \frac{1}{\epsilon_x[j] h^2} \frac{b[i,j-\frac{1}{2}]}{\epsilon_x[j-\frac{1}{2}]} + (1-\alpha)\left\{ \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} - \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} - \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} \right\} \\
&p[i+1,j-1] \rightarrow \frac{1-c-4d}{4} \frac{w^2}{k[i+1,j-1]} + (1-\alpha)\left\{\frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} \right\} \\
&p[i-1,j] \rightarrow d \frac{w^2}{k[i-1,j]} + \alpha \left\{ \frac{1}{\epsilon_z[i] h^2} \frac{b[i-\frac{1}{2},j]}{\epsilon_z[i-\frac{1}{2}]}\right\} + (1-\alpha) \left\{ \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} - \frac{b[i-\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]} - \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} \right\} \\
&p[i+1,j] \rightarrow d\frac{w^2}{k[i+1,j]} + \alpha \left\{\frac{1}{\epsilon_z[i] h^2} \frac{b[i+\frac{1}{2},j]}{\epsilon_z[i+\frac{1}{2}]}\right\} + (1-\alpha)\left\{\frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} +\frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} -\frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} -\frac{b[i+\frac{1}{2}, j-\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j-\frac{1}{2}]}\right\} \\
&p[i-1,j+1] \rightarrow \frac{1-c-4d}{4} \frac{w^2}{k[i-1,j+1]} + (1-\alpha)\left\{\frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]} + \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]}  \right\} \\
&p[i,j+1] \rightarrow d\frac{w^2}{k[i,j+1]} + \alpha \left\{\frac{1}{\epsilon_x[j] h^2} \frac{b[i,j+\frac{1}{2}]}{\epsilon_x[j+\frac{1}{2}]}\right\} + (1-\alpha) \left\{\frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]} - \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} - \frac{b[i-\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i-\frac{1}{2}]}  \right\} \\
&p[i+1,j+1] \rightarrow \frac{1-c-4d}{4} \frac{w^2}{k[i+1,j+1]} + (1-\alpha)\left\{\frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_z[i] \epsilon_z[i+\frac{1}{2}]} + \frac{b[i+\frac{1}{2}, j+\frac{1}{2}]}{4h^2 \epsilon_x[j] \epsilon_x[j+\frac{1}{2}]}  \right\}
\end{align}
$$

Outside of PML area, $\epsilon_z=\epsilon_x=1$, the above coefficients can be simplified as
To facilitate programming, we give the coefficients for each term
$$
\begin{align}
& p[i,j] \rightarrow c\frac{w^2}{k[i,j]} - \frac{a}{h^2}( b[i+\frac{1}{2},j]+ b[i-\frac{1}{2},j]+b[i,j+\frac{1}{2}]+b[i,j-\frac{1}{2}]) - \frac{1-a}{2h^2}(b[i+\frac{1}{2},j+\frac{1}{2}]+b[i-\frac{1}{2},j-\frac{1}{2}]+b[i-\frac{1}{2},j+\frac{1}{2}]+b[i+\frac{1}{2},j-\frac{1}{2}]) \\
& p[i-1,j-1] \rightarrow  \frac{1-c-4d}{4} \frac{w^2}{k[i-1,j-1]} +  \frac{1-a}{2h^2}b[i-\frac{1}{2},j-\frac{1}{2}] \\
& p[i,j-1]   \rightarrow d\frac{w^2}{k[i,j-1]} + \frac{\alpha}{h^2} b[i,j-\frac{1}{2}] \\
& p[i+1,j-1] \rightarrow \frac{1-c-4d}{4} \frac{w^2}{k[i+1,j-1]} + \frac{1-a}{2h^2}b[i+\frac{1}{2},j-\frac{1}{2}] \\
& p[i-1,j]   \rightarrow d\frac{w^2}{k[i-1,j]} + \frac{\alpha}{h^2} b[i-\frac{1}{2},j] \\
& p[i+1,j]   \rightarrow d\frac{w^2}{k[i+1,j]} + \frac{\alpha}{h^2} b[i+\frac{1}{2},j] \\
& p[i-1,j+1] \rightarrow \frac{1-c-4d}{4} \frac{w^2}{k[i-1,j+1]}+\frac{1-a}{2h^2}b[i-\frac{1}{2},j+\frac{1}{2}] \\
& p[i,j+1]   \rightarrow d\frac{w^2}{k[i,j+1]} + \frac{\alpha}{h^2} b[i,j+\frac{1}{2}] \\
& p[i+1,j+1] \rightarrow \frac{1-c-4d}{4} \frac{w^2}{k[i+1,j+1]} + \frac{1-a}{2h^2}b[i+\frac{1}{2},j+\frac{1}{2}] \\
\end{align}
$$

Use the same way of derivation, we can extend this method to forth order.

```julia

```


## The SEIR model

The SEIR model of epidemiology partitions the population into three compartments: susceptibles, S, who can catch the disease; exposed, E, that has caught the infection but is not infectious;  infectives, I, who have already caught the disease and infect susceptibles; and recovered individuals, R. Since the disease is assumed not to be fatal, the sum $N=S+I+R$ remains constant. The rate at which the susceptibles get infected is
$$
\lambda(t) = \frac{\beta I}{N}
$$
where the parameter $\beta$ is the probability of infection on contact. The infected individuals recover from the disease at a rate $\gamma$. Then, the ordinary differential equations of the SIR model are

\begin{align}
\dot S &= -\lambda(t)S \\
\dot E &= \lambda(t)S - \gamma_E E \\
\dot I &= \gamma_E E - \gamma I \\
\dot R &= \gamma I
\end{align}

Here $ 1/\zeta$ can be interpreted as the average incubation period. This example integrates the above equations to obtain what is called the **epidemic curve**: a plot of the number of susceptibles and infectives as a function of time.

```python
%%capture
## compile PyRoss for this notebook
import os
owd = os.getcwd()
os.chdir('../../')
%run setup.py install
os.chdir(owd)
```

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

```python
M = 1                  # the SIR model has no age structure
Ni = 1000*np.ones(M)   # so there is only one age group
N = np.sum(Ni)         # and the total population is the size of this age group

alpha = 0              # fraction of asymptomatic infectives
beta  = 0.2            # infection rate
gIa   = 0.1            # recovery rate of asymptomatic infectives
gIs   = 0.1            # recovery rate of symptomatic infectives
gE    = 0.04           # recovery rate of E
fsa   = 1              # the self-isolation parameter

E0  = np.array([0])
Ia0 = np.array([0])     # the SIR model has only one kind of infective
Is0 = np.array([1])     # we take these to be symptomatic
R0  = np.array([0])     # and assume there are no recovered individuals initially
S0  = N-(Ia0+Is0+R0+E0)    # so that the initial susceptibles are obtained from S + Ia + Is + R = N

# there is no contact structure
def contactMatrix(t):
    return np.identity(M)

# duration of simulation and data file
Tf = 500;  Nt=160;

# instantiate model
parameters = {'alpha':alpha, 'beta':beta, 'gIa':gIa, 'gIs':gIs, 'gE':gE, 'fsa':fsa}
model = pyross.deterministic.SEIR(parameters, M, Ni)

# simulate model
data=model.simulate(S0, E0, Ia0, Is0, contactMatrix, Tf, Nt)
```

```python
# plot the data and obtain the epidemic curve

S  = data['X'][:,0].flatten()
E  = data['X'][:,1].flatten()
Ia = data['X'][:,2].flatten()
Is = data['X'][:,3].flatten()
t = data['t']

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

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

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

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

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

plt.legend(fontsize=26); plt.grid()
plt.autoscale(enable=True, axis='x', tight=True)
```


## Decision tree for classification in plain Python

A decision tree is a **supervised** machine learning model that can be used both for **classification** and **regression**. At its core, a decision tree uses a tree structure to predict an output value for a given input example. In the tree, each path from the root node to a leaf node represents a *decision path* that ends in a predicted value.

A simple example might look as follows:

Decision trees have many advantages. For example, they are easy to understand and their decisions are easy to interpret. Also, they don't require a lot of data preparation. A more extensive list of their advantages and disadvantages can be found [here](http://scikit-learn.org/stable/modules/tree.html).

### CART training algorithm
In order to train a decision tree, various algorithms can be used. In this notebook we will focus on the *CART* algorithm (Classification and Regression Trees) for *classification*. The CART algorithm builds a *binary tree* in which every non-leaf node has exactly two children (corresponding to a yes/no answer).

Given a set of training examples and their labels, the algorithm repeatedly splits the training examples $D$ into two subsets $D_{left}, D_{right}$ using some feature $f$ and feature threshold $t_f$ such that samples with the same label are grouped together. At each node, the algorithm selects the split $\theta = (f, t_f)$ that produces the purest subsets, weighted by their size. Purity/impurity is measured using the *Gini impurity*.

So at each step, the algorithm selects the parameters $\theta$ that minimize the following cost function:

\begin{equation}
J(D, \theta) = \frac{n_{left}}{n_{total}} G_{left} + \frac{n_{right}}{n_{total}} G_{right}
\end{equation}

- $D$: remaining training examples
- $n_{total}$ : number of remaining training examples
- $\theta = (f, t_f)$: feature and feature threshold
- $n_{left}/n_{right}$: number of samples in the left/right subset
- $G_{left}/G_{right}$: Gini impurity of the left/right subset

This step is repeated recursively until the *maximum allowable depth* is reached or the current number of samples $n_{total}$ drops below some minimum number. The original equations can be found [here](http://scikit-learn.org/stable/modules/tree.html).

After building the tree, new examples can be classified by navigating through the tree, testing at each node the corresponding feature until a leaf node/prediction is reached.

### Gini Impurity

Given $K$ different classification values $k \in \{1, ..., K\}$ the Gini impurity of node $m$ is computed as follows:

\begin{equation}
G_m = 1 - \sum_{k=1}^{K} (p_{m,k})^2
\end{equation}

where $p_{m, k}$ is the fraction of training examples with class $k$ among all training examples in node $m$.

The Gini impurity can be used to evaluate how good a potential split is. A split divides a given set of training examples into two groups. Gini measures how "mixed" the resulting groups are. A perfect separation (i.e. each group contains only samples of the same class) corresponds to a Gini impurity of 0. If the resulting groups contain equally many samples of each class, the Gini impurity will reach its highest value of 0.5

### Caveats

Without regularization, decision trees are likely to overfit the training examples. This can be prevented using techniques like *pruning* or by providing a maximum allowed tree depth and/or a minimum number of samples required to split a node further.

```python
from sklearn.datasets import load_iris
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
np.random.seed(123)

%matplotlib inline
```

## Dataset

The iris dataset compromises 150 examples of 3 different species of iris flowers (Setosa, Versicolour, and Virginica). Each example is described by four attributes: sepal length (cm), sepal width (cm), petal length (cm), petal width (cm).

```python
iris = load_iris()

X, y = iris.data, iris.target
```

```python
# Split the data into a training and test set
X_train, X_test, y_train, y_test = train_test_split(X, y)

print(f'Shape X_train: {X_train.shape}')
print(f'Shape y_train: {y_train.shape}')
print(f'Shape X_test: {X_test.shape}')
print(f'Shape y_test: {y_test.shape}')
```

    Shape X_train: (112, 4)
    Shape y_train: (112,)
    Shape X_test: (38, 4)
    Shape y_test: (38,)

## Decision tree class

Parts of this code were inspired by [this tutorial](https://machinelearningmastery.com/implement-decision-tree-algorithm-scratch-python/)

```python
class DecisionTree:
    """
    Decision tree for classification
    """

    def __init__(self):
        self.root_dict = None
        self.tree_dict = None

    def split_dataset(self, X, y, feature_idx, threshold):
        """
        Splits dataset X into two subsets, according to a given feature
        and feature threshold.

        Args:
            X: 2D numpy array with data samples
            y: 1D numpy array with labels
            feature_idx: int, index of feature used for splitting the data
            threshold: float, threshold used for splitting the data

        Returns:
            splits: dict containing the left and right subsets
            and their labels
        """

        left_idx = np.where(X[:, feature_idx] < threshold)
        right_idx = np.where(X[:, feature_idx] >= threshold)

        left_subset = X[left_idx]
        y_left = y[left_idx]

        right_subset = X[right_idx]
        y_right = y[right_idx]

        splits = {
        'left': left_subset,
        'y_left': y_left,
        'right': right_subset,
        'y_right': y_right,
        }

        return splits

    def gini_impurity(self, y_left, y_right, n_left, n_right):
        """
        Computes Gini impurity of a split.

        Args:
            y_left, y_right: target values of samples in left/right subset
            n_left, n_right: number of samples in left/right subset

        Returns:
            gini_left: float, Gini impurity of left subset
            gini_right: gloat, Gini impurity of right subset
        """

        n_total = n_left + n_left

        score_left, score_right = 0, 0
        gini_left, gini_right = 0, 0

        if n_left != 0:
            for c in range(self.n_classes):
                # For each class c, compute fraction of samples with class c
                p_left = len(np.where(y_left == c)[0]) / n_left
                score_left += p_left * p_left
            gini_left = 1 - score_left

        if n_right != 0:
            for c in range(self.n_classes):
                p_right = len(np.where(y_right == c)[0]) / n_right
                score_right += p_right * p_right
            gini_right = 1 - score_right

        return gini_left, gini_right

    def get_cost(self, splits):
        """
        Computes cost of a split given the Gini impurity of
        the left and right subset and the sizes of the subsets

        Args:
            splits: dict, containing params of current split
        """
        y_left = splits['y_left']
        y_right = splits['y_right']

        n_left = len(y_left)
        n_right = len(y_right)
        n_total = n_left + n_right

        gini_left, gini_right = self.gini_impurity(y_left, y_right, n_left, n_right)
        cost = (n_left / n_total) * gini_left + (n_right / n_total) * gini_right

        return cost

    def find_best_split(self, X, y):
        """
        Finds the best feature and feature index to split dataset X into
        two groups. Checks every value of every attribute as a candidate
        split.

        Args:
            X: 2D numpy array with data samples
            y: 1D numpy array with labels

        Returns:
            best_split_params: dict, containing parameters of the best split
        """

        n_samples, n_features = X.shape

        best_feature_idx, best_threshold, best_cost, best_splits = np.inf, np.inf, np.inf, None

        for feature_idx in range(n_features):
            for i in range(n_samples):
                current_sample = X[i]
                threshold = current_sample[feature_idx]
                splits = self.split_dataset(X, y, feature_idx, threshold)
                cost = self.get_cost(splits)

                if cost < best_cost:
                    best_feature_idx = feature_idx
                    best_threshold = threshold
                    best_cost = cost
                    best_splits = splits

        best_split_params = {
            'feature_idx': best_feature_idx,
            'threshold': best_threshold,
            'cost': best_cost,
            'left': best_splits['left'],
            'y_left': best_splits['y_left'],
            'right': best_splits['right'],
            'y_right': best_splits['y_right'],
        }

        return best_split_params

    def build_tree(self, node_dict, depth, max_depth, min_samples):
        """
        Builds the decision tree in a recursive fashion.

        Args:
            node_dict: dict, representing the current node
            depth: int, depth of current node in the tree
            max_depth: int, maximum allowed tree depth
            min_samples: int, minimum number of samples needed to split a node further

        Returns:
            node_dict: dict, representing the full subtree originating from current node
        """
        left_samples = node_dict['left']
        right_samples = node_dict['right']
        y_left_samples = node_dict['y_left']
        y_right_samples = node_dict['y_right']

        if len(y_left_samples) == 0 or len(y_right_samples) == 0:
            node_dict["left_child"] = node_dict["right_child"] = self.create_terminal_node(np.append(y_left_samples, y_right_samples))
            return None

        if depth >= max_depth:
            node_dict["left_child"] = self.create_terminal_node(y_left_samples)
            node_dict["right_child"] = self.create_terminal_node(y_right_samples)
            return None

        if len(right_samples) < min_samples:
            node_dict["right_child"] = self.create_terminal_node(y_right_samples)
        else:
            node_dict["right_child"] = self.find_best_split(right_samples, y_right_samples)
            self.build_tree(node_dict["right_child"], depth+1, max_depth, min_samples)

        if len(left_samples) < min_samples:
            node_dict["left_child"] = self.create_terminal_node(y_left_samples)
        else:
            node_dict["left_child"] = self.find_best_split(left_samples, y_left_samples)
            self.build_tree(node_dict["left_child"], depth+1, max_depth, min_samples)

        return node_dict

    def create_terminal_node(self, y):
        """
        Creates a terminal node.
        Given a set of labels the most common label is computed and
        set as the classification value of the node.

        Args:
            y: 1D numpy array with labels
        Returns:
            classification: int, predicted class
        """
        classification = max(set(y), key=list(y).count)
        return classification

    def train(self, X, y, max_depth, min_samples):
        """
        Fits decision tree on a given dataset.

        Args:
            X: 2D numpy array with data samples
            y: 1D numpy array with labels
            max_depth: int, maximum allowed tree depth
            min_samples: int, minimum number of samples needed to split a node further
        """
        self.n_classes = len(set(y))
        self.root_dict = self.find_best_split(X, y)
        self.tree_dict = self.build_tree(self.root_dict, 1, max_depth, min_samples)

    def predict(self, X, node):
        """
        Predicts the class for a given input example X.

        Args:
            X: 1D numpy array, input example
            node: dict, representing trained decision tree

        Returns:
            prediction: int, predicted class
        """
        feature_idx = node['feature_idx']
        threshold = node['threshold']

        if X[feature_idx] < threshold:
            if isinstance(node['left_child'], (int, np.integer)):
                return node['left_child']
            else:
                prediction = self.predict(X, node['left_child'])
        elif X[feature_idx] >= threshold:
            if isinstance(node['right_child'], (int, np.integer)):
                return node['right_child']
            else:
                prediction = self.predict(X, node['right_child'])

        return prediction
```

## Initializing and training the decision tree

```python
tree = DecisionTree()
tree.train(X_train, y_train, max_depth=2, min_samples=1)
```

## Printing the decision tree structure

```python
def print_tree(node, depth=0):
    if isinstance(node, (int, np.integer)):
        print(f"{depth * '  '}predicted class: {iris.target_names[node]}")
    else:
        print(f"{depth * '  '}{iris.feature_names[node['feature_idx']]} < {node['threshold']}, "
             f"cost of split: {round(node['cost'], 3)}")
        print_tree(node["left_child"], depth+1)
        print_tree(node["right_child"], depth+1)
```

```python
print_tree(tree.tree_dict)
```

    petal length (cm) < 3.0, cost of split: 0.346
      sepal length (cm) < 5.4, cost of split: 0.0
        predicted class: setosa
        predicted class: setosa
      petal width (cm) < 1.8, cost of split: 0.097
        predicted class: versicolor
        predicted class: virginica

## Testing the decision tree

```python
all_predictions = []
for i in range(X_test.shape[0]):
    result = tree.predict(X_test[i], tree.tree_dict)
    all_predictions.append(y_test[i] == result)

print(f"Accuracy on test set: {sum(all_predictions) / len(all_predictions)}")
```

    Accuracy on test set: 0.9473684210526315

```python

```


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

from scipy import signal
from animation import animate_convolution

%matplotlib inline
sym.init_printing()

plt.rcParams['animation.html'] = 'jshtml'

class rect(sym.Function):

    @classmethod
    def eval(cls, arg):
        return sym.Heaviside(arg + sym.S.Half) - sym.Heaviside(arg - sym.S.Half)

def plot_signals(x_axis, x, h, ylabel, xlabel):
    plt.rcParams['figure.figsize'] = 7, 2

    p1 = sym.plot(x, (x_axis, -5, 5), show=False,
                  line_color='C0', ylabel=ylabel, xlabel=xlabel,legend=True, label = '$x$')
    p2 = sym.plot(h, (x_axis, -5, 5), show=False, line_color='C1',legend=True, label = '$h$')
    p1.extend(p2)
    p1.show()

from pole_zero_plot import pole_zero_plot as pzmapZ

def pzmapS(poles, zeros, ax=None):
    mR = -2
    MR = 2
    mI = -2
    MI = 2
    for pole in poles:
        plt.plot(complex(pole).real, complex(pole).imag, 'rx', markersize=10)
        if complex(pole).real < mR:
            mR = complex(pole).real
        if complex(pole).real > MR:
            MR = complex(pole).real
        if complex(pole).imag < mI:
            mI = complex(pole).imag
        if complex(pole).imag > MI:
            MI = complex(pole).imag

    for zero in zeros:
        plt.plot(complex(zero).real, complex(zero).imag, 'bo', markersize=10)
        if complex(zero).real < mR:
            mR = complex(zero).real
        if complex(zero).real > MR:
            MR = complex(zero).real
        if complex(zero).imag < mI:
            mI = complex(zero).imag
        if complex(zero).imag > MI:
            MI = complex(zero).imag

    if ax is None:
        ax = plt.gca()
    ax.set_aspect(1)

    plt.axis([mR*1.1, MR*1.1, mI*1.1, MI*1.1])
    plt.xlabel(r'$\Re \{ s \}$')
    plt.ylabel(r'$\Im \{ s \}$')
    plt.grid()

```

# Análisis de sistemas LTI

## Causalidad

La [propiedad de causalidad](https://en.wikipedia.org/wiki/Causality) establece que puede percibirse un efecto en la salida de un sistema únicamente después de haber sido excitado. Matemáticamente pueden encontrarse modelos de sistemas que violen la causalidad, sin embargo no pueden implementarse, pues se requeriría conocer el futuro de las señales de entrada.

A continuación se discute sobre la respuesta de un sistemas LTI ante una entrada $x(t)$ desde la operación de convolución:

\begin{equation}
y(t) = \int_{-\infty}^{\infty} x(\tau) \cdot h(t-\tau) \; d\tau
\end{equation}

```python
t, tau = sym.symbols('t tau', real=True)
x = sym.Rational(3, 4) * rect(t - 1/2)

```

### Respuesta impulsional de dos lados

Un sistema hipotético cuya respuesta impulsional sea de dos lados, es decir, que tenga valores diferentes de cero antes y después del instante $t=0$, tendrá que conocer el futuro de la señal de entrada. Observe las gráficas de convolución.

```python
h1 = sym.exp(-t) * sym.Heaviside(t) + sym.exp(t) * sym.Heaviside(-t)

plot_signals(t, x, h1, r'$h1(t)$, $x(t)$', r'$t$')

```

```python

y1 = sym.integrate(h1.subs(t,t-tau) * x.subs(t, tau), (tau, 0, t))

anim1 = animate_convolution(x, h, y, t, tau, np.arange(-2, 5, 1), np.arange(-5, 5, 1))
anim1

```

### Respuesta impulsional antes de $t=0$

Un sistema hipotético cuya respuesta impulsional se presenta antes del instante $t=0$ tendrá que conocer el futuro de la señal de entrada. Observe las gráficas de convolución.

```python
h2 = sym.exp(t) * sym.Heaviside(-t)

plot_signals(t, x, h2, r'$h2(t)$, $x(t)$', r'$t$')

```

```python
y2 = sym.integrate(h2.subs(t,t-tau) * x.subs(t, tau), (tau, 0, t))

anim2 = animate_convolution(x, h2, y2, t, tau, np.arange(-2, 5, .1), np.arange(-5, 5, 0.1))
anim2
```

### Respuesta impulsional después de $t=0$

Un sistema hipotético cuya respuesta impulsional se presenta después del instante $t=0$ tendrá que conocer el pasado de la señal de entrada. Observe las gráficas de convolución.

```python
h3 = sym.exp(-t) * sym.Heaviside(t)

plot_signals(t, x, h3, r'$h3(t)$, $x(t)$', r'$t$')
```

```python
y3 = sym.integrate(h3.subs(t,t-tau) * x.subs(t, tau), (tau, 0, t))

anim3 = animate_convolution(x, h3, y3, t, tau, np.arange(-2, 5, .1), np.arange(-5, 5, 0.02))
anim3
```

Del estudio de la transformada de Laplace se obtuvo que las señales causales están asociadas a **regiones de convergencia** tales que:

\begin{equation}
\Re \{ s \} > \Re \{ - s_0 \}
\end{equation}

Lo cual se obtiene cuando hay al menos la misma cantidad de polos que de ceros. A continuación algunos ejemplos de funciones de transferencia causales y no causales.

| Función                  | causal o no causal|
|--------------------------|-------------------|
|$\frac{1}{s+2}$           | Causal            |
|$\frac{s-1}{s+2}$         | Causal            |
|$\frac{s+1}{s^2 + 3s +2}$ | Causal            |
|$\frac{s^2 + 3s +2}{s+1}$ | No Causal         |
|$\frac{s+1}{2}$           | No Causal         |

## Causalidad en sistemas de tiempo discreto

Con base en la relación expuesta entre la transformada de Laplace y la transformada $z$, puede deducirse que la función de transferencia debe cumplir la misma condición en cuanto a la cantidad de polos y ceros. Además, el modelo como ecuación en diferencias muestra directamente la dependencia de valores futuros de las señales.

**Ejemplo 1**
Suponga la siguiente función de transferencia $F(z)=\frac{Y(z)}{X(z)}$

\begin{equation}
F(z) = \frac{Y(z)}{X(z)} = \frac{z-0.2}{z-0.5}
\end{equation}

Esta tiene la misma cantidad de polos y ceros. LLevándola a ecuación en diferencias se obtiene:

\begin{align}
\frac{Y(z)}{X(z)} &= \frac{z-0.2}{z-0.5} \\
Y(z) \cdot (z-0.5) &= X(z) \cdot (z-0.2) \\
z \cdot Y(z) - 0.5 \cdot Y(z)  &= z \cdot X(z) - 0.2 \cdot X(z) \\
y[k+1] -0.5 \cdot y[k] &= x[k+1] - 0.2 \cdot x[k]
\end{align}

Esta ecuación en diferencias puede desplazarse al instante $\kappa = k+1$

\begin{equation}
y[k+1] -0.5 \cdot y[k] = x[k+1] - 0.2 \cdot x[k] \\
y[\kappa] -0.5 \cdot y[\kappa-1] = x[\kappa] - 0.2 \cdot x[\kappa-1]
\end{equation}

De esta última puede despejarse la salida $y$ retomando la variable independiente $k$ para facilidad.

\begin{equation}
y[k] = 0.5 \cdot y[k-1] + x[k] - 0.2 \cdot x[k-1]
\end{equation}

Esta ecuación muestra que se encuentra el valor actual de salida $y[k]$ a partir del conocimiento del valor de salida en el instante anterior $y[k-1]$, el valor actual de entrada $x[k]$ y el valor anterior de entrada $x[k-1]$

**Ejemplo 2**
Suponga la siguiente función de transferencia $F(z)=\frac{Y(z)}{X(z)}$

\begin{equation}
F(z) = \frac{Y(z)}{X(z)} = \frac{z^2-0.2}{z-0.5}
\end{equation}

Esta tiene la misma cantidad de polos y ceros. LLevándola a ecuación en diferencias se obtiene:

\begin{align}
\frac{Y(z)}{X(z)} &= \frac{z^2-0.2}{z-0.5} \\
Y(z) \cdot (z-0.5) &= X(z) \cdot (z^2-0.2) \\
z \cdot Y(z) - 0.5 \cdot Y(z)  &= z^2 \cdot X(z) - 0.2 \cdot X(z) \\
y[k+1] -0.5 \cdot y[k] &= x[k+2] - 0.2 \cdot x[k]
\end{align}

Si esta ecuación en diferencias se desplaza al instante $\kappa = k+1$

\begin{equation}
y[k+1] -0.5 \cdot y[k] = x[k+2] - 0.2 \cdot x[k] \\
y[\kappa] -0.5 \cdot y[\kappa-1] = x[\kappa+1] - 0.2 \cdot x[\kappa-1]
\end{equation}

Al despejarse la salida $y$ retomando la variable independiente $k$ para facilidad.

\begin{equation}
y[k] = 0.5 \cdot y[k-1] + x[k+1] - 0.2 \cdot x[k-1]
\end{equation}

Esta ecuación muestra que se encuentra el valor actual de salida $y[k]$ a partir del conocimiento del valor de salida en el instante anterior $y[k-1]$, el valor de entrada en el instante futuro $x[k+1]$ y el valor anterior de entrada $x[k-1]$

Si esta ecuación en diferencias se desplaza al instante $\kappa = k+2$

\begin{equation}
y[k+1] -0.5 \cdot y[k] = x[k+2] - 0.2 \cdot x[k] \\
y[\kappa - 1] -0.5 \cdot y[\kappa - 2] = x[\kappa] - 0.2 \cdot x[\kappa-2]
\end{equation}

En este caso no puede despejarse la salida $y$. Además, se observa que el valor actual de entrada $x[k]$ solo puede relacionarse con valores anteriores de la salida $y$

## Estabilidad en sistemas causales

Un sistema se considera estable si tiene la propiedad de tranformar la energía de manera que es capaz de llegar a equilibrio. En términos de señales, esto se traduce en la capacidad de responder de manera acotada ante cualquier entrada acotada. Esta es la condición de estabilidad [BIBO](https://en.wikipedia.org/wiki/BIBO_stability).

Suponga un sistema tal que la señal de salida $y(t) = \mathcal{H} \{ x(t) \}$. Este sistema es estable si se cumple que:

\begin{align}
|x(t)| &< B_x \\
|y(t)| &< B_y
\end{align}

Gráficamente se representa así:

Partiendo de la respuesta impulsional del sistema se tiene que:

\begin{equation}
|y(t)| = \left| \int_{-\infty}^{\infty} x(\tau) \cdot h(t-\tau) \; d\tau \right|
\end{equation}

Considerando $B_x$ como la mayor norma posible para la entrada $x(t)$, se tiene que:

\begin{equation}
|y(t)| \leq \int_{-\infty}^{\infty} |x(\tau)| \cdot |h(t-\tau)| \; d\tau < \int_{-\infty}^{\infty} B_x \cdot |h(t-\tau)| \; d\tau
\end{equation}

Si $y(t)$ también es acotada significa que la respuesta impulsional es [absolutamente integrable](https://en.wikipedia.org/wiki/Absolutely_integrable_function), y por ende acotada.

\begin{equation}
\int_{-\infty}^{\infty} |h(t)| \; dt < B_h
\end{equation}

Lo anterior implica que un sistema es estable si la respuesta impulsional es absolutamente integrable.

Como la transformada de la señal exponencial compleja y causal es
\begin{equation}
\mathcal{L} \{ \epsilon(t) \cdot e^{-s_0 t} \} = \frac{1}{s + s_0} \qquad \forall s : \Re \{ s \} > \Re \{ - s_0 \}
\end{equation}

y las señales exponenciales son absolutamente integrables cuando la parte real del exponente es negativa, puede deducirse que un sistema es estable si los polos de la respuesta impulsional (su función de transferencia) tienen parte real negativa.

Los sistemas con polos de parte real positiva serán pues inestables.

A continuación algunos ejemplos de modelos estables e inestables.

|Función              | polos        | ceros    | Estabilidad |
|---------------------|--------------|----------|---------------------|
| \begin{equation} \frac{1}{s+1}\end{equation}      |\begin{equation} s = -1\end{equation}       | \begin{equation} s = []\end{equation}   | estable             |
| \begin{equation} \frac{1}{s-1}\end{equation}     |\begin{equation} s = 1\end{equation}        | \begin{equation} s = []\end{equation}   | inestable           |
| \begin{equation} \frac{1}{s^2+2s+1}\end{equation}|\begin{equation} s = -1 \\ s = -1\end{equation}    | \begin{equation} s = []\end{equation}   | estable             |
| \begin{equation} \frac{1}{s^2+s+1}\end{equation} |\begin{equation} s = -\frac{1}{2}+j\cdot \frac{\sqrt{3}}{2} \\ s =-\frac{1}{2}-j\cdot \frac{\sqrt{3}}{2} \end{equation}  | \begin{equation} s = []\end{equation}   | estable             |
| \begin{equation} \frac{s-1}{s^2+s+1}\end{equation} |\begin{equation} s = -\frac{1}{2}+j\cdot \frac{\sqrt{3}}{2} \\ s =-\frac{1}{2}-j\cdot \frac{\sqrt{3}}{2} \end{equation}  |\begin{equation} s = 1\end{equation}| estable             |

Suponga que $h(t) = \mathcal{L}^{-1} \{ H(s) \}$ es la respuesta impulsional de un sistema de segundo orden con función de transferencia $H(s)$

**Sistema estable**

\begin{equation}
H(s) = \frac{1}{(s-s_\infty)(s-s^*_\infty)}
\end{equation}

Suponga que $s_\infty= -1 + j$ y $s^*_\infty$ es su complejo conjugado.

```python
s = sym.symbols('s', complex=True)
t = sym.symbols('t', real=True)
s_inf = -1 + 0*sym.I

H3 = 1/((s - s_inf)*(s - sym.conjugate(s_inf)))
h3 = sym.inverse_laplace_transform(H3, s, t)
h3
```

La respuesta al impulso se grafica

```python
sym.plot(h3, (t, 0, 10), xlabel='$t$', ylabel='$h3(t)$');
```

También el mapa de polos y ceros

```python
pzmapS(sym.roots(sym.denom(H3), s), sym.roots(sym.numer(H3), s))

```

Note además que

\begin{equation}
\int_{-\infty}^{\infty} | \epsilon(t) e^{-t} \sin(t) | \; dt = \int_{0}^{\infty} | e^{-t} \sin(t) | \; dt < \int_{0}^{\infty} | e^{-t} | \; dt = 1
\end{equation}

**Sistema inestable**

\begin{equation}
H(s) = \frac{1}{(s-s_\infty)(s-s^*_\infty)}
\end{equation}

Suponga que $s_\infty= 1 + j$ y $s^*_\infty$ es su complejo conjugado.

```python
s_inf = 1 + sym.I

H4 = 1/((s - s_inf)*(s - sym.conjugate(s_inf)))
h4 = sym.inverse_laplace_transform(H4, s, t)
h4
```

La respuesta al impulso se grafica

```python
sym.plot(h4, (t, 0, 10), xlabel='$t$', ylabel='$h4(t)$');
```

También en mapa de polos y ceros

```python
pzmapS(sym.roots(sym.denom(H4), s), sym.roots(sym.numer(H4), s))

```

Note además que

\begin{equation}
\int_{-\infty}^{\infty} | \epsilon(t) e^{t} \sin(t) | \; dt = \int_{0}^{\infty} | e^{t} \sin(t) | \; dt \rightarrow \text{no converge}
\end{equation}

## Estabilidad en sistemas de tiempo discreto

Con base en la relación expuesta entre la transformada de Laplace y la transformada $z$, puede deducirse que la función de transferencia debe cumplir con la condición de ubicación de sus polos afectando la ubicación con el periodo de muestreo.

\begin{equation}
X_\text{s}(s) = X(z) \big\rvert_{z = e^{s T}}
\end{equation}

Así, los sistemas estables tienen sus polos dentro de la circunferencia de radio unitario y los inestables por fuera.

**Sistema estable**

\begin{equation}
H(z) = \frac{z}{z-0.5}
\end{equation}

Se grafica la respuesta al impulso.

```python
num = [1,0]
den =[1,-0.5]

H5 = (num, den)

t, y = signal.dimpulse((num,den,1),n=30)

plt.stem(t, np.squeeze(y))
```

y el mapa de polos y ceros

```python
pzmapZ(sym.roots(den), sym.roots(num))

```

**Sistema inestable**

\begin{equation}
H(z) = \frac{0.5}{z^2+0.2z+0.5}
\end{equation}

```python
num = [0.5]
den =[1,0.2,0.5]

H5 = (num, den)

t, y = signal.dimpulse((num,den,1),n=30)

plt.stem(t, np.squeeze(y))
```

y el mapa de polos y ceros

```python
pzmapZ(sym.roots(den), sym.roots(num))

```

## Amortiguamiento

De acuerdo con los análisis realizados a la transformada de Laplace y la transformada $z$, se presentan oscilaciones en las respuestas impulsionales solamente cuando los polos tienen parte imaginaria. El amortiguamiento es la proíedad de los sistemas para extraer la energía de sí permitiendo o no que sus elementos internos oscilen.

- Sistemas sobreamortiguados: No oscilan y sus polos tienen parte imaginaria igual a cero.
- Sistemas subamortiguados: Oscilan y sus polos tienen parte imaginaria distinta de cero.
- Sistemas críticamente amortiguados: Tiene polos reales repetidos.


```python
from sympy import *
from sympy.abc import *
init_printing()
from math import *
```

```python
theta, phi, psi
```

```python
Rz = Matrix([[cos(theta),-sin(theta),0],[sin(theta), cos(theta),0], [0,0,1]])
Rz
```

```python
Ry = Matrix([[cos(psi), 0, sin(psi)], [0, 1, 0],[-sin(psi), 0, cos(psi)]])
Ry
```

```python
Rx = Matrix([[1, 0, 0], [0, cos(phi), -sin(phi)], [0, sin(phi), cos(phi)]])
Rx
```

```python
Rz*Ry*Rx
```

```python
import numpy as np
```

```python
np.degrees(np.radians(90))
```

```python
R = (Rz*Ry*Rx).subs([(theta, np.radians(0)), (psi, np.radians(10)), (phi, np.radians(120))])
print(R)
```

    Matrix([[0.984807753012208, 0.150383733180435, -0.0868240888334651], [0, -0.500000000000000, -0.866025403784439], [-0.173648177666930, 0.852868531952443, -0.492403876506104]])

```python
R
```

```python
def deg(R):
    y = atan2(R[2,1], R[2,2])
    p = asin(-R[2,0])
    r = atan2(R[1,0],R[0,0])
    print(r, p, y)
    return np.rad2deg(r), np.degrees(p), np.degrees(y)
```

```python
deg(R)
```

```python
R = R*(Rz*Ry*Rx).subs([(theta, np.radians(120)), (psi, np.radians(0)), (phi, np.radians(0))])
```

```python
deg(R)
```

```python
360 -120
```

```python

```


# Fitting the transfer function of a coupled oscillator
This section will be undoubtedly very useful for KAGRA's suspension modeling, as the transfer functions of the suspensions all behave like oscillator. Different from that in the previous tutorials, the transfer function of osciilators has complex poles and zeros instead of simple poles and zeros. Typically, they can be represented by the model

\begin{equation}
    G(s; z_i, q_i, p_j, Q_j, k) = k\frac{\prod_i \left(\frac{1}{(2\pi z_i)^2} s^2 + \frac{1}{2\pi z_i q_i}s + 1\right)}{\prod_j \left(\frac{1}{(2\pi p_j)^2} s^2 + \frac{1}{2\pi p_j Q_j}s + 1\right)}\,,
\end{equation}
where $z_i$ and $p_j$ are the resonance frequencies of the complex zeros and poles respectively, $q_i$ and $Q_j$ are the quality factors of the complex zeros and poles respectively, and $k$ is the static gain. And because of the similarity to the simple ZPK model, we call this a complex ZPK model.

Typically it's sufficient to consider the model
\begin{equation}
    H(s; k_i, f_i, q_i) = \sum_i k_i\frac{(2\pi f_i)^2}{s^2 + \frac{2\pi f_i}{q_i} s + (2\pi f_i)^2}\,.
\end{equation}
This model has fewer parameters than $G(s)$, but it only works for second-order systems, i.e. input-output relationship of a single stage. $G(s)$ is a more general representation so this is what we will use here.

In this tutorial, we will consider the following transfer function
\begin{equation}
 F(s) = 10\frac{(2\pi)^2}{s^2 + \frac{2\pi}{10}s + (2\pi)^2} + 1\frac{(10*2\pi)^2}{s^2 + \frac{10*2\pi}{100}s + (10*2\pi)^2}\,.
\end{equation}
And, we will try to fit it with the model $G(s)$

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

s = control.tf("s")
tf = 10*(2*np.pi)**2/(s**2 + (1*2*np.pi)/10*s + (2*np.pi)**2) + 1*(10*2*np.pi)**2/(s**2 + (10*2*np.pi)/10*s + (10*2*np.pi)**2)
f = np.linspace(0.01, 100, 10000)
tf_val = tf(1j*2*np.pi*f)

plt.figure(figsize=(15, 5))
plt.subplot(121)
plt.loglog(f, abs(tf_val))
plt.grid(which="both")
plt.ylabel("Amplitude")
plt.xlabel("Frequency (Hz)")
plt.subplot(122)
plt.semilogx(f, np.angle(tf_val))
plt.grid(which="both")
plt.ylabel("Phase (rad)")
plt.xlabel("Frequency (Hz)")
plt.show()
```

Here will we show how we can obtain an inital guess for the fit.

Judging from the plot, we should expect two complex poles at around 1 and 10 Hz, and a complex zeros at around 3 Hz. We will hand fit the transfer function as an initial guess. We start from lowest frequency, so it's just a gain at around 10.

\begin{equation}
    G_0(s) = 10
\end{equation}

And then, we scan from low frequency to high frequency. The next thing we see is a peak, which correspond to a pair of complex poles.
So, we add a pair of complex poles at 1 Hz and we fine tune the Q to match the peak, to which it becomes

\begin{equation}
    G_1(s) = G_0(s)  \frac{1}{\frac{1}{(1*2\pi)^2}s^2 + \frac{1}{1*2\pi Q_1}s + 1}\,,
\end{equation}
wherer $Q_1$ is around 10.

Then, the next feature is a notch at 3 Hz, which corresponds to a pair of complex zeros.
So we add a pair of complex zeros at 3 Hz, and tune the Q value to match the depth.
So it becomes

\begin{equation}
    G_2(s) = G_1(s)  \frac{\frac{1}{(3*2\pi)^2}s^2 + \frac{1}{3*2\pi q_2}s + 1}{1}\,,
\end{equation}
where we pick $q_2=10$

The last feature is a peak at 10 Hz, so we will add a pair of complex poles at 10 Hz.
And finally, the
\begin{equation}
    G_3(s) = G_2(s)  \frac{1}{\frac{1}{(10*2\pi)^2}s^2 + \frac{1}{10*2\pi Q_3}s + 1}\,,
\end{equation}
where we pick $Q_3$ to be 20.

```python
import kontrol.curvefit

G0 = kontrol.curvefit.model.ComplexZPK(nzero_pairs=0, npole_pairs=0)
G0_args = [10]

G1 = kontrol.curvefit.model.ComplexZPK(nzero_pairs=0, npole_pairs=1)
G1_args = [1, 10, 10]

G2 = kontrol.curvefit.model.ComplexZPK(nzero_pairs=1, npole_pairs=1)
G2_args = [3, 10, 1, 10, 10]

G3 = kontrol.curvefit.model.ComplexZPK(nzero_pairs=1, npole_pairs=2)
G3_args = [3, 10, 1, 10, 10, 20, 10]

plt.figure(figsize=(15, 5))
plt.subplot(121)
plt.loglog(f, abs(tf_val), label="Original transfer function", lw=3)
plt.loglog(f, abs(G0(f, G0_args)), label="$G_0(s)$", ls="--")
plt.loglog(f, abs(G1(f, G1_args)), label="$G_1(s)$", ls="-.")
plt.loglog(f, abs(G2(f, G2_args)), label="$G_2(s)$", ls=":")
plt.loglog(f, abs(G3(f, G3_args)), label="$G_3(s)$", color="k", ls="--")
plt.legend(loc=0)
plt.grid(which="both")
plt.ylabel("Amplitude")
plt.xlabel("Frequency (Hz)")

plt.subplot(122)
plt.semilogx(f, np.angle(tf_val))
plt.semilogx(f, np.angle(G0(f, G0_args)), label="$G_0(s)$")
plt.semilogx(f, np.angle(G1(f, G1_args)), label="$G_1(s)$")
plt.semilogx(f, np.angle(G2(f, G2_args)), label="$G_2(s)$")
plt.semilogx(f, np.angle(G3(f, G3_args)), label="$G_3(s)$")
plt.legend(loc=0)
plt.grid(which="both")
plt.ylabel("Phase (rad)")
plt.xlabel("Frequency (Hz)")
plt.show()

```

As can be seen, adding zeros/poles at higher frequencies doesn't affect the features at lower frequencies. This is why we started adding features from lower to higher frequencies.

The final initial guess we obtained is a ``kontrol.curvefit.model.ComplexZPK`` model, with parameters ``[3, 10, 1, 10, 10, 20, 10]``.
The arguments are arranged in the form ``[z_i, q_i, p_j, Q_j, k]``.
We specified the number of complex zero pairs ``nzero_pairs`` and the number of complex pole pairs ``npole_pairs`` during the model declaration so the program knows how to slices the parameters.
So in this case, the first two values correspond to a pair of complex zeros at 3 Hz with Q-value of 10.
The second two values correspond to a pair of complex poles at 1 Hz with Q-value of 10, and so on.
But this initial guess is probably too good to start with, so we add some noise to it to mimic a realistic fit.

Now we can proceed to the usual curve fitting routine.

```python
import scipy.optimize
import kontrol.curvefit

a = kontrol.curvefit.CurveFit()
a.xdata = f
a.ydata = tf_val
a.model = kontrol.curvefit.model.ComplexZPK(nzero_pairs=1, npole_pairs=2)

error_func = kontrol.curvefit.error_func.tf_error
weight = 1/f
error_func_kwargs = {"weight": weight}
a.cost = kontrol.curvefit.Cost(error_func=error_func, error_func_kwargs=error_func_kwargs)
a.optimizer = scipy.optimize.minimize

np.random.seed(123)
noise = np.random.normal(0, np.array(G3_args)/10)
# print(noise)
x0 = np.array(G3_args) + noise
a.optimizer_kwargs = {"x0": x0}
res = a.fit()
```

```python
plt.figure(figsize=(15, 5))
plt.subplot(121)
plt.loglog(f, abs(tf_val), label="Data")
plt.loglog(f, abs(a.model(f, x0)), label="Initial guess", ls="--")
plt.loglog(f, abs(a.yfit), label="Fit", ls="-.")
plt.legend(loc=0)
plt.grid(which="both")
plt.ylabel("Amplitude")
plt.xlabel("Frequency (Hz)")
plt.subplot(122)
plt.semilogx(f, np.angle(tf_val), label="Data")
plt.semilogx(f, np.angle(a.model(f, x0)), label="Initial guess", ls="--")
plt.semilogx(f, np.angle(a.yfit), label="Fit", ls="-.")
plt.legend(loc=0)
plt.grid(which="both")
plt.ylabel("Phase (rad)")
plt.xlabel("Frequency (Hz)")
plt.show()
```

Again, because the parameters can be bounded, we can use global optimization without an initial guess.
Let's just demonstrate once again.

```python
# Using log scale this time.
bounds = [(np.log10(min(f)), np.log10(max(f))), (np.log10(0.5), np.log10(1e3))]*3
bounds += [(np.log10(abs(tf_val[0])/10), np.log10(abs(tf_val[0])*10))]

a.model = kontrol.curvefit.model.ComplexZPK(nzero_pairs=1, npole_pairs=2, log_args=True)
a.optimizer = scipy.optimize.differential_evolution
a.optimizer_kwargs = {"bounds": bounds, "workers": -1, "updating": "deferred"}

np.random.seed(123)
res = a.fit()
```

```python
plt.figure(figsize=(15, 5))
plt.subplot(121)
plt.loglog(f, abs(tf_val), label="Data")
# plt.loglog(f, abs(a.model(f, x0)), label="Initial guess", ls="--")
plt.loglog(f, abs(a.yfit), label="Fit", ls="-.")
plt.legend(loc=0)
plt.grid(which="both")
plt.ylabel("Amplitude")
plt.xlabel("Frequency (Hz)")
plt.subplot(122)
plt.semilogx(f, np.angle(tf_val), label="Data")
# plt.semilogx(f, np.angle(a.model(f, x0)), label="Initial guess", ls="--")
plt.semilogx(f, np.angle(a.yfit), label="Fit", ls="-.")
plt.legend(loc=0)
plt.grid(which="both")
plt.ylabel("Phase (rad)")
plt.xlabel("Frequency (Hz)")
plt.show()
```

And, it worked very well.


## Code generation

```python
# initialize sympy printing (for latex output)
from sympy import init_printing, Symbol
init_printing()

# import functions and classes for compartment models
from compartor import *
```

Here, we illustrate how to export the moment equations in LaTeX format or generate code for simulations.

We consider again the exemplary model of the paper and derive moments equations with the automated function.

```python
x = Content('x')
y = Content('y')

# Intake Distribution
pi_I = OutcomeDistribution.Poisson(Symbol('\pi_{I}(y; \lambda)'),y[0],Symbol('\lambda'))

Intake      = TransitionClass( {}        -to> [(y[0],0)],   'k_I', pi=pi_I, name='I')
Fusion      = TransitionClass( [x] + [y] -to> [x+y],        'k_F',          name='F')
Conversion  = TransitionClass( [x]       -to> [x + (-1,1)], 'k_c', x[0],    name='c')
Degradation = TransitionClass( [x]       -to> [x + (0,-1)], 'k_d', x[1],    name='d')

transitions = [ Intake, Fusion, Conversion, Degradation]
display_transition_classes(transitions)
```

$\displaystyle \begin{align} \emptyset&\overset{h_{I}}{\longrightarrow}\left[\left( {y}_{0}, \  0\right)\right] && h_{I} = \pi_{I}(y; \lambda) k_{I}\\\left[x\right] + \left[y\right]&\overset{h_{F}}{\longrightarrow}\left[x + y\right] && h_{F} = \frac{k_{F} \left(n{\left(y \right)} - \delta_{x y}\right) n{\left(x \right)}}{\delta_{x y} + 1}\\\left[x\right]&\overset{h_{c}}{\longrightarrow}\left[\left( -1, \  1\right) + x\right] && h_{c} = k_{c} n{\left(x \right)} {x}_{0}\\\left[x\right]&\overset{h_{d}}{\longrightarrow}\left[\left( 0, \  -1\right) + x\right] && h_{d} = k_{d} n{\left(x \right)} {x}_{1} \end{align}$

```python
desired_moments = [Moment(0,0), Moment(0,0)**2, Moment(1,0), Moment(1,0)**2, Moment(0,1), Moment(0,1)**2]
equations = automated_moment_equations(2, transitions, desired_moments)
display_moment_equations(equations)
```

The LaTeX source of the ODE system can be found and copy-pasted by modifying the Math Render option of the jupyter notebook. This is done by right-clicking on the system and choosing Math Settings > Math Render > Plain Source .

From the closed equations, we can also generate code to simulate the system. Currently, Python or Julia code can be generated.

For a direct code output, the user can rely on the functions:

```python
python_code = generate_python_code(equations, function_name="example")
print(python_code)
```

    # evaluate ODEs
    def example_ODEs(M, dM, parameters):
      """
      Evaluate derivatives of expected moments

      Indices in M and dM vectors are
        M[0] = Moment(0, 0)
        M[1] = Moment(0, 0)**2
        M[2] = Moment(1, 0)
        M[3] = Moment(1, 0)**2
        M[4] = Moment(0, 1)
        M[5] = Moment(0, 1)**2
        M[6] = Moment(0, 1)*Moment(1, 0)

      :param M: expected moments
      :param dM: result, the derivative dM/dt is stored here
      :param parameters: tuple of values for constants (\lambda, k_F, k_I, k_c, k_d)
      :return: dM
      """
      c4 = parameters[0] # \lambda
      c0 = parameters[1] # k_F
      c3 = parameters[2] # k_I
      c1 = parameters[3] # k_c
      c2 = parameters[4] # k_d
      # Number of Compartments (N)
      dM[0] = c3+1/2*c0*M[0]-1/2*c0*M[1]
      # N^2
      dM[1] = c3+2*c3*M[0]-1/2*c0*M[0]+3/2*c0*M[1]+c0*M[1]*M[0]-2*c0*M[1]**2/M[0]
      # Total Mass of species 1
      dM[2] = c4*c3-c1*M[2]
      # M1^2
      dM[3] = c4*c3+c3*c4**2+c1*M[2]-2*c1*M[3]+2*c4*c3*M[2]
      # Total Mass of species 2
      dM[4] = c1*M[2]-c2*M[4]
      # M2^2
      dM[5] = c1*M[2]+c2*M[4]-2*c2*M[5]+2*c1*M[6]
      # M2*M1
      dM[6] = c1*M[3]-c1*M[6]-c1*M[2]-c2*M[6]+c4*c3*M[4]
      return dM

    # initialize expected moments vector
    def example_initial(initial):
      """
      Create inital expected moments vector

      :param initial: tuple of initial values for expectations of (Moment(0, 0), Moment(1, 0), Moment(0, 1))
      :return: initial expected moments vector
      """
      import numpy as np
      M = np.zeros(7)
      # Number of Compartments (N)
      M[0] = initial[0]
      # N^2
      M[1] = M[0]**2
      # Total Mass of species 1
      M[2] = initial[1]
      # M1^2
      M[3] = M[2]**2
      # Total Mass of species 2
      M[4] = initial[2]
      # M2^2
      M[5] = M[4]**2
      # M2*M1
      M[6] = M[4]*M[2]
      return M

or, for Julia code,

```python
julia_code = generate_julia_code(equations, function_name="example")
print(julia_code)
```

    # evaluate ODEs
    function example_ODEs(dM, M, parameters, t)
      c0 = ??? # k_F, please specify!
      c1 = ??? # k_c, please specify!
      c2 = ??? # k_d, please specify!
      c3 = ??? # k_I, please specify!
      c4 = ??? # \lambda, please specify!
      # Number of Compartments (N)
      dM[1] = c3+1/2*c0*M[1]-1/2*c0*M[2]
      # N^2
      dM[2] = c3+2*c3*M[1]-1/2*c0*M[1]+3/2*c0*M[2]+c0*M[2]*M[1]-2*c0*M[2]^2/M[1]
      # Total Mass of species 1
      dM[3] = c4*c3-c1*M[3]
      # M1^2
      dM[4] = c4*c3+c3*c4^2+c1*M[3]-2*c1*M[4]+2*c4*c3*M[3]
      # Total Mass of species 2
      dM[5] = c1*M[3]-c2*M[5]
      # M2^2
      dM[6] = c1*M[3]+c2*M[5]-2*c2*M[6]+2*c1*M[7]
      # M2*M1
      dM[7] = c1*M[4]-c1*M[7]-c1*M[3]-c2*M[7]+c4*c3*M[5]
      return
    end

    # initialize expected moments vector
    function example_initial(n0)
      M=zeros[7]
      # Number of Compartments (N)
      M[1] = ??? # initial value for Moment(0, 0), please specify!
      # N^2
      M[2] = M[1]^2
      # Total Mass of species 1
      M[3] = ??? # initial value for Moment(1, 0), please specify!
      # M1^2
      M[4] = M[3]^2
      # Total Mass of species 2
      M[5] = ??? # initial value for Moment(0, 1), please specify!
      # M2^2
      M[6] = M[5]^2
      # M2*M1
      M[7] = M[5]*M[3]
      return M
    end

Both derive from a common `AbstractCodeGenerator` class, which allows the user to customize the code output with further implementation-specific details. The argument `function_name` gives sets the initial name of the two functions defined in the respective output:
* `function_name_ODEs()` implements the ODE system in the respective programming language,
* `function_name_initial()` is a supplementary function that computes the initial moments for a given initial configuration `n0` of the compartment population.

## Solving and Plotting results
We use the generated functions with `scipy.solve_ivp` to solve the system with the same initial condition and parameter values reported in the paper.

```python
from scipy.integrate import solve_ivp
import numpy as np

exec(python_code)

# vector of timepoints where to solve ODE problem
timepoints = np.linspace(0, 100, 1001)

# initial condition
# 1 compartment with 20 molecules of first species
M0 = example_initial((1, 20, 0))

# rate constants and other parameters
kI = 1.0
kF = 0.005
kc = 0.1
kd = 0.05
Lambda = 10
parameters = (Lambda, kF, kI, kc, kd)

# solve
sol = solve_ivp(
    lambda t, M: example_ODEs(M, np.zeros(7), parameters),
    (np.min(timepoints), np.max(timepoints)),
    M0,
    method='BDF',
    t_eval=timepoints)

# print solution of N
print(sol.y[1])
```

    [  1.           1.30993662   1.63969989 ... 420.1743308  420.17551396
     420.17668693]

Finally, we use `matplotlib` to plot the solution for mean and standard deviation of $N$ and of the total amounts of molecules $M^{(1,0)}$ and $M^{(0,1)}$ for the two chemical species.

```python
%matplotlib inline
import matplotlib.pyplot as plt

y = sol.y
N = y[0]
stdN = np.sqrt(y[1] - N**2)
plt.plot(timepoints, y[0], color="b")
plt.ylabel(r'$\left< N \right>$', size=12)
plt.xlabel('time', size=12)
plt.title("Expected compartment number", size=12)
plt.fill_between(timepoints, N-stdN, N+stdN, alpha=0.4, color="b")
plt.ylim((0,30))
```

```python
m1 = y[2]
stdm1 = np.sqrt(y[3] - m1**2)
m2 = y[4]
stdm2 = np.sqrt(y[5] - m2**2)
plt.plot(timepoints, m1, color="g")
plt.fill_between(timepoints, m1-stdm1, m1+stdm1, alpha=0.4, color="g",label=r'$\left< M^{(1,0)} \right>$')
plt.plot(timepoints, m2, color="orange")
plt.ylabel('molecule number', size=12)
plt.fill_between(timepoints, m2-stdm2, m2+stdm2, alpha=0.4, color="orange",label=r'$\left< M^{(0,1)} \right>$')
plt.xlabel('time', size=12)
plt.title("Expected total molecule amount", size=12)
plt.legend()
plt.ylim((0,250))
```


平面ロボットアームの運動学導出

```python
import sympy as sy
from sympy import pi, cos, sin, tan
from IPython.display import display
from sympy.printing.pycode import pycode

#q1, q2, q3, q4 = sy.symbols("q1, q2, q3, q4")  # 関節角度
t = sy.Symbol("t")
q1 = sy.Function("q1")
q2 = sy.Function("q2")
q3 = sy.Function("q3")
q4 = sy.Function("q4")
omega1 = sy.Function("omega1")
omega2 = sy.Function("omega2")
omega3 = sy.Function("omega3")
omega4 = sy.Function("omega4")

a1, a2, a3, a4 = sy.symbols("a1, a2, a3, a4")
b1, b2, b3, b4 = sy.symbols("b1, b2, b3, b4")
c1, c2, c3, c4 = sy.symbols("c1, c2, c3, c4")

l1, l2, l3, l4 = sy.symbols("l1, l2, l3, l4")  # リンク長さ
lg1, lg2, lg3, lg4 = sy.symbols("lg1, lg2, lg3, lg4")  # 重心までの長さ
m1, m2, m3, m4 = sy.symbols("m1, m2, m3, m4")  # 質量
Ig1, Ig2, Ig3, Ig4 = sy.symbols("I1, I2, I3, I4")  # 慣性モーメント
g = sy.Symbol("g")  # 重力加速度

def R(q):
    return sy.Matrix([
        [cos(q), -sin(q)],
        [sin(q), cos(q)],
    ])

def HTM(q, x, y):
    return sy.Matrix([
        [cos(q), -sin(q), x],
        [sin(q), cos(q), y],
        [0, 0, 1],
    ])

# ジョイントの位置
x1 = R(q1(t)) * sy.Matrix([[l1, 0]]).T
x2 = R(q1(t)) * sy.Matrix([[l1, 0]]).T + \
    R(q1(t) + q2(t)) * sy.Matrix([[l2, 0]]).T
x3 = R(q1(t)) * sy.Matrix([[l1, 0]]).T +\
    R(q1(t) + q2(t)) * sy.Matrix([[l2, 0]]).T +\
        R(q1(t) + q2(t) + q3(t)) * sy.Matrix([[l3, 0]]).T
x4 = R(q1(t)) * sy.Matrix([[l1, 0]]).T +\
    R(q1(t) + q2(t)) * sy.Matrix([[l2, 0]]).T +\
        R(q1(t) + q2(t) + q3(t)) * sy.Matrix([[l3, 0]]).T +\
            R(q1(t) + q2(t) + q3(t) + q4(t)) * sy.Matrix([[l4, 0]]).T

q = sy.Matrix([[q1(t), q2(t), q3(t), q4(t)]]).T

```

```python
J1 = x1.jacobian(q)
J2 = x2.jacobian(q)
J3 = x3.jacobian(q)
J4 = x4.jacobian(q)

J1_dot = sy.diff(J1, t)
J2_dot = sy.diff(J2, t)
J3_dot = sy.diff(J3, t)
J4_dot = sy.diff(J4, t)

J_all = [J1, J2, J3, J4]
J_dot_all = [J1_dot, J2_dot, J3_dot, J4_dot]

for i, J in enumerate(J_all):
    J_all[i] = J.subs([
        (sy.Derivative(q1(t),t), b1),
        (sy.Derivative(q2(t),t), b2),
        (sy.Derivative(q3(t),t), b3),
        (sy.Derivative(q4(t),t), b4),
        (q1(t), a1),
        (q2(t), a2),
        (q3(t), a3),
        (q4(t), a4),
    ])

for i, J_dot in enumerate(J_dot_all):
    J_dot_all[i] = J_dot.subs([
        (sy.Derivative(q1(t),t), b1),
        (sy.Derivative(q2(t),t), b2),
        (sy.Derivative(q3(t),t), b3),
        (sy.Derivative(q4(t),t), b4),
        (q1(t), a1),
        (q2(t), a2),
        (q3(t), a3),
        (q4(t), a4),
    ])

```

```python
f = open('sice_kinema.txt', 'w')
for i, j in enumerate([x1, x2, x3, x4]):
    s = '\nx' + str(i) + '='
    f.write(s)
    f.write(str(j))

for i, j in enumerate(J_all):
    s = '\nJ' + str(i) + '='
    f.write(s)
    f.write(str(j))
for i, j in enumerate(J_dot_all):
    s = '\nJ_dot' + str(i) + '='
    f.write(s)
    f.write(str(j))

f.close()

```

```python
print(J_dot_all[3])
```

    Matrix([[-b1*l1*cos(a1) - l2*(b1 + b2)*cos(a1 + a2) - l3*(b1 + b2 + b3)*cos(a1 + a2 + a3) - l4*(b1 + b2 + b3 + b4)*cos(a1 + a2 + a3 + a4), -l2*(b1 + b2)*cos(a1 + a2) - l3*(b1 + b2 + b3)*cos(a1 + a2 + a3) - l4*(b1 + b2 + b3 + b4)*cos(a1 + a2 + a3 + a4), -l3*(b1 + b2 + b3)*cos(a1 + a2 + a3) - l4*(b1 + b2 + b3 + b4)*cos(a1 + a2 + a3 + a4), -l4*(b1 + b2 + b3 + b4)*cos(a1 + a2 + a3 + a4)], [-b1*l1*sin(a1) - l2*(b1 + b2)*sin(a1 + a2) - l3*(b1 + b2 + b3)*sin(a1 + a2 + a3) - l4*(b1 + b2 + b3 + b4)*sin(a1 + a2 + a3 + a4), -l2*(b1 + b2)*sin(a1 + a2) - l3*(b1 + b2 + b3)*sin(a1 + a2 + a3) - l4*(b1 + b2 + b3 + b4)*sin(a1 + a2 + a3 + a4), -l3*(b1 + b2 + b3)*sin(a1 + a2 + a3) - l4*(b1 + b2 + b3 + b4)*sin(a1 + a2 + a3 + a4), -l4*(b1 + b2 + b3 + b4)*sin(a1 + a2 + a3 + a4)]])

```python
print(J4)
```

    Matrix([[-l1*sin(q1(t)) - l2*sin(q1(t) + q2(t)) - l3*sin(q1(t) + q2(t) + q3(t)) - l4*sin(q1(t) + q2(t) + q3(t) + q4(t)), -l2*sin(q1(t) + q2(t)) - l3*sin(q1(t) + q2(t) + q3(t)) - l4*sin(q1(t) + q2(t) + q3(t) + q4(t)), -l3*sin(q1(t) + q2(t) + q3(t)) - l4*sin(q1(t) + q2(t) + q3(t) + q4(t)), -l4*sin(q1(t) + q2(t) + q3(t) + q4(t))], [l1*cos(q1(t)) + l2*cos(q1(t) + q2(t)) + l3*cos(q1(t) + q2(t) + q3(t)) + l4*cos(q1(t) + q2(t) + q3(t) + q4(t)), l2*cos(q1(t) + q2(t)) + l3*cos(q1(t) + q2(t) + q3(t)) + l4*cos(q1(t) + q2(t) + q3(t) + q4(t)), l3*cos(q1(t) + q2(t) + q3(t)) + l4*cos(q1(t) + q2(t) + q3(t) + q4(t)), l4*cos(q1(t) + q2(t) + q3(t) + q4(t))]])