# Start-to-Finish Example: Setting up Polytropic [TOV](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) Initial Data, in Curvilinear Coordinates

## Authors: Zach Etienne & Phil Chang
### Formatting improvements courtesy Brandon Clark

## This module sets up initial data for a TOV star in *spherical, isotropic coordinates*, using the *Numerical* ADM Spherical to BSSN Curvilinear initial data module (numerical = BSSN $\lambda^i$'s are computed using finite-difference derivatives instead of exact expressions).

**Module Status:**  Validated

**Validation Notes:** This module has been validated to exhibit convergence to zero of the Hamiltonian constraint violation at the expected order to the exact solution (see [plots](#convergence) at bottom). Note that convergence at the surface of the star will be lower order due to the sharp drop to zero in $T^{\mu\nu}$.

### NRPy+ Source Code for this module:
* [TOV/TOV_Solver.py](../edit/TOV/TOV_Solver.py) [\[**tutorial**\]](Tutorial-ADM_Initial_Data-TOV.ipynb): Tolman-Oppenheimer-Volkoff (TOV) initial data; defines all ADM variables and nonzero $T^{\mu\nu}$ components in Spherical basis.
* [BSSN/ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py](../edit/BSSN/ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py); [\[**tutorial**\]](Tutorial-ADM_Initial_Data-Converting_Numerical_ADM_Spherical_or_Cartesian_to_BSSNCurvilinear.ipynb): *Numerical* Spherical ADM$\to$Curvilinear BSSN converter function
* [BSSN/BSSN_constraints.py](../edit/BSSN/BSSN_constraints.py); [\[**tutorial**\]](Tutorial-BSSN_constraints.ipynb): Hamiltonian constraint in BSSN curvilinear basis/coordinates

## Introduction:
Here we use NRPy+ to generate initial data for a [simple polytrope TOV star](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation).

The entire algorithm is outlined below, with NRPy+-based components highlighted in green.

1. Allocate memory for gridfunctions, including temporary storage for the RK4 time integration.
1. Set gridfunction values to initial data (**This module**).
1. Evolve the system forward in time using RK4 time integration. At each RK4 substep, do the following:
    1. Evaluate BSSN RHS expressions.
    1. Apply singular, curvilinear coordinate boundary conditions [*a la* the SENR/NRPy+ paper](https://arxiv.org/abs/1712.07658)
    1. Apply constraints on conformal 3-metric: $\det{\bar{\gamma}_{ij}}=\det{\hat{\gamma}_{ij}}$
1. At the end of each iteration in time, output the Hamiltonian constraint violation.
1. Repeat above steps at two numerical resolutions to confirm convergence to zero.

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

This notebook is organized as follows

1. [Step 1](#initializenrpy): Set core NRPy+ parameters for numerical grids and reference metric
1. [Step 2](#adm_id_tov): Set up ADM initial data for polytropic TOV Star
    1. [Step 2.a](#tov_interp): Interpolating the TOV data file as needed
    1. [Step 2.b](#source): Compute source terms $S_{ij}$, $S_{i}$, $S$, and $\rho$
    1. [Step 2.c](#jacobian): Jacobian transformation on the ADM/BSSN source terms
    1. [Step 2.d](#tensor): Rescale tensorial quantities
1. [Step 3](#adm_id_spacetime): Convert ADM spacetime quantity initial data from Spherical to BSSN Curvilinear coordinates
1. [Step 4](#validate): Validating that the TOV initial data satisfy the Hamiltonian constraint
    1. [Step 4.a](#ham_const_output): Output the Hamiltonian Constraint
    1. [Step 4.b](#apply_bcs): Apply singular, curvilinear coordinate boundary conditions
    1. [Step 4.c](#enforce3metric): Enforce conformal 3-metric $\det{\bar{\gamma}_{ij}}=\det{\hat{\gamma}_{ij}}$
1. [Step 5](#mainc): `TOV_Playground.c`: The Main C Code
1. [Step 6](#plot): Plotting the single-neutron-star initial data
1. [Step 7](#convergence): Validation: Convergence of numerical errors (Hamiltonian constraint violation) to zero
1. [Step 8](#latex_pdf_output): Output this notebook to $\LaTeX$-formatted PDF file

# Step 1: Set core NRPy+ parameters for numerical grids and reference metric \[Back to [top](#toc)\]
$$\label{initializenrpy}$$

```python
# First we import needed core NRPy+ modules
from outputC import *
import NRPy_param_funcs as par
import grid as gri
import loop as lp
import indexedexp as ixp
import finite_difference as fin
import reference_metric as rfm

thismodule = "TOV_ID_setup"

# Set spatial dimension (must be 3 for BSSN)
DIM = 3
par.set_parval_from_str("grid::DIM",DIM)

# Then we set the coordinate system for the numerical grid
par.set_parval_from_str("reference_metric::CoordSystem","Spherical")
rfm.reference_metric() # Create ReU, ReDD needed for rescaling B-L initial data, generating BSSN RHSs, etc.

import BSSN.BSSN_quantities as Bq
par.set_parval_from_str("BSSN.BSSN_quantities::EvolvedConformalFactor_cf","phi")

#################
# Next output C headers related to the numerical grids we just set up:
#################

# First output the coordinate bounds xxmin[] and xxmax[]:
with open("BSSN/xxminmax.h", "w") as file:
    file.write("const REAL xxmin[3] = {"+str(rfm.xxmin[0])+","+str(rfm.xxmin[1])+","+str(rfm.xxmin[2])+"};\n")
    file.write("const REAL xxmax[3] = {"+str(rfm.xxmax[0])+","+str(rfm.xxmax[1])+","+str(rfm.xxmax[2])+"};\n")

# Generic coordinate NRPy+ file output, Part 2: output the conversion from (x0,x1,x2) to Cartesian (x,y,z)
outputC([rfm.xxCart[0],rfm.xxCart[1],rfm.xxCart[2]],["xCart[0]","xCart[1]","xCart[2]"],
        "BSSN/xxCart.h")

# Register the Hamiltonian as a gridfunction, to be used later.
H = gri.register_gridfunctions("AUX","H")
```

    Wrote to file "BSSN/xxCart.h"

# Step 2: Set up ADM initial data for polytropic TOV Star \[Back to [top](#toc)\]
$$\label{adm_id_tov}$$

As documented [in the TOV Initial Data NRPy+ Tutorial Module](Tutorial-TOV_Initial_Data.ipynb) ([older version here](Tutorial-GRMHD_UnitConversion.ipynb)), we will now set up TOV initial data, storing the densely-sampled result to file (***Courtesy Phil Chang***).

The TOV solver uses an ODE integration routine provided by scipy, so we first make sure that scipy is installed:

```python
!pip install scipy > /dev/null
```

    DEPRECATION: A future version of pip will drop support for Python 2.7. More details about Python 2 support in pip, can be found at https://pip.pypa.io/en/latest/development/release-process/#python-2-support

Next we call the [`TOV.TOV_Solver()` function](../edit/TOV/TOV_Solver.py) ([NRPy+ Tutorial module](Tutorial-ADM_Initial_Data-TOV.ipynb)) to set up the initial data, using the default parameters for initial data. This function outputs the solution to a file named "outputTOVpolytrope.txt".

```python
############################
# Single polytrope example #
############################
import TOV.Polytropic_EOSs as ppeos

## SINGLE POLYTROPE OPTION:
# Set neos = 1 (single polytrope)
neos = 1

# Set rho_poly_tab (not needed for a single polytrope)
rho_poly_tab = []

# Set Gamma_poly_tab
Gamma_poly_tab = [2.0]

# Set K_poly_tab0
K_poly_tab0 = 1. # ZACH NOTES: CHANGED FROM 100.

# Set the eos quantities
eos = ppeos.set_up_EOS_parameters__complete_set_of_input_variables(neos,rho_poly_tab,Gamma_poly_tab,K_poly_tab0)
import TOV.TOV_Solver as TOV
TOV_Mass, TOV_RSchw, TOV_Riso = TOV.TOV_Solver(eos,
                                               outfile="outputTOVpolytrope.txt",
                                               rho_baryon_central=0.129285,
                                               verbose = True,
                                               return_M_RSchw_and_Riso=True)

## PIECEWISE POLYTROPE OPTION:
# eos = ppeos.set_up_EOS_parameters__Read_et_al_input_variables("APR4")
# import TOV.TOV_Solver as TOV
# TOV_Mass, TOV_RSchw, TOV_Riso  = TOV.TOV_Solver(eos,
#                                      outfile="outputTOVpolytrope.txt",
#                                      rho_baryon_central=2.0,
#                                      accuracy="medium",
#                                      verbose = True,
#                                      return_M_RSchw_and_Riso=True)
```

    (1256, 1256, 1256, 1256, 1256, 1256)
    Just generated a TOV star with
    * M        = 1.405030336771405e-01 ,
    * R_Schw   = 9.566044579232513e-01 ,
    * R_iso    = 8.100085557410308e-01 ,
    * M/R_Schw = 1.468768334847266e-01

## Step 2.a: Interpolating the TOV data file as needed \[Back to [top](#toc)\]
$$\label{tov_interp}$$

The TOV data file just written stored $\left(r,\rho(r),P(r),M(r),e^{\nu(r)}\right)$, where $\rho(r)$ is the total mass-energy density (cf. $\rho_{\text{baryonic}}$).

**METRIC DATA IN TERMS OF ADM QUANTITIES**

The [TOV line element](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation) in *Schwarzschild coordinates* is written (in the $-+++$ form):
$$
ds^2 = - c^2 e^\nu dt^2 + \left(1 - \frac{2GM}{rc^2}\right)^{-1} dr^2 + r^2 d\Omega^2.
$$

In *isotropic coordinates* with $G=c=1$ (i.e., the coordinate system and units we'd prefer to use), the ($-+++$ form) line element is written:
$$
ds^2 = - e^{\nu} dt^2 + e^{4\phi} \left(d\bar{r}^2 + \bar{r}^2 d\Omega^2\right),
$$
where $\phi$ here is the *conformal factor*.

The ADM 3+1 line element for this diagonal metric in isotropic spherical coordinates is given by:
$$
ds^2 = (-\alpha^2 + \beta_k \beta^k) dt^2 + \gamma_{\bar{r}\bar{r}} d\bar{r}^2 + \gamma_{\theta\theta} d\theta^2+ \gamma_{\phi\phi} d\phi^2,
$$

from which we can immediately read off the ADM quantities:
\begin{align}
\alpha &= e^{\nu(\bar{r})/2} \\
\beta^k &= 0 \\
\gamma_{\bar{r}\bar{r}} &= e^{4\phi}\\
\gamma_{\theta\theta} &= e^{4\phi} \bar{r}^2 \\
\gamma_{\phi\phi} &= e^{4\phi} \bar{r}^2 \sin^2 \theta \\
\end{align}

**STRESS-ENERGY TENSOR $T^{\mu\nu}$**

We will also need the stress-energy tensor $T^{\mu\nu}$. [As discussed here](https://en.wikipedia.org/wiki/Tolman%E2%80%93Oppenheimer%E2%80%93Volkoff_equation), the stress-energy tensor is diagonal:

\begin{align}
T^t_t &= -\rho \\
T^i_j &= P \delta^i_j \\
\text{All other components of }T^\mu_\nu &= 0.
\end{align}

Since $\beta^i=0$ the inverse metric expression simplifies to (Eq. 4.49 in [Gourgoulhon](https://arxiv.org/pdf/gr-qc/0703035.pdf)):
$$
g^{\mu\nu} = \begin{pmatrix}
-\frac{1}{\alpha^2} & \frac{\beta^i}{\alpha^2} \\
\frac{\beta^i}{\alpha^2} & \gamma^{ij} - \frac{\beta^i\beta^j}{\alpha^2}
\end{pmatrix} =
\begin{pmatrix}
-\frac{1}{\alpha^2} & 0 \\
0 & \gamma^{ij}
\end{pmatrix},
$$

and since the 3-metric is diagonal we get

\begin{align}
\gamma^{\bar{r}\bar{r}} &= e^{-4\phi}\\
\gamma^{\theta\theta} &= e^{-4\phi}\frac{1}{\bar{r}^2} \\
\gamma^{\phi\phi} &= e^{-4\phi}\frac{1}{\bar{r}^2 \sin^2 \theta}.
\end{align}

Thus raising $T^\mu_\nu$ yields a diagonal $T^{\mu\nu}$

\begin{align}
T^{tt} &= -g^{tt} \rho = \frac{1}{\alpha^2} \rho = e^{-\nu(\bar{r})} \rho \\
T^{\bar{r}\bar{r}} &= g^{\bar{r}\bar{r}} P = \frac{1}{e^{4 \phi}} P \\
T^{\theta\theta} &= g^{\theta\theta} P = \frac{1}{e^{4 \phi}\bar{r}^2} P\\
T^{\phi\phi} &= g^{\phi\phi} P = \frac{1}{e^{4\phi}\bar{r}^2 \sin^2 \theta} P
\end{align}

As all input quantities are functions of $r$, we will simply read the solution from file and interpolate it to the values of $r$ needed by the initial data.

1. First we define functions `ID_TOV_ADM_quantities()` and `ID_TOV_TUPMUNU()` that call the [1D TOV interpolator function](../edit/TOV/tov_interp.h) to evaluate the ADM spacetime quantities and $T^{\mu\nu}$, respectively, at any given point $(r,\theta,\phi)$ in the Spherical basis. All quantities are defined as above.
1. Next we will construct the BSSN/ADM source terms $\{S_{ij},S_{i},S,\rho\}$ in the Spherical basis
1. Then we will perform the Jacobian transformation on $\{S_{ij},S_{i},S,\rho\}$ to the desired (xx0,xx1,xx2) basis
1. Next we call the *Numerical* Spherical ADM$\to$Curvilinear BSSN converter function to conver the above ADM quantities to the rescaled BSSN quantities in the desired curvilinear coordinate system: [BSSN/ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py](../edit/BSSN/ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py); [\[**tutorial**\]](Tutorial-ADM_Initial_Data-Converting_Numerical_ADM_Spherical_or_Cartesian_to_BSSNCurvilinear.ipynb).

```python
with open("BSSN/ID_TOV_ADM_quantities.h", "w") as file:
    file.write("""
// This function takes as input either (x,y,z) or (r,th,ph) and outputs
//   all ADM quantities in the Cartesian or Spherical basis, respectively.
void ID_TOV_ADM_quantities(
                     const REAL xyz_or_rthph[3],

                     const ID_inputs other_inputs,

                     REAL *gammaDD00,REAL *gammaDD01,REAL *gammaDD02,REAL *gammaDD11,REAL *gammaDD12,REAL *gammaDD22,
                     REAL *KDD00,REAL *KDD01,REAL *KDD02,REAL *KDD11,REAL *KDD12,REAL *KDD22,
                     REAL *alpha,
                     REAL *betaU0,REAL *betaU1,REAL *betaU2,
                     REAL *BU0,REAL *BU1,REAL *BU2) {
      const REAL r  = xyz_or_rthph[0];
      const REAL th = xyz_or_rthph[1];
      const REAL ph = xyz_or_rthph[2];

      REAL rho,rho_baryon,P,M,expnu,exp4phi;
      TOV_interpolate_1D(r,other_inputs.Rbar,other_inputs.Rbar_idx,other_inputs.interp_stencil_size,
                         other_inputs.numlines_in_file,
                         other_inputs.r_Schw_arr,other_inputs.rho_arr,other_inputs.rho_baryon_arr,other_inputs.P_arr,other_inputs.M_arr,
                         other_inputs.expnu_arr,other_inputs.exp4phi_arr,other_inputs.rbar_arr,
                         &rho,&rho_baryon,&P,&M,&expnu,&exp4phi);

      *alpha = sqrt(expnu);
      // \gamma_{rbar rbar} = exp(4 phi)
      *gammaDD00 = exp4phi;
      // \gamma_{thth} = r^2 * exp(4 phi)
      *gammaDD11 = r*r * exp4phi;
      // \gamma_{phph} = r^2 sin^2(th) * exp(4 phi)
      *gammaDD22 = r*r*sin(th)*sin(th) * exp4phi;

      // All other quantities ARE ZERO:
      *gammaDD01 = 0.0;
      *gammaDD02 = 0.0;
      *gammaDD12 = 0.0;

      *KDD00 = 0.0;
      *KDD01 = 0.0;
      *KDD02 = 0.0;
      *KDD11 = 0.0;
      *KDD12 = 0.0;
      *KDD22 = 0.0;

      *betaU0 = 0.0;
      *betaU1 = 0.0;
      *betaU2 = 0.0;

      *BU0 = 0.0;
      *BU1 = 0.0;
      *BU2 = 0.0;
}\n""")

with open("BSSN/ID_TOV_TUPMUNU.h", "w") as file:
    file.write("""
// This function takes as input either (x,y,z) or (r,th,ph) and outputs
//   all ADM quantities in the Cartesian or Spherical basis, respectively.
void ID_TOV_TUPMUNU(const REAL xyz_or_rthph[3], const ID_inputs other_inputs,

                    REAL *T4UU00,REAL *T4UU01,REAL *T4UU02,REAL *T4UU03,
                    /**/         REAL *T4UU11,REAL *T4UU12,REAL *T4UU13,
                    /**/                      REAL *T4UU22,REAL *T4UU23,
                    /**/                                   REAL *T4UU33) {
      const REAL r  = xyz_or_rthph[0];
      const REAL th = xyz_or_rthph[1];
      const REAL ph = xyz_or_rthph[2];

      REAL rho,rho_baryon,P,M,expnu,exp4phi;
      TOV_interpolate_1D(r,other_inputs.Rbar,other_inputs.Rbar_idx,other_inputs.interp_stencil_size,
                         other_inputs.numlines_in_file,
                         other_inputs.r_Schw_arr,other_inputs.rho_arr,other_inputs.rho_baryon_arr,other_inputs.P_arr,other_inputs.M_arr,
                         other_inputs.expnu_arr,other_inputs.exp4phi_arr,other_inputs.rbar_arr,
                         &rho,&rho_baryon,&P,&M,&expnu,&exp4phi);

      //T^tt = e^(-nu) * rho
      *T4UU00 = rho / expnu;
      //T^rr = P / exp4phi
      *T4UU11 = P / exp4phi;
      //T^thth = P / (r^2 * exp4phi)
      *T4UU22 = P / (r*r * exp4phi);
      //T^phph = P / (r^2 * sin^2(th) * exp4phi)
      *T4UU33 = P / (r*r * sin(th)*sin(th) * exp4phi);

      // All other components ARE ZERO:
      *T4UU01 = 0; *T4UU02 = 0; *T4UU03 = 0;
      /**/         *T4UU12 = 0; *T4UU13 = 0;
      /**/                      *T4UU23 = 0;
}\n""")
```

## Step 2.b: Compute source terms $S_{ij}$, $S_{i}$, $S$, and $\rho$ \[Back to [top](#toc)\]
$$\label{source}$$

Compute source terms $S_{ij}$, $S_{i}$, $S$, and $\rho$, via Eqs. 10 of [Baumgarte, Montero, Cordero-Carrión, and Müller](https://arxiv.org/pdf/1211.6632.pdf):

\begin{array}
\ S_{ij} &= \gamma_{i \mu} \gamma_{j \nu} T^{\mu \nu} \\
S_{i} &= -\gamma_{i\mu} n_\nu T^{\mu\nu} \\
S &= \gamma^{ij} S_{ij} \\
\rho &= n_\mu n_\nu T^{\mu\nu},
\end{array}

`ID_TOV_TUPMUNU()` provides numerical values for $T^{\mu\nu}$, but we do not have $\gamma_{\mu \nu}$ or $n_\mu$ directly. So here we will construct the latter quantities.

First, B&S Eq. 2.27 defines $\gamma_{\mu \nu}$ as:

$$\gamma_{\mu\nu} = g_{\mu\nu} + n_\mu n_\nu$$

where

$$
n_\mu = \{-\alpha,0,0,0\}.
$$

So we will first need to construct the 4-metric based on ADM quantities. This is provided by Eq 4.47 in [Gourgoulhon](https://arxiv.org/pdf/gr-qc/0703035.pdf):
$$
g_{\mu\nu} = \begin{pmatrix}
-\alpha^2 + \beta^k \beta_k & \beta_i \\
\beta_j & \gamma_{ij}
\end{pmatrix}.
$$

```python
gammaDD = ixp.declarerank2("gammaDD", "sym01",DIM=3)
betaU   = ixp.declarerank1("betaU",DIM=3)
alpha   = sp.symbols("alpha")

# To get \gamma_{\mu \nu} = gammabar4DD[mu][nu], we'll need to construct the 4-metric, using Eq. 2.122 in B&S:

# Eq. 2.121 in B&S
betaD = ixp.zerorank1()
for i in range(DIM):
    for j in range(DIM):
        betaD[i] += gammaDD[i][j]*betaU[j]

# Now compute the beta contraction.
beta2 = sp.sympify(0)
for i in range(DIM):
    beta2 += betaU[i]*betaD[i]

# Eq. 2.122 in B&S
g4DD = ixp.zerorank2(DIM=4)
g4DD[0][0] = -alpha**2 + beta2
for i in range(DIM):
    g4DD[i+1][0] = g4DD[0][i+1] = betaD[i]
for i in range(DIM):
    for j in range(DIM):
        g4DD[i+1][j+1] = gammaDD[i][j]
```

Now let's construct $\gamma_{\mu\nu}$=`gamma4DD[mu][nu]` via $\gamma_{\mu\nu} = g_{\mu\nu} + n_\mu n_\nu$:

```python
n4D = ixp.zerorank1(DIM=4)
n4D[0] = -alpha

gamma4DD = ixp.zerorank2(DIM=4)
for mu in range(4):
    for nu in range(4):
        gamma4DD[mu][nu] = g4DD[mu][nu] + n4D[mu]*n4D[nu]
```

We now have all we need to construct the BSSN source terms in the current basis (for TOV, the Spherical basis):

\begin{array}
\ S_{ij} &= \gamma_{i \mu} \gamma_{j \nu} T^{\mu \nu} \\
S_{i} &= -\gamma_{i\mu} n_\nu T^{\mu\nu} \\
S &= \gamma^{ij} S_{ij} \\
\rho &= n_\mu n_\nu T^{\mu\nu},
\end{array}

```python
T4UU = ixp.declarerank2("T4UU", "sym01",DIM=4)

SDD = ixp.zerorank2()
SD  = ixp.zerorank1()
S   = sp.sympify(0)
rho = sp.sympify(0)

for i in range(DIM):
    for j in range(DIM):
        for mu in range(4):
            for nu in range(4):
                SDD[i][j] += gamma4DD[i+1][mu]*gamma4DD[j+1][nu] * T4UU[mu][nu]

for i in range(DIM):
    for mu in range(4):
        for nu in range(4):
            SD[i] += -gamma4DD[i+1][mu]*n4D[nu] * T4UU[mu][nu]

gammaUU, gammaDET = ixp.symm_matrix_inverter3x3(gammaDD)
for i in range(DIM):
    for j in range(DIM):
        S += gammaUU[i][j]*SDD[i][j]

for mu in range(4):
    for nu in range(4):
        rho += n4D[mu]*n4D[nu] * T4UU[mu][nu]
```

## Step 2.c: Jacobian transformation on the ADM/BSSN source terms \[Back to [top](#toc)\]
$$\label{jacobian}$$

The following discussion holds for either Spherical or Cartesian input data, so for simplicity let's just assume the data are given in Spherical coordinates.

All ADM tensors and vectors are in the Spherical coordinate basis $x^i_{\rm Sph} = (r,\theta,\phi)$, but we need them in the curvilinear coordinate basis $x^i_{\rm rfm}=$`(xx0,xx1,xx2)` set by the `"reference_metric::CoordSystem"` variable. Empirically speaking, it is far easier to write `(x(xx0,xx1,xx2),y(xx0,xx1, xx2),z(xx0,xx1,xx2))` than the inverse, so we will compute the Jacobian matrix

$$
{\rm Jac\_dUSph\_dDrfmUD[i][j]} = \frac{\partial x^i_{\rm Sph}}{\partial x^j_{\rm rfm}},
$$

via exact differentiation (courtesy SymPy), and the inverse Jacobian
$$
{\rm Jac\_dUrfm\_dDSphUD[i][j]} = \frac{\partial x^i_{\rm rfm}}{\partial x^j_{\rm Sph}},
$$

using NRPy+'s `generic_matrix_inverter3x3()` function. In terms of these, the transformation of BSSN tensors from Spherical to `"reference_metric::CoordSystem"` coordinates may be written:

\begin{align}
S^{\rm rfm}_{ij} &=
\frac{\partial x^\ell_{\rm Sph}}{\partial x^i_{\rm rfm}}
\frac{\partial x^m_{\rm Sph}}{\partial x^j_{\rm rfm}} S^{\rm Sph}_{\ell m}\\
S^{\rm rfm}_{i} &=
\frac{\partial x^\ell_{\rm Sph}}{\partial x^i_{\rm rfm}}
S^{\rm Sph}_{\ell}
\end{align}

```python
# TOV initial data are given in Spherical coordinates.
CoordType_in = "Spherical"

SSphorCartDD = ixp.zerorank2()
SSphorCartD  = ixp.zerorank1()

# Copy what we had written above, which was in Spherical coordinates, to the new tensors SSphorCartDD / SSphorCartD:
for i in range(3):
    SSphorCartD[i] = SD[i]
    for j in range(3):
        SSphorCartDD[i][j] = SDD[i][j]

# Zero out the original tensors; we're going to store the result to SD and SDD:
SDD = ixp.zerorank2()
SD  = ixp.zerorank1()

# Make sure that rfm.reference_metric() has been called.
#    We'll need the variables it defines throughout this module.
if rfm.have_already_called_reference_metric_function == False:
    print("Error. Called Convert_Spherical_ADM_to_BSSN_curvilinear() without")
    print("       first setting up reference metric, by calling rfm.reference_metric().")
    exit(1)

r_th_ph_or_Cart_xyz_oID_xx = []
if CoordType_in == "Spherical":
    r_th_ph_or_Cart_xyz_oID_xx = rfm.xxSph
elif CoordType_in == "Cartesian":
    r_th_ph_or_Cart_xyz_oID_xx = rfm.xxCart
else:
    print("Error: Can only convert ADM Cartesian or Spherical initial data to BSSN Curvilinear coords.")
    exit(1)

# Next apply Jacobian transformations to convert into the (xx0,xx1,xx2) basis

# rho and S are scalar, so no Jacobian transformations are necessary.

Jac_dUSphorCart_dDrfmUD = ixp.zerorank2()
for i in range(DIM):
    for j in range(DIM):
        Jac_dUSphorCart_dDrfmUD[i][j] = sp.diff(r_th_ph_or_Cart_xyz_oID_xx[i],rfm.xx[j])

Jac_dUrfm_dDSphorCartUD, dummyDET = ixp.generic_matrix_inverter3x3(Jac_dUSphorCart_dDrfmUD)

for i in range(DIM):
    for j in range(DIM):
        SD[i] += Jac_dUSphorCart_dDrfmUD[j][i] * SSphorCartD[j]
        for k in range(DIM):
            for l in range(DIM):
                SDD[i][j] += Jac_dUSphorCart_dDrfmUD[k][i]*Jac_dUSphorCart_dDrfmUD[l][j] * SSphorCartDD[k][l]
```

## Step 2.d: Rescale tensorial quantities \[Back to [top](#toc)\]
$$\label{tensor}$$

We rescale tensorial quantities according to the prescription described in the [BSSN in curvilinear coordinates tutorial notebook](Tutorial-BSSNCurvilinear.ipynb) (also [Ruchlin *et al.*](https://arxiv.org/pdf/1712.07658.pdf)):

Since `ReD[i]}=1/ReU[i]`, we have

\begin{align}
s_{ij} &= S_{ij} /\text{ReDD[i][j]}\\
s_{i} &= S_i \text{ReU[i]}
\end{align}

```python
# Finally rescale the tensorial quantities:
sD = ixp.zerorank1()
sDD = ixp.zerorank2()
for i in range(DIM):
    sD[i] = SD[i] * rfm.ReU[i] # ReD[i] = 1/ReU[i]
    for j in range(DIM):
        sDD[i][j] = SDD[i][j] / rfm.ReDD[i][j]
```

Next we use NRPy+ to write a C function that reads in $T^{\mu\nu}$ in the given (Spherical) basis and outputs the source terms $\{S_{ij},S_{i},S,\rho\}$ in the `(xx0,xx1,xx2)` basis:

```python
with open("BSSN/ID_TOV_BSSN_Source_Terms.h", "w") as file:
    file.write("""void ID_TOV_BSSN_Source_Terms(
                    REAL xx0xx1xx2[3],

                    const REAL gammaDD00,const REAL gammaDD01,const REAL gammaDD02,
                    /**/                 const REAL gammaDD11,const REAL gammaDD12,
                    /**/                                      const REAL gammaDD22,

                    const REAL betaU0,const REAL betaU1,const REAL betaU2,
                    const REAL alpha,

                    const REAL T4UU00,const REAL T4UU01,const REAL T4UU02,const REAL T4UU03,
                    /**/              const REAL T4UU11,const REAL T4UU12,const REAL T4UU13,
                    /**/                                const REAL T4UU22,const REAL T4UU23,
                    /**/                                                  const REAL T4UU33,

                    REAL *sDD00,REAL *sDD01,REAL *sDD02,
                    /**/        REAL *sDD11,REAL *sDD12,
                    /**/                    REAL *sDD22,

                    REAL *sD0, REAL *sD1, REAL *sD2,

                    REAL *S, REAL *rho) {

    const REAL xx0 = xx0xx1xx2[0];
    const REAL xx1 = xx0xx1xx2[1];
    const REAL xx2 = xx0xx1xx2[2];\n""")
outCparams = "preindent=1,outCfileaccess=a,outCverbose=False,includebraces=False"

outputC([sDD[0][0],sDD[0][1],sDD[0][2],sDD[1][1],sDD[1][2],sDD[2][2],
         sD[0],sD[1],sD[2], S, rho],
        ["*sDD00","*sDD01","*sDD02","*sDD11","*sDD12","*sDD22",
         "*sD0","*sD1","*sD2","*S","*rho"], "BSSN/ID_TOV_BSSN_Source_Terms.h",outCparams)

with open("BSSN/ID_TOV_BSSN_Source_Terms.h", "a") as file:
    file.write("}\n")
```

    Appended to file "BSSN/ID_TOV_BSSN_Source_Terms.h"

# Step 3: Convert ADM spacetime quantity initial data from Spherical to BSSN Curvilinear coordinates \[Back to [top](#toc)\]
$$\label{adm_id_spacetime}$$

We convert ADM spacetime quantity initial data from Spherical to BSSN Curvilinear coordinates [as documented in the corresponding tutorial notebook](Tutorial-ADM_Initial_Data-Converting_Numerical_ADM_Spherical_or_Cartesian_to_BSSNCurvilinear.ipynb)

```python
import BSSN.ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear as AtoBnum
AtoBnum.Convert_Spherical_or_Cartesian_ADM_to_BSSN_curvilinear("Spherical","ID_TOV_ADM_quantities")
```

    Appended to file "BSSN/ID_ADM_xx0xx1xx2_to_BSSN_xx0xx1xx2__ALL_BUT_LAMBDAs.h"
    Appended to file "BSSN/ID_ADM_xx0xx1xx2_to_BSSN_xx0xx1xx2__ALL_BUT_LAMBDAs.h"

# Step 4: Validating that the TOV initial data satisfy the Hamiltonian constraint \[Back to [top](#toc)\]
$$\label{validate}$$

We will validate that the TOV initial data satisfy the Hamiltonian constraint, modulo numerical finite differencing error

## Step 4.a: Output the Hamiltonian constraint \[Back to [top](#toc)\]
$$\label{ham_const_output}$$

First output the Hamiltonian constraint [as documented in the corresponding NRPy+ tutorial notebook](Tutorial-BSSN_constraints.ipynb)

```python
import BSSN.BSSN_constraints as bssncon
bssncon.output_C__Hamiltonian_h(add_T4UUmunu_source_terms=True)
```

    Generating optimized C code for Hamiltonian constraint. May take a while, depending on CoordSystem.
    Finished in 14.2037842274 seconds.
    Output C implementation of Hamiltonian constraint to BSSN/Hamiltonian.h

## Step 4.b: Apply singular, curvilinear coordinate boundary conditions \[Back to [top](#toc)\]
$$\label{apply_bcs}$$

Next apply singular, curvilinear coordinate boundary conditions [as documented in the corresponding NRPy+ tutorial notebook](Tutorial-Start_to_Finish-Curvilinear_BCs.ipynb).

```python
import CurviBoundaryConditions.CurviBoundaryConditions as cbcs
cbcs.Set_up_CurviBoundaryConditions()
```

    Wrote to file "CurviBoundaryConditions/gridfunction_defines.h"
    Wrote to file "CurviBoundaryConditions/set_parity_conditions.h"
    Wrote to file "CurviBoundaryConditions/xxCart.h"
    Wrote to file "CurviBoundaryConditions/xxminmax.h"
    Wrote to file "CurviBoundaryConditions/Cart_to_xx.h"

## Step 4.c: Enforce conformal 3-metric $\det{\bar{\gamma}_{ij}}=\det{\hat{\gamma}_{ij}}$ constraint \[Back to [top](#toc)\]
$$\label{enforce3metric}$$

Then enforce conformal 3-metric $\det{\bar{\gamma}_{ij}}=\det{\hat{\gamma}_{ij}}$ constraint (Eq. 53 of [Ruchlin, Etienne, and Baumgarte (2018)](https://arxiv.org/abs/1712.07658)), as [documented in the corresponding NRPy+ tutorial notebook](Tutorial-BSSN-Enforcing_Determinant_gammabar_equals_gammahat_Constraint.ipynb).

Applying curvilinear boundary conditions should affect the initial data at the outer boundary, and will in general cause the $\det{\bar{\gamma}_{ij}}=\det{\hat{\gamma}_{ij}}$ constraint to be violated there. Thus after we apply these boundary conditions, we must always call the routine for enforcing the $\det{\bar{\gamma}_{ij}}=\det{\hat{\gamma}_{ij}}$ constraint:

```python
import BSSN.Enforce_Detgammabar_Constraint as EGC
EGC.output_Enforce_Detgammabar_Constraint_Ccode()
```

    Output C implementation of det(gammabar) constraint to file BSSN/enforce_detgammabar_constraint.h

# Step 5: `TOV_Playground.c`: The Main C Code \[Back to [top](#toc)\]
$$\label{mainc}$$

```python
# Part P0: Set the number of ghost cells, from NRPy+'s FD_CENTDERIVS_ORDER
with open("BSSN/NGHOSTS.h", "w") as file:
    file.write("// Part P0: Set the number of ghost zones, from NRPy+'s FD_CENTDERIVS_ORDER\n")
    # Upwinding in BSSN requires that NGHOSTS = FD_CENTDERIVS_ORDER/2 + 1 <- Notice the +1.
    file.write("#define NGHOSTS   "+str(int(par.parval_from_str("finite_difference::FD_CENTDERIVS_ORDER")/2)+1)+"\n")
    file.write("#define TOV_Mass  "+str(TOV_Mass)+"\n")
    file.write("#define TOV_Riso  "+str(TOV_Riso)+"\n")
```

```python
%%writefile BSSN/TOV_Playground.c

// Part P1: Import needed header files
#include "NGHOSTS.h" // A NRPy+-generated file, which is set based on FD_CENTDERIVS_ORDER.
#include "stdio.h"
#include "stdlib.h"
#include "math.h"
#include "stdint.h" // Needed for Windows GCC 6.x compatibility

// Part P2: Add needed #define's to set data type, the IDX4() macro, and the gridfunctions
// Part P2a: set REAL=double, so that all floating point numbers are stored to at least ~16 significant digits.
#define REAL double

// Step P3: Set free parameters for the numerical grid
const REAL RMAX    = TOV_Riso*2.0;

typedef struct __ID_inputs {
    REAL Rbar;
    int Rbar_idx;
    int interp_stencil_size;
    int numlines_in_file;
    REAL *r_Schw_arr,*rho_arr,*rho_baryon_arr,*P_arr,*M_arr,*expnu_arr,*exp4phi_arr,*rbar_arr;
} ID_inputs;

// Part P4b: Declare the IDX4(gf,i,j,k) macro, which enables us to store 4-dimensions of
//           data in a 1D array. In this case, consecutive values of "i"
//           (all other indices held to a fixed value) are consecutive in memory, where
//           consecutive values of "j" (fixing all other indices) are separated by
//           Nxx_plus_2NGHOSTS[0] elements in memory. Similarly, consecutive values of
//           "k" are separated by Nxx_plus_2NGHOSTS[0]*Nxx_plus_2NGHOSTS[1] in memory, etc.
#define IDX4(g,i,j,k) \
( (i) + Nxx_plus_2NGHOSTS[0] * ( (j) + Nxx_plus_2NGHOSTS[1] * ( (k) + Nxx_plus_2NGHOSTS[2] * (g) ) ) )
#define IDX3(i,j,k) ( (i) + Nxx_plus_2NGHOSTS[0] * ( (j) + Nxx_plus_2NGHOSTS[1] * (k) ) )
// Assuming idx = IDX3(i,j,k). Much faster if idx can be reused over and over:
#define IDX4pt(g,idx)   ( (idx) + (Nxx_plus_2NGHOSTS[0]*Nxx_plus_2NGHOSTS[1]*Nxx_plus_2NGHOSTS[2]) * (g) )

// Part P4c: Set #define's for BSSN gridfunctions. C code generated above
#include "../CurviBoundaryConditions/gridfunction_defines.h"

#define LOOP_REGION(i0min,i0max, i1min,i1max, i2min,i2max) \
  for(int i2=i2min;i2<i2max;i2++) for(int i1=i1min;i1<i1max;i1++) for(int i0=i0min;i0<i0max;i0++)

// Step P5: Function for converting uniform grid coord
//         (xx[0][i0],xx[1][i1],xx[2][i2]) to
//          corresponding Cartesian coordinate.
void xxCart(REAL *xx[3],const int i0,const int i1,const int i2, REAL xCart[3]) {
    REAL xx0 = xx[0][i0];
    REAL xx1 = xx[1][i1];
    REAL xx2 = xx[2][i2];
#include "xxCart.h"
}

// Step P6: Include basic functions needed to impose curvilinear
//          parity and boundary conditions.
#include "../CurviBoundaryConditions/curvilinear_parity_and_outer_boundary_conditions.h"

// Step P7: Function for enforcing the gammabar=gammahat constraint:
#include "enforce_detgammabar_constraint.h"

// Part P8: Declare all functions for setting up UIUC initial data.
/* Routines to interpolate the TOV solution and convert to ADM & T^{munu}: */
#include "../TOV/tov_interp.h"
#include "ID_TOV_ADM_quantities.h"
#include "ID_TOV_TUPMUNU.h"

/* Next perform the basis conversion and compute all needed BSSN quantities */
#include "ID_TOV_BSSN_Source_Terms.h"
#include "ID_ADM_xx0xx1xx2_to_BSSN_xx0xx1xx2__ALL_BUT_LAMBDAs.h"
#include "ID_BSSN__ALL_BUT_LAMBDAs.h"
#include "ID_BSSN_lambdas.h"

// Part P9: Declare function for computing the Hamiltonian
//           constraint violation, which should converge to
//           zero with increasing numerical resolution.
void Hamiltonian_constraint(const int Nxx[3],const int Nxx_plus_2NGHOSTS[3],const REAL dxx[3], REAL *xx[3],
                            REAL *in_gfs, REAL *aux_gfs) {
#include "Hamiltonian.h"
}

// main() function:
// Step 0: Read command-line input, set up grid structure, allocate memory for gridfunctions, set up coordinates
// Step 1: Set up scalar wave initial data
// Step 2: Evolve scalar wave initial data forward in time using Method of Lines with RK4 algorithm,
//         applying quadratic extrapolation outer boundary conditions.
// Step 3: Output relative error between numerical and exact solution.
// Step 4: Free all allocated memory
int main(int argc, const char *argv[]) {
    // Step 0a: Read command-line input, error out if nonconformant
    if(argc != 4 || atoi(argv[1]) < NGHOSTS) {
        fprintf(stderr,"Error: Expected three command-line arguments: ./TOV_Playground Nx0 Nx1 Nx2,\n");
        fprintf(stderr,"where Nx[0,1,2] is the number of grid points in the 0, 1, and 2 directions.\n");
        fprintf(stderr,"Nx[] MUST BE larger than NGHOSTS (= %d)\n",NGHOSTS);
        exit(1);
    }
    // Step 0b: Set up numerical grid structure, first in space...
    if(atoi(argv[1])%2 != 0 || atoi(argv[2])%2 != 0 || atoi(argv[2])%2 != 0) {
        fprintf(stderr,"Error: Cannot guarantee a proper cell-centered grid if number of grid cells not set to even number.\n");
        fprintf(stderr,"       For example, in case of angular directions, proper symmetry zones will not exist.\n");
        exit(1);
    }
    const int Nxx[3] = { atoi(argv[1]), atoi(argv[2]), atoi(argv[3]) };
    const int Nxx_plus_2NGHOSTS[3] = { Nxx[0]+2*NGHOSTS, Nxx[1]+2*NGHOSTS, Nxx[2]+2*NGHOSTS };
    const int Nxx_plus_2NGHOSTS_tot = Nxx_plus_2NGHOSTS[0]*Nxx_plus_2NGHOSTS[1]*Nxx_plus_2NGHOSTS[2];
#include "xxminmax.h"

    /* TOV INPUT ROUTINE */
    // Open the data file:
    char filename[100];
    sprintf(filename,"./outputTOVpolytrope.txt");
    FILE *in1Dpolytrope = fopen(filename, "r");
    if (in1Dpolytrope == NULL) {
        fprintf(stderr,"ERROR: could not open file %s\n",filename);
        exit(1);
    }
    // Count the number of lines in the data file:
    int numlines_in_file = count_num_lines_in_file(in1Dpolytrope);
    // Allocate space for all data arrays:
    REAL *r_Schw_arr     = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *rho_arr        = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *rho_baryon_arr = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *P_arr          = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *M_arr          = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *expnu_arr      = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *exp4phi_arr    = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *rbar_arr       = (REAL *)malloc(sizeof(REAL)*numlines_in_file);

    // Read from the data file, filling in arrays
    // read_datafile__set_arrays() may be found in TOV/tov_interp.h
    if(read_datafile__set_arrays(in1Dpolytrope, r_Schw_arr,rho_arr,rho_baryon_arr,P_arr,M_arr,expnu_arr,exp4phi_arr,rbar_arr) == 1) {
        fprintf(stderr,"ERROR WHEN READING FILE %s!\n",filename);
        exit(1);
    }
    fclose(in1Dpolytrope);
    REAL Rbar = -100;
    int Rbar_idx = -100;
    for(int i=1;i<numlines_in_file;i++) {
        if(rho_arr[i-1]>0 && rho_arr[i]==0) { Rbar = rbar_arr[i-1]; Rbar_idx = i-1; }
    }
    if(Rbar<0) {
        fprintf(stderr,"Error: could not find rbar=Rbar from data file.\n");
        exit(1);
    }

    ID_inputs TOV_in;
    TOV_in.Rbar = Rbar;
    TOV_in.Rbar_idx = Rbar_idx;

    const int interp_stencil_size = 12;
    TOV_in.interp_stencil_size = interp_stencil_size;
    TOV_in.numlines_in_file = numlines_in_file;

    TOV_in.r_Schw_arr     = r_Schw_arr;
    TOV_in.rho_arr        = rho_arr;
    TOV_in.rho_baryon_arr = rho_baryon_arr;
    TOV_in.P_arr          = P_arr;
    TOV_in.M_arr          = M_arr;
    TOV_in.expnu_arr      = expnu_arr;
    TOV_in.exp4phi_arr    = exp4phi_arr;
    TOV_in.rbar_arr       = rbar_arr;
    /* END TOV INPUT ROUTINE */

    // Step 0c: Allocate memory for gridfunctions
    REAL *evol_gfs = (REAL *)malloc(sizeof(REAL) * NUM_EVOL_GFS * Nxx_plus_2NGHOSTS_tot);
    REAL *aux_gfs  = (REAL *)malloc(sizeof(REAL) * NUM_AUX_GFS * Nxx_plus_2NGHOSTS_tot);

    // Step 0d: Set up space and time coordinates
    // Step 0d.i: Set \Delta x^i on uniform grids.
    REAL dxx[3];
    for(int i=0;i<3;i++) dxx[i] = (xxmax[i] - xxmin[i]) / ((REAL)Nxx[i]);

    // Step 0d.ii: Set up uniform coordinate grids
    REAL *xx[3];
    for(int i=0;i<3;i++) {
        xx[i] = (REAL *)malloc(sizeof(REAL)*Nxx_plus_2NGHOSTS[i]);
        for(int j=0;j<Nxx_plus_2NGHOSTS[i];j++) {
            xx[i][j] = xxmin[i] + ((REAL)(j-NGHOSTS) + (1.0/2.0))*dxx[i]; // Cell-centered grid.
        }
    }

    // Step 0e: Find ghostzone mappings and parities:
    gz_map *bc_gz_map = (gz_map *)malloc(sizeof(gz_map)*Nxx_plus_2NGHOSTS_tot);
    parity_condition *bc_parity_conditions = (parity_condition *)malloc(sizeof(parity_condition)*Nxx_plus_2NGHOSTS_tot);
    set_up_bc_gz_map_and_parity_conditions(Nxx_plus_2NGHOSTS,xx,dxx,xxmin,xxmax,  bc_gz_map, bc_parity_conditions);

    // Step 1: Set up initial data to an exact solution at time=0:
    ID_BSSN__ALL_BUT_LAMBDAs(Nxx_plus_2NGHOSTS,xx,TOV_in, evol_gfs);
    apply_bcs(Nxx, Nxx_plus_2NGHOSTS, bc_gz_map,bc_parity_conditions,NUM_EVOL_GFS,evol_gf_parity, evol_gfs);
    enforce_detgammabar_constraint(Nxx_plus_2NGHOSTS, xx, evol_gfs);
    ID_BSSN_lambdas(Nxx, Nxx_plus_2NGHOSTS, xx,dxx, evol_gfs);
    apply_bcs(Nxx, Nxx_plus_2NGHOSTS, bc_gz_map,bc_parity_conditions,NUM_EVOL_GFS,evol_gf_parity, evol_gfs);
    enforce_detgammabar_constraint(Nxx_plus_2NGHOSTS, xx, evol_gfs);
{
    #pragma omp parallel for
    for(int i2=0; i2<Nxx_plus_2NGHOSTS[2]; i2++) {
        const REAL xx2 = xx[2][i2];
        for(int i1=0; i1<Nxx_plus_2NGHOSTS[1]; i1++) {
            const REAL xx1 = xx[1][i1];
            for(int i0=0; i0<Nxx_plus_2NGHOSTS[0]; i0++) {
                const REAL xx0 = xx[0][i0];
                REAL xx0xx1xx2[3] = {xx0,xx1,xx2};
                REAL gammaDD00,gammaDD01,gammaDD02,gammaDD11,gammaDD12,gammaDD22;
                REAL KDD00,KDD01,KDD02,KDD11,KDD12,KDD22;
                REAL alpha,betaU0,betaU1,betaU2;
                REAL BU0,BU1,BU2;
                REAL xyz_or_rthph[3] =  {xx0,xx1,xx2}; //FIXME

                // FIRST INTERPOLATE TO SET THE ADM AND TMUNU QUANTITIES
                ID_TOV_ADM_quantities(xyz_or_rthph, TOV_in,

                    &gammaDD00,&gammaDD01,&gammaDD02,&gammaDD11,&gammaDD12,&gammaDD22,
                    &KDD00,&KDD01,&KDD02,&KDD11,&KDD12,&KDD22,
                    &alpha,  &betaU0,&betaU1,&betaU2,  &BU0,&BU1,&BU2);

                REAL T4UU00,T4UU01,T4UU02,T4UU03,
                /**/        T4UU11,T4UU12,T4UU13,
                /**/               T4UU22,T4UU23,
                /**/                      T4UU33;

                ID_TOV_TUPMUNU(xyz_or_rthph, TOV_in,

                    &T4UU00,&T4UU01,&T4UU02,&T4UU03,
                    /**/    &T4UU11,&T4UU12,&T4UU13,
                    /**/            &T4UU22,&T4UU23,
                    /**/                    &T4UU33);

                const int idx = IDX3(i0,i1,i2);
                // THEN EVALUATE THE BSSN SOURCE TERMS
                ID_TOV_BSSN_Source_Terms(xx0xx1xx2,
                                         gammaDD00, gammaDD01, gammaDD02,
                    /**/                            gammaDD11, gammaDD12,
                    /**/                                       gammaDD22,

                    betaU0,betaU1,betaU2, alpha,

                    T4UU00, T4UU01, T4UU02, T4UU03,
                    /**/    T4UU11, T4UU12, T4UU13,
                    /**/            T4UU22, T4UU23,
                    /**/                    T4UU33,

                    &aux_gfs[IDX4pt(SDD00GF,idx)], &aux_gfs[IDX4pt(SDD01GF,idx)], &aux_gfs[IDX4pt(SDD02GF,idx)],
                    /**/                           &aux_gfs[IDX4pt(SDD11GF,idx)], &aux_gfs[IDX4pt(SDD12GF,idx)],
                    /**/                                                          &aux_gfs[IDX4pt(SDD22GF,idx)],

                    &aux_gfs[IDX4pt(SD0GF,idx)], &aux_gfs[IDX4pt(SD1GF,idx)], &aux_gfs[IDX4pt(SD2GF,idx)],

                    &aux_gfs[IDX4pt(SGF,idx)], &aux_gfs[IDX4pt(RHOGF,idx)]);
                //aux_gfs[IDX4pt(SGF,idx)] = T4UU11*gammaDD00;
            }
        }
    }
}

    // Step 1b: Apply boundary conditions, as initial data
    //          are sometimes ill-defined in ghost zones.
    //          E.g., spherical initial data might not be
    //          properly defined at points where r=-1.
//    apply_bcs(Nxx, Nxx_plus_2NGHOSTS, bc_gz_map,bc_parity_conditions,NUM_EVOL_GFS,evol_gf_parity, evol_gfs);
//    enforce_detgammabar_constraint(Nxx_plus_2NGHOSTS, xx, evol_gfs);

    // Step 2: Evaluate Hamiltonian constraint violation
    Hamiltonian_constraint(Nxx,Nxx_plus_2NGHOSTS,dxx, xx, evol_gfs, aux_gfs);

    /* Step 3: Output relative error between numerical and exact solution, */
    const int i0MIN=NGHOSTS; // In spherical, r=Delta r/2.
    const int i1mid=Nxx_plus_2NGHOSTS[1]/2;
    const int i2mid=Nxx_plus_2NGHOSTS[2]/2;
    LOOP_REGION(NGHOSTS,Nxx_plus_2NGHOSTS[0]-NGHOSTS, i1mid,i1mid+1, NGHOSTS,Nxx_plus_2NGHOSTS[2]-NGHOSTS) {
        REAL xx0 = xx[0][i0];
        REAL xx1 = xx[1][i1];
        REAL xx2 = xx[2][i2];
        REAL xCart[3];
#include "xxCart.h"
        int idx = IDX3(i0,i1,i2);
        printf("%e %e %e %e\n",xCart[0]/TOV_Mass,xCart[1]/TOV_Mass,
               log10(aux_gfs[IDX4pt(RHOGF,idx)]+1e-6),
               log10(fabs(aux_gfs[IDX4pt(HGF,idx)])));
    }

    /* Step 4: Free all allocated memory */
    free(bc_gz_map);
    free(bc_parity_conditions);
    free(rbar_arr);
    free(rho_arr);
    free(rho_baryon_arr);
    free(P_arr);
    free(M_arr);
    free(expnu_arr);

    free(aux_gfs);
    free(evol_gfs);
    for(int i=0;i<3;i++) free(xx[i]);
    return 0;
}
```

    Overwriting BSSN/TOV_Playground.c

```python
import cmdline_helper as cmd
cmd.C_compile("BSSN/TOV_Playground.c", "TOV_Playground")
cmd.delete_existing_files("out96.txt")
cmd.Execute("TOV_Playground", "96 96 96", "out96.txt")
```

    Compiling executable...
    Executing `gcc -Ofast -fopenmp -march=native -funroll-loops BSSN/TOV_Playground.c -o TOV_Playground -lm`...
    Finished executing in 4.85430407524 seconds.
    Finished compilation.
    Executing `taskset -c 0,1,2,3,4,5 ./TOV_Playground 96 96 96`...
    Finished executing in 0.844125986099 seconds.

# Step 6: Plotting the single-neutron-star initial data \[Back to [top](#toc)\]
$$\label{plot}$$

Here we plot the conformal factor of these initial data on a 2D grid, such that darker colors imply stronger gravitational fields. Hence, we see the single neutron star centered at the origin: $x/M=y/M=z/M=0$, where $M$ is an arbitrary mass scale (conventionally the [ADM mass](https://en.wikipedia.org/w/index.php?title=ADM_formalism&oldid=846335453) is chosen), and our formulation of Einstein's equations adopt $G=c=1$ [geometrized units](https://en.wikipedia.org/w/index.php?title=Geometrized_unit_system&oldid=861682626).

```python
import numpy as np
from scipy.interpolate import griddata
from pylab import savefig
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from IPython.display import Image

x96,y96,valuesCF96,valuesHam96 = np.loadtxt('out96.txt').T #Transposed for easier unpacking

bounds = 4.0
pl_xmin = -bounds
pl_xmax = +bounds
pl_ymin = -bounds
pl_ymax = +bounds

grid_x, grid_y = np.mgrid[pl_xmin:pl_xmax:100j, pl_ymin:pl_ymax:100j]
points96 = np.zeros((len(x96), 2))
for i in range(len(x96)):
    points96[i][0] = x96[i]
    points96[i][1] = y96[i]

grid96 = griddata(points96, valuesCF96, (grid_x, grid_y), method='nearest')
grid96cub = griddata(points96, valuesCF96, (grid_x, grid_y), method='cubic')

plt.clf()
plt.title("Neutron Star: log10( max(1e-6,Energy Density) )")
plt.xlabel("x/M")
plt.ylabel("y/M")

# fig, ax = plt.subplots()
# ax.plot(grid96cub.T, extent=(pl_xmin,pl_xmax, pl_ymin,pl_ymax))
# plt.close(fig)
fig96cf = plt.imshow(grid96.T, extent=(pl_xmin,pl_xmax, pl_ymin,pl_ymax))
cb = plt.colorbar(fig96cf)
savefig("BHB.png")
from IPython.display import Image
Image("BHB.png")
# #           interpolation='nearest', cmap=cm.gist_rainbow)
```

# Step 7: Validation: Convergence of numerical errors (Hamiltonian constraint violation) to zero \[Back to [top](#toc)\]
$$\label{convergence}$$

The equations behind these initial data solve Einstein's equations exactly, at a single instant in time. One reflection of this solution is that the Hamiltonian constraint violation should be exactly zero in the initial data.

However, when evaluated on numerical grids, the Hamiltonian constraint violation will *not* generally evaluate to zero due to the associated numerical derivatives not being exact. However, these numerical derivatives (finite difference derivatives in this case) should *converge* to the exact derivatives as the density of numerical sampling points approaches infinity.

In this case, all of our finite difference derivatives agree with the exact solution, with an error term that drops with the uniform gridspacing to the fourth power: $\left(\Delta x^i\right)^4$.

Here, as in the [Start-to-Finish Scalar Wave (Cartesian grids) NRPy+ tutorial](Tutorial-Start_to_Finish-ScalarWave.ipynb) and the [Start-to-Finish Scalar Wave (curvilinear grids) NRPy+ tutorial](Tutorial-Start_to_Finish-ScalarWaveCurvilinear.ipynb) we confirm this convergence.

First, let's take a look at what the numerical error looks like on the x-y plane at a given numerical resolution, plotting $\log_{10}|H|$, where $H$ is the Hamiltonian constraint violation:

```python
grid96 = griddata(points96, valuesHam96, (grid_x, grid_y), method='nearest')
grid96cub = griddata(points96, valuesHam96, (grid_x, grid_y), method='cubic')

# fig, ax = plt.subplots()

plt.clf()
plt.title("96^3 Numerical Err.: log_{10}|Ham|")
plt.xlabel("x/M")
plt.ylabel("y/M")

fig96cub = plt.imshow(grid96cub.T, extent=(pl_xmin,pl_xmax, pl_ymin,pl_ymax))
cb = plt.colorbar(fig96cub)
```

Next, we set up the same initial data but on a lower-resolution, $48\times 8\times 2$ grid (axisymmetric in the $\phi$ direction). Since the constraint violation (numerical error associated with the fourth-order-accurate, finite-difference derivatives) should converge to zero with the uniform gridspacing to the fourth power: $\left(\Delta x^i\right)^4$, we expect the constraint violation will increase (relative to the $96\times 16\times 2$ grid) by a factor of $\left(96/48\right)^4$. Here we demonstrate that indeed this order of convergence is observed as expected, *except* at the star's surface where the stress-energy tensor $T^{\mu\nu}$ sharply drops to zero.

```python
# Now rerun TOV_Playground with twice lower resolution.
cmd.delete_existing_files("out48.txt")
cmd.Execute("TOV_Playground", "48 48 48", "out48.txt")

x48,y48,valuesCF48,valuesHam48 = np.loadtxt('out48.txt').T #Transposed for easier unpacking
points48 = np.zeros((len(x48), 2))
for i in range(len(x48)):
    points48[i][0] = x48[i]
    points48[i][1] = y48[i]

grid48 = griddata(points48, valuesHam48, (grid_x, grid_y), method='cubic')

griddiff_48_minus_96 = np.zeros((100,100))
griddiff_48_minus_96_1darray = np.zeros(100*100)
gridx_1darray_yeq0 = np.zeros(100)
grid48_1darray_yeq0 = np.zeros(100)
grid96_1darray_yeq0 = np.zeros(100)
count = 0
outarray = []
for i in range(100):
    for j in range(100):
        griddiff_48_minus_96[i][j] = grid48[i][j] - grid96[i][j]
        griddiff_48_minus_96_1darray[count] = griddiff_48_minus_96[i][j]
        if j==49:
            gridx_1darray_yeq0[i] = grid_x[i][j]
            grid48_1darray_yeq0[i] = grid48[i][j] + np.log10((48./96.)**4)
            grid96_1darray_yeq0[i] = grid96[i][j]
        count = count + 1

plt.clf()
fig, ax = plt.subplots()
plt.title("Plot Demonstrating 4th-order Convergence")
plt.xlabel("x/M")
plt.ylabel("log10(Relative error)")

ax.plot(gridx_1darray_yeq0, grid96_1darray_yeq0, 'k-', label='Nr=96')
ax.plot(gridx_1darray_yeq0, grid48_1darray_yeq0, 'k--', label='Nr=48, mult by (48/96)^4')
ax.set_ylim([-12.5,1.5])

legend = ax.legend(loc='lower right', shadow=True, fontsize='x-large')
legend.get_frame().set_facecolor('C1')
plt.show()
```

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

The following code cell converts this Jupyter notebook into a proper, clickable $\LaTeX$-formatted PDF file. After the cell is successfully run, the generated PDF may be found in the root NRPy+ tutorial directory, with filename
[Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data.pdf](Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means.)

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

    [NbConvertApp] Converting notebook Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data.ipynb to latex
    [NbConvertApp] Writing 118651 bytes to Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data.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


# 2D Anisotropic Diffusion (non coercive case)

We consider the study and analysis of the eigenvalue problem associated to the following anisotropic diffusion operator
$$
    - \nabla \cdot \left( \mathbf{b} \otimes \mathbf{b} \nabla \phi \right) = \omega^2 \phi
$$
where $\mathbf{b} = (\iota, 1)$ is the normalized vector of the magnetic field.

In fact, the original problem is 3D and the iota profile is given as

$$
\iota(r) = 0.47262 + 0.32392 r + 0.49604 r^2 + 5.3991 . 10^{-6} r^3 - 3.165 . 10^{-5} r^4 + 4.7963 . 10^{-5} r^5 - 2.2824 . 10^{-5} r^6
$$

which is the LHD-$\iota$-profile.

## Formal model

```python
from sympy.core.containers import Tuple

from sympde.core     import Constant
from sympde.calculus import grad, dot
from sympde.topology import ScalarFunctionSpace
from sympde.topology import Domain
from sympde.topology import elements_of
from sympde.expr     import BilinearForm
from sympde.expr     import integral

DIM = 1
domain = Domain('\Omega', dim=DIM)

V = ScalarFunctionSpace('V', domain)

u,v = elements_of(V, names='u,v')

bx = Constant('bx')
by = Constant('by')
b  = Tuple(bx, by)

expr = integral(domain, dot(b, grad(v)) * dot(b, grad(u)))
a = BilinearForm((u,v), expr)
```

## Symbolic expression

```python
# imports from gelato
from gelato.expr     import gelatize
from gelato.printing import latex
```

We compute the GLT expression by calling the function **gelatize** for a given degree $p$

```python
p = 3

glt = gelatize(a, p)
```

```python

```

```python

```

```python

```

```python

```

```python

```

```python
# css style
from IPython.core.display import HTML
def css_styling():
    styles = open("../styles/custom.css", "r").read()
    return HTML(styles)
css_styling()
```

<link href='http://fonts.googleapis.com/css?family=Fenix' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Source+Code+Pro:300,400' rel='stylesheet' type='text/css'>
<style>
    @font-face {
        font-family: "Computer Modern";
        src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
    }
    div.cell{
        width:600px;
        margin-left:16% !important;
        margin-right:auto;
    }
    h1 {
        font-family: 'Alegreya Sans', sans-serif;
    }
    h2 {
        font-family: 'Fenix', serif;
    }
    h3{
		font-family: 'Fenix', serif;
        margin-top:12px;
        margin-bottom: 3px;
       }
	h4{
		font-family: 'Fenix', serif;
       }
    h5 {
        font-family: 'Alegreya Sans', sans-serif;
    }
    div.text_cell_render{
        font-family: 'Alegreya Sans',Computer Modern, "Helvetica Neue", Arial, Helvetica, Geneva, sans-serif;
        line-height: 135%;
        font-size: 120%;
        width:600px;
        margin-left:auto;
        margin-right:auto;
    }
    .CodeMirror{
            font-family: "Source Code Pro";
			font-size: 90%;
    }
/*    .prompt{
        display: None;
    }*/
    .text_cell_render h1 {
        font-weight: 200;
        font-size: 50pt;
		line-height: 100%;
        color:#054BCD;
        margin-bottom: 0.5em;
        margin-top: 0.5em;
        display: block;
    }
    .text_cell_render h5 {
        font-weight: 300;
        font-size: 16pt;
        color: #054BCD;
        font-style: italic;
        margin-bottom: .5em;
        margin-top: 0.5em;
        display: block;
    }

    .warning{
        color: rgb( 240, 20, 20 )
        }
</style>

```python

```


### 2018/2019 - Task List 5

[Book](https://raw.githubusercontent.com/jonesgithub/book-1/master/ML%20Machine%20Learning-A%20Probabilistic%20Perspective.pdf)

Pyro [SVI](http://pyro.ai/examples/svi_part_i.html#Setup)

Pyro [Bayesian Regression - Introduction](http://pyro.ai/examples/bayesian_regression.html)

[Video introduction to Bayesian Inference](https://www.youtube.com/watch?time_continue=3&v=DYRK0-_K2UU)

0.
 - Look at sample code
 - Watch introductory video
 - Read pyro introduction to bayesian regression
1. Select pair of parameters from Task List 3 that has bigest corellation
2. Train classic linear regresion
3. Train Bayesian lienar regresion
4. Visualize results in both cases

## Required imports

```python
%matplotlib inline
import pyro
import torch
import numpy as np
import matplotlib.pyplot as plt
import pyro.optim as optim
import pyro.distributions as dist
from torch.distributions import constraints
from tqdm import tqdm
from matplotlib import animation, rc
from IPython.display import HTML
import torch.nn as nn
from functools import partial
import pandas as pd
from pyro.contrib.autoguide import AutoDiagonalNormal
from pyro.infer import EmpiricalMarginal, SVI, Trace_ELBO, TracePredictive

import pandas as pd
import seaborn as sns
```

```python
pyro.set_rng_seed(1)
pyro.enable_validation(True)
```

## Generate some linear data

```python
a=2.0
b=3.0
noise_mean=0
noise_std=10
count=50
```

```python
base=torch.arange(0, count).type(torch.FloatTensor)
with pyro.plate("noise_plate", count):
    noise=pyro.sample("noise", pyro.distributions.Normal(noise_mean, noise_std))

x=base
y=x*a+b+noise
```

```python
x
```

    tensor([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13.,
            14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27.,
            28., 29., 30., 31., 32., 33., 34., 35., 36., 37., 38., 39., 40., 41.,
            42., 43., 44., 45., 46., 47., 48., 49.])

```python
y
```

    tensor([-12.2560,  -2.5023,   0.4602,  -7.0948,   9.9983,   6.9081,   5.2023,
              0.9090,  11.8786,  24.0372,  15.2269,  22.4854,  24.7773,  45.8711,
             33.2843,  37.6764,  28.0303,  25.3924,  45.9954,  42.9908,  51.6569,
             47.4440,  40.3709,  57.0731,  62.0168,  51.2406,  32.5444,  42.5354,
             59.6116,  54.8226,  55.0193,  63.6838,  85.7935,  68.2787,  71.6628,
             68.6296,  82.6260,  81.4151,  90.6514, 101.1539,  85.1518,  79.7581,
             85.1400,  82.5536, 106.3925,  84.3041,  61.6885,  89.5213, 110.1735,
            103.9814])

```python
plt.plot(x.numpy(), y.numpy(), "o")
```

Prepare rows of data

```python
x_data=x[:,None]
y_data=y[:,None]
```

## Model

\begin{align}
y & = a*x+b
\end{align}

```python
class RegressionModel(nn.Module):
    def __init__(self):
        # p = number of features
        super(RegressionModel, self).__init__()
        self.linear = nn.Linear(1, 1)
        self.factor = nn.Parameter(torch.tensor(1.))

    def forward(self, x):
        return self.linear(x) + (self.factor * x)
```

## Train linear regression

Learning code

```python
def train_linear_model(x, y):
    model=RegressionModel()
    loss_fn = torch.nn.MSELoss(reduction='sum')
    optim = torch.optim.Adam(model.parameters(), lr=0.0008)
    num_iterations = 2000
    for j in tqdm(range(num_iterations)):
        # run the model forward on the data
        y_pred = model(x)
        # calculate the mse loss
        loss = loss_fn(y_pred, y)
        # initialize gradients to zero
        optim.zero_grad()
        # backpropagate
        loss.backward()
        # take a gradient step
        optim.step()
        if (j + 1) % 5000 == 0:
            print("[iteration %04d] loss: %.4f" % (j + 1, loss.item()))
    # Inspect learned parameters
    print("Learned parameters:")
    for name, param in model.named_parameters():
        print(name, param.data.numpy())
    return model
```

Train

```python
trained_model = train_linear_model(x_data, y_data)
```

    100%|██████████| 2000/2000 [00:01<00:00, 1827.60it/s]

    Learned parameters:
    factor 1.53962242603302
    linear.weight [[ 0.52025026]]
    linear.bias [ 0.59336054]

```python
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
plt.plot(x_data.squeeze().numpy(), y_data.squeeze().numpy(), "bo")
y_learned=trained_model.forward(x_data)
plt.plot(x_data.squeeze().numpy(), y_learned.detach().squeeze().numpy(), 'r');
```

## Bayesian linear regresion

probabilistic model

assume some prior probabilities, get likelihood from data

```python
def create_probabilistic_model():
    regression_model=RegressionModel()
    def probabilistic_model(x_data, y_data):
        # weight and bias priors
        w_prior = pyro.distributions.Normal(torch.ones(1, 1)*3, torch.ones(1, 1)).to_event(1)
        b_prior = pyro.distributions.Normal(torch.zeros(1,1), torch.ones(1,1)).to_event(1)
        f_prior = pyro.distributions.Normal(2., 1000.)
        priors = {'linear.weight': w_prior, 'linear.bias': b_prior, 'factor': f_prior}
        scale = pyro.sample("sigma", pyro.distributions.Uniform(0., 10.))
        # lift module parameters to random variables sampled from the priors
        lifted_module = pyro.random_module("module", regression_model, priors)
        # sample a nn (which also samples w and b)
        lifted_reg_model = lifted_module()
        with pyro.plate("map", len(x_data)):
            # run the nn forward on data
            prediction_mean = lifted_reg_model(x_data).squeeze(-1)
            # condition on the observed data
            obs=pyro.sample("obs",
                        pyro.distributions.Normal(prediction_mean, scale),
                        obs=y_data)
            return prediction_mean
    return probabilistic_model
```

training code

use SVI and ELBO loss

```python
def train():
    pyro.clear_param_store()
    num_iterations=1000
    model = create_probabilistic_model()
    guide = AutoDiagonalNormal(model)
    optim = pyro.optim.Adam({"lr": 0.01})
    svi = pyro.infer.SVI(model, guide, optim, loss=pyro.infer.Trace_ELBO(), num_samples=count)
    t=tqdm(range(num_iterations))
    for j in t:
        loss = svi.step(x_data, y_data.squeeze(-1))
        t.set_postfix(loss=loss)
    return (model, svi)
```

```python
probabilistic_model, svi = train()
```

    100%|██████████| 1000/1000 [00:12<00:00, 79.80it/s, loss=220]

lets look what parameters we have

```python
for name, value in pyro.get_param_store().items():
    print(name, pyro.param(name))
```

    auto_loc tensor([0.7602, 0.9844, 0.9870, 0.9135], requires_grad=True)
    auto_scale tensor([0.7065, 0.3909, 0.3160, 1.0577], grad_fn=<AddBackward0>)

evaluation

use Delta distribution in order to force pyro to sample prediction so we will be able to get parameters of learned distribution

```python
def get_marginal(traces, sites):
    empirical_marginal=EmpiricalMarginal(traces, sites)
    return empirical_marginal._get_samples_and_weights()[0].detach().cpu().numpy()

def summary(traces, sites):
    marginal = get_marginal(traces, sites)
    site_stats = {}
    for i in range(marginal.shape[1]):
        site_name = sites[i]
        site_marginal=marginal[:, i]
        marginal_site = pd.DataFrame(site_marginal).transpose()
        describe = partial(pd.Series.describe, percentiles=[.05, 0.25, 0.5, 0.75, 0.95])
        site_stats[site_name] = marginal_site.apply(describe, axis=1) \
            [["mean", "std", "5%", "25%", "50%", "75%", "95%"]]
    return site_stats

def wrapped_model(x_data, y_data):
    model_result=probabilistic_model(x_data, y_data)
    pyro.sample("prediction", pyro.distributions.Delta(model_result))

posterior = svi.run(x_data, y_data)
```

```python
trace_pred = TracePredictive(wrapped_model,
                             posterior,
                             num_samples=count)
post_pred = trace_pred.run(x_data, None)
post_summary = summary(post_pred, sites= ['prediction', 'obs'])
mu = post_summary["prediction"]
y = post_summary["obs"]
predictions = pd.DataFrame({
    "x": x_data,
    "mu_mean": mu["mean"],
    "mu_perc_5": mu["5%"],
    "mu_perc_95": mu["95%"],
    "y_mean": y["mean"],
    "y_perc_5": y["5%"],
    "y_perc_95": y["95%"],
    "true_y": y_data,
})
```

Plot probabilities

```python
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
x_values=x_data.squeeze().numpy()
plt.plot(x_values, y_data.squeeze().numpy(), "bo")
plt.plot(x_values, predictions['mu_mean'], 'r');
plt.fill_between(x_values,
                   predictions["mu_perc_5"],
                   predictions["mu_perc_95"],
                   color='r',
                   alpha=0.5);
```

Plot outcome variables

```python
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
x_values=x_data.squeeze().numpy()
plt.plot(x_values, y_data.squeeze().numpy(), "bo")
plt.plot(x_values, predictions['y_mean'], 'r');
plt.fill_between(x_values,
                   predictions["y_perc_5"],
                   predictions["y_perc_95"],
                   color='r',
                   alpha=0.5);
```

# Solutions

```python
## Insert solution here
```

```python
# Load the datasets
abalone = pd.read_csv('abalone.csv', header=None)
abalone.columns = ["sex", "length", "diameter", "height", "whole_height", "shucked_weight", "viscera_weight", "shell_weight", "class"]
```

```python
def correlation(data):
    corr = data.corr()
    sns.heatmap(corr, cmap="PiYG",
            annot=True, center=0)

def scatter_one_pair(data, attr1, attr2):
    if type(data[attr1][0]) is not str and type(data[attr2][0]) is not str:
        if type(data['class'][0]) is not str:
            g = sns.jointplot(x=attr1, y=attr2, data=data)
            g.ax_joint.scatter(x=attr1,y=attr2, data=data, c=data['class'])
        else:
            g = sns.jointplot(x=attr1, y=attr2, data=data)

correlation(abalone)
scatter_one_pair(abalone, "length", "diameter")

x = torch.tensor(abalone.length)
y = torch.tensor(abalone.diameter)
```

```python
print(x)
print(y)

x_data=x[:,None]
y_data=y[:,None]

count = x.numel()
```

    tensor([0.4550, 0.3500, 0.5300,  ..., 0.6000, 0.6250, 0.7100])
    tensor([0.3650, 0.2650, 0.4200,  ..., 0.4750, 0.4850, 0.5550])

### Model

```python
class RegressionModel(nn.Module):
    def __init__(self):
        # p = number of features
        super(RegressionModel, self).__init__()
        self.linear = nn.Linear(1, 1)
        self.factor = nn.Parameter(torch.tensor(1.))

    def forward(self, x):
        return self.linear(x) + (self.factor * x)
```

#### Linear regression

```python
def train_linear_model(x, y):
    model=RegressionModel()
    loss_fn = torch.nn.MSELoss(reduction='sum')
    optim = torch.optim.Adam(model.parameters(), lr=0.0008)
    num_iterations = 10000
    for j in tqdm(range(num_iterations)):
        # run the model forward on the data
        y_pred = model(x)
        # calculate the mse loss
        loss = loss_fn(y_pred, y)
        # initialize gradients to zero
        optim.zero_grad()
        # backpropagate
        loss.backward()
        # take a gradient step
        optim.step()
        if (j + 1) % 5000 == 0:
            print("[iteration %04d] loss: %.4f" % (j + 1, loss.item()))
    # Inspect learned parameters
    print("Learned parameters:")
    for name, param in model.named_parameters():
        print(name, param.data.numpy())
    return model

trained_model = train_linear_model(x_data, y_data)
```

     49%|█████████████████                  | 4891/10000 [00:04<00:04, 1246.19it/s]

    [iteration 5000] loss: 1.0777

    100%|███████████████████████████████████▉| 9970/10000 [00:08<00:00, 967.04it/s]

    [iteration 10000] loss: 1.0777

    100%|██████████████████████████████████| 10000/10000 [00:08<00:00, 1204.32it/s]

    Learned parameters:
    factor 1.0641778
    linear.weight [[-0.24871705]]
    linear.bias [-0.01941371]

```python
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
plt.plot(x_data.squeeze().numpy(), y_data.squeeze().numpy(), "bo")
y_learned=trained_model.forward(x_data)
plt.plot(x_data.squeeze().numpy(), y_learned.detach().squeeze().numpy(), 'r');
```

#### Bayesian linear regression

```python
def create_probabilistic_model():
    regression_model=RegressionModel()
    def probabilistic_model(x_data, y_data):
        # weight and bias priors
        w_prior = pyro.distributions.Normal(torch.ones(1, 1)*3, torch.ones(1, 1)).to_event(1)
        b_prior = pyro.distributions.Normal(torch.zeros(1,1), torch.ones(1,1)).to_event(1)
        f_prior = pyro.distributions.Normal(2., 1000.)
        priors = {'linear.weight': w_prior, 'linear.bias': b_prior, 'factor': f_prior}
        scale = pyro.sample("sigma", pyro.distributions.Uniform(0., 10.))
        # lift module parameters to random variables sampled from the priors
        lifted_module = pyro.random_module("module", regression_model, priors)
        # sample a nn (which also samples w and b)
        lifted_reg_model = lifted_module()
        with pyro.plate("map", len(x_data)):
            # run the nn forward on data
            prediction_mean = lifted_reg_model(x_data).squeeze(-1)
            # condition on the observed data
            obs=pyro.sample("obs",
                        pyro.distributions.Normal(prediction_mean, scale),
                        obs=y_data)
            return prediction_mean
    return probabilistic_model

def train():
    pyro.clear_param_store()
    num_iterations=10000
    model = create_probabilistic_model()
    guide = AutoDiagonalNormal(model)
    optim = pyro.optim.Adam({"lr": 0.01})
    svi = pyro.infer.SVI(model, guide, optim, loss=pyro.infer.Trace_ELBO(), num_samples=count)
    t=tqdm(range(num_iterations))
    for j in t:
        loss = svi.step(x_data, y_data.squeeze(-1))
        t.set_postfix(loss=loss)
    return (model, svi)

probabilistic_model, svi = train()

for name, value in pyro.get_param_store().items():
    print(name, pyro.param(name))
```

    100%|█████████████████████| 10000/10000 [02:51<00:00, 58.17it/s, loss=-1.13e+4]

    auto_loc tensor([-6.4275,  0.3952,  0.4192, -0.0199], requires_grad=True)
    auto_scale tensor([0.0136, 0.0006, 0.0005, 0.0003], grad_fn=<AddBackward0>)

```python
def get_marginal(traces, sites):
    empirical_marginal=EmpiricalMarginal(traces, sites)
    return empirical_marginal._get_samples_and_weights()[0].detach().cpu().numpy()

def summary(traces, sites):
    marginal = get_marginal(traces, sites)
    site_stats = {}
    for i in range(marginal.shape[1]):
        site_name = sites[i]
        site_marginal=marginal[:, i]
        marginal_site = pd.DataFrame(site_marginal).transpose()
        describe = partial(pd.Series.describe, percentiles=[.05, 0.25, 0.5, 0.75, 0.95])
        site_stats[site_name] = marginal_site.apply(describe, axis=1) \
            [["mean", "std", "5%", "25%", "50%", "75%", "95%"]]
    return site_stats

def wrapped_model(x_data, y_data):
    model_result=probabilistic_model(x_data, y_data)
    pyro.sample("prediction", pyro.distributions.Delta(model_result))

posterior = svi.run(x_data, y_data)

trace_pred = TracePredictive(wrapped_model,
                             posterior,
                             num_samples=count)
post_pred = trace_pred.run(x_data, None)
post_summary = summary(post_pred, sites= ['prediction', 'obs'])
mu = post_summary["prediction"]
y = post_summary["obs"]
predictions = pd.DataFrame({
    "x": x_data,
    "mu_mean": mu["mean"],
    "mu_perc_5": mu["5%"],
    "mu_perc_95": mu["95%"],
    "y_mean": y["mean"],
    "y_perc_5": y["5%"],
    "y_perc_95": y["95%"],
    "true_y": y_data,
})
```

```python
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
x_values=x_data.squeeze().numpy()
plt.plot(x_values, y_data.squeeze().numpy(), "bo")
plt.plot(x_values, predictions['mu_mean'], 'r');
aind = np.argsort(x_values)
plt.fill_between(x_values[aind],
                   predictions["mu_perc_5"][aind],
                   predictions["mu_perc_95"][aind],
                   color='r',
                   alpha=0.5);
```

```python
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))
x_values=x_data.squeeze().numpy()
plt.plot(x_values, y_data.squeeze().numpy(), "bo")
plt.plot(x_values, predictions['y_mean'], 'r');
aind = np.argsort(x_values)
plt.fill_between(x_values[aind],
                   predictions["y_perc_5"][aind],
                   predictions["y_perc_95"][aind],
                   color='r',
                   alpha=0.5);
```

```python

```


### Simulation

```python
from IPython.display import Image
Image(filename="figure.png",width=600)
from IPython.display import display, HTML

display(HTML(data="""
<style>
    div#notebook-container    { width: 95%; }
    div#menubar-container     { width: 65%; }
    div#maintoolbar-container { width: 99%; }
</style>
"""))
```

<style>
    div#notebook-container    { width: 95%; }
    div#menubar-container     { width: 65%; }
    div#maintoolbar-container { width: 99%; }
</style>

### Equation generation

```python
import sympy as sp
import numpy as np
from IPython.display import display
sp.init_printing(use_latex='mathjax')
```

```python
# parameters
# Angular moment of inertia
J_B = 1e-2 * np.diag([1., 1., 1.])

# Gravity
g_I = np.array((-1, 0., 0.))

# Fuel consumption
alpha_m = 0.01

# Vector from thrust point to CoM
r_T_B = np.array([-1e-2, 0., 0.])

def dir_cosine(q):
        return np.matrix([
            [1 - 2 * (q[2] ** 2 + q[3] ** 2), 2 * (q[1] * q[2] +
                                                   q[0] * q[3]), 2 * (q[1] * q[3] - q[0] * q[2])],
            [2 * (q[1] * q[2] - q[0] * q[3]), 1 - 2 *
             (q[1] ** 2 + q[3] ** 2), 2 * (q[2] * q[3] + q[0] * q[1])],
            [2 * (q[1] * q[3] + q[0] * q[2]), 2 * (q[2] * q[3] -
                                                   q[0] * q[1]), 1 - 2 * (q[1] ** 2 + q[2] ** 2)]
        ])

def omega(w):
        return np.matrix([
            [0, -w[0], -w[1], -w[2]],
            [w[0], 0, w[2], -w[1]],
            [w[1], -w[2], 0, w[0]],
            [w[2], w[1], -w[0], 0],
        ])

def skew(v):
    return np.matrix([
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0]
        ])

```

```python
f = sp.zeros(14, 1)

x = sp.Matrix(sp.symbols(
    'm rx ry rz vx vy vz q0 q1 q2 q3 wx wy wz', real=True))
u = sp.Matrix(sp.symbols('ux uy uz', real=True))

g_I = sp.Matrix(g_I)
r_T_B = sp.Matrix(r_T_B)
J_B = sp.Matrix(J_B)

C_B_I = dir_cosine(x[7:11, 0])
C_I_B = C_B_I.transpose()

f[0, 0] = - alpha_m * u.norm()
f[1:4, 0] = x[4:7, 0]
f[4:7, 0] = 1 / x[0, 0] * C_I_B * u + g_I
f[7:11, 0] = 1 / 2 * omega(x[11:14, 0]) * x[7: 11, 0]
f[11:14, 0] = J_B ** -1 * \
    (skew(r_T_B) * u - skew(x[11:14, 0]) * J_B * x[11:14, 0])

```

```python
display(sp.simplify(f)) # f
```

$$\left[\begin{matrix}- 0.01 \sqrt{ux^{2} + uy^{2} + uz^{2}}\\vx\\vy\\vz\\\frac{- 1.0 m - ux \left(2 q_{2}^{2} + 2 q_{3}^{2} - 1\right) - 2 uy \left(q_{0} q_{3} - q_{1} q_{2}\right) + 2 uz \left(q_{0} q_{2} + q_{1} q_{3}\right)}{m}\\\frac{2 ux \left(q_{0} q_{3} + q_{1} q_{2}\right) - uy \left(2 q_{1}^{2} + 2 q_{3}^{2} - 1\right) - 2 uz \left(q_{0} q_{1} - q_{2} q_{3}\right)}{m}\\\frac{- 2 ux \left(q_{0} q_{2} - q_{1} q_{3}\right) + 2 uy \left(q_{0} q_{1} + q_{2} q_{3}\right) - uz \left(2 q_{1}^{2} + 2 q_{2}^{2} - 1\right)}{m}\\- 0.5 q_{1} wx - 0.5 q_{2} wy - 0.5 q_{3} wz\\0.5 q_{0} wx + 0.5 q_{2} wz - 0.5 q_{3} wy\\0.5 q_{0} wy - 0.5 q_{1} wz + 0.5 q_{3} wx\\0.5 q_{0} wz + 0.5 q_{1} wy - 0.5 q_{2} wx\\0\\1.0 uz\\- 1.0 uy\end{matrix}\right]$$

```python
display(sp.simplify(f.jacobian(x)))# A
```

$$\left[\begin{array}{cccccccccccccc}0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\frac{ux \left(2 q_{2}^{2} + 2 q_{3}^{2} - 1\right) + 2 uy \left(q_{0} q_{3} - q_{1} q_{2}\right) - 2 uz \left(q_{0} q_{2} + q_{1} q_{3}\right)}{m^{2}} & 0 & 0 & 0 & 0 & 0 & 0 & \frac{2 \left(q_{2} uz - q_{3} uy\right)}{m} & \frac{2 \left(q_{2} uy + q_{3} uz\right)}{m} & \frac{2 \left(q_{0} uz + q_{1} uy - 2 q_{2} ux\right)}{m} & \frac{2 \left(- q_{0} uy + q_{1} uz - 2 q_{3} ux\right)}{m} & 0 & 0 & 0\\\frac{- 2 ux \left(q_{0} q_{3} + q_{1} q_{2}\right) + uy \left(2 q_{1}^{2} + 2 q_{3}^{2} - 1\right) + 2 uz \left(q_{0} q_{1} - q_{2} q_{3}\right)}{m^{2}} & 0 & 0 & 0 & 0 & 0 & 0 & \frac{2 \left(- q_{1} uz + q_{3} ux\right)}{m} & \frac{2 \left(- q_{0} uz - 2 q_{1} uy + q_{2} ux\right)}{m} & \frac{2 \left(q_{1} ux + q_{3} uz\right)}{m} & \frac{2 \left(q_{0} ux + q_{2} uz - 2 q_{3} uy\right)}{m} & 0 & 0 & 0\\\frac{2 ux \left(q_{0} q_{2} - q_{1} q_{3}\right) - 2 uy \left(q_{0} q_{1} + q_{2} q_{3}\right) + uz \left(2 q_{1}^{2} + 2 q_{2}^{2} - 1\right)}{m^{2}} & 0 & 0 & 0 & 0 & 0 & 0 & \frac{2 \left(q_{1} uy - q_{2} ux\right)}{m} & \frac{2 \left(q_{0} uy - 2 q_{1} uz + q_{3} ux\right)}{m} & \frac{2 \left(- q_{0} ux - 2 q_{2} uz + q_{3} uy\right)}{m} & \frac{2 \left(q_{1} ux + q_{2} uy\right)}{m} & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & - 0.5 wx & - 0.5 wy & - 0.5 wz & - 0.5 q_{1} & - 0.5 q_{2} & - 0.5 q_{3}\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0.5 wx & 0 & 0.5 wz & - 0.5 wy & 0.5 q_{0} & - 0.5 q_{3} & 0.5 q_{2}\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0.5 wy & - 0.5 wz & 0 & 0.5 wx & 0.5 q_{3} & 0.5 q_{0} & - 0.5 q_{1}\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0.5 wz & 0.5 wy & - 0.5 wx & 0 & - 0.5 q_{2} & 0.5 q_{1} & 0.5 q_{0}\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\end{array}\right]$$

```python
sp.simplify(f.jacobian(u)) # B
```

$$\left[\begin{matrix}- \frac{0.01 ux}{\sqrt{ux^{2} + uy^{2} + uz^{2}}} & - \frac{0.01 uy}{\sqrt{ux^{2} + uy^{2} + uz^{2}}} & - \frac{0.01 uz}{\sqrt{ux^{2} + uy^{2} + uz^{2}}}\\0 & 0 & 0\\0 & 0 & 0\\0 & 0 & 0\\\frac{- 2 q_{2}^{2} - 2 q_{3}^{2} + 1}{m} & \frac{2 \left(- q_{0} q_{3} + q_{1} q_{2}\right)}{m} & \frac{2 \left(q_{0} q_{2} + q_{1} q_{3}\right)}{m}\\\frac{2 \left(q_{0} q_{3} + q_{1} q_{2}\right)}{m} & \frac{- 2 q_{1}^{2} - 2 q_{3}^{2} + 1}{m} & \frac{2 \left(- q_{0} q_{1} + q_{2} q_{3}\right)}{m}\\\frac{2 \left(- q_{0} q_{2} + q_{1} q_{3}\right)}{m} & \frac{2 \left(q_{0} q_{1} + q_{2} q_{3}\right)}{m} & \frac{- 2 q_{1}^{2} - 2 q_{2}^{2} + 1}{m}\\0 & 0 & 0\\0 & 0 & 0\\0 & 0 & 0\\0 & 0 & 0\\0 & 0 & 0\\0 & 0 & 1.0\\0 & -1.0 & 0\end{matrix}\right]$$

### Ref

- Python implementation of 'Successive Convexification for 6-DoF Mars Rocket Powered Landing with Free-Final-Time' paper
by Michael Szmuk and Behçet Açıkmeşe.

- inspired by EmbersArc/SuccessiveConvexificationFreeFinalTime: Implementation of "Successive Convexification for 6-DoF Mars Rocket Powered Landing with Free-Final-Time" https://github.com/EmbersArc/SuccessiveConvexificationFreeFinalTime


# VAE tutorial
Here we implement a VAE using Lightning.

⚡ [The video tutorial is here](...)

⚡ [Remember! we're always available for support on Slack](https://join.slack.com/t/pytorch-lightning/shared_invite/zt-f6bl2l0l-JYMK3tbAgAmGRrlNr00f1A)

---
# VAE Loss

The loss we'll be implementing for the VAE is called the ELBO.

$$
\color{red}{\min \mathbb{E}_{q}[ \log q(z|x) - \log p(z)]} - \color{blue}{\mathbb{E}_{q} \log p(x|z)}
$$

The first term is the KL-divergence and the second term the reconstruction.

For a full derivation and explanation see the next section

---
# Math

If you don't care about the derivation, please skip this part.

There are 2 ways to derive the VAE loss.
1. Using Jensen's inequality
2. Using KL divergence.

In this tutorial we'll use KL divergence.

## ELBO derivation

### What is the ELBO?
The result of this derivation is a loss called ELBO (evidence lower bound). The ELBO has two terms

Term 1: Reconstruction
$$
- \mathbb{E}_q \log p(x|z)
$$

Term 2: KL divergence
$$
\mathbb{E}_{q}  \log \frac{q(z|x)}{p(z)}
$$

Together, these two terms define the ELBO:

$$
- \mathbb{E}_{q}  \log p(x|z) + \mathbb{E}_{q}  \log \frac{q(z|x)}{p(z)}
$$

### ELBO as a bound
Minimizing KL divergence is equivalent to maximizing ELBO (with some slack).

This means that ELBO is a lowerbound on KL divergence.
Here's a rendering of what we mean by this in abstract terms

---
### How we get ELBO

Here's how we end up finding the ELBO.

First, remember we want to model a distribution over inputs, $ p(x) $. These can be images, etc...

We assume that there are factors that influence x but that we can't observe. We call these *latent* variables. We use $z$ to refer to these.

So, by using marginalization, we can rewrite this as

$$
p(x) = \int_z p(x|z)dz
$$

Marginalization is just fancy for summing over a variable.
Here's an example using a discrete probability distribution.

So, the joint model that uses $z$ variables to explain $x$ is defined by:

$$
p(x,z) = p(z|x)p(x)\\
\frac{p(x,z)}{p(x)} = p(z|x)
$$

But again, we find the $p(x)$ term at the bottom

To solve this we define $q(z|x)$ to approximate $p(z|x)$. $q$ is called the approximate posterior and $p$ the true posterior

Now this is starting to look like something we can do with machine learning. We have a source ($p$) and a target ($q$). We can treat this as an optimization problem where we parametrize $q$ (ie: a neural network) to try to get close to $p$ via gradient descent.

This is called variational inference (in general).

### KL divergence definition

The distance to minimize is called the KL divergence.
When we minimize the KL divergence, our $q$ will be an estimator of $p$.

This is the definition for the KL divergence between two distributions P, Q:

$$
D_{KL}(Q||P) = \sum Q \log \frac{Q}{P}
$$

We'll replace Q, P with our distributions $P = p(z|x)$ and $Q = q(z|x)$
$$
D_{KL}(q(z|x)||p(z|x)) = \sum q(z|x) \log \frac{q(z|x)}{p(z|x)}
$$

#### Sidenote about probability notation
I know using notation like $p(x)$ is confusing for those new to probability theory, but it is just a placeholder that stands for any distribution we want.
For example, $p(x)$ can be a Normal distribution, or a poisson, or any of the dozens of distributions available. It makes the math easier to manipulate simple letters with a few properties from probability theory

### ELBO derivation from KL(Q||P)
So, let's just play around with KL to see what comes out of the other end

\begin{align}
D_{KL}(q(z|x)||p(z|x)) &= \sum q(z|x) \log \frac{q(z|x)}{p(z|x)}\\
&= -\sum q(z|x) \log \frac{p(z|x)}{q(z|x)}\\
&= -\sum q(z|x)  \left[\log p(z|x) - \log q(z|x) \right]\\
&= -\sum q(z|x) \left[\log \frac{p(x|z)p(z)}{p(x)} - \log q(z|x) \right]\\
&= -\sum q(z|x) \left[\log p(x|z) + \log p(z) - \log p(x) - \log q(z|x) \right]\\
&= \sum q(z|x) \left[-\log p(x|z) - \log p(z) + \log p(x) + \log q(z|x) \right]\\
D_{KL}(q(z|x)||p(z|x)) - \sum q(z|x) \left[-\log p(x|z) - \log p(z) + \log q(z|x) \right] &= \log p(x)\\
-\sum q(z|x) \left[-\log p(x|z) + \log \frac{q(z|x)}{p(z)} \right] &= \\
-\mathbb{E}_{z \sim q(z|x)} \left[-\log p(x|z) + \log \frac{q(z|x)}{p(z)} \right] &=\\
-\mathbb{E}_{q} \left[-\log p(x|z) + \log \frac{q(z|x)}{p(z)} \right] &=\\
\mathbb{E}_{q} \left[\log p(x|z) - \log \frac{q(z|x)}{p(z)} \right] &=\\
D_{KL}(q(z|x)||p(z|x)) + \left[ \mathbb{E}_{q} \log p(x|z)  - \mathbb{E}_{q}  \log \frac{q(z|x)}{p(z)} \right] &= \log p(x)\\
\end{align}

### Neural network loss
To maximize the ELBO, we can minimize the negative of the ELBO.

$$
\max \mathbb{E}_{q} \log p(x|z)  - \mathbb{E}_{q}  \log \frac{q(z|x)}{p(z)} =
\min  - [\mathbb{E}_{q} \log p(x|z)  - \mathbb{E}_{q}  \log \frac{q(z|x)}{p(z)}]
$$

### Monte carlo loss
\begin{align}
-\mathbb{E}_{q} \log p(x|z)  + \mathbb{E}_{q}[ \log q(z|x) - \log p(z)]&= \\
\mathbb{E}_{q}[ \log q(z|x) - \log p(z)] - \mathbb{E}_{q} \log p(x|z)&=
\end{align}

The end result is:

\begin{align}
D_{KL}(q(z|x)||p(z|x)) + \left[ \mathbb{E}_{q} \log p(x|z)  - \mathbb{E}_{q}  \log \frac{q(z|x)}{p(z)} \right] &= \log p(x)
\end{align}

The first term is the original $D_{KL}$ we want to minimize, ($D_{KL}(q(z|x)||p(z|x))$). The second term is called the variational lower bound (ELBO).

$$
\max \mathcal{L} = \mathbb{E}_{q} \log p(x|z)  - \mathbb{E}_{q}  \log \frac{q(z|x)}{p(z)}
$$

Since $\log p(x)$ is a constant, and we want to minimize the $D_{KL}$ term, we can achieve the same by maximizing the ELBO.

So, the end goal is to minimize this function:

$$
\color{red}{\min \mathbb{E}_{q}[ \log q(z|x) - \log p(z)]} - \color{blue}{\mathbb{E}_{q} \log p(x|z)}
$$

### Other implementations
Here's where a disconnect might happen. For most tutorials, the tutorials assume that $p(z)$ (the prior) is gaussian and also assumes that $p(z|x)$ is also gaussian.

So, most tutorials end up with a KL divergence that looks like:

$$
-0.5 \sum{1 \log{(\sigma)} - \mu^2 - e^{\log{(\sigma)}}}
$$

But in our case, we won't do that because we want to retain the flexibility to modify distributions as needed.

And since we don't know the KL between all possible pairs of distributions, we'll actually just use monte-carlo sampling for this.

Recall that:
$$
\mathbb{E}_p = \frac{1}{N}\sum_{}p(x)
$$

So, anywhere we see $\mathbb{E}_q$ we will draw a sample from q and use it in the formula.

```
import torch

q = torch.distributions.Normal(2, 4)

z = q.rsample()
print(z)
```

    tensor(6.0110)

## KL divergence intuition
Here we show the intuition of minimizing the KL

```
p = torch.distributions.Normal(0, 1)
q = torch.distributions.Normal(2, 4)

log_pz = p.log_prob(z)
log_qzx = q.log_prob(z)

print('log prob pz: ', log_pz, 'prob:', torch.exp(log_pz))
print('log prob qzx: ', log_qzx, 'prob:', torch.exp(log_qzx))
```

    log prob pz:  tensor(-18.9850) prob: tensor(5.6877e-09)
    log prob qzx:  tensor(-2.8080) prob: tensor(0.0603)

```
kl_divergence = log_qzx - log_pz
kl_divergence
```

    tensor(16.1770)

now, if we manually move q closer to p, we see that this distrance has shrunk.

```
# old q = torch.distributions.Normal(2, 4)
q = torch.distributions.Normal(1, 2)

log_qzx = q.log_prob(z)
print('log prob qzx: ', log_qzx, 'prob:', torch.exp(log_qzx))
```

    log prob qzx:  tensor(-4.7508) prob: tensor(0.0086)

new kl

```
new_kl_divergence = log_qzx - log_pz
new_kl_divergence
```

    tensor(14.2341)

---
## Implementation walk-through
So, with our formula defined, let's start our implementation.

If you skipped the math and elbo derivation, remember we are implementing this loss (the ELBO)

$$
\color{red}{\min \mathbb{E}_{q}[ \log q(z|x) - \log p(z)]} - \color{blue}{\mathbb{E}_{q} \log p(x|z)}
$$

```
%%capture
! pip install pytorch-lightning
! pip install pytorch-lightning-bolts==0.2.5rc1
```

#### First 4 elements
The main 4 elements of a VAE are:
- encoder
- decoder
- mu
- sigma

mu, sigma are the parameters for the distributions in our formula.

Note: in our notation we say $q(z|x)$ where we assume $q$ to be a neural network. However, q is a neural network that DOES NOT GENERATE z. Instead it generates parameters FOR a distribution OVER z.

```
import pytorch_lightning as pl
from torch import nn
from torch.nn import functional as F
from pl_bolts.models.autoencoders.components import (
    resnet18_decoder,
    resnet18_encoder,
)

class VAE(pl.LightningModule):
    def __init__(self, enc_out_dim=512, latent_dim=256, input_height=32):
        super().__init__()

        self.save_hyperparameters()

        # encoder, decoder
        self.encoder = resnet18_encoder(False, False)
        self.decoder = resnet18_decoder(
            latent_dim=latent_dim,
            input_height=input_height,
            first_conv=False,
            maxpool1=False
        )

        # distribution parameters
        self.fc_mu = nn.Linear(enc_out_dim, latent_dim)
        self.fc_var = nn.Linear(enc_out_dim, latent_dim)
```

#### Sample z

In Lightning we define the model interactions and the core algorithm in the training_step.

For EDUCATIONAL purposes only, we are adding the training_step to the lightning module here instead of in the original class definition

The VAE algorithm goes like this:

1. Sample z

```
vae = VAE()
```

```
# DATA
# we're pretending to have an image from cifar-10 (3 channels, 32x32 pixels)
x = torch.rand(1, 3, 32, 32)

print('image shape:', x.shape)

# GET Q(z|x) PARAMETERS
# encode x to get the mu and variance parameters
x_encoded = vae.encoder(x)
mu, log_var = vae.fc_mu(x_encoded), vae.fc_var(x_encoded)

print('mu:', mu.shape)
print('log_var:', log_var.shape)

# SAMPLE Z from Q(Z|x)
std = torch.exp(log_var / 2)
q = torch.distributions.Normal(mu, std)
z = q.rsample()

print('z shape:', z.shape)
```

    image shape: torch.Size([1, 3, 32, 32])
    mu: torch.Size([1, 256])
    log_var: torch.Size([1, 256])
    z shape: torch.Size([1, 256])

```
zero = torch.zeros_like(mu)
one = torch.ones_like(std)
p = torch.distributions.Normal(zero, one)
```

#### decoding

The second part is reconstruction.

The decoder is going to take as input $z$ and it will generate PARAMETERS for a distribution from which to sample a reconstruction.

Note: Most tutorials use MNIST so the decoder generates BERNOULLI variables which can be used in a bernoulli distribution to sample a value for each pixel.

That's right, the decoder generates distribution parameters FOR EACH pixel. And since most tutorials use MNIST the outputs of the decoder will LOOK like a digit because those parameters are in the range 0-1 (which happens to be the same for black and white images).

But in our case, we'll actually use CIFAR-10 which are color images... so we need to do something more realistic:

```
x_hat = vae.decoder(z)
print(x_hat.shape)
```

    torch.Size([1, 3, 32, 32])

#### reconstruction loss

You'll likely see the reconstruction loss using MSE like this in most tutorials about autoencoders.

However, MSE DOES not apply to Variational autoencoders

```
recon_loss = F.mse_loss(x_hat, x, reduction='none')
recon_loss = recon_loss.sum(-1).sum(-1).sum(-1)
print('wrong reconstruction loss', recon_loss)
```

    wrong reconstruction loss tensor([2402.9331], grad_fn=<SumBackward1>)

Instead, the output of the decoder represents parameters for a distribution.

When the images are MNIST, the parameters (decoder output) can be considered bernoullis and thus you can use the pixel values of MNIST (because they are ones and zeros) to calculate the cross entropy (which is the "reconstruction loss").

But when dealing with color images, it is prefered to model each parameter out of the decoder as a gaussian.

```
# get the PARAMETERS for P(x|z)
x_hat = vae.decoder(z)
print('we have these many parameters for a distribution:', x_hat.shape)

# compute the probability of x under this n-dimensional distribution
log_scale = nn.Parameter(torch.Tensor([0.0]))
scale = torch.exp(log_scale)
dist = torch.distributions.Normal(x_hat, scale)
log_pxz = dist.log_prob(x)

print(log_pxz.shape)

# sum across channels, and all pixels (width, height)
# this gives the log probability of THIS particular IMAGE
log_pxz = log_pxz.sum(dim=(1, 2, 3))
print(log_pxz.shape)
print('reconstruction loss:', log_pxz.item())

recon_loss = log_pxz
```

    we have these many parameters for a distribution: torch.Size([1, 3, 32, 32])
    torch.Size([1, 3, 32, 32])
    torch.Size([1])
    reconstruction loss: -8363.24609375

#### kl divergence

Next we'll define the second term of the loss which is the KL divergence.

Now, we don't want to compute the KL divergence analytically, so we'll instead use the monte carlo estimate.

$$
= \mathbb{E}_{q}  \log \frac{q(z|x)}{p(z)}\\
= \mathbb{E}_{q} [\log q(z|x) - \log{p(z)}]
$$

To code this, we'll need to be able to get the probability of z under each of $p(z)$ and $q(z|x)$.

This means we don't need to know what the distribution is, but only sample from it.

```
# 1. define the first two probabilities (in this case Normal for both)
p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std))
q = torch.distributions.Normal(mu, std)

# 2. get the probabilities from the equation
log_qzx = q.log_prob(z)
log_pz = p.log_prob(z)

print(log_qzx.shape, log_pz.shape)

# 3. calculate the KL
kl = (log_qzx - log_pz)

# 4. this bit is a bit tricky. Since these are log probabilities
# we can sum all the individual dimensions to give us the multi-dimensional
# probability
kl = kl.sum(-1)
print(kl)
```

    torch.Size([1, 256]) torch.Size([1, 256])
    tensor([44.8220], grad_fn=<SumBackward1>)

#### ELBO

So, the full elbo is:

```
elbo = (kl - recon_loss)
print('kl:', kl.mean())
print('log likelihood (reconstruction loss):', recon_loss.mean())
print('elbo: ', elbo.mean())
```

    kl: tensor(44.8220, grad_fn=<MeanBackward0>)
    log likelihood (reconstruction loss): tensor(-4024.4458, grad_fn=<MeanBackward0>)
    elbo:  tensor(4069.2678, grad_fn=<MeanBackward0>)

# Full Implementation
So, the full implementation is as follows.

```
import pytorch_lightning as pl
from torch import nn
from torch.nn import functional as F
from pl_bolts.models.autoencoders.components import (
    resnet18_decoder,
    resnet18_encoder,
)

class VAE(pl.LightningModule):
    def __init__(self, enc_out_dim=512, latent_dim=256, input_height=32):
        super().__init__()

        self.save_hyperparameters()

        # encoder, decoder
        self.encoder = resnet18_encoder(False, False)
        self.decoder = resnet18_decoder(
            latent_dim=latent_dim,
            input_height=input_height,
            first_conv=False,
            maxpool1=False
        )

        # distribution parameters
        self.fc_mu = nn.Linear(enc_out_dim, latent_dim)
        self.fc_var = nn.Linear(enc_out_dim, latent_dim)

        # for the gaussian likelihood
        self.log_scale = nn.Parameter(torch.Tensor([0.0]))

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-4)

    def gaussian_likelihood(self, mean, logscale, sample):
        scale = torch.exp(logscale)
        dist = torch.distributions.Normal(mean, scale)
        log_pxz = dist.log_prob(sample)
        return log_pxz.sum(dim=(1, 2, 3))

    def kl_divergence(self, z, mu, std):
        # --------------------------
        # Monte carlo KL divergence
        # --------------------------
        # 1. define the first two probabilities (in this case Normal for both)
        p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std))
        q = torch.distributions.Normal(mu, std)

        # 2. get the probabilities from the equation
        log_qzx = q.log_prob(z)
        log_pz = p.log_prob(z)

        # kl
        kl = (log_qzx - log_pz)
        kl = kl.sum(-1)
        return kl

    def training_step(self, batch, batch_idx):
        x, _ = batch

        # encode x to get the mu and variance parameters
        x_encoded = self.encoder(x)
        mu, log_var = self.fc_mu(x_encoded), self.fc_var(x_encoded)

        # sample z from q
        std = torch.exp(log_var / 2)
        q = torch.distributions.Normal(mu, std)
        z = q.rsample()

        # decoded
        x_hat = vae.decoder(z)

        # reconstruction loss
        recon_loss = self.gaussian_likelihood(x_hat, self.log_scale, x)

        # kl
        kl = self.kl_divergence(z, mu, std)

        # elbo
        elbo = (kl - recon_loss)
        elbo = elbo.mean()

        self.log_dict({
            'elbo': elbo,
            'kl': kl.mean(),
            'recon_loss': recon_loss.mean(),
            'reconstruction': recon_loss.mean(),
            'kl': kl.mean(),
        })

        return elbo
```

Let's use CIFAR-10 already split up and transformed.

The Lightning Datamodule has 3 dataloaders, train, val, test

```
from pl_bolts.datamodules import CIFAR10DataModule

datamodule = CIFAR10DataModule('.')
```

Now we train!

```
pl.seed_everything(1234)

vae = VAE()
trainer = pl.Trainer(gpus=1, max_epochs=30, progress_bar_refresh_rate=10)
trainer.fit(vae, cifar10)
```

## Plot an image

```
from matplotlib.pyplot import imshow, figure
import numpy as np
from torchvision.utils import make_grid
from pl_bolts.transforms.dataset_normalizations import cifar10_normalization
figure(figsize=(8, 3), dpi=300)

# Z COMES FROM NORMAL(0, 1)
num_preds = 16
p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std))
z = p.rsample((num_preds,))

# SAMPLE IMAGES
with torch.no_grad():
    pred = vae.decoder(z.to(vae.device)).cpu()

# UNDO DATA NORMALIZATION
normalize = cifar10_normalization()
mean, std = np.array(normalize.mean), np.array(normalize.std)
img = make_grid(pred).permute(1, 2, 0).numpy() * std + mean

# PLOT IMAGES
imshow(img);
```