# Tutorial-IllinoisGRMHD: apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Authors: Leo Werneck & Zach Etienne

**This module is currently under development**

## In this tutorial module we explain two major functions within IllinoisGRMHD that are used to ensure that the results obtained throughout the simulation are Physically sound.

### Required and recommended citations:

* **(Required)** Etienne, Z. B., Paschalidis, V., Haas R., Mösta P., and Shapiro, S. L. IllinoisGRMHD: an open-source, user-friendly GRMHD code for dynamical spacetimes. Class. Quantum Grav. 32 (2015) 175009. ([arxiv:1501.07276](http://arxiv.org/abs/1501.07276)).
* **(Required)** Noble, S. C., Gammie, C. F., McKinney, J. C., Del Zanna, L. Primitive Variable Solvers for Conservative General Relativistic Magnetohydrodynamics. Astrophysical Journal, 641, 626 (2006) ([astro-ph/0512420](https://arxiv.org/abs/astro-ph/0512420)).
* **(Recommended)** Del Zanna, L., Bucciantini N., Londrillo, P. An efficient shock-capturing central-type scheme for multidimensional relativistic flows - II. Magnetohydrodynamics. A&A 400 (2) 397-413 (2003). DOI: 10.1051/0004-6361:20021641 ([astro-ph/0210618](https://arxiv.org/abs/astro-ph/0210618)).

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

This module is organized as follows

0. [Step 0](#src_dir): **Source directory creation**
1. [Step 1](#introduction): **Introduction**
1. [Step 2](#apply_tau_floor): **The `apply_tau_floor` function**
    1. [Step 2.a](#positive_definiteness_of_the_metric): *Positive-definiteness of the metric*
    1. [Step 2.b](#barbi_barb_i_barb2_and_barb): *Computing $\bar{B}^{i}$, $\bar{B}_{i}$, $\bar{B}^{2}$, and $\bar{B}$*
    1. [Step 2.c](#barbdots_hatbarbdots_and_sdots): *Computing $\bar{B}\cdot\tilde{S}$, $\hat{\bar{B}}\cdot\tilde{S}$, and $\tilde{S}^{2}$*
    1. [Step 2.d](#modifying_tau): *Modifying $\tilde{\tau}$*
        1. [Step 2.d.i](#wm_sm_and_wmin): $W_{m}$, $\tilde{S}_{m}$, and $W_{\min}$
        1. [Step 2.d.ii](#tau_min): $\tau_{\min}$
    1. [Step 2.e](#modifying_tilde_s_i): *Modifying $\tilde{S}_{i}$*
1. [Step 3](#enforce_limits_on_primitives_and_recompute_conservs): **The `IllinoisGRMHD_enforce_limits_on_primitives_and_recompute_conservs` function**
    1. [Step 3.a](#enforce_limits_on_primitives): *Enforcing physical limits on the primitive variables*
    1. [Step 3.b](#bi_b2_and_u_i): *Computing $b^{i}$, $b^{2}$, and $u_{i}$*
    1. [Step 3.c](#gupmunu): *The physical ADM inverse 4-metric, $g^{\mu\nu}$*
    1. [Step 3.d](#tupmunu): *The stress-energy tensor, $T^{\mu\nu}$*
    1. [Step 3.e](#gdnmunu): *The physical ADM 4-metric, $g_{\mu\nu}$*
    1. [Step 3.f](#tdnmunu): *The stress-energy tensor, $T_{\mu\nu}$*
    1. [Step 3.g](#recompute_conservs): *Recomputing the conservative variables*
1. [Step 4](#code_validation): **Code validation**
1. [Step 5](#latex_pdf_output): **Output this notebook to $\LaTeX$-formatted PDF file**

# Step 0: Source directory creation \[Back to [top](#toc)\]
$$\label{src_dir}$$

We will now use the [cmdline_helper.py NRPy+ module](Tutorial-Tutorial-cmdline_helper.ipynb) to create the source directory within the `IllinoisGRMHD` NRPy+ directory, if it does not exist yet.

```python
# Step 0: Creation of the IllinoisGRMHD source directory
# Step 0a: Add NRPy's directory to the path
# https://stackoverflow.com/questions/16780014/import-file-from-parent-directory
import os,sys
nrpy_dir_path = os.path.join("..","..")
if nrpy_dir_path not in sys.path:
    sys.path.append(nrpy_dir_path)

# Step 0b: Load up cmdline_helper and create the directory
import cmdline_helper as cmd
IGM_src_dir_path = os.path.join("..","src")
cmd.mkdir(IGM_src_dir_path)

# Step 0c: Create the output file path
outfile_path__tau_and_prim_limits__C = os.path.join(IGM_src_dir_path,"apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C")
```

# Step 1: Introduction \[Back to [top](#toc)\]
$$\label{introduction}$$

In this tutorial notebook we will discuss how we adjust our conservative variables given that our primitive variables are in the physical range.

For a given set of primitive variables $\left\{\rho_{b},P,v^i,B^i\right\}$ in the physical range (i.e. $\rho_{b}\geq0$, $P\geq0$ and $\epsilon\geq0$), the corresponding conservative variables $\left\{\rho_{\star},\tilde{\tau},\tilde{S}_{i},\tilde{B}^{i}\right\}$ must satisfy certain inequalities (see appendix A of [Etienne *et al.* (2012)](https://arxiv.org/pdf/1112.0568.pdf) for the full discussion). Here we provide a practical recipe to impose these inequalities approximately to reduce inversion failures, which occur mainly in regions with very low density in the artificial “atmosphere” or inside the BH horizon where high accuracy is difficult to maintain but not crucial.

# Step 2: The `apply_tau_floor` function \[Back to [top](#toc)\]
$$\label{apply_tau_floor}$$

Here we will simply declare the `apply_tau_floor()` and store some of the input in useful variables.

```python
%%writefile $outfile_path__tau_and_prim_limits__C
void eigenvalues_3by3_real_sym_matrix(CCTK_REAL & lam1, CCTK_REAL & lam2, CCTK_REAL & lam3,
                                      CCTK_REAL M11, CCTK_REAL M12, CCTK_REAL M13, CCTK_REAL M22, CCTK_REAL M23, CCTK_REAL M33);

static inline int apply_tau_floor(const int index,const CCTK_REAL tau_atm,const CCTK_REAL rho_b_atm,const CCTK_REAL Psi6threshold,CCTK_REAL *PRIMS,CCTK_REAL *METRIC,CCTK_REAL *METRIC_PHYS,CCTK_REAL *METRIC_LAP_PSI4,output_stats &stats,eos_struct &eos,  CCTK_REAL *CONSERVS) {
```

    Writing ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 2.a: Positive-definiteness of the metric \[Back to [top](#toc)\]
$$\label{positive_definiteness_of_the_metric}$$

We start by verifying if the metrix $\gamma_{ij}$ is positive definite. Notice that although we expect this to always be true, the metric may lose its positive-definiteness due to numerical error during the evolution, especially in the region deep inside the BH, near the “puncture”.

To verify whether or not the [metric is positive definite, we analyse its eigenvectors](https://en.wikipedia.org/wiki/Definiteness_of_a_matrix#Eigenvalues). If the metrix is *not* positive definite, we reset $\gamma_{ij}\to\psi^{4}\tilde{\gamma}_{ij}$, where $\tilde{\gamma}_{ij}$ corresponds to the 3D flat metric tensor.

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  //First apply the rho_star floor:

  //rho_star = alpha u0 Psi6 rho_b, alpha u0 > 1, so if rho_star < Psi6 rho_b_atm, then we are GUARANTEED that we can reset to atmosphere.
  //if(CONSERVS[RHOSTAR] < 1e4*METRIC_LAP_PSI4[PSI6]*rho_b_atm) {
  //if(CONSERVS[RHOSTAR] < 2*METRIC_LAP_PSI4[PSI6]*rho_b_atm) {

  CCTK_REAL lam1,lam2,lam3;
  eigenvalues_3by3_real_sym_matrix(lam1, lam2, lam3,METRIC[GXX], METRIC[GXY], METRIC[GXZ], METRIC[GYY], METRIC[GYZ], METRIC[GZZ]);
  if (lam1 < 0.0 || lam2 < 0.0 || lam3 < 0.0) {
    // Metric is not positive-defitive, reset the metric to be conformally-flat.
    METRIC[GXX] = 1.0;
    METRIC[GXY] = 0.0;
    METRIC[GXZ] = 0.0;
    METRIC[GYY] = 1.0;
    METRIC[GYZ] = 0.0;
    METRIC[GZZ] = 1.0;
    METRIC_PHYS[GUPXX] = METRIC_LAP_PSI4[PSIM4];
    METRIC_PHYS[GUPXY] = 0.0;
    METRIC_PHYS[GUPXZ] = 0.0;
    METRIC_PHYS[GUPYY] = METRIC_LAP_PSI4[PSIM4];
    METRIC_PHYS[GUPYZ] = 0.0;
    METRIC_PHYS[GUPZZ] = METRIC_LAP_PSI4[PSIM4];
  }
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 2.b: Computing $\bar{B}^{i}$, $\bar{B}_{i}$, $\bar{B}^{2}$, and $\bar{B}$ \[Back to [top](#toc)\]
$$\label{barbi_barb_i_barb2_and_barb}$$

We then set

$$
\boxed{\bar{B}^{i} = \frac{B^{i}}{\sqrt{4\pi}}}\ ,
$$

and

$$
\bar{B}_{i} = \gamma_{ij}\bar{B}^{j} \implies
\boxed{
\left\{
\begin{matrix}
\bar{B}_{x} = \gamma_{xx}\bar{B}^{x} + \gamma_{xy}\bar{B}^{y} + \gamma_{xz}\bar{B}^{z}\\
\bar{B}_{y} = \gamma_{yx}\bar{B}^{x} + \gamma_{yy}\bar{B}^{y} + \gamma_{yz}\bar{B}^{z}\\
\bar{B}_{z} = \gamma_{zx}\bar{B}^{x} + \gamma_{zy}\bar{B}^{y} + \gamma_{zz}\bar{B}^{z}
\end{matrix}
\right.
}\ ,
$$

then

$$
\bar{B}^{2} \equiv B_{i}B^{i} \implies \boxed{\bar{B}^{2} = B_{x}B^{x} + B_{y}B^{y} + B_{z}B^{z}}\ ,
$$

and finally

$$
\boxed{\bar{B} \equiv \sqrt{\bar{B}^{2}}}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  //Next, prepare for the tau and stilde fixes:

  CCTK_REAL Bxbar = PRIMS[BX_CENTER]*ONE_OVER_SQRT_4PI,Bybar = PRIMS[BY_CENTER]*ONE_OVER_SQRT_4PI,Bzbar = PRIMS[BZ_CENTER]*ONE_OVER_SQRT_4PI;

  CCTK_REAL Bbar_x = METRIC_PHYS[GXX]*Bxbar + METRIC_PHYS[GXY]*Bybar + METRIC_PHYS[GXZ]*Bzbar;
  CCTK_REAL Bbar_y = METRIC_PHYS[GXY]*Bxbar + METRIC_PHYS[GYY]*Bybar + METRIC_PHYS[GYZ]*Bzbar;
  CCTK_REAL Bbar_z = METRIC_PHYS[GXZ]*Bxbar + METRIC_PHYS[GYZ]*Bybar + METRIC_PHYS[GZZ]*Bzbar;
  CCTK_REAL Bbar2 = Bxbar*Bbar_x + Bybar*Bbar_y + Bzbar*Bbar_z;
  CCTK_REAL Bbar = sqrt(Bbar2);
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

The next part of the code is written to prevent [floating-point underflow](https://en.wikipedia.org/wiki/Arithmetic_underflow). We compute $\bar{B}$ in a different way. We start by evaluating

$$
\bar{B}_{\rm check} = \left|\bar{B}^{x}\right| + \left|\bar{B}^{y}\right| + \left|\bar{B}^{z}\right|\ ,
$$

and verifying whether that is a very small, positive number. Then, we determine the largest component of $\bar{B}_{\rm check}$:

$$
\bar{B}_{\max} = \max\left(\left|\bar{B}^{x}\right|,\left|\bar{B}^{y}\right|,\left|\bar{B}^{z}\right|\right)\ .
$$

Then, we rescale $\bar{B}_{i}$ and $\bar{B}^{i}$ using

$$
\left(\bar{B}^{i}\right)_{\rm tmp} \equiv \frac{\bar{B}^{i}}{\bar{B}_{\max}}\ ,\quad
\left(\bar{B}_{i}\right)_{\rm tmp} \equiv \frac{\bar{B}_{i}}{\bar{B}_{\max}}\ ,
$$

and finally recompute $\bar{B}$

$$
\bar{B} = \left[\left(\bar{B}_{i}\right)_{\rm tmp}\left(\bar{B}^{i}\right)_{\rm tmp}\right]\bar{B}_{\max}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  CCTK_REAL check_B_small = fabs(Bxbar)+fabs(Bybar)+fabs(Bzbar);
  if (check_B_small>0 && check_B_small<1.e-150) {
    // need to compute Bbar specially to prevent floating-point underflow
    CCTK_REAL Bmax = fabs(Bxbar);
    if (Bmax < fabs(Bybar)) Bmax=fabs(Bybar);
    if (Bmax < fabs(Bzbar)) Bmax=fabs(Bzbar);
    CCTK_REAL Bxtmp=Bxbar/Bmax, Bytemp=Bybar/Bmax, Bztemp=Bzbar/Bmax;
    CCTK_REAL B_xtemp=Bbar_x/Bmax, B_ytemp=Bbar_y/Bmax, B_ztemp=Bbar_z/Bmax;
    Bbar = sqrt(Bxtmp*B_xtemp + Bytemp*B_ytemp + Bztemp*B_ztemp)*Bmax;
  }
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 2.c: Computing $\bar{B}\cdot\tilde{S}$, $\hat{\bar{B}}\cdot\tilde{S}$, and $\tilde{S}^{2}$ \[Back to [top](#toc)\]
$$\label{barbdots_hatbarbdots_and_sdots}$$

Then we compute

$$
\bar{B} \cdot \tilde{S} = \bar{B}_{i}\tilde{S}^{i} = \bar{B}_{x}\tilde{S}^{x} + \bar{B}_{y}\tilde{S}^{y} + \bar{B}_{z}\tilde{S}^{z}\ .
$$

and

$$
\hat{\bar{B}}\cdot\tilde{S} \equiv \frac{\bar{B} \cdot \tilde{S}}{\bar{B}}\ .
$$

However, if $\bar{B} \ll 1$, we set $\hat{\bar{B}}\cdot\tilde{S}=0$.

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  CCTK_REAL BbardotS = Bxbar*CONSERVS[STILDEX] + Bybar*CONSERVS[STILDEY] + Bzbar*CONSERVS[STILDEZ];
  CCTK_REAL hatBbardotS = BbardotS/Bbar;
  if (Bbar<1.e-300) hatBbardotS = 0.0;

  // Limit hatBbardotS
  //CCTK_REAL max_gammav = 100.0;
  //CCTK_REAL rhob_max = CONSERVS[RHOSTAR]/METRIC_LAP_PSI4[PSI6];
  //CCTK_REAL hmax = 1.0 + 2.0*rhob_max;
  //CCTK_REAL abs_hatBbardotS_max = sqrt(SQR(max_gammav)-1.0)*CONSERVS[RHOSTAR]*hmax;
  //if (fabs(hatBbardotS) > abs_hatBbardotS_max) {
  //   CCTK_REAL fac_reduce = abs_hatBbardotS_max/fabs(hatBbardotS);
  //   CCTK_REAL hatBbardotS_max = hatBbardotS*fac_reduce;
  //   CCTK_REAL Bbar_inv = 1.0/Bbar;
  //   CCTK_REAL hat_Bbar_x = Bbar_x*Bbar_inv;
  //   CCTK_REAL hat_Bbar_y = Bbar_y*Bbar_inv;
  //   CCTK_REAL hat_Bbar_z = Bbar_z*Bbar_inv;
  //   CCTK_REAL sub_fact = hatBbardotS_max - hatBbardotS;
  //   CONSERVS[STILDEX] += sub_fact*hat_Bbar_x;
  //   CONSERVS[STILDEY] += sub_fact*hat_Bbar_y;
  //   CONSERVS[STILDEZ] += sub_fact*hat_Bbar_z;
  //   hatBbardotS = hatBbardotS_max;
  //   BbardotS *= fac_reduce;
  //   CONSERVS[STILDEX] = CONSERVS[STILDEX]; CONSERVS[STILDEY] = CONSERVS[STILDEY]; CONSERVS[STILDEZ] = CONSERVS[STILDEZ];
  //}
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

Next we compute

$$
\tilde{S}^{2} \equiv \tilde{S} \cdot \tilde{S} = \gamma^{ij}\tilde{S}_{i}\tilde{S}_{j}\ ,
$$

i.e.

$$
\boxed{
\begin{align}
\tilde{S}^{2} &= \gamma^{xx}\left(\tilde{S}_{x}\right)^{2}
                           + \gamma^{yy}\left(\tilde{S}_{y}\right)^{2}
                           + \gamma^{zz}\left(\tilde{S}_{z}\right)^{2}\\
                          &+2\left(
                             \gamma^{xy}\tilde{S}_{x}\tilde{S}_{y}
                            +\gamma^{xz}\tilde{S}_{x}\tilde{S}_{z}
                            +\gamma^{yz}\tilde{S}_{y}\tilde{S}_{z}
                           \right)
\end{align}
}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  CCTK_REAL sdots= METRIC_PHYS[GUPXX]*SQR(CONSERVS[STILDEX])+METRIC_PHYS[GUPYY]*SQR(CONSERVS[STILDEY])+METRIC_PHYS[GUPZZ]*SQR(CONSERVS[STILDEZ])+2.0*
    (METRIC_PHYS[GUPXY]*CONSERVS[STILDEX]*CONSERVS[STILDEY]+METRIC_PHYS[GUPXZ]*CONSERVS[STILDEX]*CONSERVS[STILDEZ]+METRIC_PHYS[GUPYZ]*CONSERVS[STILDEY]*CONSERVS[STILDEZ]);
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 2.d: Modifying $\tilde{\tau}$ \[Back to [top](#toc)\]
$$\label{modifying_tau}$$

### Step 2.d.i: $W_{m}$, $\tilde{S}_{m}$, and $W_{\min}$ \[Back to [top](#toc)\]
$$\label{wm_sm_and_wmin}$$

Then we compute other useful quantities, which are eqs. (A52), (A53), and (A54) in appendix A of [Etienne *et al.* (2012)](https://arxiv.org/pdf/1112.0568.pdf)

$$
\begin{align}
W_{m} &= \psi^{-6}\left[\left(\hat{\bar{B}}\cdot\tilde{S}\right)^{2}+\rho_{\star}^{2}\right]^{1/2}\ ,\\
\tilde{S}_{m}^{2} &= \frac{W_{m}^{2}\tilde{S}^{2} + \left(\bar{B}\cdot\tilde{S}\right)^{2}\left(\bar{B}^{2}+2W_{m}\right)}{\left(W_{m}+\bar{B}^{2}\right)^{2}}\ ,\\
W_{\min} &= \psi^{-6}\left(S_{m}^{2}+\rho_{\star}^{2}\right)^{1/2}\ ,\\
\end{align}
$$

respectively (notice the slightly different notation between the equations above and the one used in the paper).

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  CCTK_REAL Wm = sqrt(SQR(hatBbardotS)+ SQR(CONSERVS[RHOSTAR]))/METRIC_LAP_PSI4[PSI6];
  CCTK_REAL Sm2 = (SQR(Wm)*sdots + SQR(BbardotS)*(Bbar2+2.0*Wm))/SQR(Wm+Bbar2);
  CCTK_REAL Wmin = sqrt(Sm2 + SQR(CONSERVS[RHOSTAR]))/METRIC_LAP_PSI4[PSI6];
  CCTK_REAL sdots_fluid_max = sdots;
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

### Step 2.d.ii: $\tilde{\tau}_{\min}$ \[Back to [top](#toc)\]
$$\label{tau_min}$$

Next we evaluate

$$
\tilde{\tau}_{\min} = \tilde{\tau} - \frac{\psi^{6}}{2}\bar{B}^{2} - \frac{\bar{B}^{2}\tilde{S}^{2} - \left(\bar{B}\cdot\tilde{S}\right)^{2}}{2\psi^{6}\left(W_{\min}+\bar{B}^{2}\right)^{2}}
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  //tau fix, applicable when B==0 and B!=0:
  if(CONSERVS[TAUENERGY] < 0.5*METRIC_LAP_PSI4[PSI6]*Bbar2) {
    CONSERVS[TAUENERGY] = tau_atm+0.5*METRIC_LAP_PSI4[PSI6]*Bbar2;
    stats.failure_checker+=1000000;
  }

  CCTK_REAL tau_fluid_min = CONSERVS[TAUENERGY] - 0.5*METRIC_LAP_PSI4[PSI6]*Bbar2 - (Bbar2*sdots - SQR(BbardotS))*0.5/(METRIC_LAP_PSI4[PSI6]*SQR(Wmin+Bbar2));
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

Then we verify if $\tilde{\tau}_{\min} \geq \tilde{\tau}_{\rm atm}$. If $\tilde{\tau}_{\min} < \tilde{\tau}_{\rm atm}$, then reset $\tilde{\tau}$:

$$
\tilde{\tau} = \tilde{\tau}_{\min} + \frac{\psi^{6}}{2}\bar{B}^{2} + \frac{\bar{B}^{2}\tilde{S}^{2} - \left(\bar{B}\cdot\tilde{S}\right)^{2}}{2\psi^{6}\left(W_{\min}+\bar{B}^{2}\right)^{2}}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  //Apply Stilde fix when B==0.
  //if(PRIMS[BX_CENTER]==0 && PRIMS[BY_CENTER]==0 && PRIMS[BZ_CENTER]==0 && (METRIC_LAP_PSI4[PSI6]>30.0 || CONSERVS[RHOSTAR]/METRIC_LAP_PSI4[PSI6]<100*rho_b_atm)) {
  //if(check_B_small < 1.e-300) {

  /**********************************
   * Piecewise Polytropic EOS Patch *
   *         Computing Patm         *
   **********************************/
  /* This modification of the code trades the variable
   * "gamma_equals2" by the already defined function
   * pow().
   *
   * Also, assuming that Patm < rho_ppoly_tab[0], we skip
   * the declaration of new variables to store the
   * values of K_ppoly_tab[0] and Gamma_ppoly_tab[0]. Thus:
   *  -----------------------------------------
   * | P_{atm} = K_{0} * rho_{atm}^{Gamma_{0}} |
   *  -----------------------------------------
   */
  int polytropic_index = find_polytropic_K_and_Gamma_index(eos,rho_b_atm);
  CCTK_REAL Patm = eos.K_ppoly_tab[polytropic_index]*pow(rho_b_atm,eos.Gamma_ppoly_tab[polytropic_index]);

  if(check_B_small*check_B_small < Patm*1e-32) {
    CCTK_REAL rhot=CONSERVS[TAUENERGY]*(CONSERVS[TAUENERGY]+2.0*CONSERVS[RHOSTAR]);
    CCTK_REAL safetyfactor = 0.999999;
    //if(METRIC_LAP_PSI4[PSI6]>Psi6threshold) safetyfactor=0.99;

    if(sdots > safetyfactor*rhot) {
      CCTK_REAL rfactm1 = sqrt((safetyfactor*rhot)/sdots);
      CONSERVS[STILDEX]*=rfactm1;
      CONSERVS[STILDEY]*=rfactm1;
      CONSERVS[STILDEZ]*=rfactm1;
      stats.failure_checker+=10000000;
    }
  } else if(METRIC_LAP_PSI4[PSI6]>Psi6threshold) {
    //Apply new Stilde fix.
    if (tau_fluid_min < tau_atm*1.001) {
      tau_fluid_min = tau_atm*1.001;
      CONSERVS[TAUENERGY] = tau_fluid_min + 0.5*METRIC_LAP_PSI4[PSI6]*Bbar2 + (Bbar2*sdots - SQR(BbardotS))*0.5/(METRIC_LAP_PSI4[PSI6]*SQR(Wmin+Bbar2));
    }
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 2.e: Modifying $\tilde{S}_{i}$ \[Back to [top](#toc)\]
$$\label{modifying_tilde_s_i}$$

Then we check if $\tilde{S}^{2} \leq \tilde{\tau}_{\min}\left(\tilde{\tau}_{\min}+2\rho_{\star}\right)$. If not, we reset $\tilde{S}_{i}$

$$
\tilde{S}_{i}\to \tilde{S}_{i}\sqrt{\frac{\tilde{\tau}_{\min}\left(\tilde{\tau}_{\min}+2\rho_{\star}\right)}{\tilde{S}^{2}}}
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

    CCTK_REAL LHS = tau_fluid_min*(tau_fluid_min+2.0*CONSERVS[RHOSTAR]);
    CCTK_REAL RHS = sdots_fluid_max;

    CCTK_REAL safetyfactor = 0.999999;
    if(safetyfactor*LHS < RHS) {
      CCTK_REAL rfactm1 = sqrt((safetyfactor*LHS)/RHS);
      CONSERVS[STILDEX]*=rfactm1;
      CONSERVS[STILDEY]*=rfactm1;
      CONSERVS[STILDEZ]*=rfactm1;
      stats.failure_checker+=100000000;
    }
  }

  return 0;
}

/***********************************************************/
/***********************************************************/
/***********************************************************/
/***********************************************************/

```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

# Step 3: `The IllinoisGRMHD_ enforce_limits_on_primitives_and_recompute_conservs function` \[Back to [top](#toc)\]
$$\label{enforce_limits_on_primitives_and_recompute_conservs}$$

## Step 3.a: Enforcing physical limits on the primitive variables \[Back to [top](#toc)\]
$$\label{enforce_limits_on_primitives}$$

We start by imposing physical limits on the primitive variables $\left\{\rho_{b},P,v^{i}\right\}$, using:

1. $\rho_{b} \to \min\left(\rho_{b},\rho_{b,{\rm atm}}\right)$
1. `enforce_pressure_floor_ceiling()`: documented in the [`inlined_functions.C`](Tutorial-IllinoisGRMHD__inlined_functions.ipynb) tutorial module.
1. `impose_speed_limit_output_u0()`: documented in the [`inlined_functions.C`](Tutorial-IllinoisGRMHD__inlined_functions.ipynb) tutorial module.

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

void IllinoisGRMHD_enforce_limits_on_primitives_and_recompute_conservs(const int already_computed_physical_metric_and_inverse,CCTK_REAL *U,struct output_stats &stats,eos_struct &eos,
                                                                       CCTK_REAL *METRIC,CCTK_REAL g4dn[4][4],CCTK_REAL g4up[4][4], CCTK_REAL *TUPMUNU,CCTK_REAL *TDNMUNU,CCTK_REAL *CONSERVS) {
#ifndef ENABLE_STANDALONE_IGM_C2P_SOLVER
  DECLARE_CCTK_PARAMETERS;
#endif

  CCTK_REAL METRIC_LAP_PSI4[NUMVARS_METRIC_AUX];
  SET_LAPSE_PSI4(METRIC_LAP_PSI4,METRIC);

  // Useful debugging tool, can be used to track fixes:
  //CCTK_REAL rho_b_orig=U[RHOB],P_orig=U[PRESSURE],vx_orig=U[VX],vy_orig=U[VY],vz_orig=U[VZ];

  /***********************************************************/
  // Enforce limits on pressure, density, and v^i
  /***********************************************************/
  // Density floor:
  U[RHOB] = MAX(U[RHOB],rho_b_atm);
  // Density ceiling:
  U[RHOB] = MIN(U[RHOB],rho_b_max);

  //   Next set h, the enthalpy:
  CCTK_REAL h_enthalpy,  P_cold,eps_cold,dPcold_drho,eps_th,Gamma_cold; /* <- Note that in setting h, we need to define several
                                                                         *    other variables. Though some will be unused later
                                                                         *    in this function, they may be useful in other
                                                                         *    functions */
  compute_P_cold__eps_cold__dPcold_drho__eps_th__h__Gamma_cold(U,eos,Gamma_th,P_cold,eps_cold,dPcold_drho,eps_th,h_enthalpy,Gamma_cold);

  // Pressure floor & ceiling:
  int polytropic_index = find_polytropic_K_and_Gamma_index(eos,U[RHOB]);
  enforce_pressure_floor_ceiling(stats,eos.K_ppoly_tab[polytropic_index],P_cold,METRIC_LAP_PSI4[PSI6],Psi6threshold,U[RHOB],rho_b_atm,  U[PRESSURE]);

  // Possibly adjusted pressure, so recompute eps & h:
  CCTK_REAL eps = eps_cold + (U[PRESSURE]-P_cold)/(Gamma_th-1.0)/U[RHOB];
  h_enthalpy = 1.0 + eps + U[PRESSURE]/U[RHOB];

  CCTK_REAL uUP[4];
  impose_speed_limit_output_u0(METRIC,U,METRIC_LAP_PSI4[PSI4],METRIC_LAP_PSI4[LAPSEINV],stats, uUP[0]);
  // Compute u^i. We've already set uUP[0] in the lines above.
  for(int ii=0;ii<3;ii++) uUP[UX+ii] = uUP[0]*U[VX+ii];

  // Useful debugging tool, part 2: can be used to track fixes:
  //if(P_orig!=U[PRESSURE] || rho_b_orig!=U[RHOB] || vx_orig!=U[VX] || vy_orig!=U[VY] || vz_orig!=U[VZ]) {
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 3.b: Computing $b^{i}$, $b^{2}$, and $u_{i}$ \[Back to [top](#toc)\]
$$\label{bi_b2_and_u_i}$$

Next we compute quantities associted with $b^{i}$, using the `compute_smallba_b2_and_u_i_over_u0_psi4()` function which is documented in the [`inlined_functions.C`](Tutorial-IllinoisGRMHD__inlined_functions.ipynb) tutorial module. We also compute $u_{i}$ from the output of the `compute_smallba_b2_and_u_i_over_u0_psi4()` function.

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  /***************************************************************/
  // COMPUTE TUPMUNU, TDNMUNU, AND  CONSERVATIVES FROM PRIMITIVES
  /***************************************************************/
  // Compute b^{\mu}, b^2, and u_i/(u^0 Psi4)
  CCTK_REAL ONE_OVER_LAPSE_SQRT_4PI = METRIC_LAP_PSI4[LAPSEINV]*ONE_OVER_SQRT_4PI;
  CCTK_REAL u_x_over_u0_psi4,u_y_over_u0_psi4,u_z_over_u0_psi4;
  CCTK_REAL smallb[NUMVARS_SMALLB];
  compute_smallba_b2_and_u_i_over_u0_psi4(METRIC,METRIC_LAP_PSI4,U,uUP[0],ONE_OVER_LAPSE_SQRT_4PI,
                                          u_x_over_u0_psi4,u_y_over_u0_psi4,u_z_over_u0_psi4,smallb);
  // Compute u_i; we compute u_0 below.
  CCTK_REAL uDN[4] = { 1e200, u_x_over_u0_psi4*uUP[0]*METRIC_LAP_PSI4[PSI4],u_y_over_u0_psi4*uUP[0]*METRIC_LAP_PSI4[PSI4],u_z_over_u0_psi4*uUP[0]*METRIC_LAP_PSI4[PSI4] };
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 3.c: The physical ADM inverse 4-metric, $g^{\mu\nu}$ \[Back to [top](#toc)\]
$$\label{gupmunu}$$

Then we start evaluating the physical ADM inverse 4-metric, $g^{\mu\nu}$ (eq. 2.119 in [Numerical Relativity, by Baumgarte & Shapiro](https://books.google.com/books/about/Numerical_Relativity.html?id=dxU1OEinvRUC))

$$
\boxed{
g^{\mu\nu} =
\begin{pmatrix}
-\alpha^{-2} & \alpha^{-2}\beta^{i}\\
\alpha^{-2}\beta^{j} & \gamma^{ij} - \alpha^{-2}\beta^{i}\beta^{j}
\end{pmatrix}
}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  // Precompute some useful quantities, for later:
  CCTK_REAL alpha_sqrt_gamma=METRIC_LAP_PSI4[LAPSE]*METRIC_LAP_PSI4[PSI6];
  CCTK_REAL rho0_h_plus_b2 = (U[RHOB]*h_enthalpy + smallb[SMALLB2]);
  CCTK_REAL P_plus_half_b2 = (U[PRESSURE]+0.5*smallb[SMALLB2]);

  if(!already_computed_physical_metric_and_inverse) {
    // If you like, see Eq 2.119 in Numerical Relativity, by Baumgarte & Shapiro:
    CCTK_REAL ONE_OVER_LAPSE_SQUARED = SQR(METRIC_LAP_PSI4[LAPSEINV]);

    // g^{\mu \nu} = upper four-metric.
    //CCTK_REAL g4up[4][4];
    g4up[0][0] = -ONE_OVER_LAPSE_SQUARED;
    g4up[0][1] = ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTX];
    g4up[0][2] = ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTY];
    g4up[0][3] = ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTZ];
    g4up[1][1] = METRIC[GUPXX]*METRIC_LAP_PSI4[PSIM4] - ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTX]*METRIC[SHIFTX];
    g4up[1][2] = METRIC[GUPXY]*METRIC_LAP_PSI4[PSIM4] - ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTX]*METRIC[SHIFTY];
    g4up[1][3] = METRIC[GUPXZ]*METRIC_LAP_PSI4[PSIM4] - ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTX]*METRIC[SHIFTZ];
    g4up[2][2] = METRIC[GUPYY]*METRIC_LAP_PSI4[PSIM4] - ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTY]*METRIC[SHIFTY];
    g4up[2][3] = METRIC[GUPYZ]*METRIC_LAP_PSI4[PSIM4] - ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTY]*METRIC[SHIFTZ];
    g4up[3][3] = METRIC[GUPZZ]*METRIC_LAP_PSI4[PSIM4] - ONE_OVER_LAPSE_SQUARED*METRIC[SHIFTZ]*METRIC[SHIFTZ];
  }
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 3.d: The stress-energy tensor, $T^{\mu\nu}$ \[Back to [top](#toc)\]
$$\label{tupmunu}$$

Then we compute the stress-energy tensor, $T^{\mu\nu}$, which is given by eq. (8) in [Etienne *et al.* (2015)](https://arxiv.org/pdf/1501.07276.pdf):

$$
\boxed{
T^{\mu\nu} = \left(\rho_{0}h+b^{2}\right)u^{\mu}u^{\nu} + \left(P + \frac{b^{2}}{2}\right)g^{\mu\nu} - b^{\mu}b^{\nu}
}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  int count;
  // Next compute T^{\mu \nu}. This is very useful when computing fluxes and source terms in the GRMHD evolution equations.
  // (Eq. 33 in http://arxiv.org/pdf/astro-ph/0503420.pdf):
  // T^{mn} = (rho_0 h + b^2) u^m u^n + (P + 0.5 b^2) g^{mn} - b^m b^n, where m and n both run from 0 to 3.
  count=0; for(int ii=0;ii<4;ii++) for(int jj=ii;jj<4;jj++) { TUPMUNU[count] = rho0_h_plus_b2*uUP[ii]*uUP[jj] + P_plus_half_b2*g4up[ii][jj] - smallb[SMALLBT+ii]*smallb[SMALLBT+jj]; count++; }
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 3.e: The physical ADM 4-metric, $g_{\mu\nu}$ \[Back to [top](#toc)\]
$$\label{gdnmunu}$$

Next we wish to compute the stress-energy tensor, $T_{\mu\nu}$, which is given by

$$
T_{\mu\nu} = \left(\rho_{0}h+b^{2}\right)u_{\mu}u_{\nu} + \left(P + \frac{b^{2}}{2}\right)g_{\mu\nu} - b_{\mu}b_{\nu}\ .
$$

We start by setting up the physical ADM 4-metric, $g_{\mu\nu}$ (eq. 2.122 in [Numerical Relativity, by Baumgarte & Shapiro](https://books.google.com/books/about/Numerical_Relativity.html?id=dxU1OEinvRUC))

$$
\boxed{g_{\mu\nu} =
\begin{pmatrix}
-\alpha^{2} + \beta_{\ell}\beta^{\ell} & \beta_{i}\\
\beta_{j} & \gamma_{ij}
\end{pmatrix}}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  // Next compute T_{\mu \nu}
  // T_{mn} = (rho_0 h + b^2) u_m u_n + (P + 0.5 b^2) g_{mn} - b_m b_n, where m and n both run from 0 to 3.
  if(!already_computed_physical_metric_and_inverse) {
    CCTK_REAL LAPSE_SQUARED=SQR(METRIC_LAP_PSI4[LAPSE]);
    CCTK_REAL BETADN[4],BETAUP[4] = { 0.0, METRIC[SHIFTX],METRIC[SHIFTY],METRIC[SHIFTZ] };
    lower_4vector_output_spatial_part(METRIC_LAP_PSI4[PSI4],METRIC,BETAUP, BETADN);

    //CCTK_REAL g4dn[4][4];
    // g_{00} = - alpha^2 + gamma_{ij} beta^i beta^j = - alpha^2 beta_i beta^i
    g4dn[0][0] = -LAPSE_SQUARED + (BETAUP[1]*BETADN[1] + BETAUP[2]*BETADN[2] + BETAUP[3]*BETADN[3]);
    // g_{0i} =  gamma_{ij} beta^j = beta_i
    g4dn[0][1] = g4dn[1][0] = BETADN[1];
    g4dn[0][2] = g4dn[2][0] = BETADN[2];
    g4dn[0][3] = g4dn[3][0] = BETADN[3];
    // g_{ij} =  gamma_{ij} = Psi^4 \tilde{gamma_ij}
    g4dn[1][1] =              METRIC[GXX]*METRIC_LAP_PSI4[PSI4];
    g4dn[1][2] = g4dn[2][1] = METRIC[GXY]*METRIC_LAP_PSI4[PSI4];
    g4dn[1][3] = g4dn[3][1] = METRIC[GXZ]*METRIC_LAP_PSI4[PSI4];
    g4dn[2][2] =              METRIC[GYY]*METRIC_LAP_PSI4[PSI4];
    g4dn[2][3] = g4dn[3][2] = METRIC[GYZ]*METRIC_LAP_PSI4[PSI4];
    g4dn[3][3] =              METRIC[GZZ]*METRIC_LAP_PSI4[PSI4];
  }
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 3.f: The stress-energy tensor, $T_{\mu\nu}$ \[Back to [top](#toc)\]
$$\label{tdnmunu}$$

Then we compute $b_{\mu} = g_{\mu\nu}b^{\nu}$ and $u_{0} = g_{0\mu}u^{\mu}$.

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  CCTK_REAL smallb_lower[NUMVARS_SMALLB];
  // FIXME: This could be replaced by the function call
  //           lower_4vector_output_spatial_part(psi4,smallb,smallb_lower);
  // b_a = b^c g_{ac}
  for(int ii=0;ii<4;ii++) { smallb_lower[SMALLBT+ii]=0; for(int jj=0;jj<4;jj++) smallb_lower[SMALLBT+ii] += smallb[SMALLBT+jj]*g4dn[ii][jj]; }

  // Compute u_0, as we've already computed u_i above.
  uDN[0]=0.0; for(int jj=0;jj<4;jj++) uDN[0] += uUP[jj]*g4dn[0][jj];
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

Finally, compute

$$
\boxed{
T_{\mu\nu} = \left(\rho_{0}h+b^{2}\right)u_{\mu}u_{\nu} + \left(P + \frac{b^{2}}{2}\right)g_{\mu\nu} - b_{\mu}b_{\nu}
}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  // Compute T_{\mu \nu}
  if(update_Tmunu) {
    count=0; for(int ii=0;ii<4;ii++) for(int jj=ii;jj<4;jj++) { TDNMUNU[count] = rho0_h_plus_b2*uDN[ii]*uDN[jj] + P_plus_half_b2*g4dn[ii][jj] - smallb_lower[SMALLBT+ii]*smallb_lower[SMALLBT+jj]; count++; }
  }
```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

## Step 3.g: Recomputing the conservative variables \[Back to [top](#toc)\]
$$\label{recompute_conservs}$$

We now have all the ingredients to recompute the conservative variables

$$
\boxed{
\left\{
\begin{align}
\rho_{\star}  &= \alpha\sqrt{\gamma}\rho_{b}u^{0}\\
\tilde{S}_{i} &= \left(\rho_{\star}h + \alpha\sqrt{\gamma}u^{0}b^{2}\right)u_{i} - \alpha\sqrt{\gamma}b^{0}b_{i}\\
\tilde{\tau}  &= \alpha^{2}\sqrt{\gamma}T^{00} - \rho_{\star}
\end{align}
\right.
}\ .
$$

```python
%%writefile -a $outfile_path__tau_and_prim_limits__C

  //CCTK_VInfo(CCTK_THORNSTRING,"YAY %e",TDNMUNU[0]);

  // Finally, compute conservatives:
  CONSERVS[RHOSTAR] = alpha_sqrt_gamma * U[RHOB] * uUP[0];
  CONSERVS[STILDEX] = CONSERVS[RHOSTAR]*h_enthalpy*uDN[1] + alpha_sqrt_gamma*(uUP[0]*smallb[SMALLB2]*uDN[1] - smallb[SMALLBT]*smallb_lower[SMALLBX]);
  CONSERVS[STILDEY] = CONSERVS[RHOSTAR]*h_enthalpy*uDN[2] + alpha_sqrt_gamma*(uUP[0]*smallb[SMALLB2]*uDN[2] - smallb[SMALLBT]*smallb_lower[SMALLBY]);
  CONSERVS[STILDEZ] = CONSERVS[RHOSTAR]*h_enthalpy*uDN[3] + alpha_sqrt_gamma*(uUP[0]*smallb[SMALLB2]*uDN[3] - smallb[SMALLBT]*smallb_lower[SMALLBZ]);
  // tauL = alpha^2 sqrt(gamma) T^{00} - CONSERVS[RHOSTAR]
  CONSERVS[TAUENERGY] =  METRIC_LAP_PSI4[LAPSE]*alpha_sqrt_gamma*(rho0_h_plus_b2*SQR(uUP[0]) + P_plus_half_b2*(-SQR(METRIC_LAP_PSI4[LAPSEINV])) - SQR(smallb[SMALLBT])) - CONSERVS[RHOSTAR];
  // FIXME: We have TUPMUNU already, should replace           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ by TUPMUNU[0].
}

```

    Appending to ../src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C

# Step 4: Code validation \[Back to [top](#toc)\]
$$\label{code_validation}$$

First we download the original `IllinoisGRMHD` source code and then compare it to the source code generated by this tutorial notebook.

```python
# Verify if the code generated by this tutorial module
# matches the original IllinoisGRMHD source code

# First download the original IllinoisGRMHD source code
import urllib
from os import path

original_IGM_file_url  = "https://bitbucket.org/zach_etienne/wvuthorns/raw/5611b2f0b17135538c9d9d17c7da062abe0401b6/IllinoisGRMHD/src/apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C"
original_IGM_file_name = "apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs-original.C"
original_IGM_file_path = os.path.join(IGM_src_dir_path,original_IGM_file_name)

# Then download the original IllinoisGRMHD source code
# We try it here in a couple of ways in an attempt to keep
# the code more portable
try:
    original_IGM_file_code = urllib.request.urlopen(original_IGM_file_url).read().decode("utf-8")
    # Write down the file the original IllinoisGRMHD source code
    with open(original_IGM_file_path,"w") as file:
        file.write(original_IGM_file_code)
except:
    try:
        original_IGM_file_code = urllib.urlopen(original_IGM_file_url).read().decode("utf-8")
        # Write down the file the original IllinoisGRMHD source code
        with open(original_IGM_file_path,"w") as file:
            file.write(original_IGM_file_code)
    except:
        # If all else fails, hope wget does the job
        !wget -O $original_IGM_file_path $original_IGM_file_url

# Perform validation
Validation__tau_and_prims_limits__C  = !diff $original_IGM_file_path $outfile_path__tau_and_prim_limits__C

if Validation__tau_and_prims_limits__C == []:
    # If the validation passes, we do not need to store the original IGM source code file
    !rm $original_IGM_file_path
    print("Validation test for apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C: PASSED!")
else:
    # If the validation fails, we keep the original IGM source code file
    print("Validation test for apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C: FAILED!")
    # We also print out the difference between the code generated
    # in this tutorial module and the original IGM source code
    print("Diff:")
    for diff_line in Validation__tau_and_prims_limits__C:
        print(diff_line)
```

    Validation test for apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.C: FAILED!
    Diff:
    5,9c5
    <
    <   CCTK_REAL gamma = eos.gamma_tab[0];
    <   CCTK_REAL kpoly = eos.k_tab[0];
    <   int gamma_equals2 = 1;
    <   if (fabs(gamma-2.0) > 1.e-10) gamma_equals2 = 0;
    ---
    >
    34a31
    >
    43a41
    >
    53a52
    >
    78a78
    >
    81a82
    >
    86a88
    >
    94c96,97
    <
    ---
    >
    >
    98,105c101,119
    <   CCTK_REAL Patm;
    <   // FIXME: the operation below could be replaced by
    <   //           Patm = kpoly*fasterpow_ppm_reconstruct(rho_b_atm,gamma)
    <   if (gamma_equals2==1) {
    <     Patm = kpoly*rho_b_atm*rho_b_atm;
    <   } else {
    <     Patm = kpoly*pow(rho_b_atm,gamma);
    <   }
    ---
    >
    >   /**********************************
    >    * Piecewise Polytropic EOS Patch *
    >    *         Computing Patm         *
    >    **********************************/
    >   /* This modification of the code trades the variable
    >    * "gamma_equals2" by the already defined function
    >    * pow().
    >    *
    >    * Also, assuming that Patm < rho_ppoly_tab[0], we skip
    >    * the declaration of new variables to store the
    >    * values of K_ppoly_tab[0] and Gamma_ppoly_tab[0]. Thus:
    >    *  -----------------------------------------
    >    * | P_{atm} = K_{0} * rho_{atm}^{Gamma_{0}} |
    >    *  -----------------------------------------
    >    */
    >   int polytropic_index = find_polytropic_K_and_Gamma_index(eos,rho_b_atm);
    >   CCTK_REAL Patm = eos.K_ppoly_tab[polytropic_index]*pow(rho_b_atm,eos.Gamma_ppoly_tab[polytropic_index]);
    >
    115c129
    <       CONSERVS[STILDEZ]*=rfactm1;
    ---
    >       CONSERVS[STILDEZ]*=rfactm1;
    120c134
    <     if (tau_fluid_min < tau_atm*1.001) {
    ---
    >     if (tau_fluid_min < tau_atm*1.001) {
    123a138
    >
    136c151
    <
    ---
    >
    147a163
    >
    149a166
    > #ifndef ENABLE_STANDALONE_IGM_C2P_SOLVER
    150a168
    > #endif
    152c170
    <   CCTK_REAL METRIC_LAP_PSI4[NUMVARS_METRIC_AUX];
    ---
    >   CCTK_REAL METRIC_LAP_PSI4[NUMVARS_METRIC_AUX];
    154c172
    <
    ---
    >
    167,171c185,189
    <   CCTK_REAL h_enthalpy=0,  P_cold=0,eps_cold=0,dPcold_drho,eps_th,gamma_cold; /* <- Note that in setting h, we need to define several
    <                                                                                *    other variables. Though some will be unused later
    <                                                                                *    in this function, they may be useful in other
    <                                                                                *    functions */
    <   compute_P_cold__eps_cold__dPcold_drho__eps_th__h__gamma_cold(U,eos,P_cold,eps_cold,dPcold_drho,eps_th,h_enthalpy,gamma_cold);
    ---
    >   CCTK_REAL h_enthalpy,  P_cold,eps_cold,dPcold_drho,eps_th,Gamma_cold; /* <- Note that in setting h, we need to define several
    >                                                                          *    other variables. Though some will be unused later
    >                                                                          *    in this function, they may be useful in other
    >                                                                          *    functions */
    >   compute_P_cold__eps_cold__dPcold_drho__eps_th__h__Gamma_cold(U,eos,Gamma_th,P_cold,eps_cold,dPcold_drho,eps_th,h_enthalpy,Gamma_cold);
    174c192,193
    <   enforce_pressure_floor_ceiling(stats,eos.K_poly,P_cold,METRIC_LAP_PSI4[PSI6],Psi6threshold,U[RHOB],rho_b_atm,  U[PRESSURE]);
    ---
    >   int polytropic_index = find_polytropic_K_and_Gamma_index(eos,U[RHOB]);
    >   enforce_pressure_floor_ceiling(stats,eos.K_ppoly_tab[polytropic_index],P_cold,METRIC_LAP_PSI4[PSI6],Psi6threshold,U[RHOB],rho_b_atm,  U[PRESSURE]);
    177c196
    <   CCTK_REAL eps = eps_cold + (U[PRESSURE]-P_cold)/(gamma_th-1.0)/U[RHOB];
    ---
    >   CCTK_REAL eps = eps_cold + (U[PRESSURE]-P_cold)/(Gamma_th-1.0)/U[RHOB];
    187a207
    >
    199c219,220
    <
    ---
    >
    >
    201c222
    <   CCTK_REAL alpha_sqrt_gamma=METRIC_LAP_PSI4[LAPSE]*METRIC_LAP_PSI4[PSI6];
    ---
    >   CCTK_REAL alpha_sqrt_gamma=METRIC_LAP_PSI4[LAPSE]*METRIC_LAP_PSI4[PSI6];
    208c229
    <
    ---
    >
    222a244
    >
    228a251
    >
    235c258
    <
    ---
    >
    251a275
    >
    260a285
    >
    265a291
    >
    276a303
    >

# Step 5: 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-IllinoisGRMHD__apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.pdf](Tutorial-IllinoisGRMHD__apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means).

```python
latex_nrpy_style_path = os.path.join(nrpy_dir_path,"latex_nrpy_style.tplx")
#!jupyter nbconvert --to latex --template $latex_nrpy_style_path --log-level='WARN' Tutorial-IllinoisGRMHD__apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.ipynb
#!pdflatex -interaction=batchmode Tutorial-IllinoisGRMHD__apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.tex
#!pdflatex -interaction=batchmode Tutorial-IllinoisGRMHD__apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.tex
#!pdflatex -interaction=batchmode Tutorial-IllinoisGRMHD__apply_tau_floor__enforce_limits_on_primitives_and_recompute_conservs.tex
!rm -f Tut*.out Tut*.aux Tut*.log
```


#SHDOM single-scattering adjoint

This folder conatains the most recent (3 May 2015) single-scattering adjoint calculations from frank.

>This has the log output files (.log), radiance output files (.arad),
and optical property adjoint files (.adj) for the "true" and "guess"
cases. The first set of output is for the single scattering forward radiance calculations

    >>les0822nh15t13y135_ocaer1_w0.646_ns1true.*,

    >>les0822nh15t13y135_ocaer1_w0.646_ns1guess.*,

>and the second set of outputs is for the full multiple scattering forward radiance calculations,

    >>les0822nh15t13y135_ocaer1_w0.646_ns1true2.*,

    >>les0822nh15t13y135_ocaer1_w0.646_ns1guess2.*.

The cloud field used in these simulations is a 2D slice taken at $y_{(i_y=135)}=8.375 \mathrm{km}$ of the 320x320 high resolution simulation shown below.

The adjoint calculation is for the gradient of the misfit function,

\begin{align}
\frac{\partial \Phi(\sigma, \omega)}{\partial(\sigma)} &= \left < \Delta p, \mathcal{U}_{\sigma}[\Delta f] \right >
\end{align}

Only the first order of scattering completed so the results are approximate.  In the calculation labelled "1" the forward and adjoint calculations are only single-scattering.  In the calculation labelled "2", the Forward calculation includes multiple scatterng while the adjoint does not.

```python
%pylab inline

# Imports
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm

import pandas as pd

```

    Populating the interactive namespace from numpy and matplotlib

```python

```

```python
# Utility functions

def print_head(fname, Nlines=10, indent="\t"):
    "Print the head of the file."

    # Print a message and then the first N lines
    print("Showing head: {}".format(fname))
    for i, line in zip(range(Nlines), open(fname, 'r')):
        print(indent + line.strip())

    print('\n')

```

```python
ls
```

    adjoint-derivative-single-scattering.jpg
    adjoint-extinction-difference-and residual.jpg
    adjoint-extinction-difference-and-residual.jpg
    exploring_adjoint.ipynb
    exploring_adjoint_with_aerosol.ipynb
    les0822nh15t13y135_ocaer1_w0.646_ns1guess2.adj
    les0822nh15t13y135_ocaer1_w0.646_ns1guess2.arad
    les0822nh15t13y135_ocaer1_w0.646_ns1guess2.log
    les0822nh15t13y135_ocaer1_w0.646_ns1true2.adj
    les0822nh15t13y135_ocaer1_w0.646_ns1true2.arad
    les0822nh15t13y135_ocaer1_w0.646_ns1true2.log
    les0822nh15t13y135_w0.646_ns1guess2.adj
    les0822nh15t13y135_w0.646_ns1guess2.arad
    les0822nh15t13y135_w0.646_ns1guess2.log
    les0822nh15t13y135_w0.646_ns1true2.adj
    les0822nh15t13y135_w0.646_ns1true2.arad
    les0822nh15t13y135_w0.646_ns1true2.log

# Looking at the data files

```python
# The log file
flog = "les0822nh15t13y135_w0.646_ns1true2.log"
print_head(flog, Nlines=10)

# The radiance file
frad = "les0822nh15t13y135_w0.646_ns1true2.arad"
print_head(frad, Nlines=25)

# The adjoint file
fadj_true = "les0822nh15t13y135_w0.646_ns1true2.adj"
fadj = "les0822nh15t13y135_w0.646_ns1guess2.adj"

print_head(fadj, Nlines=35)

# The adjoint Radiance file
farad = "les0822nh15t13y135_w0.646_ns1guess2.arad"
farad_true = "les0822nh15t13y135_w0.646_ns1true2.arad"
print_head(farad, Nlines=35)
```

    Showing head: les0822nh15t13y135_w0.646_ns1true2.log

    	Name of this SHDOMADJ run:
    	les0822nh15t13y135_w0.646_ns1true2
    	Input particle properties filename (or NONE for Rayleigh only)
    	les0822nh15t13y135t.part
    	Number of input scattering tables
    	2
    	Name of each scattering table
    	water_w0.646.scat
    	aero_oc_w0.646.scat

    Showing head: les0822nh15t13y135_w0.646_ns1true2.arad
    	! Polarized Spherical Harmonic Discrete Ordinate Radiative Transfer Radiance Output
    	!  L= 15  M= 15  NLM=  256   NMU= 16  NPHI= 32  NANG=  354   NSH=   1767876
    	!  NSTOKES=1   NX= 320   NY=   1   NZ=  83    NPTS=   36044   NCELLS=   42046
    	!  PARTICLE_FILE=les0822nh15t13y135t.part
    	!  CORRELATED_K-DIST_FILE=NONE   NUM_G= 1
    	!  SOURCE_TYPE=SOLAR               DELTA-M METHOD
    	!  GRID_TYPE=EVEN (X,Y)    PROPERTY-FILE (Z)    INDEPENDENT_PIXEL=2
    	!  SURFACE_TYPE=FIXED LAMBERTIAN         HORIZ_BOUNDARY_COND=0
    	!  GROUND_ALBEDO=0.0500000  SKY_RAD= 0.00000E+00
    	!  SOLAR_FLUX= 0.100000E+01   SOLAR_MU=-0.9221000   SOLAR_AZ= 339.400
    	!  UNITS=WATTS/(M^2 MICRON STER)    WAVELENGTH=      0.65
    	!  SPLITTING_ACCURACY= 0.300E-01   SPHERICAL_HARMONIC_ACCURACY= 0.300E-02
    	!  SOLUTION_ACCURACY= 0.300E-04
    	!  MAXIMUM_ITERATIONS= 200   NUMBER_ITERATIONS= 123
    	!  RADIANCE AT Z= 15.000
    	!  320    1    1 ! NXOUT, NYOUT, NRADDIR
    	!  ID      mu       phi
    	!   1  1.0000000    0.00
    	!
    	ID  IY  IX    RADIANCE
    	1   1   1   0.205754E+00
    	1   1   2   0.196955E+00
    	1   1   3   0.171690E+00
    	1   1   4   0.196178E+00
    	1   1   5   0.211501E+00

    Showing head: les0822nh15t13y135_w0.646_ns1guess2.adj
    	! SHDOMADJ optical property adjoint output
    	!
    	X       Y       Z      Extinct  SSalbedo  dMF/dExt  dMF/dSSalb  dMF/dP(Theta_j)
    	0.0000  0.0000 -0.0087    0.00574 1.000000  0.000E+00  0.000E+00  0.000E+00
    	0.0000  0.0000  0.0095    0.00574 1.000000  0.000E+00  0.000E+00  0.000E+00
    	0.0000  0.0000  0.0315    0.00573 1.000000  0.000E+00  0.000E+00  0.000E+00
    	0.0000  0.0000  0.0574    0.00571 1.000000  0.000E+00  0.000E+00  0.000E+00
    	0.0000  0.0000  0.0878    0.00570 1.000000  0.000E+00  0.000E+00  0.000E+00
    	0.0000  0.0000  0.1245    0.00568 1.000000  1.893E-12  0.000E+00  0.000E+00
    	0.0000  0.0000  0.1650    0.00566 1.000000  9.660E-11  0.000E+00  0.000E+00
    	0.0000  0.0000  0.2050    0.00564 1.000000  3.237E-10  0.000E+00  0.000E+00
    	0.0000  0.0000  0.2450    0.00562 1.000000  5.613E-10  0.000E+00  0.000E+00
    	0.0000  0.0000  0.2850    0.00561 1.000000  7.948E-10  0.000E+00  0.000E+00
    	0.0000  0.0000  0.3250    0.00559 1.000000  8.362E-10  0.000E+00  0.000E+00
    	0.0000  0.0000  0.3650    0.00557 1.000000  6.187E-10  0.000E+00  0.000E+00
    	0.0000  0.0000  0.4050    0.00555 1.000000  3.825E-10  0.000E+00  0.000E+00
    	0.0000  0.0000  0.4450    0.00553 1.000000  1.485E-10  0.000E+00  0.000E+00
    	0.0000  0.0000  0.4850   11.03866 0.999998  7.969E-12 -6.166E-11  0.000E+00
    	0.0000  0.0000  0.5250   21.63028 0.999998  2.652E-12 -4.458E-11  0.000E+00
    	0.0000  0.0000  0.5650   30.82795 0.999997  3.485E-12 -8.711E-11  0.000E+00
    	0.0000  0.0000  0.6050   38.37085 0.999997 -3.328E-10 -7.485E-08 -4.664E-06
    	0.0000  0.0000  0.6450   40.02002 0.999997  3.989E-10 -3.205E-07 -1.682E-05
    	0.0000  0.0000  0.6850   25.36786 0.999997  3.926E-09 -3.913E-07 -1.654E-05
    	0.0000  0.0000  0.7250    0.00539 1.000000 -4.657E-08 -4.098E-10 -3.754E-09
    	0.0000  0.0000  0.7650    0.00537 1.000000 -5.320E-09 -2.011E-10 -1.842E-09
    	0.0000  0.0000  0.8050    0.00535 1.000000  1.510E-08 -7.357E-11 -6.740E-10
    	0.0000  0.0000  0.8450    0.00532 1.000000  2.142E-08 -2.975E-11 -2.726E-10
    	0.0000  0.0000  0.8850    0.00530 1.000000  2.389E-08 -1.395E-11 -1.278E-10
    	0.0000  0.0000  0.9250    0.00528 1.000000  2.497E-08 -7.086E-12 -6.492E-11
    	0.0000  0.0000  0.9650    6.33743 0.999998  1.209E-08 -4.835E-08 -3.338E-08
    	0.0000  0.0000  1.0050   21.38868 0.999998  7.411E-09 -1.245E-07 -7.732E-08
    	0.0000  0.0000  1.0450   36.03261 0.999997  1.077E-08 -3.283E-07 -4.714E-07
    	0.0000  0.0000  1.0850   25.57642 0.999997  1.458E-08 -3.145E-07 -9.886E-07
    	0.0000  0.0000  1.1250   14.30435 0.999998  1.149E-08 -1.339E-07 -6.860E-07
    	0.0000  0.0000  1.1650   30.72805 0.999997  7.772E-09 -2.360E-07 -2.281E-06

    Showing head: les0822nh15t13y135_w0.646_ns1guess2.arad
    	! Polarized Spherical Harmonic Discrete Ordinate Radiative Transfer Radiance Output
    	!  L= 15  M= 15  NLM=  256   NMU= 16  NPHI= 32  NANG=  354   NSH=   1684473
    	!  NSTOKES=1   NX= 320   NY=   1   NZ=  83    NPTS=   35125   NCELLS=   40454
    	!  PARTICLE_FILE=les0822nh15t13y135g.part
    	!  CORRELATED_K-DIST_FILE=NONE   NUM_G= 1
    	!  SOURCE_TYPE=SOLAR               DELTA-M METHOD
    	!  GRID_TYPE=EVEN (X,Y)    PROPERTY-FILE (Z)    INDEPENDENT_PIXEL=2
    	!  SURFACE_TYPE=FIXED LAMBERTIAN         HORIZ_BOUNDARY_COND=0
    	!  GROUND_ALBEDO=0.0500000  SKY_RAD= 0.00000E+00
    	!  SOLAR_FLUX= 0.100000E+01   SOLAR_MU=-0.9221000   SOLAR_AZ= 339.400
    	!  UNITS=WATTS/(M^2 MICRON STER)    WAVELENGTH=      0.65
    	!  SPLITTING_ACCURACY= 0.300E-01   SPHERICAL_HARMONIC_ACCURACY= 0.300E-02
    	!  SOLUTION_ACCURACY= 0.300E-04
    	!  MAXIMUM_ITERATIONS= 200   NUMBER_ITERATIONS= 123
    	!  RADIANCE AT Z= 15.000
    	!  320    1    1 ! NXOUT, NYOUT, NRADDIR
    	!  ID      mu       phi
    	!   1  1.0000000    0.00
    	!
    	ID  IY  IX    RADIANCE
    	1   1   1   0.197190E+00
    	1   1   2   0.188860E+00
    	1   1   3   0.165215E+00
    	1   1   4   0.187221E+00
    	1   1   5   0.201468E+00
    	1   1   6   0.204092E+00
    	1   1   7   0.193736E+00
    	1   1   8   0.170147E+00
    	1   1   9   0.133220E+00
    	1   1  10   0.107793E+00
    	1   1  11   0.666059E-01
    	1   1  12   0.487470E-01
    	1   1  13   0.455724E-01
    	1   1  14   0.446769E-01
    	1   1  15   0.981936E-01

```python
# Load the adjoint files into memory
adj_frame = pd.read_csv(fadj, quoting=3, delim_whitespace=True, skiprows=2)
adj_true_frame = pd.read_csv(fadj_true, quoting=3, delim_whitespace=True, skiprows=2)
adj_arad_frame = pd.read_csv(farad, quoting=3, delim_whitespace=True, skiprows=19)
adj_arad_true_frame = pd.read_csv(farad_true, quoting=3, delim_whitespace=True, skiprows=19)

# Get variables from the array
nx = adj_frame['X'].unique().size
nz = adj_frame["Z"].unique().size

# Get the guess case adjoint field properties
adj_x = np.array(adj_frame["X"], dtype='f8').reshape(nx, nz)
adj_z = np.array(adj_frame["Z"], dtype='f8').reshape(nx, nz)
adj_ext = np.array(adj_frame["Extinct"], dtype='f8').reshape(nx, nz)
adj_alb = np.array(adj_frame["SSalbedo"], dtype='f8').reshape(nx, nz)
adj_dMFdext = np.array(adj_frame["dMF/dExt"], dtype='f8').reshape(nx, nz)
adj_dMFdalb = np.array(adj_frame["dMF/dSSalb"], dtype='f8').reshape(nx, nz)

# Get the True case adjoint extinction field
adj_ext_true = np.array(adj_true_frame["Extinct"], dtype='f8').reshape(nx, nz)

# Define the adjoint radiance parameters
adj_arad = np.array(adj_arad_frame['RADIANCE'], dtype='f8').reshape(nx)
adj_arad_true = np.array(adj_arad_true_frame['RADIANCE'], dtype='f8').reshape(nx)

```

```python
# Check the sign of dMF/dExt for Z near the top of the atmosphere
adj_frame[abs(adj_frame["Z"]-15) < .2]
```

      X
      Y
      Z
      Extinct
      SSalbedo
      dMF/dExt
      dMF/dSSalb
      dMF/dP(Theta_j)

      82
        0.0000
       0
       15
       0.00107
       1
      -5.7630
      -0.007303
      -0.066900

      165
        0.0625
       0
       15
       0.00107
       1
      -5.4040
      -0.006903
      -0.063240

      248
        0.1250
       0
       15
       0.00107
       1
      -4.5590
      -0.005521
      -0.050580

      331
        0.1875
       0
       15
       0.00107
       1
      -6.1260
      -0.007638
      -0.069970

      414
        0.2500
       0
       15
       0.00107
       1
      -6.7710
      -0.008555
      -0.078370

      497
        0.3125
       0
       15
       0.00107
       1
      -6.5720
      -0.008433
      -0.077250

      580
        0.3750
       0
       15
       0.00107
       1
      -5.7130
      -0.007453
      -0.068280

      663
        0.4375
       0
       15
       0.00107
       1
      -4.9960
      -0.006468
      -0.059250

      746
        0.5000
       0
       15
       0.00107
       1
      -3.4130
      -0.004334
      -0.039710

      829
        0.5625
       0
       15
       0.00107
       1
      -2.9250
      -0.003888
      -0.035620

      912
        0.6250
       0
       15
       0.00107
       1
      -2.7620
      -0.003659
      -0.033520

      995
        0.6875
       0
       15
       0.00107
       1
      -2.0560
      -0.002816
      -0.025800

      1078
        0.7500
       0
       15
       0.00107
       1
      -1.9100
      -0.002647
      -0.024250

      1161
        0.8125
       0
       15
       0.00107
       1
      -1.7820
      -0.002525
      -0.023140

      1244
        0.8750
       0
       15
       0.00107
       1
      -5.3360
      -0.006965
      -0.063810

      1327
        0.9375
       0
       15
       0.00107
       1
      -6.7150
      -0.008715
      -0.079840

      1410
        1.0000
       0
       15
       0.00107
       1
      -7.2750
      -0.009091
      -0.083280

      1493
        1.0625
       0
       15
       0.00107
       1
      -7.3250
      -0.009104
      -0.083400

      1576
        1.1250
       0
       15
       0.00107
       1
      -6.3700
      -0.008180
      -0.074930

      1659
        1.1875
       0
       15
       0.00107
       1
      -5.0780
      -0.006464
      -0.059220

      1742
        1.2500
       0
       15
       0.00107
       1
      -4.0340
      -0.005292
      -0.048480

      1825
        1.3125
       0
       15
       0.00107
       1
      -4.3370
      -0.005791
      -0.053050

      1908
        1.3750
       0
       15
       0.00107
       1
      -5.9460
      -0.007700
      -0.070540

      1991
        1.4375
       0
       15
       0.00107
       1
      -6.5550
      -0.008448
      -0.077390

      2074
        1.5000
       0
       15
       0.00107
       1
      -4.7300
      -0.006188
      -0.056690

      2157
        1.5625
       0
       15
       0.00107
       1
      -4.4940
      -0.005667
      -0.051910

      2240
        1.6250
       0
       15
       0.00107
       1
      -5.1200
      -0.006369
      -0.058340

      2323
        1.6875
       0
       15
       0.00107
       1
      -5.4600
      -0.006815
      -0.062430

      2406
        1.7500
       0
       15
       0.00107
       1
      -5.8380
      -0.007253
      -0.066440

      2489
        1.8125
       0
       15
       0.00107
       1
      -4.9770
      -0.006240
      -0.057160

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      24152
       18.1250
       0
       15
       0.00107
       1
      -2.2390
      -0.003262
      -0.029880

      24235
       18.1875
       0
       15
       0.00107
       1
      -1.7660
      -0.002518
      -0.023070

      24318
       18.2500
       0
       15
       0.00107
       1
      -2.1470
      -0.003023
      -0.027690

      24401
       18.3125
       0
       15
       0.00107
       1
      -2.9900
      -0.003977
      -0.036430

      24484
       18.3750
       0
       15
       0.00107
       1
      -3.0620
      -0.003977
      -0.036430

      24567
       18.4375
       0
       15
       0.00107
       1
      -5.6620
      -0.007265
      -0.066550

      24650
       18.5000
       0
       15
       0.00107
       1
      -7.4070
      -0.009339
      -0.085550

      24733
       18.5625
       0
       15
       0.00107
       1
      -5.8450
      -0.007641
      -0.070000

      24816
       18.6250
       0
       15
       0.00107
       1
      -5.0250
      -0.006580
      -0.060280

      24899
       18.6875
       0
       15
       0.00107
       1
      -4.6920
      -0.006206
      -0.056860

      24982
       18.7500
       0
       15
       0.00107
       1
      -4.6610
      -0.006121
      -0.056070

      25065
       18.8125
       0
       15
       0.00107
       1
      -3.8760
      -0.004881
      -0.044720

      25148
       18.8750
       0
       15
       0.00107
       1
      -1.4780
      -0.002024
      -0.018540

      25231
       18.9375
       0
       15
       0.00107
       1
      -1.9090
      -0.002371
      -0.021720

      25314
       19.0000
       0
       15
       0.00107
       1
      -3.4870
      -0.004342
      -0.039770

      25397
       19.0625
       0
       15
       0.00107
       1
      -4.0000
      -0.005040
      -0.046170

      25480
       19.1250
       0
       15
       0.00107
       1
      -1.5030
      -0.001863
      -0.017060

      25563
       19.1875
       0
       15
       0.00107
       1
      -0.2800
      -0.000526
      -0.004820

      25646
       19.2500
       0
       15
       0.00107
       1
       0.3028
       0.000045
       0.000415

      25729
       19.3125
       0
       15
       0.00107
       1
      -0.3874
      -0.000854
      -0.007822

      25812
       19.3750
       0
       15
       0.00107
       1
      -1.8610
      -0.002716
      -0.024880

      25895
       19.4375
       0
       15
       0.00107
       1
      -4.9990
      -0.006609
      -0.060540

      25978
       19.5000
       0
       15
       0.00107
       1
      -6.3950
      -0.008303
      -0.076060

      26061
       19.5625
       0
       15
       0.00107
       1
      -7.0560
      -0.009055
      -0.082950

      26144
       19.6250
       0
       15
       0.00107
       1
      -6.6320
      -0.008443
      -0.077340

      26227
       19.6875
       0
       15
       0.00107
       1
      -6.8060
      -0.008879
      -0.081340

      26310
       19.7500
       0
       15
       0.00107
       1
      -6.8730
      -0.008975
      -0.082220

      26393
       19.8125
       0
       15
       0.00107
       1
      -5.3910
      -0.007234
      -0.066270

      26476
       19.8750
       0
       15
       0.00107
       1
      -5.8250
      -0.007588
      -0.069510

      26559
       19.9375
       0
       15
       0.00107
       1
      -5.8220
      -0.007439
      -0.068150

320 rows × 8 columns

```python

```

```python
# Plot uniformity parameters
H = 20 #5.5   # top of the atmosphere plots
mask_clouds = (adj_ext>= .75e-5) #* (adj_z <= 2.83)
cloud_outline = (adj_ext>= .75e-1) #* (adj_z <= 2.83)

# Scaling of the verticle coordinate

Zscale = adj_z[0,:]
DZscale = 0 * Zscale
DZscale_other = 0 * Zscale
DZscale[:-1] += 1 / (Zscale[1:] - Zscale[:-1])
DZscale_other[1:] += 1 / (Zscale[1:] - Zscale[:-1])

alpha = .99725
DZscale = alpha * DZscale + (1-alpha) * DZscale_other

# Compute the residual
adj_residual = adj_arad_true - adj_arad

```

```python
# Make a figure to plot the true extinction field
f0 = plt.figure(0, (15,8), facecolor='white')
ax0 = f0.add_axes([.1, .1, .8, .8])
contour0 = ax0.contourf(adj_x, adj_z, adj_ext_true, 100, cmap=cm.gray)
ax0.set_ybound((0,8))
ax0.set_title(" Guess for the extinction field at y=8.35")
f0.colorbar(contour0, ax=ax0)
```

```python
# Make a figure to plot the true extinction field
f0 = plt.figure(0, (15,8), facecolor='white')
ax0 = f0.add_axes([.1, .1, .8, .8])

fvals = adj_ext_true - adj_ext

contour0 = ax0.contourf(adj_x, adj_z, fvals, 100, cmap=cm.gray)
ax0.set_ybound((0,8))
ax0.set_title("Difference $\sigma_{\mathrm{true}} - \sigma_{\mathrm{guess}}$")
f0.colorbar(contour0, ax=ax0)
```

```python

```

```python
# Make a figure to plot the true single-scattering albedo field
f1 = plt.figure(1, (15,8), facecolor='white')
ax1 = f1.add_axes([.1, .1, .8, .8])

levels = np.linspace(.999995, 1, 100)
contour1 = ax1.contourf(adj_x, adj_z, adj_alb, levels=levels, cmap=cm.gray)
ax1.set_ybound((0,13))
ax1.set_title("SSAlbedo field at y=8.35")
f1.colorbar(contour1, ax=ax1)

```

```python
# Make a figure to plot the true single-scattering albedo field
f2 = plt.figure(1, (10,6), facecolor='white')
ax2 = f2.add_axes([.1, .1, .8, .8])

fvals = adj_dMFdext * mask_clouds * DZscale * adj_ext# * .1 * adj_ext
fmax = abs(fvals).max() / 20.
levels = np.linspace(-fmax, fmax, 400)

contour2 = ax2.contourf(adj_x, adj_z, fvals, levels=levels, cmap=cm.RdBu_r, extend="both") #levels=linspace(-16, 16, 100),
contour2a = ax2.contour(adj_x, adj_z, cloud_outline, 1, cmap=cm.gray) #levels=linspace(-16, 16, 100),
#contour2b = ax2.contourf(adj_x, adj_z, cloud_outline * fvals, 10 , cmap=cm.RdBu_r, linewidth=3) #levels=linspace(-16, 16, 100),

# Add the residual to the plot
ftrans = lambda y: 200 * y + .8*H
plot_overlay_origin = ftrans(0 * adj_residual)
plot_overlay_arad = ftrans(adj_residual)
plot_res_pos = ax2.fill_between(adj_x[:,0], plot_overlay_origin, plot_overlay_arad,
                             where=(plot_overlay_origin<plot_overlay_arad),
                             color='b')#, linewidth=3)
plot_res_neg = ax2.fill_between(adj_x[:,0], plot_overlay_origin, plot_overlay_arad,
                             where=(plot_overlay_origin>plot_overlay_arad),
                             color='r',)#, linewidth=3)
plot_res_line = ax2.plot(adj_x[:,0], plot_overlay_arad, 'k', linewidth=2)
plot_origin = ax2.plot(adj_x[:,0], plot_overlay_origin, 'k,', linewidth=1)
ax2.set_ybound((0,H))
ax2.set_ylabel("Z Height [km]", fontsize='x-large', fontweight='bold')
ax2.set_xlabel('X [km]', fontsize='large', fontweight='bold')
ax2.set_title("""
    Risidual ($y_{\mathrm{data}} - y_{\mathrm{model}}(\mathbf{a})$) and derivative (-$\mathrm{d}\Phi(\sigma) / \mathrm{d}\log\sigma$)
    """, fontsize='xx-large', fontweight='bold')
cb = f2.colorbar(contour2, ax=ax2)
cb.set_ticks([-fmax, 0, fmax])

f2.savefig("adjoint-derivative-single-scattering.jpg", dpi=300)

#cb.set_ticklabels(['-2', '-1', '0', '1', '2'])
```

```python
# Make a figure to plot the true single-scattering albedo field
f2 = plt.figure(1, (10,6), facecolor='white')
ax2 = f2.add_axes([.1, .1, .8, .8])

fvals = -adj_dMFdext * mask_clouds * DZscale * adj_ext# * .1 * adj_ext

fvals = - (adj_ext_true - adj_ext) #* adj_ext

fmax = abs(fvals).max() / 1.
levels = np.linspace(-fmax, fmax, 200)

contour2 = ax2.contourf(adj_x, adj_z, fvals, levels=levels, cmap=cm.RdBu_r, extend="both") #levels=linspace(-16, 16, 100),
contour2a = ax2.contour(adj_x, adj_z, cloud_outline, 1, cmap=cm.gray) #levels=linspace(-16, 16, 100),
#contour2b = ax2.contourf(adj_x, adj_z, cloud_outline * fvals, 10 , cmap=cm.RdBu_r, linewidth=3) #levels=linspace(-16, 16, 100),

# Add the residual to the plot
ftrans = lambda y: 200 * y + .8*H
plot_overlay_origin = ftrans(0 * adj_residual)
plot_overlay_arad = ftrans(adj_residual)
plot_res_pos = ax2.fill_between(adj_x[:,0], plot_overlay_origin, plot_overlay_arad,
                             where=(plot_overlay_origin<plot_overlay_arad),
                             color='b')#, linewidth=3)
plot_res_neg = ax2.fill_between(adj_x[:,0], plot_overlay_origin, plot_overlay_arad,
                             where=(plot_overlay_origin>plot_overlay_arad),
                             color='r',)#, linewidth=3)
plot_res_line = ax2.plot(adj_x[:,0], plot_overlay_arad, 'k', linewidth=2)
plot_origin = ax2.plot(adj_x[:,0], plot_overlay_origin, 'k,', linewidth=1)
ax2.set_ybound((0,H))
ax2.set_ylabel("Z Height [km]", fontsize='x-large', fontweight='bold')
ax2.set_xlabel('X [km]', fontsize='large', fontweight='bold')
ax2.set_title("""
    Risidual ($y_{\mathrm{data}} - y_{\mathrm{model}}(\mathbf{a})$) and extinction error ($\sigma_{\mathrm{true}}-\sigma_{\mathrm{guess}})$)
    """, fontsize='xx-large', fontweight='bold')
cb = f2.colorbar(contour2, ax=ax2)
cb.set_ticks([-fmax, 0, fmax])

f2.savefig("adjoint-extinction-difference-and-residual.jpg", dpi=300)

#cb.set_ticklabels(['-2', '-1', '0', '1', '2'])
```

```python

plot(Zscale, DZscale)
```


# Continuous Control

The environment for this project is [Reacher](https://github.com/udacity/deep-reinforcement-learning/tree/master/p2_continuous_control) from Unity, and it's provided in the `setup` folder. We'll implement the A2C algorithm as the synchronous version of [A3C](https://web.stanford.edu/class/psych209/Readings/MnihEtAlHassibis15NatureControlDeepRL.pdf) (but the input won't be directly from pixels). We'll also include a variation called Generalized Advantage Estimation, or [GAE](https://arxiv.org/pdf/1506.02438.pdf) for short. Results will be shown in this notebook and the best solution will be implemented in `main.py`.

> The model used to generate this gif is `final.pth` (A2C-GAE with `n = 8` steps), which was trained for 600 episodes using `main.py`.

## 1. Prepare dependencies and environment

Take a look at README.md before executing this notebook and make sure that the kernel is set to **p2_continuous_control**.

```python
!pip -q install ./setup

import sys
import platform

import numpy as np
from scipy import signal
import matplotlib.pyplot as plt

import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions.normal import Normal
from torch import nn

from setup import unityagents
from unityagents import UnityEnvironment

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
```

Unity environments contain **brains**, our interfaces for controlling agents. We'll be conrtolling the first (default) brain in the environment. It's also useful to keep information such as `state_size`, `action_size` and `num_agents`.

```python
env = None
system = platform.system()
if system == 'Linux':
    env = UnityEnvironment(file_name="setup/Reacher_Linux/Reacher.x86_64")
elif system == 'Darwin':
    env = UnityEnvironment(file_name="setup/Reacher.app")
elif system == 'Windows':
    env = UnityEnvironment(file_name="setup/Reacher_Windows_x86_64/Reacher.exe")
else:
    print('Cannot find environment for this system.')

# use the default brain
brain_name = env.brain_names[0]
brain = env.brains[brain_name]

# reset the environment
env_info = env.reset(train_mode=True)[brain_name]
num_agents = len(env_info.agents)
action_size = brain.vector_action_space_size
state_size = env_info.vector_observations.shape[1]
```

    INFO:unityagents:
    'Academy' started successfully!
    Unity Academy name: Academy
            Number of Brains: 1
            Number of External Brains : 1
            Lesson number : 0
            Reset Parameters :
    		goal_size -> 5.0
    		goal_speed -> 1.0
    Unity brain name: ReacherBrain
            Number of Visual Observations (per agent): 0
            Vector Observation space type: continuous
            Vector Observation space size (per agent): 33
            Number of stacked Vector Observation: 1
            Vector Action space type: continuous
            Vector Action space size (per agent): 4
            Vector Action descriptions: , , ,

## 2. Actor

A2C is an **actor-critic method** and as such requires two models, one to represent the policy (the actor) and one to represent the value function (the critic). These models could be represented with just one network that has two distinct final layers – this is particularly useful when there are convolutional layers at the beginning of the network, as it helps extrapolating important information from the environment – but we don't need to do that, since the input consists of only 33 values per agent. We'll instead implement two separate models.

The `Actor` class is straightforward and consists of only fully connected layers:

- `state_size` $\to 64$ followed by `relu` activations
- $64 \to 128$ and `relu`
- $128 \to 64$ and `relu`
- $64 \to \mu$, that is $64 \to$ `action_size`.

Another important aspect of the actor is that it should output a normal distribution, so $\mu$ is not enough. To output $\sigma$ I decided to create a dedicated `nn.Parameter` of size `action_size` that does not depende on the input, although can still be updated by autograd.

In this implementation, **the actor is always used in a batch**. There are 20 independent agents that share the same network, so the input is made of 20 vectors for a total size of `[num_agents, state_size]` and, similarly, the output contains 20 $\mu$ vectors that can be seen as a matrix of size `[num_agents, action_size]`. Note that $\sigma$ does not change with the number of agents, since it will be applied equally to all `torch.distributions` objects.

```python
class Actor(nn.Module):

    def __init__(self, state_size, action_size, hidden_layers=[64, 128, 64]):
        super(Actor, self).__init__()
        self.action_size = action_size

        # prepare the first hidden layer
        self.hidden_layers = nn.ModuleList([nn.Linear(state_size, hidden_layers[0])])

        # prepare the rest of the hidden layers
        A = hidden_layers[:-1]
        B = hidden_layers[1:]
        self.hidden_layers.extend([nn.Linear(a, b) for a, b in zip(A, B)])

        # the actor will output the parameters of a normal distribution,
        self.mu_layer = nn.Linear(hidden_layers[-1], action_size)
        self.sigma = nn.Parameter(torch.ones(1, action_size))

    def forward(self, state):
        for layer in self.hidden_layers:
            state = layer(state)
            state = F.relu(state)
        mu = self.mu_layer(state)
        return mu, self.sigma
```

## 3. Critic

The `Critic` class is also made of only fully connected layers:

- `state_size` $\to 64$ followed by `relu` activations
- $64 \to 128$ and `relu`
- $128 \to 64$ and `relu`
- $64 \to V(s)$, that is $64 \to 1$ and no activation.

`Critic` is very similar to the `DQN` class, except it does not process the value for state-action pairs $Q(s, a)$, but only the value of the state $V(s)$. The reason for this is that we don't care about the value of an individual state-action pair because we're going to use the critic to calculate **advantages**, and based on that we'll tweak the probability distribution that the actor generates.

As for the actor, **the critic is always used in batches** of 20 independent states, for a total size of `[num_agents, state_size]` and. Similarly, the output contains 20 $V(s)$ values that can be seen as a matrix of size `[num_agents, 1]`.

```python
class Critic(nn.Module):

    def __init__(self, state_size, hidden_layers=[64, 128, 64]):
        super(Critic, self).__init__()

        # prepare the first hidden layer
        self.hidden_layers = nn.ModuleList([nn.Linear(state_size, hidden_layers[0])])

        # prepare the rest of the hidden layers
        A = hidden_layers[:-1]
        B = hidden_layers[1:]
        self.hidden_layers.extend([nn.Linear(a, b) for a, b in zip(A, B)])

        # the critic outputs only a scalar V(s)
        self.value_layer = nn.Linear(hidden_layers[-1], 1)

    def forward(self, state):
        # connect layers to each other and put relu activations between them
        for layer in self.hidden_layers:
            state = layer(state)
            state = F.relu(state)
        value = self.value_layer(state)
        return value
```

## 4. Advantage estimation

As mentioned, the critic is used to calculate advantages. An advantage value gives us a way to evaluate an action. For example, given a tuple $(s, a, r, s')$, as an indication for the goodness of the action $a$, we can calculate the advantage $A = r + V(s') - V(s)$, which is hopefully positive, but could also be negative. Since the output of the actor is a probability distribution, we can use **gradient ascent to increase** the log probability $\nabla_\theta \ln \pi_\theta(a | s)$ of that action being taken, and if we change the sign we will **decrease** it instead.

Combining the two ideas in a single expression, we get $\nabla_\theta \ln \pi_\theta(a | s) \cdot A$. This is how we'll decide whether we'll increase or decrease the probability of the action $a$, and by how much. Of course, in order to do that, our estimate of $V$, and consequently of $A$ must be accurate. The following sections describe two ways to obtain these values, based on the fact that A2C works using $n$-step TD estimates.

### 4.1 $n$-step

The overline brackets highlights how $V$ is updated and how the advantage is estimated. In fact, when updating $V(s_{i+n-1})$ we consider the standard one-step TD error that only includes $R_{i+n}$ and $V(s_{i+n})$. However, in order to improve our next estimates, we now substitute $V(s_{i+n-1})$ with $R_{i+n} + \gamma V(s_{i+n})$, so that the update to $V(s_{i+n-2})$ will be a two-step TD error. I decided to include $G(s_{i+n+1})$ in the image to highlight that there is only one bootstrapping step in this process. The same reasoning is used to estimate the advantages. For example, if $n=3$ we have
$$
\begin{align}
\begin{split}
V(s_{i+2}) \leftarrow & R_{i+3} + \gamma V(s_{i+3}) \\
V(s_{i+1}) \leftarrow & R_{i+2} + \gamma R_{i+3} + \gamma^2 V(s_{i+3}) \\
V(s_{i}) \leftarrow & R_{i+1} + \gamma R_{i+2} + \gamma^2 R_{i+3} + \gamma^3 V(s_{i+3})
\end{split}
\qquad\qquad
\begin{split}
A(s_{i+3}) = & R_{i+3} + \gamma V(s_{i+3}) - V(s_{i+2}) \\
A(s_{i+2}) = & R_{i+2} + \gamma R_{i+3} + \gamma^2 V(s_{i+3}) - V(s_{i+1}) \\
A(s_{i+1}) = & R_{i+1} + \gamma R_{i+2} + \gamma^2 R_{i+3} + \gamma^3 V(s_{i+3}) - V(s_{i})
\end{split}
\end{align}
$$
When we take into account that the critic is updated according to the mean squared error, both equations become virtually the same thing, as they are both implemented using discounted returns.

The next two functions, `n_step` and `gae`, take more parameters than they need. Technically, both of them need `rewards, dones, future, critic_out, gamma`, with `gae` also requiring another hyperparameter `tau` that will be explained in the next markdown cell. However, `n_step` and `gae` only affect how advantages are calculated, but we're always going to compute `returns` regardless of this choice, since we need it for the critic. For this reason, I decided to also include `returns` as a parameter, so that we only need to calculate it once. The implementation of `n_step` then becomes trivial.

```python
def n_step(rewards, dones, future, returns, critic_out, gamma, tau):
    return returns - critic_out
```

### 4.2 GAE

As already mentioned, the way $V$ is updated is identical in both `n_step` and `gae`, so I left the overline brackets in the image to highlight this fact. On the other hand, the green and orange blocks indicate how the $\hat{A}$ values are computed. These values are not the actual advantages $A$ that `gae` estimates, but they help us get there – except for the last block, where $A$ is the same as $\hat{A}$. For instance, if $n = 3$ we have
$$
\begin{align}
\begin{split}
\hat{A}(s_{i+3}) &= R_{i+3} + \gamma V(s_{i+3}) - V(s_{i+2})\\
\hat{A}(s_{i+2}) &= R_{i+2} + \gamma V(s_{i+2}) - V(s_{i+1})\\
\hat{A}(s_{i+1}) &= R_{i+1} + \gamma V(s_{i+1}) - V(s_{i})
\end{split}
\qquad\qquad
\begin{split}
A(s_{i+3}) &= \hat{A}(s_{i+3})\\
A(s_{i+2}) &= \hat{A}(s_{i+2}) + \gamma \lambda A(s_{i+3})\\
A(s_{i+1}) &= \hat{A}(s_{i+1}) + \gamma \lambda A(s_{i+2})\\
\end{split}
\end{align}\\
$$
Essentially, $\hat{A}$ represents a one-step advantage (as indicated by the contents of the green and orange blocks in the image), while $A$ is a mixture of steps (or blocks in the image), controlled by $\lambda$. Indeed, if we expand $A(s_{i+2})$ we can see that it contains a little bit of both blocks
$$
A(s_{i+2}) = R_{i+2} + \gamma V(s_{i+2}) - V(s_{i+1}) + \gamma \lambda (R_{i+3} + \gamma V(s_{i+3}) - V(s_{i+2}))\\= R_{i+2} + \gamma\lambda R_{i+3} + (1 - \lambda)\gamma V(s_{i+2}) + \gamma^2\lambda V(s_{i+3}) - V(s_{i+1})
$$
Additionally, if $\lambda = 0$ the expression simplifies to a one-step TD approach, whereas with $\lambda = 1$ it becomes like a MC approach
$$
\begin{align}
\lambda = 0 \implies \quad & R_{i+2} + \gamma V(s_{i+2}) - V(s_{i+1})\\
\lambda = 1 \implies \quad & R_{i+2} + \gamma R_{i+3} + \gamma^2 V(s_{i+3}) - V(s_{i+1})
\end{align}
$$
More generally, this is called the $\lambda$-return and is often indicated with TD($\lambda$) – although we'll refer to its hyperparameter as `tau`, because `lambda` is a reserved keyword in python.

```python
def gae(rewards, dones, future, returns, critic_out, gamma, tau):

    # create a single tensor for rewards, dones and predicted_values.
    # also, dones and predicted_values have n+1 rows: predictied_values adds G or future at the end
    # whereas dones adds a row of zeros at the beginning; this simplifies access to this data in the loop
    rewards = torch.cat([r.unsqueeze(0) for r in rewards], dim=0)
    dones = torch.cat([torch.zeros([1, 20, 1]).to(device)] + [d.unsqueeze(0) for d in dones], dim=0)
    predictied_values = torch.cat((critic_out, future.unsqueeze(0)), dim=0)

    # retrieve n from a tensor's size, instead of passing it as a parameter
    n = list(critic_out.size())[0]

    # compute A_hats and As backwards, and store the As as you go
    A = torch.zeros([20, 1]).to(device)
    advantages = []
    for i in reversed(range(n)):
        A_hat = rewards[i] + gamma * predictied_values[i+1] * (1 - dones[i+1]) - predictied_values[i] * (1 - dones[i])
        A = A_hat + gamma * tau * A
        advantages.insert(0, A.unsqueeze(0))
    return torch.cat(advantages, dim=0)
```

## 5. Agent

As stated, the learning algorithm guides the actor with gradient ascent on $\nabla_\theta \ln \pi_\theta(a | s) \cdot A$ and the critic with gradient descent on the mean squared error of returns. To implement this, we can use a **temporary buffer** somewhat similar to the one found in DQN or DDPG, except it is reset after every learning step and there is no sampling involved, so the size will always be limited by `n`. Thus, we can avoid building a dedicated class for it and instead make a few python arrays. The ones that we need are `tmp_r` and `tmp_d`, respectively for storing rewards and the done flag. We do not need to store the actual states and actions, as long as we store their values and log probabilities, respectively. To do this, we use `tmp_critic_out` and `tmp_log_prob`. Finally, we'll need `tmp_ns` to store the `n`-th next state, that we'll use for bootstrapping.

There are very few hyperparameters compared to other methods. What we need is:
- `n` determines how many steps to take for one learning step; `1` is similar to other standard TD methods, while higher values make it more like a MC method
- `calc_advantages` can be either `n_step` or `gae`
- `alpha` is the learning rate of both networks
- `gamma` is the discount factor for the returns
- `tau` in `[0, 1]` will be used only if `calc_advantages` is `gae` and determines how to mix the advantages; `0` makes it a standard TD method, while `1` makes it similar to a MC method, except it is still limited by `n`

The agent has three main functions, but only two of them should be called from the outside:
- `act`, based on a state tensor `state`, choose an action. This also evaluates the state and the log probability of the action, as well as stores some information.
- `store`, receive a complete tuple and adds some data to the temporary buffer. This counts as a step towards the updates of `actor` and `critic`, and in fact calls `learn()` automatically when needed.
- `learn`, takes all the data in the buffer and calculates discounted returns, advantages using `calc_advantages`, and performs the update of the networks.

Additionally, an agent can be `reset()`.

```python
class A2CAgent:

    # -- initialization -- #

    def __init__(self, state_size, action_size, calc_advantages, n=4,
                 alpha=0.0001, gamma=0.95, tau=0.5):

        self.state_size, self.action_size = state_size, action_size
        self.alpha, self.gamma, self.tau = alpha, gamma, tau
        self.calc_advantages = calc_advantages
        self.n = n
        self.reset()

    def reset_temporary_buffer(self):
        # some values will be stored in temporary buffers to
        # avoid re-evaluating them and to speed up the execution
        self.tmp_r, self.tmp_ns, self.tmp_d = [], None, []
        self.tmp_log_prob, self.tmp_critic_out = [], []

    def reset(self):
        self.actor = Actor(self.state_size, self.action_size).to(device)
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=self.alpha)

        self.critic = Critic(self.state_size).to(device)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=self.alpha)

        self.reset_temporary_buffer()
        self.i = 0

    # -- initialization -- #

    def act(self, state):
        state = torch.FloatTensor(state).to(device)

        # evaluate the current state and choose
        #  - its value
        value = self.critic(state)

        #  - an action from its probability distribution
        mu, sigma = self.actor(state)
        dist = torch.distributions.Normal(mu, sigma)
        action = dist.sample().detach()

        # then also evaluate the log probability of choosing that action.
        # this is what we'll use for backpropagation, together with the estimated advantage
        log_prob = dist.log_prob(action)

        # store these values in a temporary buffer so that we don't need to calculate them again.
        # this is like a replay buffer in DQN, but it gets wiped after an update step.
        self.tmp_log_prob.append(log_prob.unsqueeze(0))
        self.tmp_critic_out.append(value.unsqueeze(0))

        # note: clipping the action to [-1, 1] is required by the environment, but in my tests,
        # clipping it before calculating its log probability resulted in poor performance
        return np.clip(np.array(action), -1, 1)

    def store(self, s, a, r, ns, d):
        # ignore s and a because they're already handled in act(state)
        self.tmp_r.append(torch.FloatTensor(r).unsqueeze(1).to(device))
        self.tmp_ns = torch.FloatTensor(ns).to(device)
        self.tmp_d.append(torch.FloatTensor(d).unsqueeze(1).to(device))

        # every n steps we should perform an update step
        if self.i == self.n - 1:
            self.learn()
        self.i = (self.i + 1) % self.n

    def learn(self):
        # note that this is called automatically by the agent

        # create a single tensor for the values stored in act(), namely tmp_log_prob and tmp_critic_out
        log_prob = torch.cat(self.tmp_log_prob, dim=0)
        critic_out = torch.cat(self.tmp_critic_out, dim=0)

        # calculate discounted rewards backwards, starting from G or future (the bootstrapping step)
        future = self.critic(self.tmp_ns)
        tmp_returns = []
        ret = future
        for reward, done in zip(reversed(self.tmp_r), reversed(self.tmp_d)):
            ret = reward + self.gamma * (1 - done) * ret
            tmp_returns.insert(0, ret.unsqueeze(0))
        returns = torch.cat(tmp_returns, dim=0).to(device)

        # use either n_step or gae to estimate the advantages.
        # note that not all of these parameters are guaranteed to be used by the given function.
        advantages = self.calc_advantages(self.tmp_r, self.tmp_d, future, returns, critic_out, self.gamma, self.tau)

        # train the actor by guiding the increase or decrease of log_prob with the advantages.
        # note the negative sign: pytorch executes gradient descent, so we want to invert that to gradient ascent.
        policy_loss = -torch.mean(log_prob * advantages.detach())
        self.actor_optimizer.zero_grad()
        policy_loss.backward()
        nn.utils.clip_grad_norm_(self.actor.parameters(), 5)  # the use of gradient clipping helped in my tests
        self.actor_optimizer.step()

        # train the critic on the mean squared error between its output and the discounted returns
        value_loss = torch.mean((returns.detach() - critic_out) ** 2)
        self.critic_optimizer.zero_grad()
        value_loss.backward()
        nn.utils.clip_grad_norm_(self.critic.parameters(), 5)  # the use of gradient clipping helped in my tests
        self.critic_optimizer.step()

        # finally empty the temporary buffer
        self.reset_temporary_buffer()
```

## 6. Training an agent

Creating an agent is straightforward now. We can make three agents, with some combination of `n_step`, `gae`, and `n`, to see how these parameters affect learning.

- `alpha=0.0001`
- `gamma=0.95`
- `tau=0.5`

```python
agents = [('A2Cn4', A2CAgent(state_size, action_size, n_step, n=4)),
          ('A2Cn8', A2CAgent(state_size, action_size, n_step, n=8)),
          ('A2C-GAEn8', A2CAgent(state_size, action_size, gae, n=8))]
```

What we define now is a function that connects our `agent` and `env`. Recall that we'll be using the default `brain_name` to interact with the environment. The agent is allowed to store observations and, consequently, learn every `n` steps.

```python
def execute_episode(agent, env):
    # prepare the environment
    scores = np.zeros(num_agents)
    env_info = env.reset(train_mode=True)[brain_name]

    # get the initial state
    states = env_info.vector_observations
    while True:

        # evaluate the current state
        actions = agent.act(states)

        # execute the chosen action and get the outcome
        env_info = env.step(actions)[brain_name]
        next_states = env_info.vector_observations
        rewards = env_info.rewards
        dones = env_info.local_done

        # store the experience (also automatically learn, from time to time)
        agent.store(states, actions, rewards, next_states, dones)

        # prepare for the next iteration
        states = next_states
        scores += rewards

        if np.any(dones):
            break

    # return the total rewards obtained
    return np.mean(scores)
```

We can call the function above repeatedly while **monitoring** what happens, so that we'll be able to compare different approaches. The most important indicator is going to be the average score over 100 consecutive episodes (where each episode's score is the average of its 20 independent agents): the environment is considered **solved when the agents scores +30 or more** under this metric. We'll be training each variation 3 times from scratch, for 600 episodes, and we'll average the results. This way, the learning curves will depend less on the randomness of the environment and more on the actual differences of the proposed approaches.

```python
def train(agent, env, episodes=600, repeat=3, consecutive_episodes=100, show_output=True, save_as=None):

    results = [None] * repeat
    for r in range(repeat):
        partial_results = [None] * episodes

        # reset the agent to start learning from scratch
        agent.reset()
        for i in range(episodes):

            # execute all the episodes and store the results
            score = execute_episode(agent, env)
            partial_results[i] = score

            if show_output:
                print("\r[{}] Episode: {}, Score: {}".format(r+1, i+1, score), end="")
                sys.stdout.flush()
        if show_output:
            print()

        results[r] = partial_results

        # store the trained model if requested
        if save_as is not None:
            torch.save(agent.actor.state_dict(), '{}/{}_actor.pth'.format(save_as, r+1))
            torch.save(agent.critic.state_dict(), '{}/{}_critic.pth'.format(save_as, r+1))
            np.save('{}/report_{}'.format(save_as, r+1), np.array(partial_results))

    results = np.array(results)
    if save_as is not None:
        np.save('{}/full_report'.format(save_as), np.array(results))

    # use convolutions to calculate the mean and standard deviation summarizing the training step
    mean = signal.convolve2d(results, np.ones([repeat, consecutive_episodes]) / (repeat * consecutive_episodes), mode='valid')
    v = signal.convolve2d(results, np.ones([1, consecutive_episodes]) / consecutive_episodes, mode='valid')
    std_dev = signal.convolve2d(v ** 2 - mean ** 2, np.ones([repeat, 1]) / repeat, mode='valid') ** (1/2)
    return mean.flatten(), std_dev.flatten(), results
```

```python
training_results = []
for ag in agents:
    name, agent = ag
    print('Agent: {}'.format(name))
    mean, std_dev, full_report = train(agent, env, save_as='models/{}'.format(name.replace(" ", "_")))
    training_results.append((mean, std_dev, full_report))
env.close()
```

    Agent: A2Cn4
    [1] Episode: 600, Score: 39.00949912806973
    [2] Episode: 600, Score: 39.3239991210401
    [3] Episode: 600, Score: 39.14199912510812
    Agent: A2Cn8
    [1] Episode: 600, Score: 39.081999126449226
    [2] Episode: 600, Score: 37.057999171689154
    [3] Episode: 600, Score: 39.247499122750014
    Agent: A2C-GAEn8
    [1] Episode: 600, Score: 38.7974991328083
    [2] Episode: 600, Score: 39.176499124336985
    [3] Episode: 600, Score: 38.591499137412754

## 7. Comparing the results

We can use `matplotlib` to plot the mean of the score. All proposed variants manage to reach the goal of +30, in fact they almost get +40 points. However, there are some notable differences between them.

First, we can see that, for this task, the A2C method with `n = 8` is the first one to reach the goal, solving the environment in 314 episodes, followed by the `n = 4` version and then finally by the GAE version with `n = 8`. However, we can't yet say that the best approach is A2C with `n = 8` and in fact, if we look at its stability (the orange shadow), it clearly exhibits a higher variance after 600 episodes compared to the other two variants.

This seems in line with the fact that Monte Carlo approaches usually have a higher variance, and A2C with `n = 8` is closer to MC than A2C with `n = 4`. As for the remaining variant with GAE, although it also has `n = 8`, it does not show such high variance at the end of the learning process. This may be its strenght: instead of being the fastest method, it may allow for a more flexible and easy choice of `n` (which is notably hard in some tasks) while retaining good convergence properties.

```python
def plot_results(agents, training_results):
    for agent, ts in zip(agents, training_results):
        mean, std_dev, _ = ts
        episode_solved = np.argmax(mean > 30) + 100
        print('{}: Solved after {} episodes'.format(agent[0], episode_solved))

        x = np.arange(100, 100 + mean.size)
        line = plt.plot(x, mean, label = agent[0])
        plt.fill_between(x, mean + std_dev, mean - std_dev, alpha=0.15)
        plt.axvline(episode_solved, lw=0.75, color=line[0].get_color())

    plt.axhline(30, dashes=(4, 8), lw=0.75, color='black')
    plt.xticks(np.arange(1, 1 + mean.size / 100) * 100)
    plt.yticks(np.arange(0, 9) * 5)
    plt.xlabel('Number of episodes')
    plt.title('Average score over 100 consecutive episodes')
    plt.legend()
    plt.show()

plot_results(agents, training_results)
```

## 8. Possible improvements

One of the things that is usually applied to policy gradient methods is an **entropy** term to help with exploration. In fact, given that the actor outputs probabilities, it is straightforward to add an entropy term to its loss so that it will be maximized. In particular, since we're using a normal distribution here, the ouptut of the model only consists of its parameters, $\mu$ and $\sigma$, and maximizing the entropy is essentially the same as maximizing $\sigma$, or in other words, the randomness of the agent. The entropy term should always be limited by an hyperparameters, otherwise the actor won't be able to converge to any policy.

It would be interesting to compare this solution to other policy gradient algorithms like [REINFORCE](https://link.springer.com/content/pdf/10.1007/BF00992696.pdf), [PPO](https://arxiv.org/pdf/1707.06347.pdf) or [DDPG](https://arxiv.org/pdf/1509.02971.pdf), or even to black-box optimization techniques like [ES](https://arxiv.org/pdf/1703.03864.pdf).

## 9. Conclusions

For this task, the agents proposed manage to play well every time. The models used to create the graphs in section 7 are available in the repository. The full documented code is also available and `main.py` can be executed to train or load a pre-trained A2C-GAE agent with `n = 8` and see it in action, just like it was done to get the model `final.pht`.