# Start-to-Finish validation of $\psi_4$ in curvilinear coordinates against Cartesian formulation provided by [Patrick Nelson's Weyl scalars & invariants in Cartesian coordinates module](../../Tutorial-WeylScalarsInvariants-Cartesian.ipynb)

### Author: Zach Etienne

**This module exists as a modification of [the NRPy+ $\psi_4$ in curvlinear coordinates module](../../Tutorial-Psi4.ipynb), writing all spacetime quantities in terms of ADM variables and their derivatives directly.**

## A Note on Notation

As is standard in NRPy+,

* Greek indices range from 0 to 3, inclusive, with the zeroth component denoting the temporal (time) component.
* Latin indices range from 0 to 2, inclusive, with the zeroth component denoting the first spatial component.

As a corollary, any expressions involving mixed Greek and Latin indices will need to offset one set of indices by one: A Latin index in a four-vector will be incremented and a Greek index in a three-vector will be decremented (however, the latter case does not occur in this tutorial module).

# Introduction, Table of Contents
$$\label{toc}$$

This module constructs $\psi_4$, a quantity that is immensely useful when extracting gravitational wave content from a numerical relativity simulation. $\psi_4$ is related to the gravitational wave strain via

$$
\psi_4 = \ddot{h}_+ - i \ddot{h}_\times.
$$

