# Superoperator tools

In this notebook we explore the submodules of `operator_tools` that enable easy manipulation of the various quantum channel representations.

To summarize the functionality:

- vectorization and conversions between different representations of quantum channels
- apply quantum operations
- compose quantum operations
- validate that quantum channels are physical
- project unphysical channels to physical channels

## Brief motivation and introduction

Perfect gates in **reversible classical computation** are described by permutation matrices, e.g. the [Toffoli gate](https://en.wikipedia.org/wiki/Toffoli_gate), while the input states are vectors. A noisy classical gate could be modeled as a perfect gate followed by a noise channel, e.g. [binary symmetric channel](https://en.wikipedia.org/wiki/Binary_symmetric_channel), on all the bits in the state vector.

Perfect gates in **quantum computation** are described by unitary matrices and states are described by complex vectors, e.g.

$$|\psi\rangle = U |\psi_0\rangle$$

Modeling **noisy quantum computation** often makes use of [mixed states](https://en.wikipedia.org/wiki/Density_matrix) and quantum operations or quantum noise channels.

Interestingly there are a number of ways to represent quantum noise channels, and depending on your task some can be more convenient than others. The simplest case to illustrate this point is to consider a mixed initial state $\rho$ undergoing unitary evolution

$$\rho' = U \rho U^\dagger$$

The fact that the unitary has to act on both sides of the initial state means it is a [*superoperator*](https://en.wikipedia.org/wiki/Superoperator), that is an object that can act on operators like the state matrix.

It turns out using a special matrix multiplication identity we can write this as

$$|\rho'\rangle \rangle = \mathcal U |\rho\rangle\rangle$$

where $\mathcal U = U^*\otimes U$ and $|\rho\rangle\rangle =  {\rm vec}(\rho)$. The nice thing about this is it looks like the pure state case. This is because the operator (the state) has become a vector and the superoperator (the left right action of $U$) has become an operator.

**More information**
Below we will assume that you are already an expert in these topics. If you are unfamiliar with these topics we recommend the following references

- chapter 8 of [Mike_N_Ike] which is on *Quantum noise and quantum operations*.
- chapter 3 of John Preskill's lecture notes [Physics 219/Computer Science 219](http://www.theory.caltech.edu/people/preskill/ph219/chap3_15.pdf)
- the [file](../superoperator_representations.rst) `/docs/superoperator_representations.md`
- for an intuitive but advanced treatment see [GRAPTN]

| [Mike_N_Ike] *Quantum Computation and Quantum Information*.
|     Michael A. Nielsen & Isaac L. Chuang.
|     Cambridge: Cambridge University Press (2000).

| [GRAPTN] *Tensor networks and graphical calculus for open quantum systems*.
|     Christopher Wood et al.
|     Quant. Inf. Comp. 15, 0579-0811 (2015).
|     https://arxiv.org/abs/1111.6950

## Conversion between different descriptions of quantum channels

We intentionally chose not to make quantum channels python objects with methods that would automatically transform between representations.

The functions to convert between different representations are called things like `kraus2chi`, `kraus2choi`, `pauli_liouville2choi` etc.

This assumes the user does not do silly things like input a Choi matrix to a function `chi2choi`.

```python
import numpy as np
from pyquil.gate_matrices import I, X, Y, Z, H, CNOT
```

Define some channels

```python
def amplitude_damping_kraus(p):
    Ad0 = np.asarray([[1, 0], [0, np.sqrt(1 - p)]])
    Ad1 = np.asarray([[0, np.sqrt(p)], [0, 0]])
    return [Ad0, Ad1]

def bit_flip_kraus(p):
    M0 = np.sqrt(1 - p) * I
    M1 = np.sqrt(p) * X
    return [M0, M1]
```

Define some states

```python
one_state = np.asarray([[0,0],[0,1]])
zero_state = np.asarray([[1,0],[0,0]])
rho_mixed = np.asarray([[0.9,0],[0,0.1]])
```

### vec and unvec

We can vectorize i.e. `vec` and unvec matrices.

We chose a column stacking convention so that the matrix

$$A = \begin{pmatrix} 1 & 2\\ 3 & 4\end{pmatrix}$$

becomes

$$|A\rangle\rangle = {\rm vec}(A) = \begin{pmatrix} 1\\ 3\\ 2\\ 4\end{pmatrix}$$

Let's check that

```python
from forest.benchmarking.operator_tools import vec, unvec
```

```python
A = np.asarray([[1, 2], [3, 4]])
```

```python
print(A)
print(" ")
print(vec(A))
print(" ")
print('Does the story check out? ', np.all(unvec(vec(A))==A))
```

    [[1 2]
     [3 4]]

    [[1]
     [3]
     [2]
     [4]]

    Does the story check out?  True

### Kraus to $\chi$ matrix (aka chi or process matrix)

```python
from forest.benchmarking.operator_tools import kraus2chi
```

Lets do a unitary gate first, say the Hadamard

```python
print('The Kraus operator is:\n', np.round(H,3))
print('\n')
print('The Chi matrix is:\n', kraus2chi(H))
```

    The Kraus operator is:
     [[ 0.707  0.707]
     [ 0.707 -0.707]]

    The Chi matrix is:
     [[0. +0.j 0. +0.j 0. +0.j 0. +0.j]
     [0. +0.j 0.5+0.j 0. +0.j 0.5+0.j]
     [0. +0.j 0. +0.j 0. +0.j 0. +0.j]
     [0. +0.j 0.5+0.j 0. +0.j 0.5+0.j]]

Now consider the Amplitude damping channel

```python
AD_kraus = amplitude_damping_kraus(0.1)

print('The Kraus operators are:\n', np.round(AD_kraus,3))
print('\n')
print('The Chi matrix is:\n', np.round(kraus2chi(AD_kraus),3))
```

    The Kraus operators are:
     [[[1.    0.   ]
      [0.    0.949]]

     [[0.    0.316]
      [0.    0.   ]]]

    The Chi matrix is:
     [[0.949+0.j    0.   +0.j    0.   +0.j    0.025+0.j   ]
     [0.   +0.j    0.025+0.j    0.   -0.025j 0.   +0.j   ]
     [0.   +0.j    0.   +0.025j 0.025+0.j    0.   +0.j   ]
     [0.025+0.j    0.   +0.j    0.   +0.j    0.001+0.j   ]]

### Kraus to Pauli Liouville aka the "Pauli Transfer Matrix"

```python
from forest.benchmarking.operator_tools import kraus2pauli_liouville
```

```python
Hpaulirep = kraus2pauli_liouville(H)
Hpaulirep
```

    array([[ 1.+0.j,  0.+0.j,  0.+0.j,  0.+0.j],
           [ 0.+0.j,  0.+0.j,  0.+0.j,  1.+0.j],
           [ 0.+0.j,  0.+0.j, -1.+0.j,  0.+0.j],
           [ 0.+0.j,  1.+0.j,  0.+0.j,  0.+0.j]])

We can visualize this using the tools from the plotting module.

```python
from forest.benchmarking.plotting.state_process import plot_pauli_transfer_matrix
import matplotlib.pyplot as plt
```

```python
f, (ax1) = plt.subplots(1, 1, figsize=(5, 4.2))

plot_pauli_transfer_matrix(Hpaulirep,ax=ax1)
```

The above figure is a graphical representation of:

(out operator) = H (in operator) H

Z = H X H
-Y = H Y H
X = H Z H

### Evolving states using quantum channels

In many superoperator representations evolution corresponds to multiplying the vec'ed state by the superoperator. E.g.

```python
from forest.benchmarking.operator_tools import kraus2superop

zero_state_vec = vec(zero_state)

answer_vec = np.matmul(kraus2superop([H]), zero_state_vec)

print('The vec\'ed answer is', answer_vec)
print('\n')
print('The unvec\'ed answer is\n', np.real(unvec(answer_vec)))
print('\n')
print('Let\'s compare it to the normal calculation\n', H @ zero_state @ H)
```

    The vec'ed answer is [[0.5+0.j]
     [0.5+0.j]
     [0.5+0.j]
     [0.5+0.j]]

    The unvec'ed answer is
     [[0.5 0.5]
     [0.5 0.5]]

    Let's compare it to the normal calculation
     [[0.5 0.5]
     [0.5 0.5]]

For representations with this simple application there are no inbuilt functions in forest benchmarking.

However applying a channel is more painful in the Choi and Kraus representation.

Consider the amplitude damping channel where we need to perform the following calculation to find out put of channel
$\rho_{out} = A_0 \rho A_0^\dagger + A_1 \rho A_1^\dagger$.
We provide helper functions to do these calculations.

```python
from forest.benchmarking.operator_tools import apply_kraus_ops_2_state, apply_choi_matrix_2_state, kraus2choi
```

```python
apply_kraus_ops_2_state(AD_kraus, one_state)
```

    array([[0.1, 0. ],
           [0. , 0.9]])

In the Choi representation we get the same answer:

```python
AD_choi = kraus2choi(AD_kraus)

apply_choi_matrix_2_state(AD_choi, one_state)
```

    array([[0.1, 0. ],
           [0. , 0.9]])

### Compose quantum channels

Composing channels is useful when describing larger circuits. In some representations e.g. in the superoperator  or Liouville representation it is just matrix multiplication e.g.

```python
from forest.benchmarking.operator_tools import superop2kraus, kraus2superop
```

```python
H_super = kraus2superop(H)

H_squared_super = H_super @ H_super

print('Hadamard squared as a superoperator:\n', np.round(H_squared_super,2))

print('\n As a Kraus operator:\n', np.round(superop2kraus(H_squared_super),2))
```

    Hadamard squared as a superoperator:
     [[ 1.+0.j -0.+0.j -0.+0.j  0.+0.j]
     [-0.+0.j  1.+0.j  0.+0.j -0.+0.j]
     [-0.+0.j  0.+0.j  1.+0.j -0.+0.j]
     [ 0.+0.j -0.+0.j -0.+0.j  1.+0.j]]

     As a Kraus operator:
     [[[ 1.+0.j -0.+0.j]
      [ 0.+0.j  1.+0.j]]]

Composing channels in the Kraus representation is more difficult. Consider composing two channels $\mathcal A$ (with Kraus operators $[A_0, A_1]$) and $\mathcal B$ (with Kraus operators $[B_0, B_1]$). The composition is

$$\begin{align}
\mathcal B(\mathcal A(\rho)) & = \sum_i \sum_j B_j A_i \rho A_i^\dagger B_j^\dagger
\end{align}$$

```python
from forest.benchmarking.operator_tools import compose_channel_kraus, superop2kraus
```

```python
BitFlip_kraus = bit_flip_kraus(0.2)

kraus2superop(compose_channel_kraus(AD_kraus, BitFlip_kraus))
```

    array([[0.82      +0.j, 0.        +0.j, 0.        +0.j, 0.28      +0.j],
           [0.        +0.j, 0.75894664+0.j, 0.18973666+0.j, 0.        +0.j],
           [0.        +0.j, 0.18973666+0.j, 0.75894664+0.j, 0.        +0.j],
           [0.18      +0.j, 0.        +0.j, 0.        +0.j, 0.72      +0.j]])

This is the same as if we do

```python
BitFlip_super = kraus2superop(BitFlip_kraus)
AD_super = kraus2superop(AD_kraus)

AD_super @ BitFlip_super
```

    array([[0.82      +0.j, 0.        +0.j, 0.        +0.j, 0.28      +0.j],
           [0.        +0.j, 0.75894664+0.j, 0.18973666+0.j, 0.        +0.j],
           [0.        +0.j, 0.18973666+0.j, 0.75894664+0.j, 0.        +0.j],
           [0.18      +0.j, 0.        +0.j, 0.        +0.j, 0.72      +0.j]])

We can also easily compose channels acting on independent spaces.

Consider composing the same two channels as above, $\mathcal A$ and $\mathcal B$. However this time they act on different Hilbert spaces. With respect to the tensor product structure $H_2 \otimes H_1$ the Kraus operators are $[A_0\otimes I, A_1\otimes I]$ and $[I \otimes B_0, I \otimes B_1]$.

In this case the order of the operations commutes

$$\begin{align}
\mathcal A(\mathcal B(\rho))= \mathcal B(\mathcal A(\rho)) & = \sum_i \sum_j  A_i\otimes B_j \rho A_i^\dagger\otimes B_j^\dagger
\end{align}$$

In forest benchmarking you can specify the two channels without the Identity tensored on and it will take care of it for you:

```python
from forest.benchmarking.operator_tools import tensor_channel_kraus
```

```python
np.round(tensor_channel_kraus(AD_kraus,BitFlip_kraus),3)
```

    array([[[0.894, 0.   , 0.   , 0.   ],
            [0.   , 0.894, 0.   , 0.   ],
            [0.   , 0.   , 0.849, 0.   ],
            [0.   , 0.   , 0.   , 0.849]],

           [[0.   , 0.   , 0.283, 0.   ],
            [0.   , 0.   , 0.   , 0.283],
            [0.   , 0.   , 0.   , 0.   ],
            [0.   , 0.   , 0.   , 0.   ]],

           [[0.   , 0.447, 0.   , 0.   ],
            [0.447, 0.   , 0.   , 0.   ],
            [0.   , 0.   , 0.   , 0.424],
            [0.   , 0.   , 0.424, 0.   ]],

           [[0.   , 0.   , 0.   , 0.141],
            [0.   , 0.   , 0.141, 0.   ],
            [0.   , 0.   , 0.   , 0.   ],
            [0.   , 0.   , 0.   , 0.   ]]])

### Validate quantum channels are physical

When doing process tomography sometimes the estimates returned by various estimation methods can result in unphysical processes.

The functions below can be used to check if the estimates are physical.

As a starting point, we might want to check if a process specified by Kraus operators is valid.

Unless a process is unitary you need more than one Kraus operator to be a valid quantum operation.

```python
from forest.benchmarking.operator_tools import kraus_operators_are_valid

kraus_operators_are_valid(AD_kraus[0])
```

    False

However a full set is valid:

```python
kraus_operators_are_valid(AD_kraus)
```

    True

We can also validate other properties of quantum channels such as completely positivity and trace preservation. This is done on the **Choi** representation, so you many need to convert your quantum operation to the Choi representation first.

```python
from forest.benchmarking.operator_tools import (choi_is_unitary,
                                                choi_is_unital,
                                                choi_is_trace_preserving,
                                                choi_is_completely_positive,
                                                choi_is_cptp)
```

```python
# amplitude damping is not unitary
print(choi_is_unitary(AD_choi),'\n')

# amplitude damping is not unital
print(choi_is_unital(AD_choi))
```

    False

    False

```python
# amplitude damping is trace preserving (TP)
print(choi_is_trace_preserving(AD_choi),'\n')

# amplitude damping is completely positive (CP)
print(choi_is_completely_positive(AD_choi), '\n')

# amplitude damping is CPTP
print(choi_is_cptp(AD_choi))
```

    True

    True

    True

### Project an unphysical state to the closest physical state

```python
from forest.benchmarking.operator_tools.project_state_matrix import project_state_matrix_to_physical
```

```python
# Test the method. Example from fig 1 of maximum likelihood minimum effort
# https://doi.org/10.1103/PhysRevLett.108.070502

eigs = np.diag(np.array(list(reversed([3.0/5, 1.0/2, 7.0/20, 1.0/10, -11.0/20]))))
phys = project_state_matrix_to_physical(eigs)
np.allclose(phys, np.diag([0, 0, 1.0/5, 7.0/20, 9.0/20]))
```

    True

```python
from forest.benchmarking.plotting import hinton

rho_unphys = np.random.uniform(-1, 1, (2, 2)) \
    * np.exp(1.j * np.random.uniform(-np.pi, np.pi, (2, 2)))
rho_phys = project_state_matrix_to_physical(rho_unphys)

fig, (ax1, ax2) = plt.subplots(1, 2)
hinton(rho_unphys, ax=ax1)
hinton(rho_phys, ax=ax2)
ax1.set_title('Unphysical')
ax2.set_title('Physical projection')
fig.tight_layout()
```

### Project unphysical channels to physical channels

When doing process tomography often the estimates returned by maximum likelihood estimation or linear inversion methods can result in unphysical processes.

The functions below can be used to project the unphysical estimates back to physical estimates.

```python
from forest.benchmarking.operator_tools.project_superoperators import (proj_choi_to_completely_positive,
                                                                       proj_choi_to_trace_non_increasing,
                                                                       proj_choi_to_trace_preserving,
                                                                       proj_choi_to_physical,
                                                                       proj_choi_to_unitary)
```

```python
neg_Id_choi = -kraus2choi(I)
```

```python
proj_choi_to_completely_positive(neg_Id_choi)
```

    array([[0., 0., 0., 0.],
           [0., 0., 0., 0.],
           [0., 0., 0., 0.],
           [0., 0., 0., 0.]])

```python
proj_choi_to_trace_non_increasing(neg_Id_choi)
```

    array([[-1., -0., -0., -1.],
           [-0., -0., -0., -0.],
           [-0., -0., -0., -0.],
           [-1., -0., -0., -1.]])

```python
proj_choi_to_trace_preserving(neg_Id_choi)
```

    array([[ 0.,  0., -0., -1.],
           [ 0.,  1., -0., -0.],
           [-0., -0.,  1.,  0.],
           [-1., -0.,  0.,  0.]])

```python
proj_choi_to_physical(neg_Id_choi)
```

    array([[ 0.33398438,  0.        ,  0.        , -0.33203125],
           [ 0.        ,  0.66601562,  0.        ,  0.        ],
           [ 0.        ,  0.        ,  0.66601562,  0.        ],
           [-0.33203125,  0.        ,  0.        ,  0.33398438]])

```python
# closer to identity
proj_choi_to_unitary(kraus2choi(bit_flip_kraus(0.1)))
```

    array([[1.+0.j, 0.+0.j, 0.+0.j, 1.+0.j],
           [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
           [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
           [1.+0.j, 0.+0.j, 0.+0.j, 1.+0.j]])

```python
# closer to X gate
proj_choi_to_unitary(kraus2choi(bit_flip_kraus(0.9)))
```

    array([[0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
           [0.+0.j, 1.+0.j, 1.+0.j, 0.+0.j],
           [0.+0.j, 1.+0.j, 1.+0.j, 0.+0.j],
           [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]])

### Validate operators

A lot of the work in validating the physicality of quantum channels comes down to validating properties of matrices:

```python
from forest.benchmarking.operator_tools.validate_operator import (is_square_matrix,
                                                                  is_identity_matrix,
                                                                  is_idempotent_matrix,
                                                                  is_unitary_matrix,
                                                                  is_positive_semidefinite_matrix)
```

```python
# a vector is not square
is_square_matrix(np.array([[1], [0]]))
```

    False

```python
# NBVAL_RAISES_EXCEPTION
# the line above is for testing purposes, do not remove.

# a tensor is not a matrix

tensor = np.ones(8).reshape(2,2,2)
print(tensor)

is_square_matrix(tensor)
```

```python
is_identity_matrix(X)
```

    False

```python
projector_zero = np.array([[1, 0], [0, 0]])

is_idempotent_matrix(projector_zero)
```

    True

```python
is_unitary_matrix(AD_kraus[0])
```

    False

```python
is_positive_semidefinite_matrix(I)
```

    True

```python

```


# Solving initial value problems (IVPs) in `quantecon`

## David R. Pugh

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

# comment out if you don't want plots rendered in notebook
%matplotlib inline
```

# 1. Introduction
This notebook demonstrates how to solve initial value problems (IVPs) using the [`quantecon`](http://jstac.github.io/quant-econ/) Python library. Before demonstrating how one might solve an IVP using `quantecon`, I provide formal definitions for ordinary differential equations (ODEs) and initial value problems (IVPs), as well as a short discussion of finite-difference methods that will be used to solve IVPs.

## Ordinary differential equations (ODE)
An [ordinary differential equation (ODE)](http://en.wikipedia.org/wiki/Ordinary_differential_equation) is in equation of the form

\begin{equation}
\textbf{y}'= \textbf{f}(t ,\textbf{y}) \tag{1.1}
\end{equation}

where $\textbf{f}:[t_0,\infty) \times \mathbb{R}^n\rightarrow\mathbb{R}^n$.  In the case where $n=1$, then equation 1.1 reduces to a single ODE; when $n>1$, equation 1.1 defines a system of ODEs. ODEs are one of the most basic examples of functional equations: a solution to equation 1.1 is a function $\textbf{y}(t): D \subset \mathbb{R}\rightarrow\mathbb{R}^n$. There are potentially an infinite number of solutions to the ODE defined in equation 1.1. In order to reduce the number of potentially solutions, we need to impose a bit more structure on the problem.

## Initial value problems (IVPs)
An [initial value problem (IVP)](http://en.wikipedia.org/wiki/Initial_value_problem) has the form

\begin{equation}
\textbf{y}'= \textbf{f}(t ,\textbf{y}),\ t \ge t_0,\ \textbf{y}(t_0) = \textbf{y}_0 \tag{1.2}
\end{equation}

where $\textbf{f}:[t_0,\infty) \times \mathbb{R}^n\rightarrow\mathbb{R}^n$ and the initial condition $\textbf{y}_0 \in \mathbb{R}^n$ is a given vector. Alternatively, I could also specify an initial value problem by imposing a terminal condition of the form $\textbf{y}(T) = \textbf{y}_T$. The key point is that the solution $\textbf{y}(t)$ is pinned down at one $t\in[t_0, T]$.

The solution to the IVP defined by equation 1.2 is the function $\textbf{y}(t): [t_0,T] \subset \mathbb{R}\rightarrow\mathbb{R}^n$ that satisfies the initial condition $\textbf{y}(t_0) = \textbf{y}_0$.  So long as the function $\textbf{f}$ is [reasonably well-behaved](http://en.wikipedia.org/wiki/Ordinary_differential_equation#Existence_and_uniqueness_of_solutions), the solution $\textbf{y}(t)$ exists and is unique.

## Finite-difference methods
[Finite-difference methods](http://en.wikipedia.org/wiki/Finite_difference_method) are perhaps the most commonly used class of numerical methods for approximating solutions to IVPs. The basic idea behind all finite-difference  methods is to construct a difference equation

\begin{equation}
\frac{\textbf{y}(t_i + h) - \textbf{y}_i}{h} \approx \textbf{y}'(t_i) = \textbf{f}(t_i ,\textbf{y}(t_i)) \tag{1.3}
\end{equation}

which is "similar" to the differential equation at some grid of values $t_0 < \dots < t_N$. Finite-difference methods then "solve" the original differential equation by finding for each $n=0,\dots,N$ a value $\textbf{y}_n$ that approximates the value of the solution $\textbf{y}(t_n)$.

It is important to note that finite-difference methods only approximate the solution $\textbf{y}$ at the $N$ grid points. In order to approximate $\textbf{y}$ between grid points one must resort to some form of [interpolation](http://en.wikipedia.org/wiki/Interpolation).

The literature on finite-difference methods for solving IVPs is vast and there are many excellent reference texts. Those interested in a more in-depth treatment of these topics, including formal proofs of convergence, order, and stability of the numerical methods used in this notebook, should consult (Hairer, 1993), (Butcher, 2008), (Iserles, 2009). Chapter 10 of (Judd, 1998) covers a subset of these more formal texts with a specific focus on economic applications.

# 2. Examples
The remainder of this notebook demonstrates the usage and functionality of the `quantecon.ivp` module by way of example. To get started, we need to import the `quantecon.ivp module`...

```python
from quantecon import ivp
```

## 2.1 Lotka-Volterra "Predator-Prey" model

We begin with the Lotka-Volterra model, also known as the predator-prey model, which is a pair of first order, non-linear, differential equations frequently used to describe the dynamics of biological systems in which two species interact, one a predator and the other its prey. The model was proposed independently by Alfred J. Lotka in 1925 and Vito Volterra in 1926.

\begin{align}
    \frac{du}{dt} =& au - buv \tag{2.1.1} \\
    \frac{dv}{dt} =& -cv + dbuv \tag{2.1.2}
\end{align}

where $u$ is the number of preys (for example, rabbits), $v$ is the number of predators (for example, foxes) and $a, b, c, d$ are constant parameters defining the behavior of the population.

Parameter definitions are as follows:
* $a$: the natural growing rate of prey in the absence of predators.
* $b$: the natural dying rate of prey due to predation.
* $c$: the natural dying rate of predators, in teh absence of prey.
* $d$: the factor describing how many caught prey is necessary to create a new predator.

I will use $\textbf{y}=[u, v]$ to describe the state of both populations.

### 2.1.1 Defining an instance of the `IVP` class

First, we need to create an instance of the IVP class representing the Lotka-Volterra "Predator-Prey" model. To initialize an instance of the IVP class we need to define the following...

* ``f`` : Callable of the form ``f(t, y, *f_args)``. The function ``f`` is the right hand side of the system of equations defining the model. The independent variable, ``t``, should be a ``scalar``; ``y`` is an ``ndarray`` of dependent variables with ``y.shape == (n,)``. The function `f` should return a ``scalar``, ``ndarray`` or ``list`` (but not a ``tuple``).
* ``jac`` : Callable of the form ``jac(t, y, *jac_args)``, optional(default=None). The Jacobian of the right hand side of the system of equations defining the ODE.

$$ \mathcal{J}_{i,j} = \bigg[\frac{\partial f_i}{\partial y_j}\bigg] \tag {2.1.3}$$

Most all of this information can be found in the docstring for the `ivp.IVP` class.

```python
ivp.IVP?
```

From the docstring we see that we are required to define a function describing the right-hand side of the system of differential equations that we wish to solve. While, optional, it is always a good idea to also define a function describing the Jacobian matrix of partial derivatives.

For the Lotka-Volterra model, these two functions would look as follows...

```python
def lotka_volterra_system(t, y, a, b, c, d):
    """
    Return the Lotka-Voltera system.

    Parameters
    ----------
    t : float
        Time
    y : ndarray (float, shape=(2,))
        Endogenous variables of the Lotka-Volterra system. Ordering is
        `y = [u, v]` where `u` is the number of prey and `v` is the number of
        predators.
    a : float
        Natural growth rate of prey in the absence of predators.
    b : float
        Natural death rate of prey due to predation.
    c : float
        Natural death rate of predators, due to absence of prey.
    d : float
        Factor describing how many caught prey is necessary to create a new
        predator.

    Returns
    -------
    jac : ndarray (float, shape=(2,2))
        Jacobian of the Lotka-Volterra system of ODEs.

    """
    f = np.array([ a * y[0] -   b * y[0] * y[1] ,
                  -c * y[1] + d * b * y[0] * y[1] ])
    return f

def lotka_volterra_jacobian(t, y, a, b, c, d):
    """
    Return the Lotka-Voltera Jacobian matrix.

    Parameters
    ----------
    t : float
        Time
    y : ndarray (float, shape=(2,))
        Endogenous variables of the Lotka-Volterra system. Ordering is
        `y = [u, v]` where `u` is the number of prey and `v` is the number of
        predators.
    a : float
        Natural growth rate of prey in the absence of predators.
    b : float
        Natural death rate of prey due to predation.
    c : float
        Natural death rate of predators, due to absence of prey.
    d : float
        Factor describing how many caught prey is necessary to create a new
        predator.

    Returns
    -------
    jac : ndarray (float, shape=(2,2))
        Jacobian of the Lotka-Volterra system of ODEs.

    """
    jac = np.array([[a - b * y[1],   -b * y[0]],
                    [b * d * y[1],   -c + b * d * y[0]]])
    return jac
```

We can go ahead and create our instance of the `ivp.IVP` class representing the Lotka-Volterra model using the above defined functions as follows...

```python
lotka_volterra_ivp = ivp.IVP(f=lotka_volterra_system,
                             jac=lotka_volterra_jacobian)
```

### 2.1.2 Defining model parameters

In order to simulate the model, however, we will need to supply values for the model parameters $a,b,c,d$. First, let's define some "reasonable" values.

```python
# ordering is (a, b, c, d)
lotka_volterra_params = (1.0, 0.1, 1.5, 0.75)
```

In order to add these parameter values to our model we need to pass them as arguments to the `set_f_params` and `set_jac_params` methods of the newly created instance of the `ivp.IVP` class. Check the doctrings of the methods for information on the appropriate syntax...

```python
lotka_volterra_ivp.set_f_params?
```

```python
lotka_volterra_ivp.set_jac_params?
```

From the docstring we see that both the `set_f_params` and the `set_jac_params` methods take an arbitrary number of positional arguments.

```python
lotka_volterra_ivp.set_f_params(*lotka_volterra_params)
lotka_volterra_ivp.set_jac_params(*lotka_volterra_params)
```

    <quantecon.ivp.IVP at 0x7f802afb6d30>

...and we can inspect that values of these attributes and see that the return results are the same.

```python
lotka_volterra_ivp.f_params
```

    (1.0, 0.1, 1.5, 0.75)

```python
lotka_volterra_ivp.jac_params
```

    (1.0, 0.1, 1.5, 0.75)

Alternatively, we could just directly set the `f_params` and `jac_params` attributes without needing to explicitly call either the `set_f_params` and `set_jac_params` methods!

```python
# I generally prefer to set attributes directly...
lotka_volterra_ivp.f_params = lotka_volterra_params
```

```python
# ...result is the same
lotka_volterra_ivp.f_params
```

    (1.0, 0.1, 1.5, 0.75)

### 2.1.3 Using `ivp.IVP.solve` to integrate the ODE

In order to solve a system of ODEs, the `ivp.IVP.solve` method provides an interface into the ODE integration routines provided in the  `scipy.integrate.ode` module. The method takes the following parameters...

* `t0` : float. Initial condition for the independent variable.
* `y0` : array_like (float, shape=(n,)). Initial condition for the dependent variables.
* `h` : float, optional(default=1.0). Step-size for computing the solution. Can be positive or negative depending on the desired direction of integration.
* `T` : int, optional(default=None). Terminal value for the independent variable. One of either `T` or `g` must be specified.
* `g` : Callable of the form ``g(t, y, args)``, optional(default=None). Provides a stopping condition for the integration. If specified user must also specify a stopping tolerance, `tol`.
* `tol` : float, optional (default=None). Stopping tolerance for the integration. Only required if `g` is also specifed.
* `integrator` : str, optional(default='dopri5') Must be one of 'vode', 'lsoda', 'dopri5', or 'dop853'
* `step` : bool, optional(default=False) Allows access to internal steps for those solvers that use adaptive step size routines. Currently only 'vode', 'zvode', and 'lsoda' support `step=True`.
* `relax` : bool, optional(default=False) Currently only 'vode', 'zvode', and 'lsoda' support `relax=True`.
* `**kwargs` : dict, optional(default=None). Dictionary of integrator specific keyword arguments. See the Notes section of the docstring for `scipy.integrate.ode` for a complete description of solver specific keyword arguments.

... and returns:

* `solution`: array_like (float). Simulated solution trajectory.

The above information can be found in the doctring of the `ivp.IVP.solve` method.

```python
# remember...always read the docs!
ivp.IVP.solve?
```

#### Example usage
Using `dopri5`, an embedded Runge-Kutta method of order 4(5) with adaptive step size control due to [Dormand and Prince](http://en.wikipedia.org/wiki/Dormand%E2%80%93Prince_method), integrate the model forward from an initial condition of 10 rabbits and 5 foxes for 15 years.

```python
# define the initial condition...
t0, y0 = 0, np.array([10, 5])

# ...and integrate!
solution = lotka_volterra_ivp.solve(t0, y0, h=1e-1, T=15, integrator='dopri5',
                                    atol=1e-12, rtol=1e-9)
```

#### Plotting the solution

Once we have computed the solution, we can plot it using the excellent [`matplotlib`](http://matplotlib.org/) Python library.

```python
# extract the components of the solution trajectory
t = solution[:, 0]
rabbits = solution[:, 1]
foxes = solution[:, 2]

# create the plot
fig = plt.figure(figsize=(8, 6))
plt.plot(t, rabbits, 'r.', label='Rabbits')
plt.plot(t, foxes  , 'b.', label='Foxes')
plt.grid()
plt.legend(loc=0, frameon=False, bbox_to_anchor=(1, 1))
plt.xlabel('Time', fontsize=15)
plt.ylabel('Population', fontsize=15)
plt.title('Evolution of fox and rabbit populations', fontsize=20)
plt.show()
```

Note that we have plotted the time paths of rabbit and fox populations as sequences of points rather than smooth curves. This is done to visually emphasize the fact that finite-difference methods used to approximate the solution return a discrete approximation to the true continuous solution.

### 2.1.4 Using `ivp.IVP.interpolate` to interpolate the solution

The `IVP.interpolate` method provides an interface to the parametric [B-spline interpolation](http://en.wikipedia.org/wiki/B-spline) routines in `scipy.interpolate` in order to constuct a continuous approximation to the true solution. For more details on B-spline interpolation, including some additional economic applications, see chapter 6 of (Judd, 1998). The `ivp.IVP.interpolate` method takes the following parameters...

* `traj` : array_like (float)
    Solution trajectory providing the data points for constructing the
    B-spline representation.
* `ti` : array_like (float)
    Array of values for the independent variable at which to
    interpolate the value of the B-spline.
* `k` : int, optional(default=3)
    Degree of the desired B-spline. Degree must satisfy
    :math:`1 \le k \le 5`.
* `der` : int, optional(default=0)
    The order of derivative of the spline to compute (must be less
    than or equal to `k`).
* `ext` : int, optional(default=2) Controls the value of returned elements
    for outside the original knot sequence provided by traj. For
    extrapolation, set `ext=0`; `ext=1` returns zero; `ext=2` raises a
    `ValueError`.

... and returns:

* `interp_traj`: ndarray (float). The interpolated trajectory.

#### Example usage
Approximate the solution to the Lotka-Volterra model at 1000 evenly spaced points using a 5th order B-spline interpolation and no extrapolation.

```python
# define the desired interpolation points...
ti = np.linspace(t0, solution[-1, 0], 1000)

# ...and interpolate!
interp_solution = lotka_volterra_ivp.interpolate(solution, ti, k=5, ext=2)
```

#### Plotting the interpolated solution

We can now plot the interpolated solution using `matplotlib` as follows...

```python
# extract the components of the solution
ti = interp_solution[:, 0]
rabbits = interp_solution[:, 1]
foxes = interp_solution[:, 2]

# make the plot
fig = plt.figure(figsize=(8, 6))

plt.plot(ti, rabbits, 'r-', label='Rabbits')
plt.plot(ti, foxes  , 'b-', label='Foxes')

plt.xlabel('Time', fontsize=15)
plt.ylabel('Population', fontsize=15)
plt.title('Evolution of fox and rabbit populations', fontsize=20)

plt.legend(loc='best', frameon=False, bbox_to_anchor=(1,1))
plt.grid()

plt.show()
```

Note that we have plotted the time paths of rabbit and fox populations as smooth curves. This is done to visually emphasize the fact that the B-spline interpolation methods used to approximate the solution return a continuous approximation to the true continuous solution.

### 2.1.5 Assessing accuracy using `ivp.IVP.compute_residual`

After computing a continuous approximation to the solution of our IVP, it is important to verify that the computed approximation is actually a "good" approximation. To assess the accuracy of our numerical solution we first define a residual function, $R(t)$, as the difference between the derivative of the B-spline approximation of the solution trajectory, $\hat{\textbf{y}}'(t)$, and the right-hand side of the original ODE evaluated along the approximated solution trajectory.

\begin{equation}
    \textbf{R}(t) = \hat{\textbf{y}}'(t) - f(t, \hat{\textbf{y}}(t)) \tag{2.1.4}
\end{equation}

The idea is that if our numerical approximation of the true solution is "good", then this residual function should be roughly zero everywhere within the interval of approximation. The `ivp.IVP.compute_residual` method takes the following parameters...

* `traj` : array_like (float). Solution trajectory providing the data points for constructing the B-spline representation.
* `ti` : array_like (float). Array of values for the independent variable at which to interpolate the value of the B-spline.
* `k` : int, optional(default=3). Degree of the desired B-spline. Degree must satisfy $1 \le k \le 5$.
* `ext` : int, optional(default=2). Controls the value of returned elements for outside the original knot sequence provided by `traj`. For extrapolation, set `ext=0`; `ext=1` returns zero; `ext=2` raises a `ValueError`.

... and returns:

* `residual` : array (float) Difference between the derivative of the B-spline approximation of the solution trajectory and the right-hand side of the ODE evaluated along the approximated solution trajectory.

Remember to check the docstring for more information!

```python
# life will be easier if you read the docs!
lotka_volterra_ivp.compute_residual?
```

#### Example usage
Compute the residual to the Lotka-Volterra model at 1000 evenly spaced points using a 1st order B-spline interpolation (which is equivalent to linear interpolation!).

```python
# reset original parameters
lotka_volterra_ivp.f_params = lotka_volterra_params
lotka_volterra_ivp.jac_params = lotka_volterra_params

# compute the residual
residual = lotka_volterra_ivp.compute_residual(solution, ti, k=1)
```

#### Plotting the residual

In your introductory econometrics/statistics course, your professor likely implored you to "always plot your residuals!" This maxim of data analysis is no less true in numerical analysis. However, while patterns in residuals are generally a "bad" thing in econometrics/statistics (as they suggest model mispecification, or other related problems), patterns in a residual function, $\textbf{R}(t)$, in numerical analysis are generally OK (and in certain cases actually desireable!).

In this context, what is important is that overall size of the residuals is "small".

```python
# extract the raw residuals
rabbits_residual = residual[:, 1]
foxes_residual = residual[:, 2]

# typically, normalize residual by the level of the variable
norm_rabbits_residual = np.abs(rabbits_residual) / rabbits
norm_foxes_residual = np.abs(foxes_residual) / foxes

# create the plot
fig = plt.figure(figsize=(8, 6))
plt.plot(ti, norm_rabbits_residual, 'r-', label='Rabbits')
plt.plot(ti, norm_foxes_residual**2 / foxes, 'b-', label='Foxes')
plt.axhline(np.finfo('float').eps, linestyle='dashed', color='k', label='Machine eps')
plt.xlabel('Time', fontsize=15)
plt.ylim(1e-16, 1)
plt.ylabel('Residuals (normalized)', fontsize=15)
plt.yscale('log')
plt.title('Lotka-Volterra residuals', fontsize=20)
plt.grid()
plt.legend(loc='best', frameon=False, bbox_to_anchor=(1,1))
plt.show()
```

#### Understanding determinants of accuracy
We can use [IPython widgets](https://github.com/ipython/ipywidgets) to investigate the determinants of accuracy of our approximated solution. Good candidates for exploration are...

* `h`: the step size used in computing the initial finite difference solution.
* `atol`: the absolute tolerance for the solver.
* `rtol`: the relative tolerance for the solver.
* `k`: the degree of the B-spline used in the interpolation of that finite difference solution.

```python
from ipywidgets import interact
from ipywidgets.widgets import FloatText, FloatSlider, IntSlider, Text
```

```python
# reset parameters
lotka_volterra_ivp.f_params = (1.0, 0.1, 2.0, 0.75)
lotka_volterra_ivp.jac_params = lotka_volterra_ivp.f_params
```

Now we can make use of the `@interact` decorator and the various IPython widgets to create an interactive visualization of the residual plot for the Lotka-Volterra "Predator-Prey" model.

```python
@interact(h=FloatText(value=1e0), atol=FloatText(value=1e-3),
          rtol=FloatText(value=1e-3), k=IntSlider(min=1, value=3, max=5),
          integrator=Text(value='lsoda'))
def plot_lotka_volterra_residuals(h, atol, rtol, k, integrator):
    """Plots residuals of the Lotka-Volterra system."""
    # re-compute the solution
    tmp_solution = lotka_volterra_ivp.solve(t0, y0, h=h, T=15, integrator=integrator,
                                            atol=atol, rtol=rtol)

    # re-compute the interpolated solution and residual
    tmp_ti = np.linspace(t0, tmp_solution[-1, 0], 1000)
    tmp_interp_solution = lotka_volterra_ivp.interpolate(tmp_solution, tmp_ti, k=k)
    tmp_residual = lotka_volterra_ivp.compute_residual(tmp_solution, tmp_ti, k=k)

    # extract the components of the solution
    tmp_rabbits = tmp_interp_solution[:, 1]
    tmp_foxes = tmp_interp_solution[:, 2]

    # extract the raw residuals
    tmp_rabbits_residual = tmp_residual[:, 1]
    tmp_foxes_residual = tmp_residual[:, 2]

    # typically, normalize residual by the level of the variable
    tmp_norm_rabbits_residual = np.abs(tmp_rabbits_residual) / tmp_rabbits
    tmp_norm_foxes_residual = np.abs(tmp_foxes_residual) / tmp_foxes

    # create the plot
    fig = plt.figure(figsize=(8, 6))

    plt.plot(tmp_ti, tmp_norm_rabbits_residual, 'r-', label='Rabbits')
    plt.plot(tmp_ti, tmp_norm_foxes_residual**2 / foxes, 'b-', label='Foxes')
    plt.axhline(np.finfo('float').eps, linestyle='dashed', color='k', label='Machine eps')

    plt.xlabel('Time', fontsize=15)
    plt.ylim(1e-16, 1)
    plt.ylabel('Residuals (normalized)', fontsize=15)
    plt.yscale('log')
    plt.title('Lotka-Volterra residuals', fontsize=20)

    plt.grid()
    plt.legend(loc='best', frameon=False, bbox_to_anchor=(1,1))

```

#### Sensitivity to parameters
Once we have computed and plotted an approximate solution (and verified that the approximation is a good one by plotting the residual function!), we can try and learn something about the dependence of the solution on model parameters.

```python
@interact(a=FloatSlider(min=0.0, max=5.0, step=0.5, value=1.5),
          b=FloatSlider(min=0.0, max=1.0, step=0.01, value=0.5),
          c=FloatSlider(min=0.0, max=5.0, step=0.5, value=3.5),
          d=FloatSlider(min=0.0, max=1.0, step=0.01, value=0.5))
def plot_lotka_volterra(a, b, c, d):
    """Plots trajectories of the Lotka-Volterra system."""
    # update the parameters and re-compute the solution
    lotka_volterra_ivp.f_params = (a, b, c, d)
    lotka_volterra_ivp.jac_params = (a, b, c, d)
    tmp_solution = lotka_volterra_ivp.solve(t0, y0, h=1e-1, T=15, integrator='dopri5',
                                            atol=1e-12, rtol=1e-9)
    # extract the components of the solution
    tmp_t = tmp_solution[:, 0]
    tmp_rabbits = tmp_solution[:, 1]
    tmp_foxes = tmp_solution[:, 2]

    # create the plot!
    fig = plt.figure(figsize=(8, 6))

    plt.plot(tmp_t, tmp_rabbits, 'r.', label='Rabbits')
    plt.plot(tmp_t, tmp_foxes  , 'b.', label='Foxes')

    plt.xlabel('Time', fontsize=15)
    plt.ylabel('Population', fontsize=15)
    plt.title('Evolution of fox and rabbit populations', fontsize=20)

    plt.legend(loc='best', frameon=False, bbox_to_anchor=(1,1))
    plt.grid()

```

## 2.2 The Lorenz equations

The [Lorenz equations](Lorenz Equations`:http://en.wikipedia.org/wiki/Lorenz_system) are a system of three coupled, first-order, non-linear differential equations which describe the trajectory of a particle through time. The system was originally derived by as a model of atmospheric convection, but the deceptive simplicity of the equations have made them an often-used example in fields beyond atmospheric physics.

The equations describe the evolution of the spatial variables x, y, and z, given the governing parameters $\sigma, \beta, \rho$, through the specification of the time-derivatives of the spatial variables:

\begin{align}
    \frac{dx}{dt} =& \sigma(y − x) \tag{2.2.1} \\
    \frac{dy}{dt} =& x(\rho − z) − y \tag{2.2.2} \\
    \frac{dz}{dt} =& xy − \beta z \tag{2.2.3}
\end{align}

The resulting dynamics are entirely deterministic giving a starting point $(x_0,y_0,z_0)$. Though it looks straightforward, for certain choices of the parameters, the trajectories become chaotic, and the
resulting trajectories display some surprising properties.

### 2.2.1 Incorporating SymPy
While deriving the Jacobian matrix by hand is trivial for most simple 2D or 3D systems, it can quickly become tedious and error prone for larger systems (or even some highly non-linear 2D systems!). In addition to being an important input to most ODE integrators/solvers, Jacobians are also useful for assessing the stability properties of equilibria.

An alternative approach for solving IVPs using the `ivp` module, which leverages the [SymPy](http://sympy.org/en/index.html) Python library to do the tedious computations involved in deriving the Jacobian, is as follows.

1. Define the IVP using SymPy.
2. Use SymPy routines for computing the Jacobian.
3. Wrap the symbolic expressions as callable NumPy functions.
4. Use these functions to create an instance of the `ivp.IVP` class.

The remainder of this notebook implements each of these steps to solve and analyze the Lorenz equations defined above.

#### Step 1: Defining the Lorenz equations using SymPy
We begin by defining a `sp.Matrix` instance containing the three Lorenz equations...

```python
# enables sympy LaTex printing...
sp.init_printing()
```

```python
# declare endogenous variables
t, x, y, z = sp.var('t, x, y, z')

# declare model parameters
beta, rho, sigma = sp.var('beta, rho, sigma')

# define symbolic model equations
_x_dot = sigma * (y - x)
_y_dot = x * (rho - z) - y
_z_dot = x * y - beta * z

# define symbolic system and compute the jacobian
_lorenz_system = sp.Matrix([[_x_dot], [_y_dot], [_z_dot]])
```

Let's take a check out our newly defined `_lorenz_system` and make sure it looks as expected...

```python
_lorenz_system
```

#### Step 2: Computing the Jacobian using SymPy
Once we have defined our model as a SymPy matrix, computing the Jacobian is trivial...

```python
_lorenz_jacobian = _lorenz_system.jacobian([x, y, z])
_lorenz_jacobian
```

#### Step 3: Wrap the SymPy expression to create vectorized NumPy functions
Now we wrap the SymPy matrices defining the model and the Jacobian to create vectorized NumPy functions. It is crucial that the interface for our wrapped functions matches the interface required by the `f` and `jac` parameters which we will pass to the `ivp.IVP` constructor to create an instance of the `ivp.IVP` class representing the Lorenz equations.

Recall from the `ivp.IVP` docstring that...

    f : callable `f(t, y, *f_args)`
        Right hand side of the system of equations defining the ODE. The
        independent variable, `t`, is a `scalar`; `y` is an `ndarray`
        of dependent variables with `y.shape == (n,)`. The function `f`
        should return a `scalar`, `ndarray` or `list` (but not a
        `tuple`).
    jac : callable `jac(t, y, *jac_args)`, optional(default=None)
        Jacobian of the right hand side of the system of equations defining
        the ODE.

        .. :math:

            \mathcal{J}_{i,j} = \bigg[\frac{\partial f_i}{\partial y_j}\bigg]

Thus our wrapped functions need to take a float `t` as the first argument, and array `y` as a second argument, followed by some arbitrary number of model parameters. We can handle all of this as follows...

```python
# in order to pass an array as an argument, we need to apply a change of variables
X = sp.DeferredVector('X')
change_of_vars = {'x': X[0], 'y': X[1], 'z': X[2]}
_transformed_lorenz_system = _lorenz_system.subs(change_of_vars)
_transformed_lorenz_jacobian = _transformed_lorenz_system.jacobian([X[0], X[1], X[2]])

# wrap the symbolic expressions as callable numpy funcs
_args = (t, X, beta, rho, sigma)
_f = sp.lambdify(_args, _transformed_lorenz_system,
                 modules=[{'ImmutableMatrix': np.array}, "numpy"])
_jac = sp.lambdify(_args, _transformed_lorenz_jacobian,
                   modules=[{'ImmutableMatrix': np.array}, "numpy"])
```

#### Step 4: Use these functions to create an instance of the `IVP` class
First we define functions describing the right-hand side of the ODE and the Jacobian which we need to initialize the `ivp.IVP` class...

```python
def lorenz_system(t, X, beta, rho, sigma):
    """
    Return the Lorenz system.

    Parameters
    ----------
    t : float
        Time
    X : ndarray (float, shape=(3,))
        Endogenous variables of the Lorenz system.
    beta : float
        Model parameter. Should satisfy :math:`0 < \beta`.
    rho : float
        Model parameter. Should satisfy :math:`0 < \rho`.
    sigma : float
        Model parameter. Should satisfy :math:`0 < \sigma`.

    Returns
    -------
    rhs_ode : ndarray (float, shape=(3,))
        Right hand side of the Lorenz system of ODEs.

    """
    rhs_ode = _f(t, X, beta, rho, sigma).ravel()
    return rhs_ode

def lorenz_jacobian(t, X, beta, rho, sigma):
    """
    Return the Jacobian of the Lorenz system.

    Parameters
    ----------
    t : float
        Time
    X : ndarray (float, shape=(3,))
        Endogenous variables of the Lorenz system.
    beta : float
        Model parameter. Should satisfy :math:`0 < \beta`.
    rho : float
        Model parameter. Should satisfy :math:`0 < \rho`.
    sigma : float
        Model parameter. Should satisfy :math:`0 < \sigma`.

    Returns
    -------
    jac : ndarray (float, shape=(3,3))
        Jacobian of the Lorenz system of ODEs.

    """
    jac = _jac(t, X, beta, rho, sigma)
    return jac
```

... next we define a tuple of model parameters...

```python
# parameters with ordering (beta, rho, sigma)
lorenz_params = (2.66, 28.0, 10.0)

```

... finally, we are ready to create the instance of the `ivp.IVP` class representing the Lorenz equations.

```python
# create the instance
lorenz_ivp = ivp.IVP(f=lorenz_system,
                     jac=lorenz_jacobian)

# specify the params
lorenz_ivp.f_params = lorenz_params
lorenz_ivp.jac_params = lorenz_params
```

### 2.2.2 Solving the Lorenz equations
At this point I proceed in exactly the same fashion as in the previous Lotka-Volterra equations example:

1. Solve the model using a discretized, finite-difference approximation.
2. Use the discretized approximation in conjunction with parametric B-spline interpolation to construct a continuous approximation of the true solution.
3. Compute and analyze the residual of the approximate solution.

#### Step 1. Solve the model using a discretized, finite-difference approximation
Using `dop853`, an embedded Runge-Kutta method of order 7(8) with adaptive step size control due to [Dormand and Prince](http://en.wikipedia.org/wiki/Dormand%E2%80%93Prince_method), integrate the Lorenz equations forward from an initial condition of $X_0 = (1.0, 1.0, 1.0)$ from $t=0$, to $T=100$.

```python
# declare and initial condition
t0, X0 = 0.0, np.array([1.0, 1.0, 1.0])

# solve!
solution = lorenz_ivp.solve(t0, X0, h=1e-2, T=100, integrator='dop853',
                            atol=1e-12, rtol=1e-9)
```

#### Plotting the solution time paths
We can use IPython widgets to construct a "poor man's" animation of the evolution of the Lorenz equations.

```python
@interact(T=IntSlider(min=0, value=0, max=solution.shape[0], step=5))
def plot_lorenz(T):
    """Plots the first T points in the solution trajectory of the Lorenz equations."""
    # extract the components of the solution trajectory
    t = solution[:T, 0]
    x_vals = solution[:T, 1]
    y_vals = solution[:T, 2]
    z_vals = solution[:T, 3]

    # create the plot
    fig = plt.figure(figsize=(8, 6))
    plt.plot(t, x_vals, 'r.', label='$x_t$', alpha=0.5)
    plt.plot(t, y_vals  , 'b.', label='$y_t$', alpha=0.5)
    plt.plot(t, z_vals  , 'g.', label='$z_t$', alpha=0.5)

    plt.grid()
    plt.xlabel('Time', fontsize=20)
    plt.ylabel('$x_t, y_t, z_t$', fontsize=20, rotation='horizontal')
    plt.title('Time paths of the $x,y,z$ coordinates', fontsize=25)
    plt.legend(frameon=False, bbox_to_anchor=(1.15,1))

    plt.show()
```

#### Step 2: Construct a continuous approximation to the solution.

Let's construct a continuous approximation to the solution of the Lorenz equations at 10000 evenly spaced points using a 5th order B-spline interpolation.

```python
# define the desired interpolation points...
ti = np.linspace(t0, solution[-1, 0], 1e4)

# ...and interpolate!
interp_solution = lorenz_ivp.interpolate(solution, ti, k=5, ext=2)
```

#### Plotting 2D projections of the solution in phase space
The underlying structure of the Lorenz system becomes more apparent when I plot the interpolated solution trajectories in phase space. We can plot 2D projections of the time paths of the solution in phase space...

```python
# extract the components of the solution trajectory
t = solution[:, 0]
x_vals = interp_solution[:, 1]
y_vals = interp_solution[:, 2]
z_vals = interp_solution[:, 3]

# xy phase space projection
fig, axes = plt.subplots(1, 3, figsize=(12, 6), sharex=True, sharey=True, squeeze=False)
axes[0,0].plot(x_vals, y_vals, 'r', alpha=0.5)
axes[0,0].set_xlabel('$x$', fontsize=20, rotation='horizontal')
axes[0,0].set_ylabel('$y$', fontsize=20, rotation='horizontal')
axes[0,0].set_title('$x,y$-plane', fontsize=20)
axes[0,0].grid()

# xz phase space projection
axes[0,1].plot(x_vals, z_vals  , 'b', alpha=0.5)
axes[0,1].set_xlabel('$x$', fontsize=20, rotation='horizontal')
axes[0,1].set_ylabel('$z$', fontsize=20, rotation='horizontal')
axes[0,1].set_title('$x,z$-plane', fontsize=20)
axes[0,1].grid()

# yz phase space projection
axes[0,2].plot(y_vals, z_vals  , 'g', alpha=0.5)
axes[0,2].set_xlabel('$y$', fontsize=20, rotation='horizontal')
axes[0,2].set_ylabel('$z$', fontsize=20, rotation='horizontal')
axes[0,2].set_title('$y,z$-plane', fontsize=20)
axes[0,2].grid()

plt.suptitle('Phase space projections', x=0.5, y=1.05, fontsize=25)
plt.tight_layout()
plt.show()
```

#### Step 3: Compute the residuals to assese the accuracy of our solution

Finally, to assess the accuracy of our solution we need to compute and plot the solution residuals at 10000 evenly spaced points using a 5th order B-spline interpolation.

```python
# compute the residual
ti = np.linspace(0, solution[-1,0], 10000)
residual = lorenz_ivp.compute_residual(solution, ti, k=5)
```

Again, we want to confirm that the residuals are "small" everywhere. Patterns, if they exists, are not cause for concern.

```python
# extract the raw residuals
x_residuals = residual[:, 1]
y_residuals = residual[:, 2]
z_residuals = residual[:, 3]

# typically, normalize residual by the level of the variable
norm_x_residuals = np.abs(x_residuals) / x_vals
norm_y_residuals = np.abs(y_residuals) / y_vals
norm_z_residuals = np.abs(z_residuals) / z_vals

# create the plot
fig = plt.figure(figsize=(8, 6))
plt.plot(ti, norm_x_residuals, 'r-', label='$x(t)$', alpha=0.5)
plt.plot(ti, norm_y_residuals, 'b-', label='$y(t)$', alpha=0.5)
plt.plot(ti, norm_z_residuals, 'g-', label='$z(t)$', alpha=0.5)
plt.axhline(np.finfo('float').eps, linestyle='dashed', color='k', label='Machine eps')
plt.xlabel('Time', fontsize=15)
plt.ylim(1e-16, 1)
plt.ylabel('Residuals', fontsize=15)
plt.yscale('log')
plt.title('Lorenz equations residuals', fontsize=20)
plt.grid()
plt.legend(loc='best', frameon=False)
plt.show()
```

```python

```


# Replication Carneiro & Heckman & Vytlacil 2011

This notebook replicates results from

Carneiro, P., Heckman, J. J., & Vytlacil, E. J. (2011). [Estimating marginal returns to education.](https://pubs.aeaweb.org/doi/pdfplus/10.1257/aer.101.6.2754)
*American Economic Review, 101*(6), 2754-81.

For this purpose we rely on *grmpy*. It is an open-source Python package for the simulation and estimation of the generalized Roy model. In addition it serves as a teaching tool to promote the conceptual framework of the generalized Roy model, illustrate a variety of issues in the econometrics of policy evaluation, and showcase basic software engineering practices.

# Motivation

The authors emphasize that the evaluation of marginal returns to policy interventions is one of the central tasks for evaluating the optimal size of social programs. They state that this is an relatively simple task if at least one of the following conditions is satisfied:

**1.) The effect of a policy intervention is homogeneous.**

**2.) The effect of a policy intervention is heterogeneous given observable variables but agent select not into treatment based on unobservable returns**

Under condition 1 or 2 it is possible to estimate the mean marginal return to schooling via a simplified mincer framework. Possible selection problems (the decision to attend schooling isn't exogenous) can be solved by applying methods like Instrumental variables.
Nevertheless the application of empirically tests for this conditions provide evidence for heterogeneous marginal returns. In addition it seems like agents would select into treatment based on at least partly knowledge about idiosyncratic returns. Therefore the authors make use of a local instrumental variable approach to estimate the marginal return to college attendance within the framework of the generalized Roy model.

# Theoretical Framework

Since the current version of grmpy is not capable of estimating non-parametric versions of the Roy models, our
replication of Carneiro et al. (2011) focuses on reproducing the results for the linear-in-parameter version of the model. It is characterized by the following set of equations:

\begin{align*}
 &\textbf{Potential Outcomes} & & \textbf{Choice}  &\\
 & Y_1 = \delta_1 X + U_{1} &  &D^{*} = \gamma Z - V&\\
 & Y_0 = \delta_0 X + U_{0} &  &D_i = \mathbf{1}\{D^{*} >0\} =  \mathbf{1}\{\Phi(\gamma Z) > u_D\}&\\
&&& \text{with $u_D = \Phi(V)$}  &\\
&&&&\\
&\textbf{Distributional Characteristics}&&&\\
&\{U_{1}, U_{0}, V\} \sim \mathcal{N}\left(0, \Sigma\right)&&\Sigma =  \begin{bmatrix}
    \sigma_1^{2} & \sigma_{1,0} & \sigma_{1,V} \\
    \sigma_{1,0} & \sigma_0^{2} & \sigma_{0,V} \\
    \sigma_{1,V} & \sigma_{0,V} & \sigma_V^{2} \\
  \end{bmatrix}&\\
&&&&\\
& \textbf{Observed Outcome} &&&\\
& Y = D Y_1 + (1-D) Y_0 &&&
\end{align*}

$Y_1$ and $Y_0$  are objective outcomes associated with each potential treatment state $D$ and realized after the treatment decision.
From the perspective of the econometrician, $(X,Z)$ are observable while $(U_1,U_0,V)$ are not. $X$ are the observed determinants of potential outcomes $(Y_1,Y_0)$, and $Z$ are the observed determinants of the latent indicator variable $D^{*}$. It is assumed that the unobservable parameters $(U_1,U_0,V)$ follow a joint normal distribution with mean zero and covariance matrix $\Sigma$.

$Y$ illustrates the observable outcome of the individual. Additionally its functional form emphasizes a crucial issue in the econometrics of policy evaluation which is referred to as the **Evaluation Problem**. The problem arises because the econometrician is limited to observe either $Y_1$ or $Y_0$ for the individual dependent on the selection decision $D$. Thus, the effect of treatment cannot be determined on an individual level. If the treatment choice $D$ depends on the potential outcomes, then there is also a selection problem (**Essential Heterogeneity**). If that is the case, then the treated and untreated sub populations differ not only in their treatment status but in other characteristics as well. A naive comparison of the treated and untreated leads to misleading conclusions.

# The Marginal Treatment Effect

The concept of the Marginal Treatment Effect, henceforth $B^{MTE}$, was first introduced by Björklund and Moffit (1987). It is defined as the gain from treatment to an indifferent individual given a specific degree of unobservable desire for selecting into treatment. Adding our functional form assumptions from above allows us to express $B^{MTE}$ via the following equation

\begin{align}
B^{MTE}(x,u_D) &= E\left[Y_1 - Y_0 | X=x, U_D = u_D\right]\\
& = \left(\delta_1 - \delta_0\right)x + \left(\sigma_{1,V} - \sigma_{0,V}\right) \phi^{-1}(u_D)\\
\end{align}

Note that $B^{MTE}$ does not assign explicitly one value as an effect, but rather provides a continuum of effects along the distribution of the unobservable parameter $V$. Additionally it provides the framework for deriving a broad variation of different treatment effects by applying effect specific weights along the distribution of $V$ on the marginal effect.

Next we import necessary packages:

```python
%load_ext autoreload

import pandas as pd
import grmpy

from auxiliary.auxiliary import process_data
from auxiliary.auxiliary import plot_est_mte
```

Due to data privace issues, we are not able to merge the data provided by the authors so that they fully coincide with the original data set. Therefore our replication setup makes use of a mock data set. For this purpose we randomly merge the individual specific data with the local characteristics.

```python
basic = pd.read_stata('data/basicvariables.dta')
local = pd.read_stata('data/localvariables.dta')
df = pd.concat([basic, local], axis = 1)
process_data(df,'data/aer-replication-mock')
```

In the next step we have to create a initialization file that fully coincides with the setup by Carneiro et. al. (2011). The file is based on the information that the authors provide in their [online appendix](https://assets.aeaweb.org/assets/production/articles-attachments/aer/data/oct2011/20061111_app.pdf).

We then conduct an estimation based on the initialization file.

```python
rslt = grmpy.fit('files/replication.grmpy.yml')
```

Next we plot the Marginal Treatment Effect based on our estimation results and compare it with the original values.

```python
mte = plot_est_mte(rslt, 'files/replication.grmpy.yml')
```

# Results

The figure above presents $B^{MTE}$ as well as the associated 90% confidence interval. The y-axis shows the the marginal return for one year of college attendance evaluated on $x=E\left[X\right]$. The x-axis shows $u_D$, the corresponding quantile within the distribution of $V$ in which the individual's is located. Note that $V$ enters the choice equation with a negative sign. This means that agents with a relatively low value of $u_D$ are more likely to select into treatment compared to individuals with a high values. Accordingly the decreasing curve is a sign for selection on gains.

The graph also illustrates that the return to college attendance is subjected to a significant degree of heterogeneity along the distribution of the unobservable $V$. Individuals which are more likely to attend college (low values of $u_D$) are are able to realize returns up to $28\%$ per year. In contrary college attendance would actually harm individuals with a high values of $u_D$ because they face returns up to $-15\%$ per year.

# References

- Abbring, J. H. & Heckman, J. J. (2007). [Econometric Evaluation of Social Programs, Part III: Distributional Treatment Effects, Dynamic Treatment Tffects, Dynamic Discrete Choice, and General Equilibrium Policy Evaluation.](https://www.sciencedirect.com/science/article/pii/S1573441207060722) *Handbook of Econometrics, Volume 6B*, Chapter 72, p. 5145-5303. Elsevier Science.

- Carneiro, P., Heckman, J. J., & Vytlacil, E. J. (2011). [Estimating marginal returns to education.](https://pubs.aeaweb.org/doi/pdfplus/10.1257/aer.101.6.2754)
*American Economic Review, 101(6)*, 2754-81.

- Heckman, J. J., Schmierer D., & Urzua S. (2010). [Testing the Correlated Random Coefficient Model.](https://www.sciencedirect.com/science/article/pii/S0304407610000084) *Journal of Econometrics, 158(2)*, 177-203

- Heckman, J. J., & Vytlacil, E. J. (2005). [Structural Equations, Treatment Effects, and
Econometric Policy Evaluation.](https://www.jstor.org/stable/3598865?) *Econometrica, 73(3), 669-738.

- Heckman, J. J., Vytlacil., E. J. (2007). [Econometric Evaluation of Social Programs, Part I: Causal Models, Structural Models and Econometric Policy Evaluation.](https://www.sciencedirect.com/science/article/pii/S1573441207060709) In *Handbook of Econometrics*, Volume 6B, chapter 70, pages 4779–4874. Elsevier Science, 2007.

- Heckman, J. J. & Vytlacil, E. J. (2007). [Econometric Evaluation of Social Programs, Part II: Using the Marginal Treatment Effect to Organize Alternative Econometric Estimators to Evaluate Social Programs, and to Forecast their Effects in New Environments.](https://www.sciencedirect.com/science/article/pii/S1573441207060710) In *Handbook of Econometrics*, Volume 6B, chapter 71, pages 4875–5143. Elsevier Science.

```python

```


```python
# General import

import numpy as np
import scipy.sparse as sparse
from scipy.integrate import ode
import time
import matplotlib.pyplot as plt
```

```python
# pyMPC import

from pyMPC.mpc import MPCController
```

## System dynamics ##

The system to be controlled is an inverted pendulum on a cart (see next Figure).

The system is governed by the following differential equations:

\begin{equation}
 \begin{aligned}
 (M+m)\ddot p + ml\ddot\phi \cos\phi - ml \dot \phi ^2 \sin \phi + b\dot p &= F \\
 l \ddot \phi + \ddot p \cos \phi - g \sin\phi &= -f_\phi\dot \phi
\end{aligned}
\end{equation}

Introducing the state vector $x=[p\; \dot p\; \phi\; \dot \phi]$ and the input $u=F$, the system dynamics are described in state-space by a set of an nonlinear ordinary differential equations: $\dot x = f(x,u)$ with

\begin{equation}
\begin{split}
  f(x,u) &=
 \begin{bmatrix}
 x_2\\
 \frac{-mg \sin x_3\cos x_3 + mlx_4^3\sin x_3 + f_\phi m x_4 \cos x_3 - bx_2 + u }{M+(1-\cos^2 x_3)m}\\
 x_3\\
 \frac{(M+m)(g \sin x_3 - f_\phi x_4) - (lm x_4^2 \sin x_3 - bx_2 + u)\cos x_3}{l(M+(1-\cos^2 x_3)m)}
 \end{bmatrix}\\
 \end{split}
 \end{equation}

For MPC control design, the system is linearized about the upright (unstable) equilibrium point, i.e., about the point $x_{eq} = [0, \; 0\;, 0,\; 0]^\top$.
The linearized system has form $\dot  x =  A_c x + B_c u$ with

\begin{equation}
  A =
 \begin{bmatrix}
 0& 1& 0& 0\\
 0& -\frac{b}{M}& -g\frac{m}{M}& f_\theta\frac{m}{M}\\
 0&0&0&1\\
 0&\frac{b}{Ml}& \frac{g(M+m)}{Ml}&-\frac{(M+m)f_\theta}{M l}
 \end{bmatrix},\qquad B=
 \begin{bmatrix}
 0\\
 \frac{1}{M}\\
 0\\
 -\frac{1}{Ml}&
 \end{bmatrix}
 \end{equation}

Next, the system is discretized with sampling time $T_s = 50\;\text{ms}$. Here we just use a Forward Euler dsicretization scheme for the sake of simplicity.

```python
# Constants #

M = 0.5
m = 0.2
b = 0.1
ftheta = 0.1
l = 0.3
g = 9.81

Ts = 50e-3
```

```python
# System dynamics: \dot x = f_ODE(t,x,u)

def f_ODE(t,x,u):
    F = u
    v = x[1]
    theta = x[2]
    omega = x[3]
    der = np.zeros(4)
    der[0] = v
    der[1] = (m * l * np.sin(theta) * omega ** 2 - m * g * np.sin(theta) * np.cos(theta) + m * ftheta * np.cos(theta) * omega + F - b * v) / (M + m * (1 - np.cos(theta) ** 2))
    der[2] = omega
    der[3] = ((M + m) * (g * np.sin(theta) - ftheta * omega) - m * l * omega ** 2 * np.sin(theta) * np.cos(theta) - (F - b * v) * np.cos(theta)) / (l * (M + m * (1 - np.cos(theta) ** 2)))
    return der
```

```python
# Linearized System Matrices

# Continuous-time system matrices, linearized about the upright, unstable equilibrium point
Ac =np.array([[0,       1,          0,                  0],
              [0,       -b/M,       -(g*m)/M,           (ftheta*m)/M],
              [0,       0,          0,                  1],
              [0,       b/(M*l),    (M*g + g*m)/(M*l),  -(M*ftheta + ftheta*m)/(M*l)]])

Bc = np.array([
    [0.0],
    [1.0/M],
    [0.0],
    [-1/(M*l)]
])

[nx, nu] = Bc.shape # number of states and number or inputs

# Simple forward euler discretization
Ad = np.eye(nx) + Ac*Ts
Bd = Bc*Ts
```

```python
# MPC reference input and states (set-points)

xref = np.array([0.3, 0.0, 0.0, 0.0]) # reference state
uref = np.array([0.0])    # reference input
uminus1 = np.array([0.0])     # input at time step negative one - used to penalize the first delta u at time instant 0. Could be the same as uref.

# Constraints
xmin = np.array([-100.0, -100, -100, -100])
xmax = np.array([100.0,   100.0, 100, 100])

umin = np.array([-20])
umax = np.array([20])

Dumin = np.array([-5])
Dumax = np.array([5])
```

```python
# MPC constraints
xmin = np.array([-100.0, -100, -100, -100])
xmax = np.array([100.0,   100.0, 100, 100])

umin = np.array([-20])
umax = np.array([20])

Dumin = np.array([-5])
Dumax = np.array([5])
```

```python
# MPC cost function weights

Qx = sparse.diags([0.3, 0, 1.0, 0])   # Quadratic cost for states x0, x1, ..., x_N-1
QxN = sparse.diags([0.3, 0, 1.0, 0])  # Quadratic cost for xN
Qu = 0.0 * sparse.eye(1)        # Quadratic cost for u0, u1, ...., u_N-1
QDu = 0.01 * sparse.eye(1)       # Quadratic cost for Du0, Du1, ...., Du_N-1
```

```python
# Initialize simulation system

phi0 = 15*2*np.pi/360
x0 = np.array([0, 0, phi0, 0]) # initial state
t0 = 0
system_dyn = ode(f_ODE).set_integrator('vode', method='bdf')
system_dyn.set_initial_value(x0, t0)
system_dyn.set_f_params(0.0)
```

    <scipy.integrate._ode.ode at 0x2024e539b70>

```python
Np = 40
```

```python
# Initialize and setup MPC controller

K = MPCController(Ad,Bd,Np=Np, x0=x0,xref=xref,uminus1=uminus1,
                  Qx=Qx, QxN=QxN, Qu=Qu,QDu=QDu,
                  xmin=xmin,xmax=xmax,umin=umin,umax=umax,Dumin=Dumin,Dumax=Dumax)
K.setup() # this initializes the QP problem for the first step
```

```python
# Simulate in closed loop. Use MPC model as real system

# Simulate in closed loop
[nx, nu] = Bd.shape # number of states and number or inputs
len_sim = 10 # simulation length (s)
nsim = int(len_sim/Ts) # simulation length(timesteps)
xsim = np.zeros((nsim,nx))
usim = np.zeros((nsim,nu))
tsim = np.arange(0,nsim)*Ts

time_start = time.time()

t_step = t0
uMPC = uminus1
for i in range(nsim):
    xsim[i,:] = system_dyn.y

    # MPC update and step. Could be in just one function call
    K.update(system_dyn.y, uMPC) # update with measurement
    uMPC = K.output() # MPC step (u_k value)
    usim[i,:] = uMPC

    # System simulation
    system_dyn.set_f_params(uMPC) # set current input value
    system_dyn.integrate(t_step + Ts)

    # Time update
    t_step += Ts
time_sim = time.time() - time_start

```

```python
# Plot results

fig,axes = plt.subplots(3,1, figsize=(10,10))
axes[0].plot(tsim, xsim[:,0], "k", label='p')
axes[0].plot(tsim, xref[0]*np.ones(np.shape(tsim)), "r--", label="p_ref")
axes[0].set_title("Position (m)")

axes[1].plot(tsim, xsim[:,2]*360/2/np.pi, label="phi")
axes[1].plot(tsim, xref[2]*360/2/np.pi*np.ones(np.shape(tsim)), "r--", label="phi_ref")
axes[1].set_title("Angle (deg)")

axes[2].plot(tsim, usim[:,0], label="u")
axes[2].plot(tsim, uref*np.ones(np.shape(tsim)), "r--", label="u_ref")
axes[2].set_title("Force (N)")

for ax in axes:
    ax.grid(True)
    ax.legend()
```


# Implementation of a Devito viscoacoustic equations

## This tutorial is contributed by SENAI CIMATEC (2020)

This tutorial is based on:

**Linear inversion in layered viscoacoustic media using a time‐domain method** (1994)

Joakim O. Blanch and William W. Symes

SEG Technical Program Expanded Abstracts

https://doi.org/10.1190/1.1822695

**True-amplitude prestack depth migration** (2007)

Feng Deng and George A. McMechan

GEOPHYSICS Technical Papers

https://doi.org/10.1190/1.2714334

**Attenuation compensation for least-squares reverse time migration using the viscoacoustic-wave equation** (2014)

Gaurav Dutta and Gerard T. Schuster

GEOPHYSICS Technical Papers

https://doi.org/10.1190/geo2013-0414.1

**Multiscale viscoacoustic waveform inversion with the second generation wavelet transform and adaptive time–space domain finite-difference method** (2014)

Zhiming Ren, Yang Liu,and Qunshan Zhang

Geophysical Journal International, Volume 197, Issue 2, 1 May 2014, Pages 948–974

https://doi.org/10.1093/gji/ggu024

**Viscoacoustic prestack reverse time migration based on the optimal time-space domain high-order finite-difference method** (2014)

Yan Zhao, Yang Liu, and Zhi-Ming Ren

Appl. Geophys. 11, 50–62.

https://doi.org/10.1007/s11770-014-0414-8

**A stable and efficient approach of Q reverse time migration** (2018)

Yan Zhao, Ningbo Mao, and Zhiming Ren

GEOPHYSICS Technical Papers

https://doi.org/10.1190/geo2018-0022.1

## Introduction

The conversion of mechanical energy to heat, occurs during the propagation of seismic waves on the subsurface, due to the viscosity of the rocks. The presence of oil and gas in these rocks causes seismic attenuations. Thus, associated effects, such as dispersion and dissipation, can significantly affect the amplitudes, as well as the phase of the seismic pulse.  However, in the seismic exploration, the subsurface has still been considered as an ideal elastic/acoustic medium, that is, disregarding its mitigating effect. In practice, the propagation of seismic waves on the subsurface is in many ways different from propagation in an ideal solid.

For example, some subsurface rocks have anisotropic properties, are heterogeneous, porous and so on. The acoustic/elastic wave equation is not sensitive enough to describe propagation in these more complicated mediums. Generally, the viscosity of materials in the subsurface causes energy dissipation and consequently a decrease in amplitude, in addition to modifying the frequency content of the waves. This phenomenon of energy dissipation of the wave is called seismic absorption or attenuation.

The goal of this tutorial is to perform a seismic modeling taking into account the viscosity of the medium, so that it is possible to more accurately simulate the seismic data and consequently build images with better resolution in the processing of this data, in addition to extracting more detailed information on rocky materials through seismic inversion.

This tutorial follow three main viscoacoustic approaches in time-space domain:

- Blanch and Symes (1995) / Dutta and Schuster (2014)

- Ren et al. (2014)

- Deng and McMechan (2007)

Table of symbols

| Symbol         | Description
| :---         | :---
|$f$                                  |Frequency                                            |
|$f_o$                                |Reference frequency                                  |
|$\omega$                             |Angular frenquency                                   |
|$\omega_0$                           |Angular Reference Frequency                          |
|$v$                                  |Velocity model                                       |
|$v_0$                                |Reference velocity at $\omega_0$                     |
|$\kappa$                             |Bulk modulus                                         |
|$g$                                  |Absorption coefficient                               |
|$\tau$                               |Relaxation time                                      |
|$\tau_\sigma$                        |Stress relaxation parameter                          |
|$\tau_\epsilon$                      |Strain relaxation parameter                          |
|$Q$                                  |Quality factor                                       |
|$\eta$                               |Viscosity                                            |
|$\rho$                               |Density                                              |
|$\nabla$                             |Nabla operator                                       |
|$P({\bf x},t)$                       |Pressure field                                       |
|$r({\bf x},t)$                       |Memory variable                                      |
|${\bf v}({\bf x},t)$                 |Particle velocity                                    |
|$S({\bf x}_s,t)$                     |Source                                               |

# Seismic modelling with Devito

Before start with the viscoacoustic approaches we will describe a setup of seismic modelling with Devito in a simple 2D case. We will create a physical model of our domain and define a single source and an according set of receivers to model for the forward model. But first, we initialize some basic utilities.

```python
import numpy as np
import sympy as sp

import matplotlib.pyplot as plt

from devito import *
from examples.seismic.source import RickerSource, WaveletSource, TimeAxis
from examples.seismic import ModelViscoacoustic, plot_image, setup_geometry, plot_velocity
```

```python
nx = 300
nz = 300

# Define a physical size
shape = (nx, nz)
spacing = (20., 20.)
origin = (0., 0.)
nlayers = 3
nbl = 50
space_order = 8
dtype = np.float32

# Model physical parameters:
vp = np.zeros(shape)
qp = np.zeros(shape)
rho = np.zeros(shape)

# Define a velocity profile. The velocity is in km/s
vp_top = 1.5
vp_bottom = 3.5

# Define a velocity profile in km/s
v = np.empty(shape, dtype=dtype)
v[:] = vp_top  # Top velocity (background)
vp_i = np.linspace(vp_top, vp_bottom, nlayers)
for i in range(1, nlayers):
    v[..., i*int(shape[-1] / nlayers):] = vp_i[i]  # Bottom velocity

qp[:] = 3.516*((v[:]*1000.)**2.2)*10**(-6) # Li's empirical formula

rho[:] = 0.31*(v[:]*1000.)**0.25 # Gardner's relation
```

```python
#NBVAL_IGNORE_OUTPUT
model = ModelViscoacoustic(space_order=space_order, vp=v, qp=qp, b=1/rho,
                           origin=origin, shape=shape, spacing=spacing,
                           nbl=nbl)
```

    Operator `initdamp` run in 0.01 s
    Operator `pad_vp` run in 0.01 s
    Operator `pad_b` run in 0.01 s
    Operator `pad_qp` run in 0.01 s

```python
#NBVAL_IGNORE_OUTPUT
aspect_ratio = model.shape[0]/model.shape[1]

plt_options_model = {'cmap': 'jet', 'extent': [model.origin[0], model.origin[0] + model.domain_size[0],
                                               model.origin[1] + model.domain_size[1], model.origin[1]]}
fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))

slices = [slice(model.nbl, -model.nbl), slice(model.nbl, -model.nbl)]

img1 = ax[0].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, **plt_options_model)
fig.colorbar(img1, ax=ax[0])
ax[0].set_title(r"V (km/s)", fontsize=20)
ax[0].set_xlabel('X (m)', fontsize=20)
ax[0].set_ylabel('Depth (m)', fontsize=20)
ax[0].set_aspect('auto')

img2 = ax[1].imshow(np.transpose(qp), vmin=15, vmax=220, **plt_options_model)
fig.colorbar(img2, ax=ax[1])
ax[1].set_title("Q", fontsize=20)
ax[1].set_xlabel('X (m)', fontsize=20)
ax[1].set_ylabel('Depth (m)', fontsize=20)
ax[1].set_aspect('auto')

img3 = ax[2].imshow(np.transpose(rho), vmin=1.9, vmax=2.4, **plt_options_model)
fig.colorbar(img3, ax=ax[2])
ax[2].set_title(r"Density $\rho$ (g/cm^3)", fontsize=20)
ax[2].set_xlabel('X (m)', fontsize=20)
ax[2].set_ylabel('Depth (m)', fontsize=20)
ax[2].set_aspect('auto')

plt.tight_layout()
```

```python
f0 = 0.005 # peak/dominant frequency
b = model.b
rho = 1./b

# velocity model
vp = model.vp

t_s = (sp.sqrt(1.+1./model.qp**2)-1./model.qp)/f0
t_ep = 1./(f0**2*t_s)
tt = (t_ep/t_s) - 1.

s = model.grid.stepping_dim.spacing
damp = model.damp
```

```python
# Time step in ms and time range:
t0, tn = 0., 2000.
dt = model.critical_dt
time_range = TimeAxis(start=t0, stop=tn, step=dt)
```

```python
from examples.seismic import Receiver

def src_rec(p, model):
    src = RickerSource(name='src', grid=model.grid, f0=f0, time_range=time_range)
    src.coordinates.data[0, :] = np.array(model.domain_size) * .5
    src.coordinates.data[0, -1] = 8.

    # Create symbol for receivers
    rec = Receiver(name='rec', grid=model.grid, npoint=shape[0], time_range=time_range)

    # Prescribe even spacing for receivers along the x-axis
    rec.coordinates.data[:, 0] = np.linspace(0, model.domain_size[0], num=shape[0])
    rec.coordinates.data[:, 1] = 8.

    src_term = src.inject(field=p.forward, expr=(s*src))
    rec_term = rec.interpolate(expr=p)

    return src_term + rec_term, src, rec
```

Auxiliary functions for plotting data:

```python
def plot_receiver(rec):
    rec_plot = rec.resample(num=1001)
    scale_for_plot = np.diag(np.linspace(1.0, 2.5, 1001)**2.0)
    # Pressure (txx + tzz) data at sea surface
    extent = [rec_plot.coordinates.data[0, 0], rec_plot.coordinates.data[-1, 0], 1e-3*tn, t0]
    aspect = rec_plot.coordinates.data[-1, 0]/(1e-3*tn)/.5

    plt.figure(figsize=(10, 10))
    plt.imshow(np.dot(scale_for_plot, rec_plot.data), vmin=-.01, vmax=.01, cmap="seismic",
           interpolation='lanczos', extent=extent, aspect=aspect)
    plt.ylabel("Time (s)", fontsize=20)
    plt.xlabel("Receiver position (m)", fontsize=20)
```

```python
def plot_v_and_p(model, v, p):

    slices = [slice(model.nbl, -model.nbl), slice(model.nbl, -model.nbl)]
    scale = .5*1e-3

    plt_options_model = {'extent': [model.origin[0] , model.origin[0] + model.domain_size[0],
                                    model.origin[1] + model.domain_size[1], model.origin[1]]}

    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 7))

    ax[0].imshow(np.transpose(v[0].data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[0].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[0].set_aspect('auto')
    ax[0].set_xlabel('X (m)', fontsize=20)
    ax[0].set_ylabel('Depth (m)', fontsize=20)
    ax[0].set_title(r"$v_{x}$", fontsize=20)

    ax[1].imshow(np.transpose(v[1].data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[1].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[1].set_aspect('auto')
    ax[1].set_xlabel('X (m)', fontsize=20)
    ax[1].set_title(r"$v_{z}$", fontsize=20)

    ax[2].imshow(np.transpose(p.data[0][slices]), vmin=-scale, vmax=scale, cmap="RdGy", **plt_options_model)
    ax[2].imshow(np.transpose(model.vp.data[slices]), vmin=1.5, vmax=3.5, cmap="jet", alpha=.5, **plt_options_model)
    ax[2].set_aspect('auto')
    ax[2].set_xlabel('X (m)', fontsize=20)
    ax[2].set_title(r"$P$", fontsize=20)
```

## Blanch and Symes (1995) / Dutta and Schuster (2014)

The equations of motion for a viscoacoustic medium can be written as:

\begin{equation}
\left\{
\begin{array}{lcl}
\frac{\partial P}{\partial t} + \kappa (\tau + 1)(\nabla \cdot {\bf v}) + r = S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho}\nabla{P} = 0 \\
\frac{\partial r}{\partial t} + \frac{1}{\tau_{\sigma}} [r + \tau \kappa (\nabla \cdot {\bf v})] = 0.
\end{array}
\right.
\label{first-order-dutta}
\end{equation}

Where $\tau = \tau_{\epsilon}/\tau_{\sigma} -1$ represents the magnitude of  $Q$. $\tau_{\epsilon}$ and $\tau_{\sigma}$ are, respectively, the stress and strain relaxation parameters, given by:

\begin{equation}
\tau_\sigma = \frac{\sqrt{Q^2+1}-1}{2 \pi f_0 Q}
\end{equation}
and
\begin{equation}
\tau_\epsilon= \frac{\sqrt{Q^2+1}+1}{2\pi f_0 Q}
\end{equation}

```python
# Stencil created from Blanch and Symes (1995) / Dutta and Schuster (2014)
def blanch_symes(model, p, r, v):

    # Bulk modulus
    bm = rho * (vp * vp)

    # Define PDE to v
    pde_v = v - s * b * grad(p)
    u_v = Eq(v.forward, damp * pde_v)

    # Define PDE to r
    pde_r = r - s * (1. / t_s) * r - s * (1. / t_s) * tt * bm * div(v.forward)
    u_r = Eq(r.forward, damp * pde_r)

    # Define PDE to p
    pde_p = p - s * bm * (tt + 1.) * div(v.forward) - s * r.forward
    u_p = Eq(p.forward, damp * pde_p)

    return [u_v, u_r, u_p]
```

```python
# Seismic Modelling from Blanch and Symes (1995) / Dutta and Schuster (2014) viscoacoustic wave equation.
def modelling_blanch_symes(model):

    # Create symbols for particle velocity, pressure field, memory variable, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    r = TimeFunction(name="r", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = blanch_symes(model, p, r, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_blanch_symes(model)
```

    Operator `Kernel` run in 0.52 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 16, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.87797, atol=1e-3, rtol=0)
```

## Ren et al. (2014)

The viscoacoustic wave equation in time domain is written as:

\begin{equation}
\frac{\partial^{2}P}{\partial{t^2}} - v^{2}\nabla^{2}{P} - \eta\nabla^{2}\left(\frac{\partial P}{\partial t}\right) = S({\bf x}_{s}, t),
\end{equation}

where $\eta = \frac{v^2}{\omega_{0}Q}$ represents the viscosity of medium.

Considering the variable density $\rho$, the equation can be rewritten as:

\begin{equation}
\frac{\partial^{2}P}{\partial{t^2}} - \kappa \nabla \cdot \frac{1}{\rho} \nabla{P} - \eta \rho \nabla \cdot \frac{1}{\rho} \nabla \left(\frac{\partial{P}}{\partial{t}}\right) = S({\bf x}_{s}, t).
\end{equation}

The equation can be written using a first order formulation, given by:

\begin{equation}
\left\{
\begin{array}{ll}
\frac{\partial P}{\partial t} + \kappa \nabla \cdot {\bf v} - \eta \rho \nabla \cdot \frac{1}{\rho} \nabla{P} = S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho} \nabla{P} = 0
\end{array}
\right.
\end{equation}

```python
# Stencil created from Ren et al. (2014) viscoacoustic wave equation.
def ren(model, p, v):

    # Angular frequency
    w = 2. * np.pi * f0

    # Define PDE to v
    pde_v = v - s * b * grad(p)
    u_v = Eq(v.forward, damp * pde_v)

    # Define PDE to p
    pde_p = p - s * vp * vp * rho * div(v.forward) + s * \
            ((vp * vp * rho) / (w * model.qp)) * div(b * grad(p, shift=.5), shift=-.5)

    u_p = Eq(p.forward, damp * pde_p)

    return [u_v, u_p]
```

```python
# Seismic Modelling from Ren et al. (2014) viscoacoustic wave equation.
def modelling_ren(model):

    # Create symbols for particle velocity, pressure field, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = ren(model, p, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_ren(model)
```

    Operator `Kernel` run in 0.85 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 15, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.0639238, atol=1e-3, rtol=0)
```

## Deng and McMechan (2007)

The viscoacoustic wave equation for the propagating pressure $P$ in the time-space domain:

\begin{equation}
\frac{1}{v^2}\frac{\partial^{2}P}{\partial{t^2}} - \nabla^{2}P + \frac{g}{v}\frac{\partial P}{\partial{t}}  = S({\bf x}_{s}, t),
\label{eq-deng}
\end{equation}

where $g$ is the absorption coefficient, given by:

\begin{equation}
g = \frac{2\pi f_{0}}{vQ},
\end{equation}

The equation can be written using a first order formulation, given by:

\begin{equation}
\left\{
\begin{array}{lcl}
\frac{\partial P}{\partial t} + \kappa (\nabla \cdot {\bf v}) + \frac{2\pi f_{0}}{Q}P= S({\bf x}_{s}, t) \\
\frac{\partial {\bf v}}{\partial t} + \frac{1}{\rho}\nabla{P} = 0 \\
\end{array}
\right.
\end{equation}

```python
# Stencil created from Deng and McMechan (2007) viscoacoustic wave equation.
def deng_mcmechan(model, p, v):

    # Angular frequency
    w = 2. * np.pi * f0

    # Define PDE to v
    pde_v = v - s * b * grad(p)
    u_v = Eq(v.forward, damp * pde_v)

    # Define PDE to p
    pde_p = p - s * vp * vp * rho * div(v.forward) - s * (w / model.qp) * p
    u_p = Eq(p.forward, damp * pde_p)

    return [u_v, u_p]
```

```python
# Seismic Modelling from Deng and McMechan (2007) viscoacoustic wave equation.
def modelling_deng_mcmechan(model):

    # Create symbols for particle velocity, pressure field, source and receivers

    v = VectorTimeFunction(name="v", grid=model.grid, time_order=1, space_order=space_order)

    p = TimeFunction(name="p", grid=model.grid, time_order=1, space_order=space_order,
                     staggered=NODE)

    # define the source injection and create interpolation expression for receivers

    src_rec_expr, src, rec = src_rec(p, model)

    eqn = deng_mcmechan(model, p, v)

    op = Operator(eqn + src_rec_expr, subs=model.spacing_map)

    op(time=time_range.num-1, dt=dt, src=src, rec=rec)

    return rec, v, p
```

```python
#NBVAL_IGNORE_OUTPUT
rec, v, p = modelling_deng_mcmechan(model)
```

    Operator `Kernel` run in 0.40 s

```python
#NBVAL_IGNORE_OUTPUT
plot_receiver(rec)
```

```python
assert np.isclose(np.linalg.norm(rec.data), 16, rtol=10)
```

```python
#NBVAL_IGNORE_OUTPUT
plot_v_and_p(model, v, p)
```

```python
assert np.isclose(norm(v[0]), 1.1323929, atol=1e-3, rtol=0)
```

# More references

[1] https://academic.oup.com/gji/article/197/2/948/616510

[2] https://link.springer.com/article/10.1007/s11770-014-0414-8

[3] https://janth.home.xs4all.nl/Software/fdelmodcManual.pdf


__MECA654, March 2018, Polytech Annecy-Chambéry__
[This interactive notebook is on Github](https://github.com/a1rb4Ck/MECA654)

## Exercise 12 - Kinematics composition

TODO

```python
# MECA654, March 2018, Pierre Nagorny
# Polytech Annecy-Chambéry, SYMME lab, Savoie Mont Blanc University

from IPython.display import display, HTML, IFrame, Math
from IPython.core.interactiveshell import InteractiveShell
from IPython.display import YouTubeVideo
import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt  # Loading Pyplot
from scipy.interpolate import interp1d
# Setting Matplotlib render backend
%matplotlib nbagg
#%pylab inline
from sympy import *  # Loading Sympy for symbolic calculus
from sympy.physics.vector import *  # Everything needed for vector calculs, Reference frame, dynamic
from sympy.physics.vector import init_vprinting  # To print time varying vectors

init_printing()
init_vprinting(pretty_print=True)

# Display full output in Jupyter Notebook
InteractiveShell.ast_node_interactivity = "all"
```

```python
# Define used symbols for SymPy, constant variables
R, e, t = symbols('R e t')  # the time t is a symbol
# Define dynamic symbols, variables function of time t
d, x_2, theta_1, theta_3, OM = dynamicsymbols('d x_2 theta_1 theta_3 OM')
theta_1dot = dynamicsymbols('theta_1', 1)  # First order derivative of theta_1
ddot = dynamicsymbols('d', 1)  # First order derivative of theta_2

# Our reference frames for each body:

B0 = ReferenceFrame('B0', indices=['i', 'j', 'k'])
B1 = ReferenceFrame('B1', indices=['i', 'j', 'k'])

# We define space transformation relations in our problem:
B1.orient(B0, 'Axis', [theta_1, B0.z])  # Rotation of axis (z_0), same as (z_1)

# We can view the Direct Cosinus Matrix between both reference frames !
# B0.dcm(B1)
```


```python
from sympy import Matrix, zeros, MatAdd, MatMul
from sympy.abc import x, y, z
import numpy as np
from lmi_sdp import NonLinearExpressionError, NonLinearMatrixError, \
    lin_expr_coeffs, lm_sym_to_coeffs, lm_coeffs_to_sym, lm_sym_expanded
```

### test_lin_expr_coeffs()

```python
e = 1.2 + 3*x - 4.5*y + z
coeffs, const = lin_expr_coeffs(e, [x, y, z])
assert coeffs == [3.0, -4.5, 1.0]
assert const == 1.2
```

### test_lin_expr_coeffs_exceptions()

```python
except_ok = False
try:
    lin_expr_coeffs(1.2 + x + y*z, [x, y, z])
except NonLinearExpressionError:
    except_ok = True
assert except_ok

except_ok = False
try:
    lin_expr_coeffs(1.2 + x*y, [x])
except NonLinearExpressionError:
    except_ok = True
assert except_ok
```

### test_lm_sym_to_coeffs()

```python
m = Matrix([[1.2, x], [3.4*y, 1.2 + 3*x - 4.5*y + z]])
coeffs = lm_sym_to_coeffs(m, [x, y, z])
assert len(coeffs) == 2
assert len(coeffs[0]) == 3
assert (coeffs[0][0] == np.array([[0.0, 1.0], [0.0, 3.0]])).all()
assert (coeffs[0][1] == np.array([[0.0, 0.0], [3.4, -4.5]])).all()
assert (coeffs[0][2] == np.array([[0.0, 0.0], [0.0, 1.0]])).all()
assert (coeffs[1] == np.array([[1.2, 0.0], [0.0, 1.2]])).all()

assert lm_sym_to_coeffs(Matrix([0.0]), [x, y, z]) == \
    ([np.array([[0.0]]), np.array([[0.0]]), np.array([[0.0]])],
     np.array([[0.0]]))
```

### test_lm_sym_to_coeffs_sparse()

```python
try:
    import scipy
except ImportError:  # pragma: no cover
    pass
else:
    m = Matrix([[1.2, x], [3.4*y, 1.2 + 3*x - 4.5*y + z]])
    coeffs = lm_sym_to_coeffs(m, [x, y, z], sparse=True)
    assert len(coeffs) == 2
    assert len(coeffs[0]) == 3
    assert (coeffs[0][0].toarray() ==
            np.array([[0.0, 1.0], [0.0, 3.0]])).all()
    assert (coeffs[0][1].toarray() ==
            np.array([[0.0, 0.0], [3.4, -4.5]])).all()
    assert (coeffs[0][2].toarray() ==
            np.array([[0.0, 0.0], [0.0, 1.0]])).all()
    assert (coeffs[1].toarray() == np.array([[1.2, 0.0], [0.0, 1.2]])).all()
```

### test_lm_sym_to_coeffs_exceptions()

```python
except_ok = False
try:
    lm_sym_to_coeffs(Matrix([1.2 + x + y*z]), [x, y, z])
except NonLinearMatrixError:
    except_ok = True
assert except_ok

except_ok = False
try:
    lm_sym_to_coeffs(Matrix([1.2 + x*y]), [x])
except NonLinearMatrixError:
    except_ok = True
assert except_ok
```

### test_lm_coeffs_to_sym()

```python
var_coeffs = [None]*3
var_coeffs[0] = np.array([[0.0, 1.0], [0.0, 3.0]])
var_coeffs[1] = np.array([[0.0, 0.0], [3.4, -4.5]])
var_coeffs[2] = np.array([[0.0, 0.0], [0.0, 1.0]])
consts = np.array([[1.2, 0.0], [0.0, 1.2]])
coeffs = (var_coeffs, consts)
m = Matrix([[1.2, x], [3.4*y, 1.2 + 3*x - 4.5*y + z]])
assert lm_coeffs_to_sym(coeffs, [x, y, z]) - m == zeros(2)
```

### test_lm_sym_expanded()

```python
m = Matrix([[0, x], [3.4*y, 3*x - 4.5*y + z]])
c = Matrix([[1.2, 0], [0, 1.2]])
cx = MatMul(Matrix([[0.0, 1.0], [0.0, 3.0]]), x)
cy = MatMul(Matrix([[0.0, 0.0], [3.4, -4.5]]), y)
cz = MatMul(Matrix([[0.0, 0.0], [0.0, 1.0]]), z)
cc = Matrix([[1.2, 0.0], [0.0, 1.2]])
assert MatAdd(cx, cy, cz, cc) == lm_sym_expanded(m+c, [x, y, z])
assert MatAdd(cx, cy, cz) == lm_sym_expanded(m, [x, y, z])
assert cc == lm_sym_expanded(c, [x, y, z])
```