We construct $\psi_4$ from the standard ADM spatial metric $\gamma_{ij}$ and extrinsic curvature $K_{ij}$, and their derivatives. The full expression is given by Eq. 5.1 in [Baker, Campanelli, Lousto (2001)](https://arxiv.org/pdf/gr-qc/0104063.pdf):

\begin{align}
\psi_4 &= \left[ {R}_{ijkl}+2K_{i[k}K_{l]j}\right]
{n}^i\bar{m}^j{n}^k\bar{m}^l  \\
& -8\left[ K_{j[k,l]}+{\Gamma }_{j[k}^pK_{l]p}\right]
{n}^{[0}\bar{m}^{j]}{n}^k\bar{m}^l \\
& +4\left[ {R}_{jl}-K_{jp}K_l^p+KK_{jl}\right]
{n}^{[0}\bar{m}^{j]}{n}^{[0}\bar{m}^{l]},
\end{align}

Note that $\psi_4$ is imaginary, with the imaginary components originating from the tetrad vector $m^\mu$. This module does not specify a tetrad; instead it only constructs the above expression leaving $m^\mu$ and $n^\mu$ unspecified. The [next module on tetrads defines these tetrad quantities](Tutorial-Psi4_tetrads.ipynb) (currently only a quasi-Kinnersley tetrad is supported).

**This tutorial module is organized as follows:**

1. [Step 1](#initializenrpy): Initialize needed NRPy+ modules
1. [Step 2](#riemann): Constructing the 3-Riemann tensor $R_{ik\ell m}$
1. [Step 3](#termone): Constructing the rank-4 tensor in Term 1 of $\psi_4$: $R_{ijkl} + 2 K_{i[k} K_{l]j}$
1. [Step 4](#termtwo): Constructing the rank-3 tensor in Term 2 of $\psi_4$: $-8 \left(K_{j[k,l]} + \Gamma^{p}_{j[k} K_{l]p} \right)$
1. [Step 5](#termthree): Constructing the rank-2 tensor in term 3 of $\psi_4$: $+4 \left(R_{jl} - K_{jp} K^p_l + K K_{jl} \right)$
1. [Step 6](#psifour): Constructing $\psi_4$ through contractions of the above terms with arbitrary tetrad vectors $n^\mu$ and $m^\mu$
1. [Step 7](#code_validation): Code Validation against BSSN.Psi4 NRPy+ module
1. [Step 8](#latex_pdf_output): Output this module to $\LaTeX$-formatted PDF

# Step 1: Initialize core NRPy+ modules \[Back to [top](#toc)\]
$$\label{initializenrpy}$$

Let's start by importing all the needed modules from NRPy+:

```python
# Step 1.a: import all needed modules from NRPy+:
import sympy as sp
from outputC import *
import NRPy_param_funcs as par
import indexedexp as ixp
import grid as gri
import finite_difference as fin
import reference_metric as rfm

# Step 1.b: Set the coordinate system for the numerical grid
par.set_parval_from_str("reference_metric::CoordSystem","Cartesian")

# Step 1.c: Given the chosen coordinate system, set up
#           corresponding reference metric and needed
#           reference metric quantities
# The following function call sets up the reference metric
#    and related quantities, including rescaling matrices ReDD,
#    ReU, and hatted quantities.
rfm.reference_metric()

# Step 1.d: Set spatial dimension (must be 3 for BSSN, as BSSN is
#           a 3+1-dimensional decomposition of the general
#           relativistic field equations)
DIM = 3

# Step 1.e: Import all ADM quantities as written in terms of BSSN quantities
# import BSSN.ADM_in_terms_of_BSSN as AB
# AB.ADM_in_terms_of_BSSN()
```

# Step 2: Constructing the 3-Riemann tensor $R_{ik\ell m}$ \[Back to [top](#toc)\]
$$\label{riemann}$$

Analogously to Christoffel symbols, the Riemann tensor is a measure of the curvature of an $N$-dimensional manifold. Thus the 3-Riemann tensor is not simply a projection of the 4-Riemann tensor (see e.g., Eq. 2.7 of [Campanelli *et al* (1998)](https://arxiv.org/pdf/gr-qc/9803058.pdf) for the relation between 4-Riemann and 3-Riemann), as $N$-dimensional Riemann tensors are meant to define a notion of curvature given only the associated $N$-dimensional metric.

So, given the ADM 3-metric, the Riemann tensor in arbitrary dimension is given by the 3-dimensional version of Eq. 1.19 in Baumgarte & Shapiro's *Numerical Relativity*. I.e.,

$$
R^i_{jkl} = \partial_k \Gamma^{i}_{jl} - \partial_l \Gamma^{i}_{jk} + \Gamma^i_{mk} \Gamma^m_{jl} - \Gamma^{i}_{ml} \Gamma^{m}_{jk},
$$
where $\Gamma^i_{jk}$ is the Christoffel symbol associated with the 3-metric $\gamma_{ij}$:

$$
\Gamma^l_{ij} = \frac{1}{2} \gamma^{lk} \left(\gamma_{ki,j} + \gamma_{kj,i} - \gamma_{ij,k} \right)
$$

Notice that this equation for the Riemann tensor is equivalent to the equation given in the Wikipedia article on [Formulas in Riemannian geometry](https://en.wikipedia.org/w/index.php?title=List_of_formulas_in_Riemannian_geometry&oldid=882667524):

$$
R^\ell{}_{ijk}=
\partial_j \Gamma^\ell{}_{ik}-\partial_k\Gamma^\ell{}_{ij}
+\Gamma^\ell{}_{js}\Gamma_{ik}^s-\Gamma^\ell{}_{ks}\Gamma^s{}_{ij},
$$
with the replacements $i\to \ell$, $j\to i$, $k\to j$, $l\to k$, and $s\to m$. Wikipedia also provides a simpler form in terms of second-derivatives of three-metric itself (using the definition of Christoffel symbol), so that we need not define derivatives of the Christoffel symbol:

$$
R_{ik\ell m}=\frac{1}{2}\left(
\gamma_{im,k\ell}
+ \gamma_{k\ell,im}
- \gamma_{i\ell,km}
- \gamma_{km,i\ell} \right)
+\gamma_{np} \left(
\Gamma^n{}_{k\ell} \Gamma^p{}_{im} -
\Gamma^n{}_{km} \Gamma^p{}_{i\ell} \right).
$$

First we construct the term on the left:

```python
# Step 2: Construct the (rank-4) Riemann curvature tensor associated with the ADM 3-metric:
RDDDD = ixp.zerorank4()
gammaDD = ixp.register_gridfunctions_for_single_rank2("AUX","gammaDD", "sym01") # The AUX or EVOL designation is *not*
                                                                                # used in diagnostic modules.
kDD = ixp.register_gridfunctions_for_single_rank2("AUX","kDD", "sym01")
gammaDD_dD = ixp.declarerank3("gammaDD_dD","sym01")
gammaDD_dDD = ixp.declarerank4("gammaDD_dDD","sym01_sym23")

# gammaDD_dDD = AB.gammaDD_dDD

for i in range(DIM):
    for k in range(DIM):
        for l in range(DIM):
            for m in range(DIM):
                RDDDD[i][k][l][m] = sp.Rational(1,2) * \
                    (gammaDD_dDD[i][m][k][l] + gammaDD_dDD[k][l][i][m] - gammaDD_dDD[i][l][k][m] - gammaDD_dDD[k][m][i][l])
```

... then we add the term on the right:

```python
# ... then we add the term on the right:
# Define the Christoffel symbols
GammaUDD = ixp.zerorank3(DIM)
gammaUU,gammadetdummy = ixp.symm_matrix_inverter3x3(gammaDD)
for i in range(DIM):
    for k in range(DIM):
        for l in range(DIM):
            for m in range(DIM):
                GammaUDD[i][k][l] += (sp.Rational(1,2))*gammaUU[i][m]*\
                                     (gammaDD_dD[m][k][l] + gammaDD_dD[m][l][k] - gammaDD_dD[k][l][m])

for i in range(DIM):
    for k in range(DIM):
        for l in range(DIM):
            for m in range(DIM):
                for n in range(DIM):
                    for p in range(DIM):
                        RDDDD[i][k][l][m] += gammaDD[n][p] * \
                            (GammaUDD[n][k][l]*GammaUDD[p][i][m] - GammaUDD[n][k][m]*GammaUDD[p][i][l])
```

# Step 3: Constructing the rank-4 tensor in Term 1 of $\psi_4$: $R_{ijkl} + 2 K_{i[k} K_{l]j}$ \[Back to [top](#toc)\]
$$\label{termone}$$

Following Eq. 5.1 in [Baker, Campanelli, Lousto (2001)](https://arxiv.org/pdf/gr-qc/0104063.pdf), the rank-4 tensor in the first term of $\psi_4$ is given by

$$
R_{ijkl} + 2 K_{i[k} K_{l]j} = R_{ijkl} + K_{ik} K_{lj} - K_{il} K_{kj}
$$

```python
# Step 3: Construct the (rank-4) tensor in term 1 of psi_4 (referring to Eq 5.1 in
#   Baker, Campanelli, Lousto (2001); https://arxiv.org/pdf/gr-qc/0104063.pdf
rank4term1 = ixp.zerorank4()
# kDD = AB.kDD

for i in range(DIM):
    for j in range(DIM):
        for k in range(DIM):
            for l in range(DIM):
                rank4term1[i][j][k][l] = RDDDD[i][j][k][l] + kDD[i][k]*kDD[l][j] - kDD[i][l]*kDD[k][j]
```

# Step 4: Constructing the rank-3 tensor in Term 2 of $\psi_4$: $-8 \left(K_{j[k,l]} + \Gamma^{p}_{j[k} K_{l]p} \right)$ \[Back to [top](#toc)\]
$$\label{termtwo}$$

Following Eq. 5.1 in [Baker, Campanelli, Lousto (2001)](https://arxiv.org/pdf/gr-qc/0104063.pdf), the rank-3 tensor in the second term of $\psi_4$ is given by

$$
-8 \left(K_{j[k,l]} + \Gamma^{p}_{j[k} K_{l]p} \right)
$$
First let's construct the first term in this sum: $K_{j[k,l]} = \frac{1}{2} (K_{jk,l} - K_{jl,k})$:

```python
# Step 4: Construct the (rank-3) tensor in term 2 of psi_4 (referring to Eq 5.1 in
#   Baker, Campanelli, Lousto (2001); https://arxiv.org/pdf/gr-qc/0104063.pdf
rank3term2 = ixp.zerorank3()
# kDD_dD = AB.kDD_dD
kDD_dD = ixp.declarerank3("kDD_dD","sym01")

for j in range(DIM):
    for k in range(DIM):
        for l in range(DIM):
            rank3term2[j][k][l] = sp.Rational(1,2)*(kDD_dD[j][k][l] - kDD_dD[j][l][k])
```

... then we construct the second term in this sum: $\Gamma^{p}_{j[k} K_{l]p} = \frac{1}{2} (\Gamma^{p}_{jk} K_{lp}-\Gamma^{p}_{jl} K_{kp})$:

```python
# ... then we construct the second term in this sum:
#  \Gamma^{p}_{j[k} K_{l]p} = \frac{1}{2} (\Gamma^{p}_{jk} K_{lp}-\Gamma^{p}_{jl} K_{kp}):
for j in range(DIM):
    for k in range(DIM):
        for l in range(DIM):
            for p in range(DIM):
                rank3term2[j][k][l] += sp.Rational(1,2)*(GammaUDD[p][j][k]*kDD[l][p] - GammaUDD[p][j][l]*kDD[k][p])
```

Finally, we multiply the term by $-8$:

```python
# Finally, we multiply the term by $-8$:
for j in range(DIM):
    for k in range(DIM):
        for l in range(DIM):
            rank3term2[j][k][l] *= sp.sympify(-8)
```

# Step 5: Constructing the rank-2 tensor in term 3 of $\psi_4$: $+4 \left(R_{jl} - K_{jp} K^p_l + K K_{jl} \right)$ \[Back to [top](#toc)\]
$$\label{termthree}$$

Following Eq. 5.1 in [Baker, Campanelli, Lousto (2001)](https://arxiv.org/pdf/gr-qc/0104063.pdf), the rank-2 tensor in the third term of $\psi_4$ is given by

$$
+4 \left(R_{jl} - K_{jp} K^p_l + K K_{jl} \right),
$$
where
\begin{align}
R_{jl} &= R^i_{jil} \\
&= \gamma^{im} R_{ijml} \\
K &= K^i_i \\
&= \gamma^{im} K_{im}
\end{align}

Let's build the components of this term: $R_{jl}$, $K^p_l$, and $K$, as defined above:

```python
# Step 5: Construct the (rank-2) tensor in term 3 of psi_4 (referring to Eq 5.1 in
#   Baker, Campanelli, Lousto (2001); https://arxiv.org/pdf/gr-qc/0104063.pdf

# Step 5.1: Construct 3-Ricci tensor R_{ij} = gamma^{im} R_{ijml}
RDD = ixp.zerorank2()
for j in range(DIM):
    for l in range(DIM):
        for i in range(DIM):
            for m in range(DIM):
                RDD[j][l] += gammaUU[i][m]*RDDDD[i][j][m][l]

# Step 5.2: Construct K^p_l = gamma^{pi} K_{il}
KUD = ixp.zerorank2()
for p in range(DIM):
    for l in range(DIM):
        for i in range(DIM):
            KUD[p][l] += gammaUU[p][i]*kDD[i][l]

# Step 5.3: Construct trK = gamma^{ij} K_{ij}
trK = sp.sympify(0)
for i in range(DIM):
    for j in range(DIM):
        trK += gammaUU[i][j]*kDD[i][j]
```

Next we put these terms together to construct the entire term:
$$
+4 \left(R_{jl} - K_{jp} K^p_l + K K_{jl} \right),
$$

```python
# Next we put these terms together to construct the entire term in parentheses:
# +4 \left(R_{jl} - K_{jp} K^p_l + K K_{jl} \right),
rank2term3 = ixp.zerorank2()
for j in range(DIM):
    for l in range(DIM):
        rank2term3[j][l] = RDD[j][l] + trK*kDD[j][l]
        for p in range(DIM):
            rank2term3[j][l] += - kDD[j][p]*KUD[p][l]
# Finally we multiply by +4:
for j in range(DIM):
    for l in range(DIM):
        rank2term3[j][l] *= sp.sympify(4)
```

# Step 6: Constructing $\psi_4$ through contractions of the above terms with an arbitrary tetrad vectors $m^\mu$ and $n^\mu$ \[Back to [top](#toc)\]
$$\label{psifour}$$

Eq. 5.1 in [Baker, Campanelli, Lousto (2001)](https://arxiv.org/pdf/gr-qc/0104063.pdf) writes $\psi_4$ (which is complex) as the contraction of each of the above terms with products of tetrad vectors:

\begin{align}
\psi_4 &= \left[ {R}_{ijkl}+2K_{i[k}K_{l]j}\right]
{n}^i\bar{m}^j{n}^k\bar{m}^l  \\
& -8\left[ K_{j[k,l]}+{\Gamma }_{j[k}^pK_{l]p}\right]
{n}^{[0}\bar{m}^{j]}{n}^k\bar{m}^l \\
& +4\left[ {R}_{jl}-K_{jp}K_l^p+KK_{jl}\right]
{n}^{[0}\bar{m}^{j]}{n}^{[0}\bar{m}^{l]},
\end{align}
where $\bar{m}^\mu$ is the complex conjugate of $m^\mu$, and $n^\mu$ is real. The third term is given by
\begin{align}
{n}^{[0}\bar{m}^{j]}{n}^{[0}\bar{m}^{l]}
&= \frac{1}{2}({n}^{0}\bar{m}^{j} - {n}^{j}\bar{m}^{0} )\frac{1}{2}({n}^{0}\bar{m}^{l} - {n}^{l}\bar{m}^{0} )\\
&= \frac{1}{4}({n}^{0}\bar{m}^{j} - {n}^{j}\bar{m}^{0} )({n}^{0}\bar{m}^{l} - {n}^{l}\bar{m}^{0} )\\
&= \frac{1}{4}({n}^{0}\bar{m}^{j}{n}^{0}\bar{m}^{l} - {n}^{j}\bar{m}^{0}{n}^{0}\bar{m}^{l} - {n}^{0}\bar{m}^{j}{n}^{l}\bar{m}^{0} +  {n}^{j}\bar{m}^{0}{n}^{l}\bar{m}^{0})
\end{align}

Only $m^\mu$ is complex, so we can separate the real and imaginary parts of $\psi_4$ by hand, defining $M^\mu$ to now be the real part of $\bar{m}^\mu$ and $\mathcal{M}^\mu$ to be the imaginary part. All of the above products are of the form ${n}^\mu\bar{m}^\nu{n}^\eta\bar{m}^\delta$, so let's evalute the real and imaginary parts of this product once, for all such terms:

\begin{align}
{n}^\mu\bar{m}^\nu{n}^\eta\bar{m}^\delta
&= {n}^\mu(M^\nu - i \mathcal{M}^\nu){n}^\eta(M^\delta - i \mathcal{M}^\delta) \\
&= \left({n}^\mu M^\nu {n}^\eta M^\delta -
{n}^\mu \mathcal{M}^\nu {n}^\eta \mathcal{M}^\delta \right)+
i \left(
-{n}^\mu M^\nu {n}^\eta \mathcal{M}^\delta
-{n}^\mu \mathcal{M}^\nu {n}^\eta M^\delta
\right)
\end{align}

```python
# mre4U = ixp.declarerank1("mre4U",DIM=4)
# mim4U = ixp.declarerank1("mim4U",DIM=4)
# n4U   = ixp.declarerank1("n4U"  ,DIM=4)

import BSSN.Psi4_tetrads as P4t
P4t.Psi4_tetrads()
mre4U = P4t.mre4U
mim4U = P4t.mim4U
n4U = P4t.n4U

def tetrad_product__Real_psi4(n,Mre,Mim,  mu,nu,eta,delta):
    return +n[mu]*Mre[nu]*n[eta]*Mre[delta] - n[mu]*Mim[nu]*n[eta]*Mim[delta]

def tetrad_product__Imag_psi4(n,Mre,Mim,  mu,nu,eta,delta):
    return -n[mu]*Mre[nu]*n[eta]*Mim[delta] - n[mu]*Mim[nu]*n[eta]*Mre[delta]

psi4_re = sp.sympify(0)
psi4_im = sp.sympify(0)
# First term:
for i in range(DIM):
    for j in range(DIM):
        for k in range(DIM):
            for l in range(DIM):
                psi4_re += rank4term1[i][j][k][l]*tetrad_product__Real_psi4(n4U,mre4U,mim4U, i+1,j+1,k+1,l+1)
                psi4_im += rank4term1[i][j][k][l]*tetrad_product__Imag_psi4(n4U,mre4U,mim4U, i+1,j+1,k+1,l+1)

# Second term:
for j in range(DIM):
    for k in range(DIM):
        for l in range(DIM):
            psi4_re += rank3term2[j][k][l] * \
                       sp.Rational(1,2)*(+tetrad_product__Real_psi4(n4U,mre4U,mim4U, 0,j+1,k+1,l+1)
                                         -tetrad_product__Real_psi4(n4U,mre4U,mim4U, j+1,0,k+1,l+1) )
            psi4_im += rank3term2[j][k][l] * \
                       sp.Rational(1,2)*(+tetrad_product__Imag_psi4(n4U,mre4U,mim4U, 0,j+1,k+1,l+1)
                                         -tetrad_product__Imag_psi4(n4U,mre4U,mim4U, j+1,0,k+1,l+1) )
# Third term:
for j in range(DIM):
    for l in range(DIM):
        psi4_re += rank2term3[j][l] * \
                       (sp.Rational(1,4)*(+tetrad_product__Real_psi4(n4U,mre4U,mim4U, 0,j+1,0,l+1)
                                          -tetrad_product__Real_psi4(n4U,mre4U,mim4U, j+1,0,0,l+1)
                                          -tetrad_product__Real_psi4(n4U,mre4U,mim4U, 0,j+1,l+1,0)
                                          +tetrad_product__Real_psi4(n4U,mre4U,mim4U, j+1,0,l+1,0)))
        psi4_im += rank2term3[j][l] * \
                       (sp.Rational(1,4)*(+tetrad_product__Imag_psi4(n4U,mre4U,mim4U, 0,j+1,0,l+1)
                                          -tetrad_product__Imag_psi4(n4U,mre4U,mim4U, j+1,0,0,l+1)
                                          -tetrad_product__Imag_psi4(n4U,mre4U,mim4U, 0,j+1,l+1,0)
                                          +tetrad_product__Imag_psi4(n4U,mre4U,mim4U, j+1,0,l+1,0)))
```

    initialize_param() minor warning: Did nothing; already initialized parameter reference_metric::M_PI
    initialize_param() minor warning: Did nothing; already initialized parameter reference_metric::xmin
    initialize_param() minor warning: Did nothing; already initialized parameter reference_metric::xmax
    initialize_param() minor warning: Did nothing; already initialized parameter reference_metric::ymin
    initialize_param() minor warning: Did nothing; already initialized parameter reference_metric::ymax
    initialize_param() minor warning: Did nothing; already initialized parameter reference_metric::zmin
    initialize_param() minor warning: Did nothing; already initialized parameter reference_metric::zmax

# Step 6: Code validation against BSSN.Psi4 NRPy+ module \[Back to [top](#toc)\]
$$\label{code_validation}$$

As a code validation check, we verify agreement in the SymPy expressions for the RHSs of the BSSN equations between
1. this tutorial and
2. the NRPy+ BSSN.Psi4 module.

By default, we compare all quantities in Spherical coordinates, though other coordinate systems may be chosen.

```python
outCparams = "preindent=1,outCfileaccess=w,outCverbose=False,includebraces=False"
print("STARTING NEW")
fin.FD_outputC("Psi4_new.h", lhrh(lhs="psi4_real", rhs=psi4_im), outCparams)
print("FINISHED NEW")

gri.glb_gridfcs_list = []

import WeylScal4NRPy.WeylScalars_Cartesian as W4
W4.WeylScalars_Cartesian()
print("STARTING OLD")
fin.FD_outputC("Psi4_old.h", lhrh(lhs="psi4_real", rhs=W4.psi4i), outCparams)
print("FINISHED OLD")

# print("FullSimplify["+str(sp.mathematica_code(psi4_re-W4.psi4r))+"]")
# with open("math.txt","w") as file:
#     file.write("FullSimplify["+str(sp.mathematica_code(psi4_re-W4.psi4r))+"]")

# # Call the BSSN_RHSs() function from within the
# #          BSSN/BSSN_RHSs.py module,
# #          which should do exactly the same as in Steps 1-16 above.
# print("vvv Ignore the minor warnings below. vvv")
# import BSSN.Psi4 as BP4
# BP4.Psi4()
# print("^^^ Ignore the minor warnings above. ^^^\n")

# print("Consistency check between this tutorial and BSSN.Psi4 NRPy+ module: ALL SHOULD BE ZERO.")

# print("psi4_im - BP4.psi4_im = " + str(psi4_im - BP4.psi4_im))
# print("psi4_re - BP4.psi4_re = " + str(psi4_re - BP4.psi4_re))
```

    STARTING NEW
    Wrote to file "Psi4_new.h"
    FINISHED NEW
    STARTING OLD
    Wrote to file "Psi4_old.h"
    FINISHED OLD

```python
!gcc -O2 psi4_tester.c -o psi4_tester -lm
!./psi4_tester 4 4 4
```

    GOOD: 9.858968807130199e-16 : 1.678026634904849e-08 1.678026634904847e-08

# Step 7: Output this module to $\LaTeX$-formatted PDF file \[Back to [top](#toc)\]
$$\label{latex_pdf_output}$$

```python
!jupyter nbconvert --to latex --template latex_nrpy_style.tplx Tutorial-Psi4.ipynb
!pdflatex -interaction=batchmode Tutorial-Psi4.tex
!pdflatex -interaction=batchmode Tutorial-Psi4.tex
!pdflatex -interaction=batchmode Tutorial-Psi4.tex
!rm -f Tut*.out Tut*.aux Tut*.log
```

    [NbConvertApp] Converting notebook Tutorial-Psi4.ipynb to latex
    [NbConvertApp] Writing 61056 bytes to Tutorial-Psi4.tex
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode
    This is pdfTeX, Version 3.14159265-2.6-1.40.18 (TeX Live 2017/Debian) (preloaded format=pdflatex)
     restricted \write18 enabled.
    entering extended mode


```python
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.ensemble import RandomForestClassifier
from sklearn import metrics
```

# Árvores de Decisão e Florestas Aleatórias

Neste artigo, será introduzido o conceito do métodos de Aprendizado de Máquina para classificação, denominado como árvore de decisão.

O artigo assume que o leitor já conheça os conceitos de variáveis aleatórias e esperança de variáveis aleatórias.

****

# Entropia

Antes de entender o que são árvores de decisão, precisamos introduzir o conceito de Entropia.

A entropia é uma métrica para a incerteza de determinada variável aleatória. Sendo $X$ uma variável aleatória, com espaço de estados $\mathcal{S}$, e função de densidade de probabilidade $P(X = x), x \in \mathcal{S}$, temos que a entropia de $X$, denotada por $H(X)$, é:

\begin{equation}
H(X) = - \sum_{x \in \mathcal{S}} P(X = x) log_2(P(X = x))
\end{equation}

Por convenção $0\cdot log_2(0) = 0$, o que é facilmente verificado pelo fato de que $\lim_{x \to 0} xlog_2(x) = 0$

\begin{equation}
\lim_{x \to 0} xlog_2(x) = 0
\end{equation}

Por padrão, a unidade de medida da entropia são bits, a unidade de medida depende inteiramente na base do logarítmo que usamos para calcular a entropia, caso usassemos um $log_{10}$ ou $log_3$, a entropia seria medida em dits e trits, respectivamente.

Observação: Note que a definição de entropia $H(X)$, é similar a esperança da função $g(X) = \frac{1}{log_2(P(X = X))}$ de determinada variável aleatória $X$.

\begin{equation}
H(X) = E\left[\frac{1}{log_2(P(X = X))}\right]
\end{equation}

```python
#Função para calcular de determinada distribuição de valores
def calcularEntropia(prob_function):
    entropia = 0
    for i in range(len(prob_function)):
        if(prob_function[i] != 0):
            entropia += prob_function[i]*np.log2(prob_function[i])
    return -entropia

#Para calculo da entropia, usaremos um ensaio de Bernoulli onde a nossa variável aleatória assume os valores, 0 e 1
#com as probabilidades (1 - p) e p, respectivamente.
#Note que o ensaio de Bernoulli, nos da uma distribuição diferente para cada valor de p, consequentemente, a entropia
#para cada uma dessas distribuições será diferente.
x = np.linspace(0,1, 20)
y = [calcularEntropia([1 - p, p]) for p in x]
plt.plot(x, y, '-bo')
plt.xlabel("Parâmetro p para o ensaio de Bernoulli")
plt.ylabel("Entropia")
plt.show
```

Acima vemos um exemplo do gráfico da entropia, para os diferentes valores do parâmetro p em um ensaio de Bernoulli. É possível perceber que quando $p = 0$ ou $p = 1$, a entropia é 0, o que intuitivamente, faz total sentido, uma vez que, nesses casos o ensaio de Bernoulli só tem 0 ou 1 como resultado, ou seja, não existe incerteza sobre a distribuição.

Além disso, temos que a entropia é maxima quando $p = 0.5$, que é exatamente quando os valores 0 e 1 estão igualmente distribuidos, ou seja, a incerteza sobre qual dos dois valores teremos como resultado será máxima.

# Informação Mútua

Além do conceito de Entropia, precisamos também entender como funciona a Informação Mútua:

\begin{equation}
I(X,Y) = H(X) - H(X|Y)
\end{equation}

Perceba que agora introduzimos uma espécie de Entropia Condicional $H(X|Y)$, antes de falarmos sobre isso, vamos entender a intuição por trás da Informação Mútua, estamos calculando o quanto foi reduzido na entropia de X, quando a condicionamos a variável aleatória Y. Em outras palavras o quanto condicionar X a Y, diminuiu a incerteza de X.

A partir disso, podemos perceber que Informação Mútua é uma medida de dependência entre variáveis aleatórias, em outras palavras se temos uma grande redução na incerteza de X, quando condicionamos o valor de X a Y, podemos afirmar que Y explica muito bem X.

***

A Entropia Condicional $H(X|Y)$, vem da definição de Entropia, apresentada anteriormente, pela esperança de uma função de variáveis aleatória:

\begin{equation}
H(X|Y) = E\left[\frac{1}{log_2(P(X = x | Y = y))} \right] = - E\left[log_2(P(X = x | Y = y)) \right]
\end{equation}
\begin{equation}
= - \sum_{x \in \mathcal{X}}\sum_{y \in \mathcal{Y}} P(X = x, Y = y)\cdot log_2(P(X = x | Y = y))
\end{equation}

Sendo $ \mathcal{Y}$ e $\mathcal{X}$, os espaços de estados das variáveis aleatórias Y e X respectivamente.

Agora que conhecemos a Entropia Condicional, podemos chegar em uma fórmula melhor para a Informação Mútua:

\begin{equation}
I(X,Y) = H(X) - H(X|Y) = - \sum_{x \in \mathcal{S}} P(X = x) log_2(P(X = x))  - \left(- \sum_{x \in \mathcal{X}}\sum_{y \in \mathcal{Y}} P(X = x, Y = y)\cdot log_2(P(X = x | Y = y)) \right)
\end{equation}
\begin{equation}
= - \sum_{x \in \mathcal{X}} P(X = x) log_2(P(X = x))  + \sum_{x \in \mathcal{X}}\sum_{y \in \mathcal{Y}} P(X = x, Y = y)\cdot log_2(P(X = x | Y = y))
\end{equation}
\begin{equation}
= - \sum_{x \in \mathcal{X}}\sum_{y \in \mathcal{Y}} P(X = x, Y = y) \cdot log_2(P(X = x))  + \sum_{x \in \mathcal{X}}\sum_{y \in \mathcal{Y}} P(X = x, Y = y)\cdot log_2(P(X = x | Y = y))
\end{equation}
\begin{equation}
= \sum_{x \in \mathcal{X}}\sum_{y \in \mathcal{Y}} P(X = x, Y = y) \cdot log_2\left(\frac{P(X = x | Y = y)}{P(X = x)}\right)
\end{equation}
\begin{equation}
I(X, Y)= \sum_{x \in \mathcal{X}}\sum_{y \in \mathcal{Y}} P(X = x, Y = y) \cdot log_2\left(\frac{P(X = x, Y = y)}{P(X = x)P(Y = y)}\right)
\end{equation}

Onde a segunda para a terceira igualdade, vem do fato que:

\begin{equation}
P(X = x) = \sum_{y \in \mathcal{Y}} P(X = x|Y = y)P(Y=y) = \sum_{y \in \mathcal{Y}} P(X = x,Y = y)
\end{equation}
O resultado acima é a aplicação da Lei da Probabilidade Total.

Sendo assim, chegamos num resultado fechado para a informação mútua:

\begin{equation}
    I(X, Y)= \sum_{x \in \mathcal{X}}\sum_{y \in \mathcal{Y}} P(X = x, Y = y) \cdot log_2\left(\frac{P(X = x, Y = y)}{P(X = x)P(Y = y)}\right)
\end{equation}

# Árvores de decisão
Uma árvore de decisão é uma estrutura de relação entre as categorias. Cada um do nós de uma Árvore de Decisão representa uma espécie de "teste" em determinada feature, as relações entre cada um dos nós ("ramos" da árvore) representam determinado valor que aquele feature pode assumir.

---

Para melhor entendimento, usaremos o seguinte dataset como exemplo, ele contém dados sobre como o tempo estava em determinado dia e se pessoas jogaram tênis naquele dia em específico ou não.

A coluna Outlook contém os valores Ensolarado, Nublado e Chuvoso.

A coluna Temperature contém os valores Quente, Morno e Frio.

A coluna Humidity contém os valores Alta e Normal.

A coluna Wind contém os valores Forte e Fraco.

A coluna PlayTennis contém os valores Sim e Não, indicando se tennis foi ou não jogado naquele dia.

```python
col_names = ['Outlook','Temperature','Humidity','Wind','PlayTennis']

exemplo = pd.read_csv("exemplo_arvore.csv", header=None, names=col_names)
exemplo.head()
```

      Outlook
      Temperature
      Humidity
      Wind
      PlayTennis

      0
      Sunny
      Hot
      High
      Weak
      No

      1
      Sunny
      Hot
      High
      Strong
      No

      2
      Overcast
      Hot
      High
      Weak
      Yes

      3
      Rain
      Mild
      High
      Weak
      Yes

      4
      Rain
      Cool
      Normal
      Weak
      Yes

Pode-se observar uma possível árvore de decisão para este dataset na figura abaixo.
 Observe que estamos decidindo se iremos jogar tênis ou não, baseado nas informações meteorológicas que temos.

A árvore de decisão acima, nos diz que:

Caso o tempo esteja ensolarado, jogaremos tênis se a umidade estiver normal, caso esteja alta, não jogamos.

Caso o tempo esteja Nublado, jogaremos tênis.

Caso o tempo esteja Chuvoso, jogaremos tênis se o vento estiver fraco, caso esteja forte, não jogamos.

Observe que quando queremos modelar um problema, que quando a coluna target, no exemplo acima 'PlayTennis', pode assumir somente dois valores, temos uma função booleana.

No caso acima podemos mapear:

No $\rightarrow$ False

Yes $\rightarrow$ True

E assim, construir uma tabela verdade para a árvore de decisão acima:

    Outlook = Sunny
    Outlook = Overcast
    Outlook = Rain
    Humidity = High
    Humidity = Normal
    Wind = Strong
    Wind = Weak
    PlayTennis

    True
    False
    False
    True
    False
    False
    False
    False

    True
    False
    False
    False
    True
    False
    False
    True

    False
    True
    False
    False
    False
    False
    False
    True

    False
    False
    True
    False
    False
    True
    False
    False

    False
    False
    True
    False
    False
    False
    True
    True

Muitas linhas da tabela acima foram omitidas, caso o contrário, tal tabela seria gigantesca. Mais especificamente, como temos 7 colunas usadas para classificação, e cada uma delas pode assumir dois valores True ou False. Temos $2^7 = 128$ possíveis combinações de valores, ou seja, essa tabela verdade completa, teria 128 linhas.

# O problema com a melhor Árvore de Decisão
Tendo definido como funciona uma árvore de decisão, é fácil observar que modificando as features que escolhemos para determinados nós, temos diferentes árvores.

A partir dessa observação, podemos ver que existem várias árvores de decisão para o mesmo problema, no exemplo acima, a árvore de decisão é construida usando 7 atributos ('Outlook = Sunny', 'Outlook = Overcast', 'Outlook = Rain', 'Humidity = High', 'Humidity = Normal', 'Wind = Strong' e 'Wind = Weak') e vimos que a tabela verdade para esta árvore teria $2^7$ linhas.

Observe que dado determinado número de atributos, digamos $n$, para uma única árvore de decisão $T$, a tabela verdade de $T$ teria $2^n$ linhas.

Intuitivamente se pegássemos somente a coluna da classificação 'PlayTennis', teríamos um vetor de $2^n$ elementos, como observado abaixo:

    PlayTennis

    False

    True

    True

    False

    True

    ...

Diferentes árvores de decisão, que resolvem o mesmo problema, usam os mesmos atributos para classificação, a única coisa que muda nessas árvores, é a própria classificação. Ou seja, para todas as diferentes árvores de decisão possíveis, temos a mesma tabela verdade, com exceção da coluna target, no nosso exemplo a coluna 'PlayTennis'.

Ou seja, a quantidade total de árvores de decisão para determinado problema, é a combinação total de valores que cada um dos elementos da coluna target na tabela verdade pode assumir. Como o vetor gerado pela coluna target é um vetor de $2^n$ elementos (sendo $n$ o número de atributos), temos que cada um dos elementos desse vetor pode assumir um valor (True ou False), sendo assim, o total de combinações possíveis é $2^{2^n}$.

Sendo assim, no exemplo acima, como estamos usando 7 atributos, teriamos $2^{2^7} = 3.4 \cdot 10^{38}$ possíveis árvores de decisão que podem ser construídas (e nem consideramos todos os possíveis atributos para se jogar tennis).

Para determinar qual a melhor árvore de decisão, precisariamos iterar por todas as $3.4 \cdot 10^{38}$ possíveis árvores e ver qual a que melhor classifica o nosso target 'PlayTennis'. É, definitivamente inviável iterar por tantas árvores, e por isso precisamos de uma espécie heurística que nos ajudar a estimar uma árvore que seja boa o suficiente.

# Como encontrar uma árvore "boa"

Note que usamos o termo "boa", entre aspas, pois estamos montando um algoritmo para encontrar uma árvore que é boa e não ótima (como visto anteriormente a árvore ótima é um problema intratável).

A ideia é colocar os atributos mais importantes mais próximos ao topo da árvore, enquanto atributos menos importantes ficam nas camadas mais baixas. Por "atributos mais importantes" nos referimos aos atributos que classificam melhor o nosso target.

Por exemplo, no problema acima, vemos que o atributo 'Outlook = Overcast', classifica muito bem o 'PlayTennis', visto que, em todos os casos onde 'Outlook = Overcast' temos 'PlayTennis' como True, em outras palavras, condicionar 'PlayTennis' à 'Outlook = Overcast', reduz a incerteza (ou entropia) em 'PlayTennis'.

Para estruturar o algoritmo, ao invés de escolhermos os atributos que melhor classificam o target, escolheremos as features, como visto anteriormente, temos que as features são os nós das árvores e os atributos as arestas entre os nós.

Além disso, temos uma ótima maneira de determinar quanto de incerteza em uma variável aleatória é reduzido, quando a condicionamos a outra, e esta é a Informação Mútua.

No exemplo de 'PlayTennis', para escolher a feature que melhor classifica 'PlayTennis', calcularíamos I(PlayTennis, Outlook), I(PlayTennis, Humidity), I(PlayTennis, Wind) e I(PlayTennis, Temperature) e escolheríamos a feature que máxima a Informação Mútua.

Depois de decidida a primeira feature que usaremos para classificar o target, usaremos tal feature, para quebrar o nosso dataset em $n$ datasets, onde $n$ é o número de atributos na feature escolhida.

Aplicando essa ideia ao exemplo de 'PlayTennis', supondo que feature escolhida seja 'Humidity', que tem dois atributos associados a ela ('High' ou 'Normal'), construiremos dois dataframes que juntos formam o dataframe original, em um deles teremos todas as linhas do dataframe original em que 'Humidity = High' e no outro teremos todas as linhas do dataframe original em que 'Humidity = Normal'.

Usando os dataframes gerados, a partir da quebra do dataframe original, realizaremos o mesmo processo de classificação, escolhendo a melhor feature que classifica 'PlayTennis' em cada um dos dataframes gerados. Note que assim que o dataframe original é quebrado em $n$ dataframes, cada um desses novos dataframes é um novo problema de árvore de decisão, com menos exemplos (ou linhas) e uma feature a menos.

Realizamos tal processo de quebra de dataframes, até que não tenhamos mais features para usar na classificação, ou até que a entropia em target seja 0, ou seja, todos os valores na coluna target são iguais.

# A árvore de decisão do scikit-learn

Vamos usar o método 'DecisionTreeClassifier' da biblioteca scikit-learn no python, para classificar o nosso dataset de 'PlayTennis'.

```python
X = exemplo.loc[:, exemplo.columns != 'PlayTennis']
y = exemplo.PlayTennis
```

Os dados do nosso exemplo são categóricos e o classificador de árvore de decisão do scikit não funciona para este tipo de dado. Precisamos preparar os dados antes construir nosso classificador. Para isso vamos usar a função do pandas "pd.get_dummies" que converte os atributos categóricos em atributos com valores 0 e 1.

Por exemplo, o atributo Humidity que pode receber os valores High e Normal, passa a ser representado pelos atributos Humidity_High e Humidity_Normal, que podem receber os valores 0 e 1.

Essas variáveis criadas são chamadas de variáveis dummies, uma variável categorica com $n$ atributos, da origem à $n$ variáveis dummy, uma para cada atributo.

```python
X_dummy = pd.get_dummies(X) #transformando o dataset
X_dummy.head()
```

      Outlook_Overcast
      Outlook_Rain
      Outlook_Sunny
      Temperature_Cool
      Temperature_Hot
      Temperature_Mild
      Humidity_High
      Humidity_Normal
      Wind_Strong
      Wind_Weak

      0
      0
      0
      1
      0
      1
      0
      1
      0
      0
      1

      1
      0
      0
      1
      0
      1
      0
      1
      0
      1
      0

      2
      1
      0
      0
      0
      1
      0
      1
      0
      0
      1

      3
      0
      1
      0
      0
      0
      1
      1
      0
      0
      1

      4
      0
      1
      0
      1
      0
      0
      0
      1
      0
      1

E na seguinte célula, criamos a classe para o classificador de Árvores de Decisão do scikit-learn, além de usar o método 'fit' para 'treinar' a árvore aos dados disponibilizados

```python
# Criação do classificador de árvore de decisão
clf = DecisionTreeClassifier(criterion="entropy")

# Usamos o método fit para construir o classificador a partir do nosso conjunto de treinamento
clf = clf.fit(X_dummy, y)
```

```python
clf.predict(X_dummy)
```

Podemos usar o método 'plot_tree', também da biblioteca sklearn para vizualizar a árvore.

```python
fig = plt.figure(figsize=(25,20))
_ = plot_tree(clf, feature_names=X_dummy.columns, class_names=list(y.unique()), filled=True)

```

# Florestas Aleatórias

Como dito anteriormente, as Árvores de Decisão são construidas com base em uma heurística de qual é a melhor árvore. Uma vez que essa construção é por meio de uma heurística, não temos garantia de que essa árvore é, de fato, a melhor árvore, na verdade, existem outras heurística para a construção de uma Árvore de Decisão, o que foi apresentado anteriormente, é apenas uma delas.

A partir dessa ideia de que podemos ter diversas Árvores de Decisão que sejam capazes de classificar o target, surge a noção de Florestas Aleatórias.

---

A ideia por trás das Florestas Aleatórias, está em gerar diferentes Árvores de Decisão, que conferem diferentes classificações do target, e assim fazer-las "votar" na melhor classificação.

Para começar a construção de uma Floresta Aleatória, inicialmente decidimos a quantidade de árvores que estarão presentes naquela floresta.

Após decidir quantas árvores irão compor a floresta, geraremos um dataset a partir do original, selecionando amostras do dataset original aleatoriamente, as quais estarão presentes nesse novo dataset. Essa técnica se chama Bagging .
 Um dataset para cada Árvore na Floresta será gerado a partir do Bagging, estes serão usados para treinar a árvore, ao qual tal dataset está relacionado.

Existem varias maneiras de se escolher diferentes heurísticas para gerar as Árvores de Decisão em uma Floresta Aleatória. O padrão da biblioteca 'sklearn', é selecionar a cada iteração um número limitado de features que podem compor o nó da árvore na iteração corrente.

No sklearn, é possível definir uma constante $n$, de maneira que a cada iteração o algoritmo pode escolher somente $n$ features selecionadas aleatoriamente, ou definir, uma função do número de features disponíveis.

Por exemplo, sendo $n$ o número de features disponíveis em uma iteração, o padrão do 'sklearn', é selecionar $\sqrt{n}$ features que podem ser escolhidas para aquele nó.

# A Floresta Aleatória do scikit-learn

Vamos usar o método 'RandomForestClassifier' da biblioteca scikit-learn no python, para classificar o nosso dataset de 'PlayTennis'.

```python
# Criação do classificador de floresta aleatória.
#O parâmetro n_estimators diz quantas árvores serão geradas.
#O parâmetro bootstrap está dizendo para o algoritmo realizar o processo de Bagging para cada árvore gerada.
#O parâmetro max_features diz o máximo de features a cada execução, no caso, o máximo será a raiz quadrada
#do número de features disponível
#O parâmetro criterion diz respeito ao critério ultilizado para construir a árvore, no caso Entropia/Informação Mútua.
clf = RandomForestClassifier(n_estimators=100, bootstrap=True, max_features="sqrt", criterion="entropy")

# Usamos o método fit para construir o classificador a partir do nosso conjunto de treinamento
clf = clf.fit(X_dummy, y)
```

```python
clf.predict(X_dummy)
```

A floresta gerada é composta por 100 árvores (parâmetro que pode ser modificado em 'n_estimators'), por isso não é possível vizualizar todas as árvores na floresta, dessa forma, escolheremos uma única árvore para mostrar.

```python
fig = plt.figure(figsize=(25,20))
_ = plot_tree(clf.estimators_[25], feature_names=X_dummy.columns, class_names=list(y.unique()), filled=True)
```


```
%load_ext autoreload
%autoreload 2

from __future__ import division

import numpy as np
import matplotlib.pyplot as plt

import pydy
import sympy as sym
import sympy.physics.mechanics as me

from scipy.linalg import solve, lstsq
from scipy.integrate import odeint

import time
import functools

me.Vector.simp = False  # to increase computation speed

from sympy.physics.vector import init_vprinting
init_vprinting()

from matplotlib import animation
from matplotlib.patches import Ellipse, Rectangle

from IPython.html import widgets
from IPython.html.widgets import interact, interactive

from pydy_code_gen.code import generate_ode_function

# from loading import load_cython
# filename_prefix = 'multibody_system'
# rhs = load_cython(filename_prefix, len(q), True)
```

```
from IPython.display import Image
Image(filename='snake-schematic.png')
```

```
n = 3

# dynamic variables
q = me.dynamicsymbols('q:' + str(n + 2))
u = me.dynamicsymbols('u:' + str(n + 2))
T = me.dynamicsymbols('T:' + str(n + 1))

# variables for bodies
g, t = sym.symbols('g t')
m = sym.symbols('m:' + str(n))
l = sym.symbols('l:' + str(n))
Izz = sym.symbols('I_zz_:' + str(n))

# frames
frames = [me.ReferenceFrame('fr_N')]
for i in range(n):
    name = 'fr_' + str(i)
    frame = frames[i].orientnew(name, 'Axis', [q[i], frames[i].z])
    frame.set_ang_vel(frames[i], u[i] * frames[i].z)
    frames.append(frame)

# origin of inertial frame
joints = [me.Point('N_or')]
joints[0].set_vel(frames[0], 0)

# joints
for i in range(n + 1):
    name = 'jt_' + str(i)
    if i == 0:
        joint = joints[i].locatenew(name, q[-2] * frames[0].x + q[-1] * frames[0].y)
        joint.set_vel(frames[0], u[-2] * frames[0].x + u[-1] * frames[0].y)
    else:
        joint = joints[i].locatenew(name, l[i - 1] * frames[i].x)
        joint.v2pt_theory(joints[i], frames[0], frames[i])
    joints.append(joint)

# mass centers
mc = []
for i in range(n):
    name = 'mc_' + str(i)
    mcent = joints[i + 1].locatenew(name, l[i] / 2 * frames[i + 1].x)
    mcent.v2pt_theory(joints[i + 1], frames[0], frames[i + 1])
    mc.append(mcent)

# inertia dyads
dyads = []
for i in range(n):
    dyad = (me.inertia(frames[i + 1], 0, 0, Izz[i]), mc[i])
    dyads.append(dyad)

# bodies
bodies = []
for i in range(n):
    name = 'bd_' + str(i)
    body = me.RigidBody(name, mc[i], frames[i + 1], m[i], dyads[i])
    bodies.append(body)

# kinematic differential equations
kd = []
for i in range(n + 2):
    kd.append(q[i].diff(t) - u[i])

forces = []
for i in range(n):
    forces.append((mc[i], -m[i] * g * frames[0].z))

torques = []
for i in range(n):
    torques.append((frames[i + 1], (T[i] - T[i + 1]) * frames[0].z))

fl = forces + torques
```

```
import nchain
dynamic, symbol, setup, fbd = nchain.formulate_nchain_parameters(n)
eom = nchain.make_kane_eom(dynamic, setup, fbd)
```

```
# equations of motion using Kane's method
kane = me.KanesMethod(frame=frames[0],
                      q_ind=q,
                      u_ind=u,
                      kd_eqs=kd,
                      q_dependent=[],
                      configuration_constraints=[],
                      u_dependent=[],
                      velocity_constraints=[])

now = time.time()
(fr, frstar) = kane.kanes_equations(fl, bodies)
kanezero = fr + frstar
print('elapsed time: {0:.3f} sec'.format(time.time() - now))
```

    elapsed time: 1.560 sec

```
now = time.time()
kz = sym.trigsimp(kanezero)
print('elapsed time: {0:.3f} sec'.format(time.time() - now))
```

    elapsed time: 72.597 sec

```
tt = kz[1]
exp = str(me.mlatex(tt))
#print(exp)  # copy this output
```

```
#print(str(me.mlatex(kz)))
```

```
if False:
    com_vec = 1 / (m[0] + m[1]) * (m[0] * mc[0].pos_from(joints[0]) + m[1] * mc[1].pos_from(joints[0]))
    com = joints[0].locatenew('com', com_vec)

    #print me.mprint(com.pos_from(joints[2]))

    com.v2pt_theory(joints[0], frames[0], frames[0])

    #print me.angular_momentum(com, frames[0], *bodies).express(frames[0])
```

```
constants = [g]
for i in range(n):
    constants = constants + [m[i], l[i], Izz[i]]

now = time.time()
mass_matrix = kane.mass_matrix_full
forcing_vector = kane.forcing_full
print('Equation rearranging:  {0:.3f} sec'.format(time.time() - now))

now = time.time()
right_hand_side = generate_ode_function(mass_matrix, forcing_vector, constants,
                                        q, u, T, generator='cython')
print('Code generation:       {0:.3f} sec'.format(time.time() - now))
```

    Equation rearranging:  0.006 sec
    Code generation:       4.295 sec

```
def heavy(trq):
    ih = np.where(trq > 0)[0]
    il = np.where(trq < 0)[0]

    th = trq[ih].mean()
    tl = trq[il].mean()

    trq[ih] = th
    trq[il] = tl

    return trq
```

```
def shape_space(x, t):
    trq = []
    for i in range(n + 1):
        trq.append(np.interp(t, ts, torques[:, i]))
    return np.array(trq)

# actual snake values
masses = 40.5 * np.ones(n) / n / 1000  # kg
#masses = np.r_[50, 30] / 1000
lengths = 68.6 * np.ones(n) / n / 100  # m
#lengths = np.r_[20, 45] / 100

# masses = np.r_[5, 10, 10, 10, 10] / 1000
# lengths = np.r_[5, 10, 15, 15, 15] / 100
width = 2.2 / 100                      # m
#mom_inertias = masses * (lengths**2 + width**2) / 12
mom_inertias = masses / 4 * width**2 + masses / 12 * lengths**2

constant_vals = [9.81]
for i in range(n):
    constant_vals = constant_vals + [masses[i], lengths[i], mom_inertias[i]]

# initial conditions
qstart = np.zeros(n + 2)
qstart[0] = np.pi / 2
ustart = np.zeros(n + 2)
ustart[0] = 0
x0 = np.hstack((qstart, ustart))

npts = 2000
ts, dt = linspace(0, 10, npts, retstep=True)

torques = np.zeros((npts, n + 1))
for i in range(n + 1):
    #torques[:, i] = .0025 * np.cos(10 * ts + np.pi/2 * i) #* np.exp(-2 * ts)
    trq = .005 * np.cos(20 * ts + np.pi/2 * i) #* np.exp(-2 * ts)
    torques[:, i] = heavy(trq)

#torques[:, 1] = 0
#torques[:, 2] = 0
#torques[:, 4] = 0
torques[:, 0] = 0  # .5 * torques[:, 1]
torques[:, n] = 0

args = {'constants': np.array(constant_vals),
        'specified': shape_space}

now = time.time()
y = odeint(right_hand_side, x0, ts, args=(args,))
print('Simulation time: {0:.2f} sec'.format(time.time() - now))
```

    Simulation time: 1.64 sec

```
fig, ax = plt.subplots()
ax.plot(ts, torques);

fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True)
iq = n + 2

lines = ax1.plot(ts, y[:, :iq])
#lab = xlabel('Time [sec]')
leg = ax1.legend(q)

lines = ax2.plot(ts, y[:, iq:])
#lab = xlabel('Time [sec]')
leg = ax2.legend(u)

fig, ax = plt.subplots()
ax.plot(y[:, 1], y[:, 0])
```

    [<matplotlib.lines.Line2D at 0x7453d90>]

```
plt.close('all')
```

```
def organize_data(y, n, ls, ms):
    """Organize the data so it can be animated with ellipses.
    """

    npts = y.shape[0]
    angs = y[:, :n]
    rads = angs.cumsum(axis=1)
    degs = np.rad2deg(rads)

    jtx, jty = y[:, n:n + 2].T
    jtx = jtx.reshape(-1, 1)
    jty = jty.reshape(-1, 1)

    mcx, mcy = [], []
    for i in range(n):
        dx = ls[i] * np.cos(rads[:, i])
        dy = ls[i] * np.sin(rads[:, i])

        mcx.append((jtx[:, i] + dx / 2).flatten())
        mcy.append((jty[:, i] + dy / 2).flatten())

        x = (jtx[:, i] + dx).reshape(-1, 1)
        y = (jty[:, i] + dy).reshape(-1, 1)

        jtx = np.hstack((jtx, x))
        jty = np.hstack((jty, y))

    mcx = np.array(mcx).T
    mcy = np.array(mcy).T

    # center of mass
    mms = np.array(ms)
    comx = (mcx * mms).sum(axis=1) / mms.sum()
    comy = (mcy * mms).sum(axis=1) / mms.sum()

    return rads, degs, jtx, jty, mcx, mcy, comx, comy

def n_movie_maker(n, lengths, ts, y, limx=(-1, 1), limy=(-1, 1), hst=50):

    rads, degs, jtx, jty, mcx, mcy, comx, comy = organize_data(y, n, lengths, masses)

    ws = .05 * lengths
    ws = width * np.ones(n)
    ii = 0
    st = 0

    fig, ax = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(9, 9))
    ax.axis('image')
    ax.set_xlim(limx)
    ax.set_ylim(limy)

    fig.set_facecolor('w')
    fig.tight_layout()

    time_fmt = '{0:2.3} sec'
    time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes, fontsize='large')

    ellipse_params = dict(alpha=.5, color='green')
    patches = []
    for i in range(n):
        e = Ellipse(xy=(mcx[ii, i], mcy[ii, i]),
                    width=ws[i], height=1.05 * lengths[i],
                    angle=degs[ii, i] - 90, **ellipse_params)
        ax.add_patch(e)
        patches.append(e)

    head, = ax.plot(jtx[ii, 0], jty[ii, 0], 'go', alpha=.6, markersize=7)
    jlines, mclines = [], []
    for i in range(n):
        jline, = ax.plot(jtx[st:ii, i + 1], jty[st:ii, i + 1], 'b--', alpha=.5)
        mcline, = ax.plot(mcx[st:ii, i], mcy[st:ii, i], 'r', alpha=.5)
        jlines.append(jline)
        mclines.append(mcline)

    coml, = ax.plot(comx[st:ii], comy[st:ii], 'k-', lw=2, alpha=.9)

    def init():
        ii = 0
        head.set_data(jtx[ii, 0], jty[ii, 0])
        for i in range(n):
            patches[i].set_alpha(.3)
        time_text.set_text('')
        return [head, time_text] + patches

    def animate(ii):

        st = ii - hst
        if st < 0:
            st = 0

        head.set_data(jtx[ii, 0], jty[ii, 0])
        for i in range(n):
            patches[i].set_visible(True)
            patches[i].set_alpha(.5)
            patches[i].center = (mcx[ii, i], mcy[ii, i])
            patches[i].angle = degs[ii, i] - 90
            jlines[i].set_data(jtx[st:ii, i + 1], jty[st:ii, i + 1])
            mclines[i].set_data(mcx[st:ii, i], mcy[st:ii, i])

        time_text.set_text(time_fmt.format(ts[ii]))

        coml.set_data(comx[st:ii], comy[st:ii])

        return patches + [head, coml, time_text] + jlines + mclines

    anim = animation.FuncAnimation(fig, animate, frames=len(ts),
                interval=ts[-1] / len(ts) * 10000, blit=True, repeat=False, init_func=init)

    return anim

anim = n_movie_maker(n, lengths, ts, y, limx=(-.3, .3), limy=(-.05, .8), hst=100)
```

```
#anim.save('3a-snake.mp4')
```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```

```
Hj0 = me.angular_momentum(joints[1], frames[0], *bodies).express(frames[0]).simplify()
```

```
Hj0.to_matrix(frames[0])[2]
```

```
#sym.solve(Hj0.to_matrix(frames[0])[2], q[1])
```

```
sym.trigsimp(Hj0)
```

```
##
# this is like the lizard
##
n = 2

# dynamic variables
q = me.dynamicsymbols('q:' + str(n))
u = me.dynamicsymbols('u:' + str(n))
T = me.dynamicsymbols('T:' + str(n + 1))

# variables for bodies
g, t = sym.symbols('g t')
m = sym.symbols('m:' + str(n))
l = sym.symbols('l:' + str(n))
Izz = sym.symbols('I_zz_:' + str(n))

# frames
frames = [me.ReferenceFrame('fr_N')]
for i in range(n):
    name = 'fr_' + str(i)
    frame = frames[i].orientnew(name, 'Axis', [q[i], frames[0].z])
    frame.set_ang_vel(frames[i], u[i] * frames[0].z)
    frames.append(frame)

# origin of inertial frame
joints = [me.Point('N_or')]
joints[0].set_vel(frames[0], 0)

# joints
for i in range(n + 1):
    name = 'jt_' + str(i)
    joint = joints[i].locatenew(name, l[i - 1] * frames[i].x)
    joint.v2pt_theory(joints[i], frames[0], frames[i])
    joints.append(joint)

# mass centers
mc = []
for i in range(n):
    name = 'mc_' + str(i)
    mcent = joints[i + 1].locatenew(name, l[i] / 2 * frames[i + 1].x)
    mcent.v2pt_theory(joints[i + 1], frames[0], frames[i + 1])
    mc.append(mcent)

# inertia dyads
dyads = []
for i in range(n):
    dyad = (me.inertia(frames[i + 1], 0, 0, Izz[i]), mc[i])
    dyads.append(dyad)

# bodies
bodies = []
for i in range(n):
    name = 'bd_' + str(i)
    body = me.RigidBody(name, mc[i], frames[i + 1], m[i], dyads[i])
    bodies.append(body)

# kinematic differential equations
kd = []
for i in range(n):
    kd.append(q[i].diff(t) - u[i])

forces = []
for i in range(n):
    forces.append((mc[i], -m[i] * g * frames[0].z))

torques = []
for i in range(n):
    torques.append((frames[i + 1], (T[i] - T[i + 1]) * frames[0].z))

fl = forces + torques
```

```
L = me.Lagrangian(frames[0], *bodies)
LM = me.LagrangesMethod(L, q, forcelist=fl, frame=frames[0])

now = time.time()
langeqn = LM.form_lagranges_equations()
print('Lagrangian time: {0:.2f} sec'.format(time.time() - now))
```

    Lagrangian time: 0.34 sec

```
eom = nchain.make_lagrange_eom(dynamic, setup, fbd)
```

```
sym.trigsimp(eom['L'])
```

```
lz = sym.trigsimp(L)
```

```
lz
```

```
le = sym.trigsimp(langeqn)
```

```
#print(str(me.mlatex(le)))
```

```
le[0]
```

```
kz[0]
```

```
le[1]
```

```
kz[1]
```

```
le[2]
```

```
kz[2]
```

```
sym.trigsimp(forcing_vector[6])
```

```
mass_matrix[6, :]
```

```
LM.mass_matrix_full
```

```
forcing_vector_lang = LM.forcing_full
```

```
sym.trigsimp(forcing_vector_lang[6])
```

```
sym.trigsimp(LM.forcing_full)
```

```
#anim.save('two-link-flower.mp4')
```

```
rads, degs, jtx, jty, mcx, mcy, comx, comy = organize_data(y, n, lengths, masses)
```

```
plt.close('all')
```

```
animation.FuncAnimation??
```

```
mcx.shape
```

```
jtx.shape
```

```
jtx.shape
```

```

```

```
def shape_space(x, t):

    torques = []
    for i in range(1, n - 1):
        torques.append(i)
    torques = [0] + torques + [0]
    #tk1 = np.interp(t, ts, t1)
    #tk2 = np.interp(t, ts, t2)

    return np.array(torques)

args = {'constants': np.array(constant_vals),
        'specified': shape_space}

qstart = np.pi/2
ustart = (.3, 0.0, .0, 0, 0)
#ustart = np.zeros(5)
x0 = np.hstack((qstart, ustart))

npts = 1001
ts, dt = linspace(0,  10, npts, retstep=True)

# t1 = 60 * np.ones(npts)
# t2 = -30 * np.ones(npts)

t1 = .025 * np.cos(10 * ts)
t2 = -.025 * np.cos(10 * ts + np.pi/4)

# t1 *= 0
# t2 *= 0

# t1[t1 > 0] = t1.max()
# t2[tt2 > 0] = t2.max()

# t1[theta1 < 0] = t1.min()
# t2[theta2 < 0] = t2.min()

now = time.time()
y1 = odeint(right_hand_side, x0, ts, args=(args,))
y1 = y1.T
print('Elapsed time: {0:.2f} sec'.format(time.time() - now))

# organize the data
th0, th1, th2 = y1[:3]
th1 = th1 + th0
th2 = th2 + th1

dg0, dg1, dg2 = np.rad2deg(y1[:3])
dg1 = dg1 + dg0
dg2 = dg2 + dg1

j0 = np.c_[0, 0]
j0 = y1[3:5].T
j1 = j0 + np.c_[ls   * np.cos(th0), ls   * np.sin(th0)]
j2 = j1 + np.c_[ls   * np.cos(th1), ls   * np.sin(th1)]
j3 = j2 + np.c_[ls   * np.cos(th2), ls   * np.sin(th2)]
c0 = j0 + np.c_[ls/2 * np.cos(th0), ls/2 * np.sin(th0)]
c1 = j1 + np.c_[ls/2 * np.cos(th1), ls/2 * np.sin(th1)]
c2 = j2 + np.c_[ls/2 * np.cos(th2), ls/2 * np.sin(th2)]

# center of mass
mms = np.array([ms, ms, ms])
cx = np.c_[c0[:, 0], c1[:, 0], c2[:, 0]]
cy = np.c_[c0[:, 1], c1[:, 1], c2[:, 1]]
com = np.c_[(cx * mms).sum(axis=1), (cy * mms).sum(axis=1)] / mms.sum()
```

```
# number of links (not currently used)
n = 3

# dynamic variables
q0, q1, q2 = me.dynamicsymbols('q0 q1 q2')
u0, u1, u2 = me.dynamicsymbols('u0 u1 u2')

# location of the head
q3, q4 = me.dynamicsymbols('q3 q4')
u3, u4 = me.dynamicsymbols('u3 u4')

# time varying torques
T0, T1, T2, T3 = sym.symbols('T_0 T_1 T_2 T_3')

# variables for bodies
g, t = sym.symbols('g t')
m0, m1, m2 = sym.symbols('m_0 m_1 m_2')
l0, l1, l2 = sym.symbols('l_0 l_1 l_2')
Izz0, Izz1, Izz2 = sym.symbols('I_zz_0 I_zz_1 I_zz_2')

# inertial reference frame
frN = me.ReferenceFrame('fr_N')

# body-fixed frames
fr0 = frN.orientnew('fr_0', 'Axis', [q0, frN.z])
fr1 = fr0.orientnew('fr_1', 'Axis', [q1, fr0.z])
fr2 = fr1.orientnew('fr_2', 'Axis', [q2, fr1.z])

# angular velocities of frames
fr0.set_ang_vel(frN, u0 * frN.z)
fr1.set_ang_vel(fr0, u1 * fr0.z)
fr2.set_ang_vel(fr1, u2 * fr1.z)

# origin of inertial frame
Nor = me.Point('N_or')

# joints
jt0 = Nor.locatenew('jt_0', q3 * frN.x + q4 * frN.y)
jt1 = jt0.locatenew('jt_1', l0 * fr0.x)
jt2 = jt1.locatenew('jt_2', l1 * fr1.x)
jt3 = jt2.locatenew('jt_3', l2 * fr2.x)

# mass centers
mc0 = jt0.locatenew('mc_0', l0 / 2 * fr0.x)
mc1 = jt1.locatenew('mc_1', l1 / 2 * fr1.x)
mc2 = jt2.locatenew('mc_2', l2 / 2 * fr2.x)

# constrained point velocities
Nor.set_vel(frN, 0)
jt0.set_vel(frN, u3 * frN.x + u4 * frN.y)

# unconstrained point velocities
jt1.v2pt_theory(jt0, frN, fr0)
jt2.v2pt_theory(jt1, frN, fr1)
jt3.v2pt_theory(jt2, frN, fr2)

mc0.v2pt_theory(jt0, frN, fr0)
mc1.v2pt_theory(jt1, frN, fr1)
mc2.v2pt_theory(jt2, frN, fr2)

# inertia dyads
in0 = (me.inertia(fr0, 0, 0, Izz0), mc0)
in1 = (me.inertia(fr1, 0, 0, Izz1), mc1)
in2 = (me.inertia(fr2, 0, 0, Izz2), mc2)

# bodies
bd0 = me.RigidBody('bd_0', mc0, fr0, m0, in0)
bd1 = me.RigidBody('bd_1', mc1, fr1, m1, in1)
bd2 = me.RigidBody('bd_2', mc2, fr2, m2, in2)

# kinematic differential equations
kd = [q0.diff(t) - u0,
      q1.diff(t) - u1,
      q2.diff(t) - u2,
      q3.diff(t) - u3,
      q4.diff(t) - u4]

# forces and toques
fl = [(mc0, -m0 * g * frN.z),
      (mc1, -m1 * g * frN.z),
      (mc2, -m2 * g * frN.z),
      (fr0, (T0 - T1) * frN.z),
      (fr1, (T1 - T2) * frN.z),
      (fr2, (T2 - T3) * frN.z)]

# bodies
bl = [bd0, bd1, bd2]

# q_ind and u_ind
coords = [q0, q1, q2, q3, q4]
speeds = [u0, u1, u2, u3, u4]

# equations of motion using Kane's method
kane = me.KanesMethod(frame=frN,
                      q_ind=coords,
                      u_ind=speeds,
                      kd_eqs=kd,
                      q_dependent=[],
                      configuration_constraints=[],
                      u_dependent=[],
                      velocity_constraints=[])

(fr, frstar) = kane.kanes_equations(fl, bl)
kanezero = fr + frstar

# these take to long...
# solve Kane's equations for the derivatives of the speeds
#eom = sym.solvers.solve(kanezero, q0d, q1d, q2d)

# add the kinematical differential equations to get the equations of motion
#eom.update(kane.kindiffdict())

# print the results
#me.mprint(eom)

#rhs = kane.mass_matrix.inv() * kane.forcing
#me.mprint(rhs)

#j3pos = jt3.pos_from(jt0).express(frN).simplify()
#sym.trigsimp(kanezero)
```

```
kz = sym.trigsimp(kanezero)
```

```
sym.trigsimp(kanezero[1])
```

```
kz[4]
```

```
from pydy_code_gen.code import generate_ode_function

# parameter values
ls = .25                          # segment lengths
ms = .15                          # segment masses
ws_ratio = .1                     # width is a function of length
ws = ls * ws_ratio                # segment widths
Izzs = ms * (ls**2 + ws**2) / 12  # flat plate about com

coords = [q0, q1, q2, q3, q4]
speeds = [u0, u1, u2, u3, u4]
constants = [g, m0, l0, Izz0, m1, l1, Izz1, m2, l2, Izz2]
constant_vals = [9.81, ms, ls, Izzs, ms, ls, Izzs, ms, ls, Izzs]
specified = [T0, T1, T2, T3]
mass_matrix = kane.mass_matrix_full
forcing_vector = kane.forcing_full

right_hand_side = generate_ode_function(mass_matrix, forcing_vector, constants,
                                        coords, speeds, specified, generator='cython')
```

```
def shape_space(x, t):

    tk1 = np.interp(t, ts, t1)
    tk2 = np.interp(t, ts, t2)

    return np.r_[0, tk1, tk2, 0]

args = {'constants': np.array(constant_vals),
        'specified': shape_space}

qstart = (0, 0, 0, 0, 0)
ustart = (.3, 0.0, .0, 0, 0)
#ustart = np.zeros(5)
x0 = np.hstack((qstart, ustart))

npts = 1001
ts, dt = linspace(0,  10, npts, retstep=True)

# t1 = 60 * np.ones(npts)
# t2 = -30 * np.ones(npts)

t1 = .025 * np.cos(10 * ts)
t2 = -.025 * np.cos(10 * ts + np.pi/4)

# t1 *= 0
# t2 *= 0

# t1[t1 > 0] = t1.max()
# t2[tt2 > 0] = t2.max()

# t1[theta1 < 0] = t1.min()
# t2[theta2 < 0] = t2.min()

now = time.time()
y1 = odeint(right_hand_side, x0, ts, args=(args,))
y1 = y1.T
print('Elapsed time: {0:.2f} sec'.format(time.time() - now))

# organize the data
th0, th1, th2 = y1[:3]
th1 = th1 + th0
th2 = th2 + th1

dg0, dg1, dg2 = np.rad2deg(y1[:3])
dg1 = dg1 + dg0
dg2 = dg2 + dg1

j0 = np.c_[0, 0]
j0 = y1[3:5].T
j1 = j0 + np.c_[ls   * np.cos(th0), ls   * np.sin(th0)]
j2 = j1 + np.c_[ls   * np.cos(th1), ls   * np.sin(th1)]
j3 = j2 + np.c_[ls   * np.cos(th2), ls   * np.sin(th2)]
c0 = j0 + np.c_[ls/2 * np.cos(th0), ls/2 * np.sin(th0)]
c1 = j1 + np.c_[ls/2 * np.cos(th1), ls/2 * np.sin(th1)]
c2 = j2 + np.c_[ls/2 * np.cos(th2), ls/2 * np.sin(th2)]

# center of mass
mms = np.array([ms, ms, ms])
cx = np.c_[c0[:, 0], c1[:, 0], c2[:, 0]]
cy = np.c_[c0[:, 1], c1[:, 1], c2[:, 1]]
com = np.c_[(cx * mms).sum(axis=1), (cy * mms).sum(axis=1)] / mms.sum()
```

```
fig, ax = plt.subplots()
ax.plot(ts, t1)
ax.plot(ts, t2)

fig, ax = plt.subplots()
ax.plot(np.rad2deg(y1[1]), np.rad2deg(y1[2]))
```

```
plt.close('all')
```

```
fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(6, 6))

ax1.plot(y1[0])
ax1.plot(y1[1])
ax1.plot(y1[2])

ax2.plot(y1[5])
ax2.plot(y1[6])
ax2.plot(y1[7])

ax2.set_xlabel('iteration')
ax1.set_ylabel('q', rotation='horizontal')
ax2.set_ylabel('u', rotation='horizontal')

fig.tight_layout()

#display(fig)
```

```
#%matplotlib inline

def movie_maker():

    ii = 0
    st = 0
    hst = 250

    fig, ax = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(9, 9))

    ax.axis('image')
    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)

    fig.set_facecolor('w')
    fig.tight_layout()

    time_text = ax.text(0.05, 0.9, '{0:.3} sec'.format(ts[ii]), transform=ax.transAxes)

    ellipse_params = dict(alpha=.5, edgecolor='None', color='green')
    e0 = Ellipse(xy=c0[ii], width=ws, height=1.05 * ls, angle=dg0[ii] - 90, **ellipse_params)
    e1 = Ellipse(xy=c1[ii], width=ws, height=1.05 * ls, angle=dg1[ii] - 90, **ellipse_params)
    e2 = Ellipse(xy=c2[ii], width=ws, height=1.05 * ls, angle=dg2[ii] - 90, **ellipse_params)
    patches = [e0, e1, e2]
    map(ax.add_patch, patches)

    j0l, = ax.plot(j0[ii, 0], j0[ii, 1], 'go', alpha=.6, markersize=7)
    j1l, = ax.plot(j1[st:ii, 0], j1[st:ii, 1], 'b--')
    j2l, = ax.plot(j2[st:ii, 0], j2[st:ii, 1], 'g--')
    j3l, = ax.plot(j3[st:ii, 0], j3[st:ii, 1], 'r--')

    c0l, = ax.plot(c0[st:ii, 0], c0[st:ii, 1], 'b')
    c1l, = ax.plot(c1[st:ii, 0], c1[st:ii, 1], 'g')
    c2l, = ax.plot(c2[st:ii, 0], c2[st:ii, 1], 'r')

    coml, = ax.plot(com[st:ii, 0], com[st:ii, 1], 'k-', lw=2, alpha=.9)

    def animate(ii):

        st = ii - hst
        if st < 0:
            st = 0

        e0.center = c0[ii]
        e0.angle = dg0[ii] - 90

        e1.center = c1[ii]
        e1.angle = dg1[ii] - 90

        e2.center = c2[ii]
        e2.angle = dg2[ii] - 90

        j0l.set_data(j0[ii, 0], j0[ii, 1])
        j1l.set_data(j1[st:ii, 0], j1[st:ii, 1])
        j2l.set_data(j2[st:ii, 0], j2[st:ii, 1])
        j3l.set_data(j3[st:ii, 0], j3[st:ii, 1])

        c0l.set_data(c0[st:ii, 0], c0[st:ii, 1])
        c1l.set_data(c1[st:ii, 0], c1[st:ii, 1])
        c2l.set_data(c2[st:ii, 0], c2[st:ii, 1])

        coml.set_data(com[st:ii, 0], com[st:ii, 1])

        return e0, e1, e2, j0l, j1l, j2l, j3l, c0l, c1l, c2l

    anim = animation.FuncAnimation(fig, animate, frames=len(ts),
                interval=ts[-1] / len(ts) * 1000, blit=True, repeat=False)

    return anim

movie_maker()
```

```
plt.close('all')
```

```
fig, ax = plt.subplots()

ax.plot(j0[:, 0], j0[:, 1], 'r')
ax.plot(j3[:, 0], j3[:, 1], 'g')
ax.plot(com[:, 0], com[:, 1], 'k')
```

```
sym.angular_momentum?
```

```

```

```

```

```

```

```
M = 3 * ms
1 / M * np.sum()
```

```
fig, ax = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(6, 6))

fig.tight_layout()

# display the current time
time_text = ax.text(0.05, 0.9, 'time here', transform=ax.transAxes)

# create a rectangular cart
ell = [Ellipse(xy=(0, 0), width=.05*ls, height=1.05 * ls, angle=100 * np.random.random(),
                    alpha=.5, edgecolor='None', color='green')
            for i in range(n)]
[ax.add_patch(ellis) for ellis in ell]

lim = 3 * ls
ax.set_xlim(-lim, lim)
ax.set_ylim(-lim, lim)

th0, th1, th2 = y1[:3]
dg0, dg1, dg2 = np.rad2deg(y1[:3])

j0 = np.c_[0, 0]
c0 = j0 + np.c_[ls/2 * np.cos(th0), ls/2 * np.sin(th0)]

j1 = j0 + np.c_[ls   * np.cos(th0), ls   * np.sin(th0)]
c1 = j1 + np.c_[ls/2 * np.cos(th1), ls/2 * np.sin(th1)]

j2 = j1 + np.c_[ls   * np.cos(th1), ls   * np.sin(th1)]
c2 = j2 + np.c_[ls/2 * np.cos(th2), ls/2 * np.sin(th2)]

j3 = j2 + np.c_[ls   * np.cos(th2), ls   * np.sin(th2)]

ii = 100

ell[0].center = c0[ii]
ell[1].center = c1[ii]
ell[2].center = c2[ii]

ell[0].angle = dg0[ii] - 90
ell[1].angle = dg1[ii] - 90
ell[2].angle = dg2[ii] - 90

ax.axis('image')
ax.set_xlim(-.7, .7)
ax.set_ylim(-.7, .7)

ax.plot(j1[:ii, 0], j1[:ii, 1], 'b--')
ax.plot(j2[:ii, 0], j2[:ii, 1], 'g--')
ax.plot(j3[:ii, 0], j3[:ii, 1], 'r--')

ax.plot(c0[:ii, 0], c0[:ii, 1], 'b')
ax.plot(c1[:ii, 0], c1[:ii, 1], 'g')
ax.plot(c2[:ii, 0], c2[:ii, 1], 'r')
```

```
-dg2[:20] + 90
```

```
fig, ax = plt.subplots(figsize=(7, 7))
ax.plot(j3[:, 0], j3[:, 1], alpha=.4)
ax.plot(j2[:, 0], j2[:, 1], alpha=.7)
ax.plot(j1[:, 0], j1[:, 1], alpha=.7)
ax.plot(j0[:, 0], j0[:, 1], alpha=.7)

ax.axis('equal')

fig.tight_layout()
```

```
    # the number of pendulum bobs
    numpoints = states.shape[1] / 2

    # first set up the figure, the axis, and the plot elements we want to animate
    fig = plt.figure()

    # some dimesions
    cart_width = 0.4
    cart_height = 0.2

    # set the limits based on the motion
    xmin = around(states[:, 0].min() - cart_width / 2.0, 1)
    xmax = around(states[:, 0].max() + cart_width / 2.0, 1)

    # create the axes
    ax = plt.axes(xlim=(xmin, xmax), ylim=(-1.1, 1.1), aspect='equal')

    # display the current time
    time_text = ax.text(0.04, 0.9, '', transform=ax.transAxes)

    # create a rectangular cart
    rect = Rectangle([states[0, 0] - cart_width / 2.0, -cart_height / 2],
        cart_width, cart_height, fill=True, color='red', ec='black')
    ax.add_patch(rect)

    # blank line for the pendulum
    line, = ax.plot([], [], lw=2, marker='o', markersize=6)

    # initialization function: plot the background of each frame
    def init():
        time_text.set_text('')
        rect.set_xy((0.0, 0.0))
        line.set_data([], [])
        return time_text, rect, line,

    # animation function: update the objects
    def animate(i):
        time_text.set_text('time = {:2.2f}'.format(t[i]))
        rect.set_xy((states[i, 0] - cart_width / 2.0, -cart_height / 2))
        x = hstack((states[i, 0], zeros((numpoints - 1))))
        y = zeros((numpoints))
        for j in arange(1, numpoints):
            x[j] = x[j - 1] + length * cos(states[i, j])
            y[j] = y[j - 1] + length * sin(states[i, j])
        line.set_data(x, y)
        return time_text, rect, line,

    # call the animator function
    anim = animation.FuncAnimation(fig, animate, frames=len(t), init_func=init,
            interval=t[-1] / len(t) * 1000, blit=True, repeat=False)

    # save the animation if a filename is given
    if filename is not None:
        anim.save(filename, fps=30, codec='libx264')

##############

from matplotlib import animation
ani = animation.FuncAnimation(fig, updatefig, frames=1000, interval=75,
                              repeat=True, blit=True)

# video save information
metadata = dict(author='Isaac J. Yeaton',
                description='Standard and Otsu threholded images.')
Writer = animation.writers['avconv']
writer = Writer(fps=15)  # , extra_args=['-vcodec', 'libx264'], metadata=metadata)

if True:
    #ani.save('/home/isaac/Desktop/tmp1.mp4', writer=writer)
    ani.save('/home/isaac/Desktop/processed_shadowgraph_fields.mp4',
             writer=writer)
```

```

```

```
# check that the partial evaulation works
xx = x0
#xx = y.T[696]
arguments = np.hstack((param_vals, xx))

mm = M_func(*arguments)
ff = F_func(*arguments)

# this does not work...
# aargs = dict(zip(param_names, dummy_symbols, arguments))
# mm = M_func(*aargs)
# ff = F_func(*aargs)

print solve(mm, ff).flatten()
print lstsq(mm, ff)[0].flatten()

#part_args = dict(zip(param_names, param_vals))

print solve(mfunc(*xx), ffunc(*xx)).flatten()
```

```
frame0 = I.orientnew('frame0', 'Axis', [q[0], I.z])

mc0 = O.locatenew('mc0', q[1] * frame0.x)
mc0.v2pt_theory(O, I, frame0)

in0 = me.inertia(frame0, 0, 0, Izz[0])
body0 = me.RigidBody('body0', mc0, frame0, m[0], (in0, mc0))

frame1 = frame0.orientnew('frame1', 'Axis', [q[2], I.z])
frame2 = frame1.orientnew('frame2', 'Axis', [q[3], I.z])
frame1.set_ang_vel(I, u[2] * I.z)
frame2.set_ang_vel(I, u[3] * I.z)

in1 = me.inertia(frame1, 0, 0, Izz[1])
in2 = me.inertia(frame2, 0, 0, Izz[2])

# mc1 = mc0.locatenew('mc1', l[0] * frame0.x +
#                            l[1] * sym.cos(q[2]) * frame1.x +
#                            l[1] * sym.sin(q[2]) * frame1.y)
# mc2 = mc1.locatenew('mc2', l[1] * frame1.x +
#                            l[2] * sym.cos(q[3]) * frame2.x +
#                            l[2] * sym.sin(q[3]) * frame2.y)

# after looking at the n-pendulum example
mc1 = mc0.locatenew('mc1', l[0] * frame0.x + l[1] * frame1.x)
mc2 = mc1.locatenew('mc2', l[1] * frame1.x + l[2] * frame2.x)
mc1.v2pt_theory(mc0, I, frame1)
mc2.v2pt_theory(mc1, I, frame2)

body1 = me.RigidBody('body1', mc1, frame1, m[1], (in1, mc1))
body2 = me.RigidBody('body2', mc2, frame2, m[2], (in2, mc2))

# each one might need four entries...
frames = [I, frame0, frame1, frame2]
points = [O, mc0, mc1, mc2]
bodies = [body0, body1, body2]
forces = [0, 0, 0]
kindiffs = [qq.diff(t) - uu for qq, uu in zip(q, u)]

# need to figure out forces and how to setup me.KanesMethod

# frames = [I]                              # List to hold the n + 1 frames
# points = [P0]                             # List to hold the n + 1 points
# particles = [Pa0]                         # List to hold the n + 1 particles
# forces = [(P0, f * I.x - m[0] * g * I.y)] # List to hold the n + 1 applied forces, including the input force, f
# kindiffs = [q[0].diff(t) - u[0]]          # List to hold kinematic ODE's
# for i in range(n):
#     Bi = I.orientnew('B' + str(i), 'Axis', [q[i + 1], I.z])   # Create a new frame
#     Bi.set_ang_vel(I, u[i + 1] * I.z)                         # Set angular velocity
#     frames.append(Bi)                                         # Add it to the frames list

#     Pi = points[-1].locatenew('P' + str(i + 1), l[i] * Bi.x)  # Create a new point
#     Pi.v2pt_theory(points[-1], I, Bi)                         # Set the velocity
#     points.append(Pi)                                         # Add it to the points list

#     Pai = Particle('Pa' + str(i + 1), Pi, m[i + 1])           # Create a new particle
#     particles.append(Pai)                                     # Add it to the particles list

#     forces.append((Pi, -m[i + 1] * g * I.y))                  # Set the force applied at the point

#     kindiffs.append(q[i + 1].diff(t) - u[i + 1])              # Define the kinematic ODE:  dq_i / dt - u_i = 0
```


```python
# Imports
import torch
cuda = torch.cuda.is_available()
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import sys
sys.path.append("../../semi-supervised")
sys.path.append("../../../gen-vqeeg/experiments/experiment1/")
sys.path.append("../../../gen-vqeeg/experiments/")
```

# Deep Generative Model

In this notebook we show how you can use the deep "generative model for semi-supervised learning" as presented in [[Kingma 2014]](https://arxiv.org/abs/1406.5298). The paper posits three different models, though we are just interested in two of these: the M2 model and the M1+M2 model.

The M1 model is just a variational autoencoder, so we refer to the previous notebook for more information on this. The M2 model however is an extension to the VAE to include label information for a semi-supervised objective. The structure is shown below (left: inference model, right: generative model).

The point of the generative model is to seperate the partially observed label information $y$ from the latent variable $z$ in order to learn a representation that seperates these two variables. We can use this model for semi-supervised learning as the inference model must also infer the label from the data $x$ along with the latent variable $z$.

```python
from models import DeepGenerativeModel, StackedDeepGenerativeModel

y_dim = 10
z_dim = 32
h_dim = [256, 128]

model = DeepGenerativeModel([784, y_dim, z_dim, h_dim])
model
```

    ../../semi-supervised/models/vae.py:114: UserWarning: nn.init.xavier_normal is now deprecated in favor of nn.init.xavier_normal_.
      init.xavier_normal(m.weight.data)
    ../../semi-supervised/models/dgm.py:50: UserWarning: nn.init.xavier_normal is now deprecated in favor of nn.init.xavier_normal_.
      init.xavier_normal(m.weight.data)

    DeepGenerativeModel(
      (encoder): Encoder(
        (hidden): ModuleList(
          (0): Linear(in_features=794, out_features=256, bias=True)
          (1): Linear(in_features=256, out_features=128, bias=True)
        )
        (sample): GaussianSample(
          (mu): Linear(in_features=128, out_features=32, bias=True)
          (log_var): Linear(in_features=128, out_features=32, bias=True)
        )
      )
      (decoder): Decoder(
        (hidden): ModuleList(
          (0): Linear(in_features=42, out_features=128, bias=True)
          (1): Linear(in_features=128, out_features=256, bias=True)
        )
        (reconstruction): Linear(in_features=256, out_features=784, bias=True)
        (output_activation): Sigmoid()
      )
      (classifier): Classifier(
        (dense): Linear(in_features=784, out_features=256, bias=True)
        (logits): Linear(in_features=256, out_features=10, bias=True)
      )
    )

```python
print(model.encoder.hidden[0])
print(model.decoder.hidden[0])
```

    Linear(in_features=794, out_features=256, bias=True)
    Linear(in_features=42, out_features=128, bias=True)

Notice how theres now a classifier associated with model. This classifier will just be a simple model that takes the size of the first layer encoder network. We also have a larger input space on both the encoder and decoder to make room for label information, in this case 10 labels.

## Training

Recall the ELBO from the VAE formulation, we want to construct a similar ELBO when we include labelled data $y$. In the case that we have labels, the ELBO has a simple formulation that is similar to the one for the VAE. The difference here is that we must also have a prior over labels $p(y)$, which we choose to be uniform over the different classes.

\begin{align}
\log p(x, y) &= \log \int q(z|x, y) \frac{p(x, y, z)}{q(z|xy)} \ dz \geq \int q(z|x, y) \log \frac{p(x, y, z)}{q(z|xy)} \ dz\\
&= \int q(z|x, y) [ \log p(x|z,y) + \log p(y) ] \ dz + \int q(z|x, y) \log \frac{p(z)}{q(z|xy)} \ dz\\
&= \mathbb{E}_{q(z|x, y)} [ \log p(x|z,y) + \log p(y) ] - KL(p(z)||q(z|xy)) = - \mathcal{L}(x, y)
\end{align}

In the case when the labels are not observed, we can instead integrate over all of the labels to achieve the same effect.

\begin{align}
\log p(x) &= \log \sum_{y} \int q(z,y|x) \frac{p(x, y, z)}{q(z,y|x)} \ dz \geq \sum_{y} q(y|x) \int q(z|x, y) \log \frac{p(x, y, z)}{q(z,y|x)} \ dz\\
&= \sum_{y} q(y|x) \int q(z|x, y) \log \frac{p(x, y, z)}{q(z,y|x)} \ dz + \sum_{y} q(y|x) \log q(y|x) \int q(z|x, y) \ dz\\
&= \sum_{y} q(y|x) (- \mathcal{L}(x,y)) + \mathcal{H}(q(y|x)) = - \mathcal{U}(x)
\end{align}

Notice how in both cases we need to compute the labelled bound, but in the unlabelled case we need to do it $n$ times where $n$ is the number of classes. In this model, we do not learn directly from the labelled class, as there is no cross entropy term between $y$ and our model output $q(y|x)$. We therefore add an auxiliary loss to arrive at the final loss objective.

$$\mathcal{J}^{\alpha} = \sum_{(x_l, y_l)}\mathcal{L}(x_l, y_l) + \alpha \cdot \mathbb{E}_{x_l, y_l}[- \log q(y_l|x_l)] + \sum_{(x_u)}\mathcal{U}(x_u)$$

Where $l, u$ denotes labelled and unlabelled data respectively and $\alpha$ is a hyperparameter that denotes the reliance of labelled data.

```python
from datautils import get_mnist

# Only use 10 labelled examples per class
# The rest of the data is unlabelled.
labelled, unlabelled, validation = get_mnist(location="./", batch_size=64, labels_per_class=10)
alpha = 0.1 * len(unlabelled) / len(labelled)

def binary_cross_entropy(r, x):
    return -torch.sum(x * torch.log(r + 1e-8) + (1 - x) * torch.log(1 - r + 1e-8), dim=-1)

optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, betas=(0.9, 0.999))
```

```python
from dynaconf import settings
from dataloading1 import EEGContinuousClipsDataset
from dataloading_labeled1 import EEGContinuousClipsDatasetLabeled
```

```python
%load_ext autoreload
%autoreload 2
```

```python
# TODO: figure out how to import dynaconfig settings
EEGML_EEGDBS='/mnt/data1/eegdbs'
EEGML_SEC01_DATADIR = "/mnt/data1/eegdbs/SEC-0.1"
EEGML_TEMPLE_TUEG_V100 = "/mnt/data1/eegdbs/TUH/temple/v1.0.0"
EEGML_TEMPLE_TUEG_v120 = "/mnt/data2/eegdbs2/TUH/tuh_eeg/v1.2.0"
EEGML_TEMPLE_SZ_V151 = "/mnt/data1/eegdbs/TUH/temple/tuh_eeg_seizure/v1.5.1"
EEGML_STEV_NEONATES = "/mnt/data1/eegdbs/stevenson_neonatal_eeg"
# temple artifacts corpus (1) eye movement (EYEM), (2) chewing (CHEW), (3) shivering (SHIV), (4) electrode pop, electrode static, and lead artifacts (ELPP), and (5) muscle artifacts (MUSC)
EEGML_TUH_TUAR_v100 = "/mnt/data2/eegdbs2/TUH/tuh_eeg_artifact/v1.0.0"
EEGML_TUH_TUAR_V100 = "/mnt/data2/eegdbs2/TUH/tuh_eeg_artifact/v1.0.0"

EEGML_BONN="/mnt/data2/eegdbs2/bonn-ieeg-data"

VISION_GOT10K="/mnt/data2/GOT-10k"
```

```python
class PartiallyLabeledDataSet:

    def __init__(
        self,
        dataset,
        prob_label=0.1,
        prob_one_chan_bad=0.5,
        num_channels_select=2,
        noise_amplitude=50.0,
    ):
        self.prob_label = prob_label
        self.prob_one_chan_bad = prob_one_chan_bad
        self.num_channels_select = num_channels_select
        self.dataset = dataset
        self.fs = float(dataset.eegs[0].sample_frequency)
        self.clip_length = dataset.clip_length  # seconds
        self.clip_shape = dataset[0].shape
        self.noise_amplitude = noise_amplitude

    def __getitem__(self, idx):
        if np.random.uniform() < self.prob_label:  # then we have a label
            hasLabel = True
        else:
            hasLabel = False

        hasBad = np.random.uniform() < self.prob_one_chan_bad
#         channels = np.random.randint(0, 22, size=self.num_channels_select)
        channels = np.random.randint(0, 10, size=self.num_channels_select)
        # truth_arr = torch.rand(batch_size) > self.prob_one_chan_bad
        try:
            if hasBad:
                y = np.ones(self.num_channels_select)
                noisy_channels_indx = np.random.randint(0, 2, size=1)
                y[noisy_channels_indx] = 0
                noise = self.noise_amplitude * np.random.uniform(
                    -0.5, 0.5, size=self.clip_shape[1]
                )
                arr = self.dataset[idx][channels]
                arr[noisy_channels_indx] += noise
                if hasLabel:
                    return arr, y
                else:
                    return arr, np.zeros(self.num_channels_select)
            else:
                if hasLabel:
                    return self.dataset[idx][channels], np.ones(self.num_channels_select)
                else:
                    return self.dataset[idx][channels], np.zeros(self.num_channels_select)
        except ValueError:
            return self.__getitem__(idx+1)

    def __len__(self):
        return self.dataset.length
        # we will have a bad channel
```

```python

```

```python
def gen_sin(fs, length):
    f0 = 1
    t = np.arange(length)
    sinusoid = np.sin(2*np.pi*t*(f0/fs))
    sinusoid = np.expand_dims(sinusoid, axis=0)
    return sinusoid

def custom_norm(x):
    abs_mean = (np.abs(x).mean())
    if abs_mean == 0:
        return x
    x = x/(abs_mean)
    x = (np.tanh(x) + 1)/2
    return x

def custom_norm_batch(batch):
    out = []
    for x in batch:
        out.append(custom_norm(x))
    out = np.vstack(out)
    return torch.from_numpy(out).type(torch.DoubleTensor)

class ConditionalyLabeledDataSet:

    def __init__(
        self,
        dataset,
        has_label,
        prob_one_chan_bad=.5,
        num_channels_select=1,
        noise_amplitude=1.0,
    ):
        self.prob_one_chan_bad = prob_one_chan_bad
        self.num_channels_select = num_channels_select
        self.dataset = dataset
        self.fs = float(dataset.eegs[0].sample_frequency)
        self.clip_length = dataset.clip_length  # seconds
        self.clip_shape = dataset[0].shape
        self.noise_amplitude = noise_amplitude
        self.has_label = has_label
        self.sin_noise = gen_sin(60, self.clip_shape[1])

    def convert_one_hot(self, y):
        one_hot_y = []
        noisy = np.zeros(10)
        noisy[1] = 1
        clean = np.zeros(10)
        clean[0] = 1

        for label in y:
            if label == 0:
                one_hot_y.append(noisy)
            else:
                one_hot_y.append(clean)
        return np.asarray(one_hot_y)

    def __getitem__(self, idx):
        hasLabel = self.has_label

        hasBad = np.random.uniform() < self.prob_one_chan_bad
#         channels = np.random.randint(0, 22, size=self.num_channels_select)
        channels = np.random.randint(0, 18, size=self.num_channels_select)
        # truth_arr = torch.rand(batch_size) > self.prob_one_chan_bad
        try:
            if hasBad:
                y = np.ones(self.num_channels_select)
                noisy_channels_indx = np.random.randint(0, self.num_channels_select, size=1)
                y[noisy_channels_indx] = 0
                noise = self.noise_amplitude * self.sin_noise
                arr = custom_norm_batch(self.dataset[idx][channels]).numpy()
                arr[noisy_channels_indx] += noise
                if hasLabel:
                    return custom_norm_batch(arr), self.convert_one_hot(y)
                else:
                    return custom_norm_batch(arr)
            else:
                if hasLabel:
                    return custom_norm_batch(self.dataset[idx][channels]), self.convert_one_hot(np.ones(self.num_channels_select))
                else:
                    return custom_norm_batch(self.dataset[idx][channels])
        except FileNotFoundError: #ValueError:
            return self.__getitem__(idx+1)

    def __len__(self):
        return self.dataset.length
        # we will have a bad channel
```

```python
from dataloading1 import (
    EEGContinuousClipsDataset,
    reasonable_normal_eegs_sample_5s_clip,
    plot_tuh_tcp,
)
cliplist = reasonable_normal_eegs_sample_5s_clip
dset_labeled = EEGContinuousClipsDataset(
    EEGML_TUH_TUAR_V100 + "/eeg.h5.files.csv", EEGML_TUH_TUAR_V100, clip_length=4, files_range=(0, 40), normalize=False
)
dset_unlabeled = EEGContinuousClipsDataset(
    EEGML_TUH_TUAR_V100 + "/eeg.h5.files.csv", EEGML_TUH_TUAR_V100, clip_length=4, files_range=(40,110), normalize=False)

# partially_labelled_dataset = PartiallyLabeledDataSet(dset)
# partially_labelled_loader = torch.utils.data.DataLoader(partially_labelled_dataset, batch_size=64, shuffle=True)
labelled_dataset = ConditionalyLabeledDataSet(dset_labeled, has_label=True)
unlabelled_dataset = ConditionalyLabeledDataSet(dset_unlabeled, has_label=False)

unlabelled = torch.utils.data.DataLoader(unlabelled_dataset, batch_size=64*2, shuffle=True)
labelled = torch.utils.data.DataLoader(labelled_dataset, batch_size=64*2, shuffle=True)
```

    length filenames 40
    length filenames 70

```python
### OLD VERSION (uncurated set)
'''
unlabelled_dataset = EEGContinuousClipsDataset(
        EEGML_TUH_TUAR_v100 + "/eeg.h5.files.csv", EEGML_TUH_TUAR_v100, clip_length=4
)
unlabelled = torch.utils.data.DataLoader(unlabelled_dataset, batch_size=64, shuffle=True)

labelled_dataset = EEGContinuousClipsDatasetLabeled(
        EEGML_TUH_TUAR_v100 + "/eeg.h5.files.csv", EEGML_TUH_TUAR_v100, clip_length=4
)
labelled = torch.utils.data.DataLoader(labelled_dataset, batch_size=64, shuffle=True)
print(len(unlabelled_dataset))
print(len(labelled_dataset))
'''
```

    '\nunlabelled_dataset = EEGContinuousClipsDataset(\n        EEGML_TUH_TUAR_v100 + "/eeg.h5.files.csv", EEGML_TUH_TUAR_v100, clip_length=4\n)\nunlabelled = torch.utils.data.DataLoader(unlabelled_dataset, batch_size=64, shuffle=True)\n\nlabelled_dataset = EEGContinuousClipsDatasetLabeled(\n        EEGML_TUH_TUAR_v100 + "/eeg.h5.files.csv", EEGML_TUH_TUAR_v100, clip_length=4\n)\nlabelled = torch.utils.data.DataLoader(labelled_dataset, batch_size=64, shuffle=True)\nprint(len(unlabelled_dataset))\nprint(len(labelled_dataset))\n'

```python
print("# Labelled batches", len(labelled))
print("# Unlabelled batches", len(unlabelled))
```

    # Labelled batches 93
    # Unlabelled batches 150

```python
from itertools import cycle
from inference import SVI, ImportanceWeightedSampler

# You can use importance weighted samples [Burda, 2015] to get a better estimate
# on the log-likelihood.
alpha = 0.1 * len(unlabelled_dataset) / len(labelled_dataset)
# alpha = 0.1 * (.5) # .1 is prob of label

def binary_cross_entropy(r, x):
    return -torch.sum(x * torch.log(r + 1e-8) + (1 - x) * torch.log(1 - r + 1e-8), dim=-1)

optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, betas=(0.9, 0.999))
sampler = ImportanceWeightedSampler(mc=1, iw=1)

if cuda: model = model.cuda()
elbo = SVI(model, likelihood=binary_cross_entropy, sampler=sampler)
```

The library is conventially packed with the `SVI` method that does all of the work of calculating the lower bound for both labelled and unlabelled data depending on whether the label is given. It also manages to perform the enumeration of all the labels.

Remember that the labels have to be in a *one-hot encoded* format in order to work with SVI.

```python
from torch.autograd import Variable
for epoch in range(30):
    model.train()
    total_loss, accuracy = (0, 0)
#     labelled = iter(labelled)
#     unlabelled = iter(unlabelled)
#     for i in range(num_batches):
#         print("i", i)
#         x, y = next(labelled)
#         u = next(unlabelled)
#     for (x, y), (u) in zip(cycle(labelled), unlabelled):
    for i, ((x, y), (u)) in enumerate(zip(cycle(labelled), unlabelled)):
        x, y, u = Variable((x).type(torch.FloatTensor)), Variable(y.type(torch.FloatTensor)), Variable((u).type(torch.FloatTensor))

        x = x[:, 0, :784]
        u = u[:, 0, :784]
        y = y[:, 0, :10]
#         y = y[:, :10]

        if cuda:
            # They need to be on the same device and be synchronized.
            x, y = x.cuda(device=0), y.cuda(device=0)
            u = u.cuda(device=0)

        L = -elbo(x, y)
        U = -elbo(u)

        # Add auxiliary classification loss q(y|x)
        logits = model.classify(x)

        # Regular cross entropy
        classication_loss = torch.sum(y * torch.log(logits + 1e-8), dim=1).mean()

        J_alpha = L - alpha * classication_loss + U

        J_alpha.backward()
        optimizer.step()
        optimizer.zero_grad()

        total_loss += J_alpha.item()
#         total_loss += J_alpha.data[0]
        accuracy += torch.mean((torch.max(logits, 1)[1].data == torch.max(y, 1)[1].data).float())

    if epoch % 1 == 0:
#         model.eval()
        m = len(unlabelled)
        print("Epoch: {}".format(epoch))
        print("[Train]\t\t J_a: {:.2f}, accuracy: {:.2f}".format(total_loss / m, accuracy / m))

#         total_loss, accuracy = (0, 0)
#         for x, y in validation:
#             x, y = Variable(x), Variable(y)

#             if cuda:
#                 x, y = x.cuda(device=0), y.cuda(device=0)

#             L = -elbo(x, y)
#             U = -elbo(x)

#             logits = model.classify(x)
#             classication_loss = -torch.sum(y * torch.log(logits + 1e-8), dim=1).mean()

#             J_alpha = L + alpha * classication_loss + U

# #             total_loss += J_alpha.data[0]
#             total_loss += J_alpha.item()

#             _, pred_idx = torch.max(logits, 1)
#             _, lab_idx = torch.max(y, 1)
#             accuracy += torch.mean((torch.max(logits, 1)[1].data == torch.max(y, 1)[1].data).float())

#         m = len(validation)
#         print("[Validation]\t J_a: {:.2f}, accuracy: {:.2f}".format(total_loss / m, accuracy / m))
```

## Conditional generation

When the model is done training you can generate samples conditionally given some normal distributed noise $z$ and a label $y$.

*The model below has only trained for 10 iterations, so the perfomance is not representative*.

```python
i += 1
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = [20, 10]
sys.path.append("../../../visualize/")
import stacklineplot
sample = np.vstack((labelled_dataset[i][0][0], labelled_dataset[i+1][0][0], labelled_dataset[i+2][0][0], labelled_dataset[i+3][0][0]))
print(labelled_dataset[i][1][0, :10])
print(labelled_dataset[i+1][1][0, :10])
print(labelled_dataset[i+2][1][0, :10])
print(labelled_dataset[i+3][1][0, :10])
stacklineplot.stackplot(sample, yscale=.3)
```

### Normal Signal

```python
from utils import onehot

model.eval()

z = Variable(torch.randn(16, 32)).cuda()
# z = Variable(torch.zeros(16, 32)).cuda()
# z = Variable(torch.arange(0, 512).view(16, 32).type(torch.FloatTensor)).cuda() / (32* 6) -1

# Generate a batch of 8s
y = Variable(onehot(10)(0).repeat(16, 1)).cuda()
# y = Variable(torch.zeros(16,10)).cuda()

x_mu = model.sample(z, y)
samples = x_mu.data.view(-1, 28 * 28).cpu().numpy()
ylabels=[str(x) for x in range(16)]
stacklineplot.stackplot(samples[:14], ylabels=ylabels, yscale=1)
```

### 60hz Noise

```python
z = Variable(torch.randn(16, 32)).cuda()
# z = Variable(torch.zeros(16, 32)).cuda()
# z = Variable(torch.arange(0, 512).view(16, 32).type(torch.FloatTensor)).cuda() / (32* 6) -1

# Generate a batch of 8s
y = Variable(onehot(10)(1).repeat(16, 1)).cuda()
# y = Variable(torch.zeros(16,10)).cuda()

x_mu = model.sample(z, y)
samples = x_mu.data.view(-1, 28 * 28).cpu().numpy()
ylabels=[str(x) for x in range(16)]
stacklineplot.stackplot(samples[:14], ylabels=ylabels, yscale=1)
```

```python
f, axarr = plt.subplots(1, 16, figsize=(18, 12))

samples = x_mu.data.view(-1, 28, 28).cpu().numpy()

for i, ax in enumerate(axarr.flat):
    ax.imshow(samples[i])
    ax.axis("off")
```

## Stacked Deep Generative Model

The M1+M2 model also described in the same paper is an M1 model (VAE) with an M2 model stacked on top of it. That means that we train the a VAE end-to-end on the data given dataset, then we use the learned encoder as a feature extractor and feed the data transformed by the M1 encoder into the M2 model.

We approach is somewhat similar to restricted boltzmann machines (RBMs) in the sense that we perform a layerwise training of the whole model by first training a level-1 feature extractor and stacking another model on top of this. The stacked model is therefore also more modular, but cannot be trained end-to-end, which is a downside.

```python
from models import VariationalAutoencoder, StackedDeepGenerativeModel

features = VariationalAutoencoder([784, z_dim, h_dim])
model = StackedDeepGenerativeModel([784, y_dim, z_dim, h_dim], features)
```

    ../../semi-supervised/models/vae.py:114: UserWarning: nn.init.xavier_normal is now deprecated in favor of nn.init.xavier_normal_.
      init.xavier_normal(m.weight.data)
    ../../semi-supervised/models/dgm.py:50: UserWarning: nn.init.xavier_normal is now deprecated in favor of nn.init.xavier_normal_.
      init.xavier_normal(m.weight.data)

Typically, you would want to load a pretrained feature extractor VAE instead.

```python
features = torch.load("./your-pretrained-vae.pt")
```

```python
from torch.autograd import Variable

for epoch in range(10):
    model.train()
    total_loss, accuracy = (0, 0)
    for (x, y), (u, _) in zip(cycle(labelled), unlabelled):
        # Wrap in variables
        x, y, u = Variable(x), Variable(y), Variable(u)

        if cuda:
            # They need to be on the same device and be synchronized.
            x, y = x.cuda(device=0), y.cuda(device=0)
            u = u.cuda(device=0)

        L = -elbo(x, y)
        U = -elbo(u)

        # Add auxiliary classification loss q(y|x)
        logits = model.classify(x)

        # Regular cross entropy
        classication_loss = torch.sum(y * torch.log(logits + 1e-8), dim=1).mean()

        J_alpha = L - alpha * classication_loss + U

        J_alpha.backward()
        optimizer.step()
        optimizer.zero_grad()

        total_loss += J_alpha.data[0]
        accuracy += torch.mean((torch.max(logits, 1)[1].data == torch.max(y, 1)[1].data).float())

    if epoch % 1 == 0:
        model.eval()
        m = len(unlabelled)
        print("Epoch: {}".format(epoch))
        print("[Train]\t\t J_a: {:.2f}, accuracy: {:.2f}".format(total_loss / m, accuracy / m))

        total_loss, accuracy = (0, 0)
        for x, y in validation:
            x, y = Variable(x), Variable(y)

            if cuda:
                x, y = x.cuda(device=0), y.cuda(device=0)

            L = -elbo(x, y)
            U = -elbo(x)

            logits = model.classify(x)
            classication_loss = -torch.sum(y * torch.log(logits + 1e-8), dim=1).mean()

            J_alpha = L + alpha * classication_loss + U

            total_loss += J_alpha.data[0]

            _, pred_idx = torch.max(logits, 1)
            _, lab_idx = torch.max(y, 1)
            accuracy += torch.mean((torch.max(logits, 1)[1].data == torch.max(y, 1)[1].data).float())

        m = len(validation)
        print("[Validation]\t J_a: {:.2f}, accuracy: {:.2f}".format(total_loss / m, accuracy / m))
```

## Additional tips

You can change the built-in classifier with your own, for example if you want a CNN as a classifier you can do the following.

```python
import torch.nn as nn
import torch.nn.functional as F

class ConvolutionalClassifier(nn.Module):
    def __init__(self):
        super(ConvolutionalClassifier, self).__init__()
        self.conv1 = nn.Conv2d(1, 64, kernel_size=3)
        self.conv2 = nn.Conv2d(64, 32, kernel_size=3)
        self.pool = nn.MaxPool2d(kernel_size=4)

        size = int((28 - 3) + 1)//4
        size = int((size - 3) + 1)//4

        self.fc1 = nn.Linear(32*size**2, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        batch, *_ = x.size()
        x = x.view(-1, 1, 28, 28)
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = F.relu(x.view(batch, -1))
        x = self.fc1(x)
        x = self.fc2(x)
        return F.softmax(x, dim=-1)

classifier = ConvolutionalClassifier()
model.classifier = classifier
```

```python
model
```

    DeepGenerativeModel(
      (encoder): Encoder(
        (hidden): ModuleList(
          (0): Linear(in_features=794, out_features=256)
          (1): Linear(in_features=256, out_features=128)
        )
        (sample): GaussianSample(
          (mu): Linear(in_features=128, out_features=32)
          (log_var): Linear(in_features=128, out_features=32)
        )
      )
      (decoder): Decoder(
        (hidden): ModuleList(
          (0): Linear(in_features=42, out_features=128)
          (1): Linear(in_features=128, out_features=256)
        )
        (reconstruction): Linear(in_features=256, out_features=784)
        (output_activation): Sigmoid()
      )
      (classifier): ConvolutionalClassifier(
        (conv1): Conv2d (1, 64, kernel_size=(3, 3), stride=(1, 1))
        (conv2): Conv2d (64, 32, kernel_size=(3, 3), stride=(1, 1))
        (pool): MaxPool2d(kernel_size=(4, 4), stride=(4, 4), dilation=(1, 1))
        (fc1): Linear(in_features=32, out_features=50)
        (fc2): Linear(in_features=50, out_features=10)
      )
    )

```python
x = labelled_dataset[0][1]
for i in range(1, len(labelled_dataset)):
    x += labelled_dataset[i][1]
```

```python
x
```

    tensor([506.,   0.,   0.,   0.,   6.,   0.,   6.,   0.,   0.,   0.,   0.,   0.,
              0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,
              0.,   0.,   0.], dtype=torch.float64)

```python

```

examples/notebooks/Deep Generative Model EEG.ipynb