# Tutorial-IllinoisGRMHD: Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

## Authors: Leo Werneck & Zach Etienne

**This module is currently under development**

## In this tutorial module we explain the construction of the right-hand side of the evolution equations of $\left[\sqrt{\gamma}\Phi\right]$ and add gauge terms to the right-hand side of the evolution equations of $A_{i}$

### 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)).

If using the version of `IllinoisGRMHD` with piecewise polytropic *or* tabulated (coming soon!) EOS support, then the following citation is also required:

* **(Required)** Etienne, Z. B., Werneck, L., Paschalidis, V., Haas R., Mösta P., and Shapiro, S. L., *IllinoisGRMHD github repository* (2019). Source Code URL: https://github.com/zachetienne/nrpytutorial/tree/master/IllinoisGRMHD/.

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

This module is organized as follows

0. [Step 0](#src_dir): **Source directory creation**
1. [Step 1](#lorenz_psi6phi_rhs__add_gauge_terms_to_a_i_rhs__c): **`Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C`**
    1. [Step 1.a](#interpolations): *Performing all necessary interpolations*
        1. [Step 1.a.i](#interpolation_algorithm): The interpolation algorithm
        1. [Step 1.a.ii](#interp_and_alpha_phi_minus_betaj_A_j): Interpolating gridfunctions and computing $\left(\alpha\Phi-\beta^{j}A_{j}\right)$
    1. [Step 1.b](#partial_t_a_i_gauge): *Computing $\partial_{t}A_{i}^{\rm gauge}$*
    1. [Step 1.c](#shift_advection_terms): *Computing $\partial_{j}\beta^{j}\left[\sqrt{\gamma}\Phi\right]$*
    1. [Step 1.d](#partial_j_alpha_psi6_aj): *Computing $-\partial_{j}\left(\alpha\sqrt{\gamma}A^{j}\right)-\xi\alpha\left[\sqrt{\gamma}\Phi\right]$*
    1. [Step 1.e](#fct_avg): *The `avg()` function*
1. [Step 2](#code_validation): **Code validation**
1. [Step 3](#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
outdir = os.path.join("..","src")
cmd.mkdir(outdir)
```

# Step 1: The `Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C` file \[Back to [top](#toc)\]
$$\label{lorenz_psi6phi_rhs__add_gauge_terms_to_a_i_rhs__c}$$

In the `Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C` file we compute the gauge terms in $\partial_{t}A_{i}$, as well as the right-hand side of $\left[\sqrt{\gamma}\Phi\right]$, according to equations (16) and (17) of [Etienne *et al*.](https://arxiv.org/pdf/1501.07276.pdf):

$$
\begin{align}
\partial_{t}A_{i}^{\rm gauge} &\equiv \partial_{t}A_{i} - \epsilon_{ijk}v^{j}\tilde{B}^{k} = -\partial_{i}\left(\alpha\Phi-\beta^{j}A_{j}\right)\ ,\\
\partial_{t}\left[\sqrt{\gamma}\Phi\right] &= -\partial_{j}\left(\alpha\sqrt{\gamma}A^{j} - \beta^{j}\left[\sqrt{\gamma}\Phi\right]\right) - \xi\alpha\left[\sqrt{\gamma}\Phi\right]\ .
\end{align}
$$

```python
%%writefile $outdir/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C
static inline CCTK_REAL avg(CCTK_REAL f[PLUS2+1][PLUS2+1][PLUS2+1],int imin,int imax, int jmin,int jmax, int kmin,int kmax);

static void Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs(const cGH *cctkGH,const int *cctk_lsh,const int *cctk_nghostzones,CCTK_REAL *dX,CCTK_REAL **in_vars,CCTK_REAL *psi6phi,
                                                           CCTK_REAL *shiftx_iphjphkph,CCTK_REAL *shifty_iphjphkph,CCTK_REAL *shiftz_iphjphkph,
                                                           CCTK_REAL *alpha_iphjphkph,CCTK_REAL *alpha_Phi_minus_betaj_A_j_iphjphkph,CCTK_REAL *alpha_sqrtg_Ax_interp,
                                                           CCTK_REAL *alpha_sqrtg_Ay_interp,CCTK_REAL *alpha_sqrtg_Az_interp,
                                                           CCTK_REAL *psi6phi_rhs,CCTK_REAL *Ax_rhs,CCTK_REAL *Ay_rhs,CCTK_REAL *Az_rhs) {
  DECLARE_CCTK_PARAMETERS;
  /* Compute \partial_t psi6phi = -\partial_i (  \alpha psi^6 A^i - psi6phi \beta^i)
   *    (Eq 13 of http://arxiv.org/pdf/1110.4633.pdf), using Lorenz gauge.
   * Note that the RHS consists of a shift advection term on psi6phi and
   *    a term depending on the vector potential.
   * psi6phi is defined at (i+1/2,j+1/2,k+1/2), but instead of reconstructing
   *    to compute the RHS of \partial_t psi6phi, we instead use standard
   *    interpolations.
   */

  CCTK_REAL dXm1=1.0/dX[0];
  CCTK_REAL dYm1=1.0/dX[1];
  CCTK_REAL dZm1=1.0/dX[2];
```

## Step 1.a: Performing all necessary interpolations \[Back to [top](#toc)\]
$$\label{interpolations}$$

### Step 1.a.i: The interpolation algorithm \[Back to [top](#toc)\]
$$\label{interpolation_algorithm}$$

It is important to notice the different staggerrings. We are ultimately interested in the RHS of $\left[\sqrt{\gamma}\Phi\right]_{i,j,k}=\left[\psi^{6}\Phi\right]_{i,j,k}$, which is actually located at $\left(i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}\right)$. Remember the staggerings:

|               Quantitity               |              Actual location on the grid               |
|:--------------------------------------:|:------------------------------------------------------:|
|            $\alpha_{i,j,k}$            |                  $\left(i,j,k\right)$                  |
|           $\psi^{6}_{i,j,k}$           |                  $\left(i,j,k\right)$                  |
|      $\left(A_{x}\right)_{i,j,k}$      |      $\left(i,j+\frac{1}{2},k+\frac{1}{2}\right)$      |
|      $\left(A_{y}\right)_{i,j,k}$      |      $\left(i+\frac{1}{2},j,k+\frac{1}{2}\right)$      |
|      $\left(A_{z}\right)_{i,j,k}$      |      $\left(i+\frac{1}{2},j+\frac{1}{2},k\right)$      |
|$\left[\sqrt{\gamma}\Phi\right]_{i,j,k}$|$\left(i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}\right)$|

To obtain the staggerings we use an interpolation routine which averages the value of the gridfunctions around the point of interest. For metric quantities, such as $\alpha$ and $\psi$, it is very easy to understand what is happening, so we will use them as the our main example.

Because we are solving the equations in 3 spatial dimensions, our numerical grid can be seen as a collection of small "cubes" whose sizes are $\left(dx,dy,dz\right)$ (notice that, while we could have $dx\neq dy\neq dz$, we will refer to the grid as cubes to make it easier to describe it). Then, the metric quantities, which are unstaggered, are defined on the *vertices* of the cubes. On the other hand, $\left[\sqrt{\gamma}\Phi\right]$ is defined at the *centers* of the cubes. Thus, to obtain metric quantities on the center of the cube, we use a *averaging interpolation algorithm* which  uses all 8 vertices of the cube.

As a concrete example, we obtain the lapse at the grid location $\left(i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}\right)$ by averaging

$$
\alpha_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} =
\frac{
\alpha_{i,j,k}+
\alpha_{i+1,j,k}+
\alpha_{i,j+1,k}+
\alpha_{i,j,k+1}+
\alpha_{i+1,j+1,k}+
\alpha_{i+1,j,k+1}+
\alpha_{i,j+1,k+1}+
\alpha_{i+1,j+1,k+1}
}
{8}\ .
$$

However, our algorithm is not restricted to using 8 points when performing such an average. For example, suppose we need the quantity $\left(A_{x}\right)_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}}$, but we normaly only have $\left(A_{x}\right)_{i,j+\frac{1}{2},k+\frac{1}{2}}$. The average algorithm would then be

$$
\left(A_{x}\right)_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} = \frac{\left(A_{x}\right)_{i,j+\frac{1}{2},k+\frac{1}{2}} + \left(A_{x}\right)_{i+1,j+\frac{1}{2},k+\frac{1}{2}}}{2}\ .
$$

### Step 1.a.ii: Interpolating gridfunctions and computing $\left(\alpha\Phi-\beta^{j}A_{j}\right)$ \[Back to [top](#toc)\]
$$\label{interp_and_alpha_phi_minus_betaj_A_j}$$

Now, to compute $\partial_{j}\left(\alpha\psi^{6}A^{j}\right)$ we need different interpolations for different quantitites. We will focus on the $x$-direction to give a concrete example.

Consider, then, the term

$$
\partial_{x}\left(\alpha\psi^{6}A^{x}\right)\ .
$$

As we know, the second-order, centered finite differences approximation of this equation is

$$
\left[\partial_{x}\left(\alpha\psi^{6}A^{x}\right)\right]_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} = \frac{\left(\alpha\psi^{6}A^{x}\right)_{i+1,j+\frac{1}{2},k+\frac{1}{2}} - \left(\alpha\psi^{6}A^{x}\right)_{i,j+\frac{1}{2},k+\frac{1}{2}}}{dx}\ ,
$$

where the subscripts indicate ***the actual gridpoint location***. The $A_{x}$ staggering, then, already makes it available at the needed points for the derivative. However, we need to compute $A^{x} = g^{xi}A_{i} = \gamma^{xx}A_{x} + \gamma^{xy}A_{y} + \gamma^{xz}A_{z}$, which means that to compute $A^{x}$ alone we need to interpolate $\left(\gamma^{xx},\gamma^{xy},\gamma^{xz},A_{y},A_{z}\right)$ to the grid locations $\left(i,j+\frac{1}{2},k+\frac{1}{2}\right)$ and $\left(i+1,j+\frac{1}{2},k+\frac{1}{2}\right)$. On top of that, to be able to compute the derivative, we also need $\alpha$ and $\psi^{6}$ at those locations. Finally, once all interpolations have been performed, we can then add the appropriate terms to the RHSs of $\partial_{t}\left[\sqrt{\gamma}\Phi\right]$ and $\partial_{t}A_{i}$.

The outline of the algorithm is as follows:

1. Read in gridfunctions at point $(i,j,k)$ and points around that, so that we can perform interpolations.
1. Interpolate $\alpha$ to the point $\left(i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}\right)$
1. $A^{x}$ term:
    1. Interpolate $\bar{\gamma}^{xx},\bar{\gamma}^{xy},\bar{\gamma}^{xz},\alpha,\psi$ to the point $\left(i,j+\frac{1}{2},k+\frac{1}{2}\right)$.
    1. Interpolate $A_{x}$, $A_{y}$, $A_{z}$ to the point $\left(i,j+\frac{1}{2},k+\frac{1}{2}\right)$.
    1. Compute $A^{x}$ at the point $\left(i,j+\frac{1}{2},k+\frac{1}{2}\right)$.
    1. Compute $\alpha\psi^{6}A^{x}$ at the point $\left(i,j+\frac{1}{2},k+\frac{1}{2}\right)$.
1. $A^{y}$ term:
    1. Interpolate $\bar{\gamma}^{xy},\bar{\gamma}^{yy},\bar{\gamma}^{yz},\alpha,\psi$ to the point $\left(i+\frac{1}{2},j,k+\frac{1}{2}\right)$.
    1. Interpolate $A_{x}$, $A_{y}$, $A_{z}$ to the point $\left(i+\frac{1}{2},j,k+\frac{1}{2}\right)$.
    1. Compute $A^{y}$ at the point $\left(i+\frac{1}{2},j,k+\frac{1}{2}\right)$.
    1. Compute $\alpha\psi^{6}A^{y}$ at the point $\left(i+\frac{1}{2},j,k+\frac{1}{2}\right)$.
1. $A^{z}$ term:
    1. Interpolate $\bar{\gamma}^{xz},\bar{\gamma}^{yz},\bar{\gamma}^{zz},\alpha,\psi$ to the point $\left(i+\frac{1}{2},j+\frac{1}{2},k\right)$.
    1. Interpolate $A_{x}$, $A_{y}$, $A_{z}$ to the point $\left(i+\frac{1}{2},j+\frac{1}{2},k\right)$.
    1. Compute $A^{z}$ at the point $\left(i+\frac{1}{2},j+\frac{1}{2},k\right)$.
    1. Compute $\alpha\psi^{6}A^{z}$ at the point $\left(i+\frac{1}{2},j+\frac{1}{2},k\right)$.
1. Interpolate $A_{x},A_{y},A_{z}$ to the point $\left(i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}\right)$.
1. Interpolate $\beta^{x},\beta^{y},\beta^{z}$ to the point $\left(i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}\right)$.
1. Compute $\left(\alpha\Phi-\beta^{j}A_{j}\right)$ at $\left(i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}\right)$.

```python
%%writefile -a $outdir/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

  // The stencil here is {-1,1},{-1,1},{-1,1} for x,y,z directions, respectively.
  //     Note that ALL input variables are defined at ALL gridpoints, so no
  //     worries about ghostzones.
#pragma omp parallel for
  for(int k=1;k<cctk_lsh[2]-1;k++) for(int j=1;j<cctk_lsh[1]-1;j++) for(int i=1;i<cctk_lsh[0]-1;i++) {
        int index=CCTK_GFINDEX3D(cctkGH,i,j,k);
        CCTK_REAL INTERP_VARS[MAXNUMINTERP][PLUS2+1][PLUS2+1][PLUS2+1];

        // First compute \partial_j \alpha \sqrt{\gamma} A^j (RHS of \partial_i psi6phi)
        // FIXME: Would be much cheaper & easier to unstagger A_i, raise, then interpolate A^i.
        //        However, we keep it this way to be completely compatible with the original
        //        Illinois GRMHD thorn, called mhd_evolve.
        //
        //Step 1) j=x: Need to raise A_i, but to do that, we must have all variables at the same gridpoints:
        // The goal is to compute \partial_j (\alpha \sqrt{\gamma} A^j) at (i+1/2,j+1/2,k+1/2)
        //    We do this by first interpolating (RHS1x) = (\alpha \sqrt{\gamma} A^x) at
        //    (i,j+1/2,k+1/2)and (i+1,j+1/2,k+1/2), then taking \partial_x (RHS1x) =
        //    [ RHS1x(i+1,j+1/2,k+1/2) - RHS1x(i,j+1/2,k+1/2) ]/dX.
        // First bring gup's, psi, and alpha to (i,j+1/2,k+1/2):
        int num_vars_to_interp;
        int vars_to_interpolate[MAXNUMINTERP] = {INTERP_GAMMATILDEUPXX,INTERP_GAMMATILDEUPXY,INTERP_GAMMATILDEUPXZ,INTERP_GAMMATILDEUPYY,INTERP_GAMMATILDEUPYZ,INTERP_GAMMATILDEUPZZ,INTERP_LAPM1,INTERP_PSI,INTERP_SHIFTX,INTERP_SHIFTY,INTERP_SHIFTZ};
        num_vars_to_interp = 11;
        // We may set interp_limits to be more general than we need.
        int interp_limits[6] = {-1,1,-1,1,-1,1}; SET_INDEX_ARRAYS_3DBLOCK(interp_limits);
        //SET_INDEX_ARRAYS_3DBLOCK(interp_limits);
        for(int ww=0;ww<num_vars_to_interp;ww++) {
          int whichvar=vars_to_interpolate[ww];
          // Read in variable at interp. stencil points from main memory, store in INTERP_VARS.
          for(int kk=PLUS0;kk<=PLUS1;kk++) for(int jj=PLUS0;jj<=PLUS1;jj++) for(int ii=PLUS0;ii<=PLUS1;ii++) {
                INTERP_VARS[whichvar][kk][jj][ii] = in_vars[whichvar][index_arr_3DB[kk][jj][ii]]; }
        }

        // Next set \alpha at (i+1/2,j+1/2,k+1/2). Will come in handy when computing damping term later.
        alpha_iphjphkph[index] = avg(INTERP_VARS[INTERP_LAPM1] , PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS1)+1.0;

        //A_x needs a stencil s.t. interp_limits={0,1,-1,1,-1,1}:
        for(int kk=MINUS1;kk<=PLUS1;kk++) for(int jj=MINUS1;jj<=PLUS1;jj++) for(int ii=PLUS0;ii<=PLUS1;ii++) {
              INTERP_VARS[INTERP_AX][kk][jj][ii] = in_vars[INTERP_AX][index_arr_3DB[kk][jj][ii]]; }
        //A_y needs a stencil s.t. interp_limits={-1,1,0,1,-1,1}:
        for(int kk=MINUS1;kk<=PLUS1;kk++) for(int jj=PLUS0;jj<=PLUS1;jj++) for(int ii=MINUS1;ii<=PLUS1;ii++) {
              INTERP_VARS[INTERP_AY][kk][jj][ii] = in_vars[INTERP_AY][index_arr_3DB[kk][jj][ii]]; }
        //A_z needs a stencil s.t. interp_limits={-1,1,-1,1,0,1}:
        for(int kk=PLUS0;kk<=PLUS1;kk++) for(int jj=MINUS1;jj<=PLUS1;jj++) for(int ii=MINUS1;ii<=PLUS1;ii++) {
              INTERP_VARS[INTERP_AZ][kk][jj][ii] = in_vars[INTERP_AZ][index_arr_3DB[kk][jj][ii]]; }

        // FIRST DO A^X TERM (interpolate to (i,j+1/2,k+1/2) )
        // \alpha \sqrt{\gamma} A^x = \alpha psi^6 A^x (RHS of \partial_i psi6phi)
        // Note that gupij is \tilde{\gamma}^{ij}, so we need to multiply by \psi^{-4}.
        CCTK_REAL gupxx_jphkph = avg(INTERP_VARS[INTERP_GAMMATILDEUPXX], PLUS0,PLUS0, PLUS0,PLUS1, PLUS0,PLUS1);
        CCTK_REAL gupxy_jphkph = avg(INTERP_VARS[INTERP_GAMMATILDEUPXY], PLUS0,PLUS0, PLUS0,PLUS1, PLUS0,PLUS1);
        CCTK_REAL gupxz_jphkph = avg(INTERP_VARS[INTERP_GAMMATILDEUPXZ], PLUS0,PLUS0, PLUS0,PLUS1, PLUS0,PLUS1);

        for(int kk=PLUS0;kk<=PLUS1;kk++) for(int jj=PLUS0;jj<=PLUS1;jj++) for(int ii=PLUS0;ii<=PLUS1;ii++) {
              CCTK_REAL Psi2 = INTERP_VARS[INTERP_PSI][kk][jj][ii]*INTERP_VARS[INTERP_PSI][kk][jj][ii];
              CCTK_REAL alpha = INTERP_VARS[INTERP_LAPM1][kk][jj][ii]+1.0;
              INTERP_VARS[INTERP_LAPSE_PSI2][kk][jj][ii]=alpha*Psi2;
              INTERP_VARS[INTERP_LAPSE_OVER_PSI6][kk][jj][ii]=alpha/(Psi2*Psi2*Psi2);
            }

        CCTK_REAL lapse_Psi2_jphkph = avg(INTERP_VARS[INTERP_LAPSE_PSI2], PLUS0,PLUS0, PLUS0,PLUS1, PLUS0,PLUS1);

        CCTK_REAL A_x_jphkph   = avg(INTERP_VARS[INTERP_AX], PLUS0,PLUS0, PLUS0,PLUS0, PLUS0,PLUS0); // @ (i,j+1/2,k+1/2)
        CCTK_REAL A_y_jphkph   = avg(INTERP_VARS[INTERP_AY],MINUS1,PLUS0, PLUS0,PLUS1, PLUS0,PLUS0); // @ (i+1/2,j,k+1/2)
        CCTK_REAL A_z_jphkph   = avg(INTERP_VARS[INTERP_AZ],MINUS1,PLUS0, PLUS0,PLUS0, PLUS0,PLUS1); // @ (i+1/2,j+1/2,k)

        alpha_sqrtg_Ax_interp[index] = lapse_Psi2_jphkph*
          ( gupxx_jphkph*A_x_jphkph + gupxy_jphkph*A_y_jphkph + gupxz_jphkph*A_z_jphkph );

        // DO A^Y TERM (interpolate to (i+1/2,j,k+1/2) )
        // \alpha \sqrt{\gamma} A^y = \alpha psi^6 A^y (RHS of \partial_i psi6phi)
        // Note that gupij is \tilde{\gamma}^{ij}, so we need to multiply by \psi^{-4}.
        CCTK_REAL gupxy_iphkph = avg(INTERP_VARS[INTERP_GAMMATILDEUPXY], PLUS0,PLUS1, PLUS0,PLUS0, PLUS0,PLUS1);
        CCTK_REAL gupyy_iphkph = avg(INTERP_VARS[INTERP_GAMMATILDEUPYY], PLUS0,PLUS1, PLUS0,PLUS0, PLUS0,PLUS1);
        CCTK_REAL gupyz_iphkph = avg(INTERP_VARS[INTERP_GAMMATILDEUPYZ], PLUS0,PLUS1, PLUS0,PLUS0, PLUS0,PLUS1);

        CCTK_REAL lapse_Psi2_iphkph = avg(INTERP_VARS[INTERP_LAPSE_PSI2], PLUS0,PLUS1, PLUS0,PLUS0, PLUS0,PLUS1);
        //CCTK_REAL lapse_iphkph = avg(INTERP_VARS[INTERP_LAPM1], PLUS0,PLUS1, PLUS0,PLUS0, PLUS0,PLUS1)+1.0;
        //CCTK_REAL psi_iphkph   = avg(INTERP_VARS[INTERP_PSI  ], PLUS0,PLUS1, PLUS0,PLUS0, PLUS0,PLUS1);

        CCTK_REAL A_x_iphkph   = avg(INTERP_VARS[INTERP_AX], PLUS0,PLUS1,MINUS1,PLUS0, PLUS0,PLUS0); // @ (i,j+1/2,k+1/2)
        CCTK_REAL A_y_iphkph   = avg(INTERP_VARS[INTERP_AY], PLUS0,PLUS0, PLUS0,PLUS0, PLUS0,PLUS0); // @ (i+1/2,j,k+1/2)
        CCTK_REAL A_z_iphkph   = avg(INTERP_VARS[INTERP_AZ], PLUS0,PLUS0,MINUS1,PLUS0, PLUS0,PLUS1); // @ (i+1/2,j+1/2,k)

        alpha_sqrtg_Ay_interp[index] = lapse_Psi2_iphkph*
          ( gupxy_iphkph*A_x_iphkph + gupyy_iphkph*A_y_iphkph + gupyz_iphkph*A_z_iphkph );

        // DO A^Z TERM (interpolate to (i+1/2,j+1/2,k) )
        // \alpha \sqrt{\gamma} A^z = \alpha psi^6 A^z (RHS of \partial_i psi6phi)
        // Note that gupij is \tilde{\gamma}^{ij}, so we need to multiply by \psi^{-4}.
        CCTK_REAL gupxz_iphjph = avg(INTERP_VARS[INTERP_GAMMATILDEUPXZ], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS0);
        CCTK_REAL gupyz_iphjph = avg(INTERP_VARS[INTERP_GAMMATILDEUPYZ], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS0);
        CCTK_REAL gupzz_iphjph = avg(INTERP_VARS[INTERP_GAMMATILDEUPZZ], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS0);
        //CCTK_REAL lapse_iphjph = avg(INTERP_VARS[INTERP_LAPM1], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS0)+1.0;
        //CCTK_REAL psi_iphjph   = avg(INTERP_VARS[INTERP_PSI  ], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS0);

        CCTK_REAL lapse_Psi2_iphjph = avg(INTERP_VARS[INTERP_LAPSE_PSI2], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS0);

        CCTK_REAL A_x_iphjph   = avg(INTERP_VARS[INTERP_AX], PLUS0,PLUS1, PLUS0,PLUS0,MINUS1,PLUS0); // @ (i,j+1/2,k+1/2)
        CCTK_REAL A_y_iphjph   = avg(INTERP_VARS[INTERP_AY], PLUS0,PLUS0, PLUS0,PLUS1,MINUS1,PLUS0); // @ (i+1/2,j,k+1/2)
        CCTK_REAL A_z_iphjph   = avg(INTERP_VARS[INTERP_AZ], PLUS0,PLUS0, PLUS0,PLUS0, PLUS0,PLUS0); // @ (i+1/2,j+1/2,k)

        alpha_sqrtg_Az_interp[index] = lapse_Psi2_iphjph*
          ( gupxz_iphjph*A_x_iphjph + gupyz_iphjph*A_y_iphjph + gupzz_iphjph*A_z_iphjph );

        // Next set \alpha \Phi - \beta^j A_j at (i+1/2,j+1/2,k+1/2):
        //   We add a "L" suffix to shifti_iphjphkph to denote "LOCAL", as we set
        //      shifti_iphjphkph[] gridfunction below.
        CCTK_REAL shiftx_iphjphkphL = avg(INTERP_VARS[INTERP_SHIFTX], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS1);
        CCTK_REAL shifty_iphjphkphL = avg(INTERP_VARS[INTERP_SHIFTY], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS1);
        CCTK_REAL shiftz_iphjphkphL = avg(INTERP_VARS[INTERP_SHIFTZ], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS1);
        CCTK_REAL lapse_over_Psi6_iphjphkphL = avg(INTERP_VARS[INTERP_LAPSE_OVER_PSI6], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS1);
        //CCTK_REAL psi_iphjphkph = avg(INTERP_VARS[INTERP_PSI  ], PLUS0,PLUS1, PLUS0,PLUS1, PLUS0,PLUS1);
        //CCTK_REAL psi2_iphjphkph= psi_iphjphkph*psi_iphjphkph;
        //CCTK_REAL psi6_iphjphkph= psi2_iphjphkph*psi2_iphjphkph*psi2_iphjphkph;
        CCTK_REAL A_x_iphjphkph = avg(INTERP_VARS[INTERP_AX], PLUS0,PLUS1, PLUS0,PLUS0, PLUS0,PLUS0); // @ (i,j+1/2,k+1/2)
        CCTK_REAL A_y_iphjphkph = avg(INTERP_VARS[INTERP_AY], PLUS0,PLUS0, PLUS0,PLUS1, PLUS0,PLUS0); // @ (i+1/2,j,k+1/2)
        CCTK_REAL A_z_iphjphkph = avg(INTERP_VARS[INTERP_AZ], PLUS0,PLUS0, PLUS0,PLUS0, PLUS0,PLUS1); // @ (i+1/2,j+1/2,k)

        alpha_Phi_minus_betaj_A_j_iphjphkph[index] = psi6phi[index]*lapse_over_Psi6_iphjphkphL
          - (shiftx_iphjphkphL*A_x_iphjphkph + shifty_iphjphkphL*A_y_iphjphkph + shiftz_iphjphkphL*A_z_iphjphkph);

        // Finally, save shifti_iphjphkph, for \partial_j \beta^j psi6phi
        shiftx_iphjphkph[index]=shiftx_iphjphkphL;
        shifty_iphjphkph[index]=shifty_iphjphkphL;
        shiftz_iphjphkph[index]=shiftz_iphjphkphL;
      }

```

### Step 1.b: Computing $\partial_{t}A_{i}^{\rm gauge}$ \[Back to [top](#toc)\]
$$\label{partial_t_a_i_gauge}$$

Now that we have access to the gridfunctions at all the necessary gridpoints, we proceed to the computation of the gauge terms in $\partial_{t}A_{i}$, i.e.

$$
\partial_{t}A_{i}^{\rm gauge} = -\partial_{i}\left(\alpha\Phi-\beta^{j}A_{j}\right)\ ,
$$

or, more explicitly:

$$
\begin{align}
\left(\partial_{t}A_{x}^{\rm gauge}\right)_{i,j+\frac{1}{2},k+\frac{1}{2}}
&=
\frac{
\left(\alpha\Phi-\beta^{j}A_{j}\right)_{i-\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} - \left(\alpha\Phi-\beta^{j}A_{j}\right)_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}}
}{dx}\ ,\\
\left(\partial_{t}A_{y}^{\rm gauge}\right)_{i+\frac{1}{2},j,k+\frac{1}{2}}
&=
\frac{
\left(\alpha\Phi-\beta^{j}A_{j}\right)_{i+\frac{1}{2},j-\frac{1}{2},k+\frac{1}{2}} - \left(\alpha\Phi-\beta^{j}A_{j}\right)_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}}
}{dy}\ ,\\
\left(\partial_{t}A_{z}^{\rm gauge}\right)_{i+\frac{1}{2},j+\frac{1}{2},k}
&=
\frac{
\left(\alpha\Phi-\beta^{j}A_{j}\right)_{i-\frac{1}{2},j+\frac{1}{2},k-\frac{1}{2}} - \left(\alpha\Phi-\beta^{j}A_{j}\right)_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}}
}{dz}\ .
\end{align}
$$

```python
%%writefile -a $outdir/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

  // This loop requires two additional ghostzones in every direction. Hence the following loop definition:
#pragma omp parallel for
  for(int k=cctk_nghostzones[2];k<cctk_lsh[2]-cctk_nghostzones[2];k++) for(int j=cctk_nghostzones[1];j<cctk_lsh[1]-cctk_nghostzones[1];j++) for(int i=cctk_nghostzones[0];i<cctk_lsh[0]-cctk_nghostzones[0];i++) {
        int index = CCTK_GFINDEX3D(cctkGH,i,j,k);

        // \partial_t A_i = [reconstructed stuff] + [gauge stuff],
        //    where [gauge stuff] = -\partial_i (\alpha \Phi - \beta^j A_j)
        CCTK_REAL alpha_Phi_minus_betaj_A_j_iphjphkphL = alpha_Phi_minus_betaj_A_j_iphjphkph[index];
        // - partial_i -> - (A_{i} - A_{i-1})/dX = (A_{i-1} - A_{i})/dX, for Ax
        Ax_rhs[index] += dXm1*(alpha_Phi_minus_betaj_A_j_iphjphkph[CCTK_GFINDEX3D(cctkGH,i-1,j,k)] - alpha_Phi_minus_betaj_A_j_iphjphkphL);
        Ay_rhs[index] += dYm1*(alpha_Phi_minus_betaj_A_j_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j-1,k)] - alpha_Phi_minus_betaj_A_j_iphjphkphL);
        Az_rhs[index] += dZm1*(alpha_Phi_minus_betaj_A_j_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j,k-1)] - alpha_Phi_minus_betaj_A_j_iphjphkphL);
```

    Appending to ../src/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

### Step 1.c: Computing $\partial_{j}\beta^{j}\left[\sqrt{\gamma}\Phi\right]$ \[Back to [top](#toc)\]
$$\label{shift_advection_terms}$$

We now compute the shift-advection terms, $\partial_{j}\beta^{j}\left[\sqrt{\gamma}\Phi\right]$. For these terms we use [forwards/backwards finite difference stencils](https://en.wikipedia.org/wiki/Finite_difference_coefficient). The sign of $\beta^{i}$ determines whether we will use forwards or backwards finite difference, according to the following criterion:

| $\beta^{i}$ | Finite difference |
|:-----------:|:-----------------:|
|    $>0$     |     Backward      |
|    $<0$     |      Forward      |

For example, assume $\beta^{x}<0$. Then the term which contains the derivative in the $x$-direction would read

$$
\left[\partial_{x}\beta^{x}\left(\sqrt{\gamma}\Phi\right)\right]_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} =
\frac{
\left[\beta^{x}\left(\sqrt{\gamma}\Phi\right)\right]_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} -
4\left[\beta^{x}\left(\sqrt{\gamma}\Phi\right)\right]_{i-\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} +
3\left[\beta^{x}\left(\sqrt{\gamma}\Phi\right)\right]_{i-\frac{3}{2},j+\frac{1}{2},k+\frac{1}{2}}
}
{2dx}\ .
$$

Similarly, if $\beta^{y}>0$, then the term which contains the derivative in the $y$-direction would read

$$
\left[\partial_{y}\beta^{y}\left(\sqrt{\gamma}\Phi\right)\right]_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} =
\frac{
-3\left[\beta^{y}\left(\sqrt{\gamma}\Phi\right)\right]_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}} +
4\left[\beta^{y}\left(\sqrt{\gamma}\Phi\right)\right]_{i+\frac{1}{2},j+\frac{3}{2},k+\frac{1}{2}} -
\left[\beta^{y}\left(\sqrt{\gamma}\Phi\right)\right]_{i+\frac{1}{2},j+\frac{5}{2},k+\frac{1}{2}}
}
{2dy}\ .
$$

```python
%%writefile -a $outdir/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

        // \partial_t psi6phi = [shift advection term] + \partial_j (\alpha \sqrt{\gamma} A^j)
        // Here we compute [shift advection term] = \partial_j (\beta^j psi6phi)
        // Cache misses are likely more expensive than branch mispredictions here,
        //       which is why we use if() statements and array lookups inside the if()'s.
        CCTK_REAL psi6phi_rhsL=0.0;
        CCTK_REAL psi6phiL=psi6phi[index];
        CCTK_REAL shiftx_iphjphkphL=shiftx_iphjphkph[index];
        CCTK_REAL shifty_iphjphkphL=shifty_iphjphkph[index];
        CCTK_REAL shiftz_iphjphkphL=shiftz_iphjphkph[index];

        // \partial_x (\beta^x psi6phi) :
        if(shiftx_iphjphkphL < 0.0) {
          psi6phi_rhsL+=0.5*dXm1*(+    shiftx_iphjphkph[CCTK_GFINDEX3D(cctkGH,i-2,j,k)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i-2,j,k)]
                                  -4.0*shiftx_iphjphkph[CCTK_GFINDEX3D(cctkGH,i-1,j,k)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i-1,j,k)]
                                  +3.0*shiftx_iphjphkphL*                               psi6phiL);
        } else {
          psi6phi_rhsL+=0.5*dXm1*(-    shiftx_iphjphkph[CCTK_GFINDEX3D(cctkGH,i+2,j,k)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i+2,j,k)]
                                  +4.0*shiftx_iphjphkph[CCTK_GFINDEX3D(cctkGH,i+1,j,k)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i+1,j,k)]
                                  -3.0*shiftx_iphjphkphL*                               psi6phiL);
        }

        // \partial_y (\beta^y psi6phi) :
        if(shifty_iphjphkphL < 0.0) {
          psi6phi_rhsL+=0.5*dYm1*(+    shifty_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j-2,k)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i,j-2,k)]
                                  -4.0*shifty_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j-1,k)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i,j-1,k)]
                                  +3.0*shifty_iphjphkphL*                               psi6phiL);
        } else {
          psi6phi_rhsL+=0.5*dYm1*(-    shifty_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j+2,k)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i,j+2,k)]
                                  +4.0*shifty_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j+1,k)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i,j+1,k)]
                                  -3.0*shifty_iphjphkphL*                               psi6phiL);
        }

        // \partial_z (\beta^z psi6phi) :
        if(shiftz_iphjphkphL < 0.0) {
          psi6phi_rhsL+=0.5*dZm1*(+    shiftz_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j,k-2)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i,j,k-2)]
                                  -4.0*shiftz_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j,k-1)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i,j,k-1)]
                                  +3.0*shiftz_iphjphkphL*                               psi6phiL);
        } else {
          psi6phi_rhsL+=0.5*dZm1*(-    shiftz_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j,k+2)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i,j,k+2)]
                                  +4.0*shiftz_iphjphkph[CCTK_GFINDEX3D(cctkGH,i,j,k+1)]*psi6phi[CCTK_GFINDEX3D(cctkGH,i,j,k+1)]
                                  -3.0*shiftz_iphjphkphL*                               psi6phiL);
        }
```

    Appending to ../src/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

### Step 1.d: Computing $-\partial_{j}\left(\alpha\sqrt{\gamma}A^{j}\right)-\xi\alpha\left[\sqrt{\gamma}\Phi\right]$ \[Back to [top](#toc)\]
$$\label{partial_j_alpha_psi6_aj}$$

Now we have the simple task of computing the remaining terms on the RHS of $\partial_{t}\left[\sqrt{\gamma}\Phi\right]$. For the derivative term, $\partial_{j}\left(\alpha\sqrt{\gamma}A^{j}\right)$, we can now use ordinary centered finite differences:

$$
\begin{align}
-\left[\partial_{j}\left(\alpha\sqrt{\gamma}A^{j}\right)\right]_{i+\frac{1}{2},j+\frac{1}{2},k+\frac{1}{2}}
&= \frac{\left(\alpha\sqrt{\gamma}A^{x}\right)_{i,j+\frac{1}{2},k+\frac{1}{2}}-\left(\alpha\sqrt{\gamma}A^{x}\right)_{i+1,j+\frac{1}{2},k+\frac{1}{2}}}{dx} \\
&+ \frac{\left(\alpha\sqrt{\gamma}A^{y}\right)_{i+\frac{1}{2},j,k+\frac{1}{2}}-\left(\alpha\sqrt{\gamma}A^{y}\right)_{i+\frac{1}{2},j+1,k+\frac{1}{2}}}{dy} \\
&+ \frac{\left(\alpha\sqrt{\gamma}A^{z}\right)_{i+\frac{1}{2},j+\frac{1}{2},k}-\left(\alpha\sqrt{\gamma}A^{z}\right)_{i+\frac{1}{2},j+\frac{1}{2},k+1}}{dz}
\end{align}
$$

The [*generalized Lorenz gauge*](https://arxiv.org/pdf/1207.3354.pdf) term, $\xi\alpha\left[\sqrt{\gamma}\Phi\right]$, is then trivially implemented.

```python
%%writefile -a $outdir/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

        // Next we add \partial_j (\alpha \sqrt{\gamma} A^j) to \partial_t psi6phi:
        psi6phi_rhsL+=dXm1*(alpha_sqrtg_Ax_interp[index] - alpha_sqrtg_Ax_interp[CCTK_GFINDEX3D(cctkGH,i+1,j,k)])
          +           dYm1*(alpha_sqrtg_Ay_interp[index] - alpha_sqrtg_Ay_interp[CCTK_GFINDEX3D(cctkGH,i,j+1,k)])
          +           dZm1*(alpha_sqrtg_Az_interp[index] - alpha_sqrtg_Az_interp[CCTK_GFINDEX3D(cctkGH,i,j,k+1)]);

        // *GENERALIZED* LORENZ GAUGE:
        // Finally, add damping factor to \partial_t psi6phi
        //subtract lambda * alpha psi^6 Phi
        psi6phi_rhsL+=-damp_lorenz*alpha_iphjphkph[index]*psi6phiL;

        psi6phi_rhs[index] = psi6phi_rhsL;
      }
}

```

    Appending to ../src/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

### Step 1.e: The `avg()` function \[Back to [top](#toc)\]
$$\label{fct_avg}$$

This is the implementation of the algorithm we discussed in [step 1.a.i](#interpolation_algorithm).

```python
%%writefile -a $outdir/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

static inline CCTK_REAL avg(CCTK_REAL f[PLUS2+1][PLUS2+1][PLUS2+1],int imin,int imax, int jmin,int jmax, int kmin,int kmax) {
  CCTK_REAL retval=0.0,num_in_sum=0.0;
  for(int kk=kmin;kk<=kmax;kk++) for(int jj=jmin;jj<=jmax;jj++) for(int ii=imin;ii<=imax;ii++) {
        retval+=f[kk][jj][ii]; num_in_sum++;
      }
  return retval/num_in_sum;
}
```

    Appending to ../src/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C

# Step 2: 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/Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C"
# original_IGM_file_name = "Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs-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__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs__C  = !diff $original_IGM_file_path $outfile_path__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs__C

# if Validation__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs__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 Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.C: PASSED!")
# else:
#     # If the validation fails, we keep the original IGM source code file
#     print("Validation test for Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.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__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs__C:
#         print(diff_line)
```

# Step 3: 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__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.pdf](Tutorial-IllinoisGRMHD__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.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__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.ipynb
#!pdflatex -interaction=batchmode Tutorial-IllinoisGRMHD__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.tex
#!pdflatex -interaction=batchmode Tutorial-IllinoisGRMHD__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.tex
#!pdflatex -interaction=batchmode Tutorial-IllinoisGRMHD__Lorenz_psi6phi_rhs__add_gauge_terms_to_A_i_rhs.tex
!rm -f Tut*.out Tut*.aux Tut*.log
```


##### Copyright 2020 The TensorFlow Authors.

```
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
```

# 量子卷积神经网络

  在 TensorFlow.org 上查看
  在 Google Colab 中运行
  在 GitHub 上查看源代码
  下载笔记本

本教程介绍如何实现一个简化的量子卷积神经网络 (QCNN)，即对同样具有*平移不变性*的经典卷积神经网络的提议量子模拟。

本示例演示如何检测量子数据源的某些属性，例如设备的量子传感器或复杂模拟。量子数据源是可能有或可能没有激发的簇态，后者是 QCNN 将学习检测的对象（论文中使用的数据集是 SPT 阶段分类）。

## 设置

```
!pip install tensorflow==2.4.1
```

安装 TensorFlow Quantum：

```
!pip install tensorflow-quantum
```

```
# Update package resources to account for version changes.
import importlib, pkg_resources
importlib.reload(pkg_resources)
```

现在，导入 TensorFlow 和模块依赖项：

```
import tensorflow as tf
import tensorflow_quantum as tfq

import cirq
import sympy
import numpy as np

# visualization tools
%matplotlib inline
import matplotlib.pyplot as plt
from cirq.contrib.svg import SVGCircuit
```

## 1. 构建 QCNN

### 1.1 在 TensorFlow 计算图中组装电路

TensorFlow Quantum (TFQ) 提供了专为计算图中的电路构造而设计的层类。一个示例是从 `tf.keras.Layer` 继承的 `tfq.layers.AddCircuit` 层。此层可以追加或附加到电路的输入批次，如下图所示。

下面的代码段使用了此层：

```
qubit = cirq.GridQubit(0, 0)

# Define some circuits.
circuit1 = cirq.Circuit(cirq.X(qubit))
circuit2 = cirq.Circuit(cirq.H(qubit))

# Convert to a tensor.
input_circuit_tensor = tfq.convert_to_tensor([circuit1, circuit2])

# Define a circuit that we want to append
y_circuit = cirq.Circuit(cirq.Y(qubit))

# Instantiate our layer
y_appender = tfq.layers.AddCircuit()

# Run our circuit tensor through the layer and save the output.
output_circuit_tensor = y_appender(input_circuit_tensor, append=y_circuit)
```

检查输入张量：

```
print(tfq.from_tensor(input_circuit_tensor))
```

检查输出张量：

```
print(tfq.from_tensor(output_circuit_tensor))
```

虽然不使用 `tfq.layers.AddCircuit` 也可以运行下面的示例，但这是一个理解如何将复杂的功能嵌入 TensorFlow 计算图的好机会。

### 1.2 问题概述

您将准备*簇态*，并训练一个量子分类器来检测它是否处于“激发”状态。簇态是高度纠缠的，不过，这对经典计算机而言并非难以解决的问题。为了让您更清楚地理解，我们使用的这一数据集比论文中使用的更简单。

对于此分类任务，您将实现一个类似深度 MERA 的 QCNN 架构，因为：

1. 就像 QCNN 一样，环上的簇态具有平移不变性。
2. 簇态是高度纠缠的。

这种架构在减少纠缠方面应该会很有效，通过读出一个量子位来获得分类。

根据定义，“激发”簇态是指将 `cirq.rx` 门应用到其任何量子位的簇态。Qconv 和 QPool 在本教程的后续部分讨论。

### 1.3 为 TensorFlow 构建块

使用 TensorFlow Quantum 解决此问题的一种方式是实现以下几点：

1. 模型的输入是一个电路张量，空电路或表明激发的特定量子位上的 X 门。
2. 使用 `tfq.layers.AddCircuit` 层构造模型的其他量子组件。
3. 使用 `tfq.layers.PQC` 层进行推断。它会读取 $\langle \hat{Z} \rangle$，并将其与激发态的标签 1 或非激发态的标签 -1 进行比较。

### 1.4 数据

在构建模型之前，您可以生成数据。在本例中，它将是簇态的激发（原论文使用的是一个更复杂的数据集）。激发通过 `cirq.rx` 门表示。我们将足够大的旋转视为激发，并使用 `1` 进行标记，不够大的旋转则使用 `-1` 标记，我们将其视为未激发。

```
def generate_data(qubits):
    """Generate training and testing data."""
    n_rounds = 20  # Produces n_rounds * n_qubits datapoints.
    excitations = []
    labels = []
    for n in range(n_rounds):
        for bit in qubits:
            rng = np.random.uniform(-np.pi, np.pi)
            excitations.append(cirq.Circuit(cirq.rx(rng)(bit)))
            labels.append(1 if (-np.pi / 2) <= rng <= (np.pi / 2) else -1)

    split_ind = int(len(excitations) * 0.7)
    train_excitations = excitations[:split_ind]
    test_excitations = excitations[split_ind:]

    train_labels = labels[:split_ind]
    test_labels = labels[split_ind:]

    return tfq.convert_to_tensor(train_excitations), np.array(train_labels), \
        tfq.convert_to_tensor(test_excitations), np.array(test_labels)
```

您可以看到，就像使用常规的机器学习一样，您创建了一个用于对模型进行基准测试的训练和测试集。利用以下代码段，您可以快速查看某些数据点：

```
sample_points, sample_labels, _, __ = generate_data(cirq.GridQubit.rect(1, 4))
print('Input:', tfq.from_tensor(sample_points)[0], 'Output:', sample_labels[0])
print('Input:', tfq.from_tensor(sample_points)[1], 'Output:', sample_labels[1])
```

### 1.5 定义层

现在，我们在 TensorFlow 中定义上图中显示的层。

#### 1.5.1 簇态

第一步是使用 Cirq（Google 为量子电路编程提供的框架）定义簇态。由于这是模型的一个静态部分，因此使用 `tfq.layers.AddCircuit` 功能将其嵌入。

```
def cluster_state_circuit(bits):
    """Return a cluster state on the qubits in `bits`."""
    circuit = cirq.Circuit()
    circuit.append(cirq.H.on_each(bits))
    for this_bit, next_bit in zip(bits, bits[1:] + [bits[0]]):
        circuit.append(cirq.CZ(this_bit, next_bit))
    return circuit
```

显示 cirq.GridQubit 的矩形的一个簇态电路：

```
SVGCircuit(cluster_state_circuit(cirq.GridQubit.rect(1, 4)))
```

#### 1.5.2 QCNN 层

按照 Cong 和 Lukin QCNN 论文定义组成模型的层。这需要具备几个前提条件：

- Tucci 论文中提出的单或双量子位参数化酉矩阵。
- 一个通用的参数化双量子位池化运算。

```
def one_qubit_unitary(bit, symbols):
    """Make a Cirq circuit enacting a rotation of the bloch sphere about the X,
    Y and Z axis, that depends on the values in `symbols`.
    """
    return cirq.Circuit(
        cirq.X(bit)**symbols[0],
        cirq.Y(bit)**symbols[1],
        cirq.Z(bit)**symbols[2])

def two_qubit_unitary(bits, symbols):
    """Make a Cirq circuit that creates an arbitrary two qubit unitary."""
    circuit = cirq.Circuit()
    circuit += one_qubit_unitary(bits[0], symbols[0:3])
    circuit += one_qubit_unitary(bits[1], symbols[3:6])
    circuit += [cirq.ZZ(*bits)**symbols[6]]
    circuit += [cirq.YY(*bits)**symbols[7]]
    circuit += [cirq.XX(*bits)**symbols[8]]
    circuit += one_qubit_unitary(bits[0], symbols[9:12])
    circuit += one_qubit_unitary(bits[1], symbols[12:])
    return circuit

def two_qubit_pool(source_qubit, sink_qubit, symbols):
    """Make a Cirq circuit to do a parameterized 'pooling' operation, which
    attempts to reduce entanglement down from two qubits to just one."""
    pool_circuit = cirq.Circuit()
    sink_basis_selector = one_qubit_unitary(sink_qubit, symbols[0:3])
    source_basis_selector = one_qubit_unitary(source_qubit, symbols[3:6])
    pool_circuit.append(sink_basis_selector)
    pool_circuit.append(source_basis_selector)
    pool_circuit.append(cirq.CNOT(control=source_qubit, target=sink_qubit))
    pool_circuit.append(sink_basis_selector**-1)
    return pool_circuit
```

要查看您创建的对象，请打印出单量子位酉电路：

```
SVGCircuit(one_qubit_unitary(cirq.GridQubit(0, 0), sympy.symbols('x0:3')))
```

以及双量子位酉电路：

```
SVGCircuit(two_qubit_unitary(cirq.GridQubit.rect(1, 2), sympy.symbols('x0:15')))
```

以及双量子位池化电路：

```
SVGCircuit(two_qubit_pool(*cirq.GridQubit.rect(1, 2), sympy.symbols('x0:6')))
```

##### 1.5.2.1 量子卷积

按照 Cong 和 Lukin 的论文所述，将一维量子卷积定义为对每对步长为 1 的相邻量子位的双量子位参数化酉的应用。

```
def quantum_conv_circuit(bits, symbols):
    """Quantum Convolution Layer following the above diagram.
    Return a Cirq circuit with the cascade of `two_qubit_unitary` applied
    to all pairs of qubits in `bits` as in the diagram above.
    """
    circuit = cirq.Circuit()
    for first, second in zip(bits[0::2], bits[1::2]):
        circuit += two_qubit_unitary([first, second], symbols)
    for first, second in zip(bits[1::2], bits[2::2] + [bits[0]]):
        circuit += two_qubit_unitary([first, second], symbols)
    return circuit
```

显示（高度水平的）电路：

```
SVGCircuit(
    quantum_conv_circuit(cirq.GridQubit.rect(1, 8), sympy.symbols('x0:15')))
```

##### 1.5.2.2 量子池化

量子池化层使用上面定义的双量子位池从 $N$ 个量子位池化为 $\frac{N}{2}$ 个量子位。

```
def quantum_pool_circuit(source_bits, sink_bits, symbols):
    """A layer that specifies a quantum pooling operation.
    A Quantum pool tries to learn to pool the relevant information from two
    qubits onto 1.
    """
    circuit = cirq.Circuit()
    for source, sink in zip(source_bits, sink_bits):
        circuit += two_qubit_pool(source, sink, symbols)
    return circuit
```

检查池化组件电路：

```
test_bits = cirq.GridQubit.rect(1, 8)

SVGCircuit(
    quantum_pool_circuit(test_bits[:4], test_bits[4:], sympy.symbols('x0:6')))
```

### 1.6 模型定义

现在，使用定义的层构造纯量子 CNN。首先创建八个量子位，再将其池化为一个量子位，然后测量 $\langle \hat{Z} \rangle$。

```
def create_model_circuit(qubits):
    """Create sequence of alternating convolution and pooling operators
    which gradually shrink over time."""
    model_circuit = cirq.Circuit()
    symbols = sympy.symbols('qconv0:63')
    # Cirq uses sympy.Symbols to map learnable variables. TensorFlow Quantum
    # scans incoming circuits and replaces these with TensorFlow variables.
    model_circuit += quantum_conv_circuit(qubits, symbols[0:15])
    model_circuit += quantum_pool_circuit(qubits[:4], qubits[4:],
                                          symbols[15:21])
    model_circuit += quantum_conv_circuit(qubits[4:], symbols[21:36])
    model_circuit += quantum_pool_circuit(qubits[4:6], qubits[6:],
                                          symbols[36:42])
    model_circuit += quantum_conv_circuit(qubits[6:], symbols[42:57])
    model_circuit += quantum_pool_circuit([qubits[6]], [qubits[7]],
                                          symbols[57:63])
    return model_circuit

# Create our qubits and readout operators in Cirq.
cluster_state_bits = cirq.GridQubit.rect(1, 8)
readout_operators = cirq.Z(cluster_state_bits[-1])

# Build a sequential model enacting the logic in 1.3 of this notebook.
# Here you are making the static cluster state prep as a part of the AddCircuit and the
# "quantum datapoints" are coming in the form of excitation
excitation_input = tf.keras.Input(shape=(), dtype=tf.dtypes.string)
cluster_state = tfq.layers.AddCircuit()(
    excitation_input, prepend=cluster_state_circuit(cluster_state_bits))

quantum_model = tfq.layers.PQC(create_model_circuit(cluster_state_bits),
                               readout_operators)(cluster_state)

qcnn_model = tf.keras.Model(inputs=[excitation_input], outputs=[quantum_model])

# Show the keras plot of the model
tf.keras.utils.plot_model(qcnn_model,
                          show_shapes=True,
                          show_layer_names=False,
                          dpi=70)
```

### 1.7 训练模型

在整个批次上训练模型以简化此示例。

```
# Generate some training data.
train_excitations, train_labels, test_excitations, test_labels = generate_data(
    cluster_state_bits)

# Custom accuracy metric.
@tf.function
def custom_accuracy(y_true, y_pred):
    y_true = tf.squeeze(y_true)
    y_pred = tf.map_fn(lambda x: 1.0 if x >= 0 else -1.0, y_pred)
    return tf.keras.backend.mean(tf.keras.backend.equal(y_true, y_pred))

qcnn_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
                   loss=tf.losses.mse,
                   metrics=[custom_accuracy])

history = qcnn_model.fit(x=train_excitations,
                         y=train_labels,
                         batch_size=16,
                         epochs=25,
                         verbose=1,
                         validation_data=(test_excitations, test_labels))
```

```
plt.plot(history.history['loss'][1:], label='Training')
plt.plot(history.history['val_loss'][1:], label='Validation')
plt.title('Training a Quantum CNN to Detect Excited Cluster States')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()
```

## 2. 混合模型

您不必使用量子卷积将八个量子位池化为一个量子位，您可以执行一到两轮的量子卷积，然后将结果馈送到经典神经网络中。本部分探讨量子-经典混合模型。

### 2.1 使用单个量子滤波器的混合模型

应用一层量子卷积，在后跟密集连接的神经网络的所有位上读出 $\langle \hat{Z}_n \rangle$。

#### 2.1.1 模型定义

```
# 1-local operators to read out
readouts = [cirq.Z(bit) for bit in cluster_state_bits[4:]]

def multi_readout_model_circuit(qubits):
    """Make a model circuit with less quantum pool and conv operations."""
    model_circuit = cirq.Circuit()
    symbols = sympy.symbols('qconv0:21')
    model_circuit += quantum_conv_circuit(qubits, symbols[0:15])
    model_circuit += quantum_pool_circuit(qubits[:4], qubits[4:],
                                          symbols[15:21])
    return model_circuit

# Build a model enacting the logic in 2.1 of this notebook.
excitation_input_dual = tf.keras.Input(shape=(), dtype=tf.dtypes.string)

cluster_state_dual = tfq.layers.AddCircuit()(
    excitation_input_dual, prepend=cluster_state_circuit(cluster_state_bits))

quantum_model_dual = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_dual)

d1_dual = tf.keras.layers.Dense(8)(quantum_model_dual)

d2_dual = tf.keras.layers.Dense(1)(d1_dual)

hybrid_model = tf.keras.Model(inputs=[excitation_input_dual], outputs=[d2_dual])

# Display the model architecture
tf.keras.utils.plot_model(hybrid_model,
                          show_shapes=True,
                          show_layer_names=False,
                          dpi=70)
```

#### 2.1.2 训练模型

```
hybrid_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
                     loss=tf.losses.mse,
                     metrics=[custom_accuracy])

hybrid_history = hybrid_model.fit(x=train_excitations,
                                  y=train_labels,
                                  batch_size=16,
                                  epochs=25,
                                  verbose=1,
                                  validation_data=(test_excitations,
                                                   test_labels))
```

```
plt.plot(history.history['val_custom_accuracy'], label='QCNN')
plt.plot(hybrid_history.history['val_custom_accuracy'], label='Hybrid CNN')
plt.title('Quantum vs Hybrid CNN performance')
plt.xlabel('Epochs')
plt.legend()
plt.ylabel('Validation Accuracy')
plt.show()
```

如您所见，在适当的经典模型的帮助下，混合模型通常比纯量子版本收敛得更快。

### 2.2 使用多个量子滤波器的混合卷积

现在，我们来试试使用多个量子卷积和一个经典神经网络的架构，将其组合在一起。

#### 2.2.1 模型定义

```
excitation_input_multi = tf.keras.Input(shape=(), dtype=tf.dtypes.string)

cluster_state_multi = tfq.layers.AddCircuit()(
    excitation_input_multi, prepend=cluster_state_circuit(cluster_state_bits))

# apply 3 different filters and measure expectation values

quantum_model_multi1 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

quantum_model_multi2 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

quantum_model_multi3 = tfq.layers.PQC(
    multi_readout_model_circuit(cluster_state_bits),
    readouts)(cluster_state_multi)

# concatenate outputs and feed into a small classical NN
concat_out = tf.keras.layers.concatenate(
    [quantum_model_multi1, quantum_model_multi2, quantum_model_multi3])

dense_1 = tf.keras.layers.Dense(8)(concat_out)

dense_2 = tf.keras.layers.Dense(1)(dense_1)

multi_qconv_model = tf.keras.Model(inputs=[excitation_input_multi],
                                   outputs=[dense_2])

# Display the model architecture
tf.keras.utils.plot_model(multi_qconv_model,
                          show_shapes=True,
                          show_layer_names=True,
                          dpi=70)
```

#### 2.2.2 训练模型

```
multi_qconv_model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),
    loss=tf.losses.mse,
    metrics=[custom_accuracy])

multi_qconv_history = multi_qconv_model.fit(x=train_excitations,
                                            y=train_labels,
                                            batch_size=16,
                                            epochs=25,
                                            verbose=1,
                                            validation_data=(test_excitations,
                                                             test_labels))
```

```
plt.plot(history.history['val_custom_accuracy'][:25], label='QCNN')
plt.plot(hybrid_history.history['val_custom_accuracy'][:25], label='Hybrid CNN')
plt.plot(multi_qconv_history.history['val_custom_accuracy'][:25],
         label='Hybrid CNN \n Multiple Quantum Filters')
plt.title('Quantum vs Hybrid CNN performance')
plt.xlabel('Epochs')
plt.legend()
plt.ylabel('Validation Accuracy')
plt.show()
```


```python
#we may need some code in the ../python directory and/or matplotlib styles
import sys
sys.path.append('../python/')

#matplotlib for plotting
import matplotlib as mpl
from matplotlib import pyplot as plt
plt.style.use('mplstyles/stylelib/standard.mplstyle')

#other computational libraries
import numpy as np
```

```python
npegs=27
nboxes=36
pn = np.zeros((npegs,))

pn[0]=36-9
pn[3]=9

pn=pn/36

print(pn)
```

    [0.75 0.   0.   0.25 0.   0.   0.   0.   0.   0.   0.   0.   0.   0.
     0.   0.   0.   0.   0.   0.   0.   0.   0.   0.   0.   0.   0.  ]

```python
def fairness(v):
    retval=0
    for i,p in enumerate(v):
        if (i>0) & (p>0):
            retval+=p*np.log(p)
    return -10*retval

print(fairness(pn))
```

    3.4657359027997265

```python
#looking for a data structure to represent pegs and grid
grid = np.zeros((6,6))

print(grid)
```

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

```python
grid[0,0] = 3
grid[0,1] = 3
grid[0,2] = 3

grid[1,0] = 3
grid[1,1] = 3
grid[1,2] = 3

grid[2,0] = 3
grid[2,1] = 3
grid[2,2] = 3
```

```python
print(grid)
```

    [[3. 3. 3. 0. 0. 0.]
     [3. 3. 3. 0. 0. 0.]
     [3. 3. 3. 0. 0. 0.]
     [0. 0. 0. 0. 0. 0.]
     [0. 0. 0. 0. 0. 0.]
     [0. 0. 0. 0. 0. 0.]]

```python
#find out how many energy pegs total in the grid structure
print(np.sum(grid))
```

    27.0

```python
#print out every individual box value
for v in grid:
    for n in v:
      print(n)
```

    3.0
    3.0
    3.0
    0.0
    0.0
    0.0
    3.0
    3.0
    3.0
    0.0
    0.0
    0.0
    3.0
    3.0
    3.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0
    0.0

```python
#get total number of grid spaces
print(grid.size)
```

    36

Explanation of Micro/Macrostates is here.

\begin{equation}
S = k\log{\Omega}
\end{equation}

```python

```


# Trax : Ungraded Lecture Notebook

In this notebook you'll get to know about the Trax framework and learn about some of its basic building blocks.

## Background

### Why Trax and not TensorFlow or PyTorch?

TensorFlow and PyTorch are both extensive frameworks that can do almost anything in deep learning. They offer a lot of flexibility, but that often means verbosity of syntax and extra time to code.

Trax is much more concise. It runs on a TensorFlow backend but allows you to train models with 1 line commands. Trax also runs end to end, allowing you to get data, model and train all with a single terse statements. This means you can focus on learning, instead of spending hours on the idiosyncrasies of big framework implementation.

### Why not Keras then?

Keras is now part of Tensorflow itself from 2.0 onwards. Also, trax is good for implementing new state of the art algorithms like Transformers, Reformers, BERT because it is actively maintained by Google Brain Team for advanced deep learning tasks. It runs smoothly on CPUs,GPUs and TPUs as well with comparatively lesser modifications in code.

### How to Code in Trax
Building models in Trax relies on 2 key concepts:- **layers** and **combinators**.
Trax layers are simple objects that process data and perform computations. They can be chained together into composite layers using Trax combinators, allowing you to build layers and models of any complexity.

### Trax, JAX, TensorFlow and Tensor2Tensor

You already know that Trax uses Tensorflow as a backend, but it also uses the JAX library to speed up computation too. You can view JAX as an enhanced and optimized version of numpy.

**Watch out for assignments which import `import trax.fastmath.numpy as np`. If you see this line, remember that when calling `np` you are really calling Trax’s version of numpy that is compatible with JAX.**

As a result of this, where you used to encounter the type `numpy.ndarray` now you will find the type `jax.interpreters.xla.DeviceArray`.

Tensor2Tensor is another name you might have heard. It started as an end to end solution much like how Trax is designed, but it grew unwieldy and complicated. So you can view Trax as the new improved version that operates much faster and simpler.

### Resources

- Trax source code can be found on Github: [Trax](https://github.com/google/trax)
- JAX library: [JAX](https://jax.readthedocs.io/en/latest/index.html)

## Installing Trax

Trax has dependencies on JAX and some libraries like JAX which are yet to be supported in [Windows](https://github.com/google/jax/blob/1bc5896ee4eab5d7bb4ec6f161d8b2abb30557be/README.md#installation) but work well in Ubuntu and MacOS. We would suggest that if you are working on Windows, try to install Trax on WSL2.

Official maintained documentation - [trax-ml](https://trax-ml.readthedocs.io/en/latest/) not to be confused with this [TraX](https://trax.readthedocs.io/en/latest/index.html)

```python
!pip install trax==1.3.1 # Use this version for this notebook
```

    Requirement already satisfied: trax==1.3.1 in /opt/conda/lib/python3.7/site-packages (1.3.1)
    Requirement already satisfied: scipy in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (1.2.1)
    Requirement already satisfied: t5 in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (0.6.1)
    Requirement already satisfied: gin-config in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (0.3.0)
    Requirement already satisfied: numpy in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (1.18.5)
    Requirement already satisfied: gym in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (0.17.2)
    Requirement already satisfied: jax in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (0.1.72)
    Requirement already satisfied: jaxlib in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (0.1.51)
    Requirement already satisfied: six in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (1.14.0)
    Requirement already satisfied: absl-py in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (0.9.0)
    Requirement already satisfied: tensorflow-text in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (2.2.1)
    Requirement already satisfied: tensorflow-datasets in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (2.1.0)
    Requirement already satisfied: tensor2tensor in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (1.15.7)
    Requirement already satisfied: funcsigs in /opt/conda/lib/python3.7/site-packages (from trax==1.3.1) (1.0.2)
    Requirement already satisfied: torch in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (1.5.1)
    Requirement already satisfied: nltk in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (3.5)
    Requirement already satisfied: babel in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (2.8.0)
    Requirement already satisfied: rouge-score in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (0.0.4)
    Requirement already satisfied: sentencepiece in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (0.1.91)
    Requirement already satisfied: tfds-nightly in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (3.1.0.dev202007020105)
    Requirement already satisfied: scikit-learn in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (0.20.3)
    Requirement already satisfied: mesh-tensorflow[transformer]>=0.1.13 in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (0.1.16)
    Requirement already satisfied: pandas in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (0.24.2)
    Requirement already satisfied: sacrebleu in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (1.4.10)
    Requirement already satisfied: transformers>=2.7.0 in /opt/conda/lib/python3.7/site-packages (from t5->trax==1.3.1) (3.0.0)
    Requirement already satisfied: pyglet<=1.5.0,>=1.4.0 in /opt/conda/lib/python3.7/site-packages (from gym->trax==1.3.1) (1.5.0)
    Requirement already satisfied: cloudpickle<1.4.0,>=1.2.0 in /opt/conda/lib/python3.7/site-packages (from gym->trax==1.3.1) (1.3.0)
    Requirement already satisfied: opt-einsum in /opt/conda/lib/python3.7/site-packages (from jax->trax==1.3.1) (3.2.1)
    Requirement already satisfied: tensorflow<2.3,>=2.2.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-text->trax==1.3.1) (2.2.0)
    Requirement already satisfied: attrs>=18.1.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (19.1.0)
    Requirement already satisfied: termcolor in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (1.1.0)
    Requirement already satisfied: dill in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (0.2.9)
    Requirement already satisfied: requests>=2.19.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (2.21.0)
    Requirement already satisfied: protobuf>=3.6.1 in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (3.7.0)
    Requirement already satisfied: future in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (0.18.2)
    Requirement already satisfied: wrapt in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (1.12.1)
    Requirement already satisfied: tqdm in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (4.47.0)
    Requirement already satisfied: promise in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (2.3)
    Requirement already satisfied: tensorflow-metadata in /opt/conda/lib/python3.7/site-packages (from tensorflow-datasets->trax==1.3.1) (0.22.2)
    Requirement already satisfied: tensorflow-gan in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (2.0.0)
    Requirement already satisfied: flask in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (1.1.2)
    Requirement already satisfied: tf-slim in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (1.1.0)
    Requirement already satisfied: opencv-python in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (4.2.0.34)
    Requirement already satisfied: gevent in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (20.6.2)
    Requirement already satisfied: kfac in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (0.2.0)
    Requirement already satisfied: tensorflow-probability==0.7.0 in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (0.7.0)
    Requirement already satisfied: oauth2client in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (4.1.3)
    Requirement already satisfied: gunicorn in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (20.0.4)
    Requirement already satisfied: dopamine-rl in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (3.0.1)
    Requirement already satisfied: pypng in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (0.0.20)
    Requirement already satisfied: tensorflow-addons in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (0.10.0)
    Requirement already satisfied: h5py in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (2.9.0)
    Requirement already satisfied: sympy in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (1.3)
    Requirement already satisfied: google-api-python-client in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (1.9.3)
    Requirement already satisfied: bz2file in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (0.98)
    Requirement already satisfied: Pillow in /opt/conda/lib/python3.7/site-packages (from tensor2tensor->trax==1.3.1) (5.4.1)
    Requirement already satisfied: joblib in /opt/conda/lib/python3.7/site-packages (from nltk->t5->trax==1.3.1) (0.16.0)
    Requirement already satisfied: click in /opt/conda/lib/python3.7/site-packages (from nltk->t5->trax==1.3.1) (7.0)
    Requirement already satisfied: regex in /opt/conda/lib/python3.7/site-packages (from nltk->t5->trax==1.3.1) (2020.6.8)
    Requirement already satisfied: pytz>=2015.7 in /opt/conda/lib/python3.7/site-packages (from babel->t5->trax==1.3.1) (2018.9)
    Requirement already satisfied: python-dateutil>=2.5.0 in /opt/conda/lib/python3.7/site-packages (from pandas->t5->trax==1.3.1) (2.8.0)
    Requirement already satisfied: portalocker in /opt/conda/lib/python3.7/site-packages (from sacrebleu->t5->trax==1.3.1) (1.7.0)
    Requirement already satisfied: packaging in /opt/conda/lib/python3.7/site-packages (from transformers>=2.7.0->t5->trax==1.3.1) (19.0)
    Requirement already satisfied: sacremoses in /opt/conda/lib/python3.7/site-packages (from transformers>=2.7.0->t5->trax==1.3.1) (0.0.43)
    Requirement already satisfied: tokenizers==0.8.0-rc4 in /opt/conda/lib/python3.7/site-packages (from transformers>=2.7.0->t5->trax==1.3.1) (0.8.0rc4)
    Requirement already satisfied: filelock in /opt/conda/lib/python3.7/site-packages (from transformers>=2.7.0->t5->trax==1.3.1) (3.0.12)
    Requirement already satisfied: astunparse==1.6.3 in /opt/conda/lib/python3.7/site-packages (from tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (1.6.3)
    Requirement already satisfied: tensorboard<2.3.0,>=2.2.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (2.2.2)
    Requirement already satisfied: google-pasta>=0.1.8 in /opt/conda/lib/python3.7/site-packages (from tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (0.2.0)
    Requirement already satisfied: gast==0.3.3 in /opt/conda/lib/python3.7/site-packages (from tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (0.3.3)
    Requirement already satisfied: tensorflow-estimator<2.3.0,>=2.2.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (2.2.0)
    Requirement already satisfied: keras-preprocessing>=1.1.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (1.1.2)
    Requirement already satisfied: grpcio>=1.8.6 in /opt/conda/lib/python3.7/site-packages (from tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (1.30.0)
    Requirement already satisfied: wheel>=0.26; python_version >= "3" in /opt/conda/lib/python3.7/site-packages (from tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (0.33.1)
    Requirement already satisfied: urllib3<1.25,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->tensorflow-datasets->trax==1.3.1) (1.24.1)
    Requirement already satisfied: chardet<3.1.0,>=3.0.2 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->tensorflow-datasets->trax==1.3.1) (3.0.4)
    Requirement already satisfied: idna<2.9,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->tensorflow-datasets->trax==1.3.1) (2.8)
    Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests>=2.19.0->tensorflow-datasets->trax==1.3.1) (2019.3.9)
    Requirement already satisfied: setuptools in /opt/conda/lib/python3.7/site-packages (from protobuf>=3.6.1->tensorflow-datasets->trax==1.3.1) (41.2.0)
    Requirement already satisfied: googleapis-common-protos in /opt/conda/lib/python3.7/site-packages (from tensorflow-metadata->tensorflow-datasets->trax==1.3.1) (1.52.0)
    Requirement already satisfied: tensorflow-hub>=0.2 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gan->tensor2tensor->trax==1.3.1) (0.8.0)
    Requirement already satisfied: itsdangerous>=0.24 in /opt/conda/lib/python3.7/site-packages (from flask->tensor2tensor->trax==1.3.1) (1.1.0)
    Requirement already satisfied: Jinja2>=2.10.1 in /opt/conda/lib/python3.7/site-packages (from flask->tensor2tensor->trax==1.3.1) (2.11.2)
    Requirement already satisfied: Werkzeug>=0.15 in /opt/conda/lib/python3.7/site-packages (from flask->tensor2tensor->trax==1.3.1) (1.0.1)
    Requirement already satisfied: zope.interface in /opt/conda/lib/python3.7/site-packages (from gevent->tensor2tensor->trax==1.3.1) (5.1.0)
    Requirement already satisfied: zope.event in /opt/conda/lib/python3.7/site-packages (from gevent->tensor2tensor->trax==1.3.1) (4.4)
    Requirement already satisfied: greenlet>=0.4.16; platform_python_implementation == "CPython" in /opt/conda/lib/python3.7/site-packages (from gevent->tensor2tensor->trax==1.3.1) (0.4.16)
    Requirement already satisfied: decorator in /opt/conda/lib/python3.7/site-packages (from tensorflow-probability==0.7.0->tensor2tensor->trax==1.3.1) (4.3.2)
    Requirement already satisfied: httplib2>=0.9.1 in /opt/conda/lib/python3.7/site-packages (from oauth2client->tensor2tensor->trax==1.3.1) (0.18.1)
    Requirement already satisfied: pyasn1>=0.1.7 in /opt/conda/lib/python3.7/site-packages (from oauth2client->tensor2tensor->trax==1.3.1) (0.4.8)
    Requirement already satisfied: rsa>=3.1.4 in /opt/conda/lib/python3.7/site-packages (from oauth2client->tensor2tensor->trax==1.3.1) (4.6)
    Requirement already satisfied: pyasn1-modules>=0.0.5 in /opt/conda/lib/python3.7/site-packages (from oauth2client->tensor2tensor->trax==1.3.1) (0.2.8)
    Requirement already satisfied: typeguard>=2.7 in /opt/conda/lib/python3.7/site-packages (from tensorflow-addons->tensor2tensor->trax==1.3.1) (2.9.1)
    Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.7/site-packages (from sympy->tensor2tensor->trax==1.3.1) (1.1.0)
    Requirement already satisfied: google-api-core<2dev,>=1.18.0 in /opt/conda/lib/python3.7/site-packages (from google-api-python-client->tensor2tensor->trax==1.3.1) (1.21.0)
    Requirement already satisfied: google-auth>=1.16.0 in /opt/conda/lib/python3.7/site-packages (from google-api-python-client->tensor2tensor->trax==1.3.1) (1.18.0)
    Requirement already satisfied: google-auth-httplib2>=0.0.3 in /opt/conda/lib/python3.7/site-packages (from google-api-python-client->tensor2tensor->trax==1.3.1) (0.0.3)
    Requirement already satisfied: uritemplate<4dev,>=3.0.0 in /opt/conda/lib/python3.7/site-packages (from google-api-python-client->tensor2tensor->trax==1.3.1) (3.0.1)
    Requirement already satisfied: pyparsing>=2.0.2 in /opt/conda/lib/python3.7/site-packages (from packaging->transformers>=2.7.0->t5->trax==1.3.1) (2.3.1)
    Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (0.4.1)
    Requirement already satisfied: markdown>=2.6.8 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (3.2.2)
    Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (1.7.0)
    Requirement already satisfied: MarkupSafe>=0.23 in /opt/conda/lib/python3.7/site-packages (from Jinja2>=2.10.1->flask->tensor2tensor->trax==1.3.1) (1.1.1)
    Requirement already satisfied: cachetools<5.0,>=2.0.0 in /opt/conda/lib/python3.7/site-packages (from google-auth>=1.16.0->google-api-python-client->tensor2tensor->trax==1.3.1) (4.1.1)
    Requirement already satisfied: requests-oauthlib>=0.7.0 in /opt/conda/lib/python3.7/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (1.3.0)
    Requirement already satisfied: importlib-metadata; python_version < "3.8" in /opt/conda/lib/python3.7/site-packages (from markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (1.7.0)
    Requirement already satisfied: oauthlib>=3.0.0 in /opt/conda/lib/python3.7/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (3.1.0)
    Requirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.7/site-packages (from importlib-metadata; python_version < "3.8"->markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow<2.3,>=2.2.0->tensorflow-text->trax==1.3.1) (3.1.0)
    WARNING: You are using pip version 20.1.1; however, version 20.2 is available.
    You should consider upgrading via the '/opt/conda/bin/python -m pip install --upgrade pip' command.

## Imports

```python
import numpy as np  # regular ol' numpy

from trax import layers as tl  # core building block
from trax import shapes  # data signatures: dimensionality and type
from trax import fastmath  # uses jax, offers numpy on steroids
```

    INFO:tensorflow:tokens_length=568 inputs_length=512 targets_length=114 noise_density=0.15 mean_noise_span_length=3.0

```python
# Trax version 1.3.1 or better
!pip list | grep trax
```

    trax                     1.3.1
    WARNING: You are using pip version 20.1.1; however, version 20.2 is available.
    You should consider upgrading via the '/opt/conda/bin/python -m pip install --upgrade pip' command.

## Layers
Layers are the core building blocks in Trax or as mentioned in the lectures, they are the base classes.

They take inputs, compute functions/custom calculations and return outputs.

You can also inspect layer properties. Let me show you some examples.

### Relu Layer
First I'll show you how to build a relu activation function as a layer. A layer like this is one of the simplest types. Notice there is no object initialization so it works just like a math function.

**Note: Activation functions are also layers in Trax, which might look odd if you have been using other frameworks for a longer time.**

```python
# Layers
# Create a relu trax layer
relu = tl.Relu()

# Inspect properties
print("-- Properties --")
print("name :", relu.name)
print("expected inputs :", relu.n_in)
print("promised outputs :", relu.n_out, "\n")

# Inputs
x = np.array([-2, -1, 0, 1, 2])
print("-- Inputs --")
print("x :", x, "\n")

# Outputs
y = relu(x)
print("-- Outputs --")
print("y :", y)
```

    -- Properties --
    name : Relu
    expected inputs : 1
    promised outputs : 1

    -- Inputs --
    x : [-2 -1  0  1  2]

    -- Outputs --
    y : [0 0 0 1 2]

### Concatenate Layer
Now I'll show you how to build a layer that takes 2 inputs. Notice the change in the expected inputs property from 1 to 2.

```python
# Create a concatenate trax layer
concat = tl.Concatenate()
print("-- Properties --")
print("name :", concat.name)
print("expected inputs :", concat.n_in)
print("promised outputs :", concat.n_out, "\n")

# Inputs
x1 = np.array([-10, -20, -30])
x2 = x1 / -10
print("-- Inputs --")
print("x1 :", x1)
print("x2 :", x2, "\n")

# Outputs
y = concat([x1, x2])
print("-- Outputs --")
print("y :", y)
```

    -- Properties --
    name : Concatenate
    expected inputs : 2
    promised outputs : 1

    -- Inputs --
    x1 : [-10 -20 -30]
    x2 : [1. 2. 3.]

    -- Outputs --
    y : [-10. -20. -30.   1.   2.   3.]

## Layers are Configurable
You can change the default settings of layers. For example, you can change the expected inputs for a concatenate layer from 2 to 3 using the optional parameter `n_items`.

```python
# Configure a concatenate layer
concat_3 = tl.Concatenate(n_items=3)  # configure the layer's expected inputs
print("-- Properties --")
print("name :", concat_3.name)
print("expected inputs :", concat_3.n_in)
print("promised outputs :", concat_3.n_out, "\n")

# Inputs
x1 = np.array([-10, -20, -30])
x2 = x1 / -10
x3 = x2 * 0.99
print("-- Inputs --")
print("x1 :", x1)
print("x2 :", x2)
print("x3 :", x3, "\n")

# Outputs
y = concat_3([x1, x2, x3])
print("-- Outputs --")
print("y :", y)
```

    -- Properties --
    name : Concatenate
    expected inputs : 3
    promised outputs : 1

    -- Inputs --
    x1 : [-10 -20 -30]
    x2 : [1. 2. 3.]
    x3 : [0.99 1.98 2.97]

    -- Outputs --
    y : [-10.   -20.   -30.     1.     2.     3.     0.99   1.98   2.97]

**Note: At any point,if you want to refer the function help/ look up the [documentation](https://trax-ml.readthedocs.io/en/latest/) or use help function.**

```python
#help(tl.Concatenate) #Uncomment this to see the function docstring with explaination
```

## Layers can have Weights
Some layer types include mutable weights and biases that are used in computation and training. Layers of this type require initialization before use.

For example the `LayerNorm` layer calculates normalized data, that is also scaled by weights and biases. During initialization you pass the data shape and data type of the inputs, so the layer can initialize compatible arrays of weights and biases.

```python
# Uncomment any of them to see information regarding the function
# help(tl.LayerNorm)
# help(shapes.signature)

```

```python
# Layer initialization
norm = tl.LayerNorm()
# You first must know what the input data will look like
x = np.array([0, 1, 2, 3], dtype="float")

# Use the input data signature to get shape and type for initializing weights and biases
norm.init(shapes.signature(x)) # We need to convert the input datatype from usual tuple to trax ShapeDtype

print("Normal shape:",x.shape, "Data Type:",type(x.shape))
print("Shapes Trax:",shapes.signature(x),"Data Type:",type(shapes.signature(x)))

# Inspect properties
print("-- Properties --")
print("name :", norm.name)
print("expected inputs :", norm.n_in)
print("promised outputs :", norm.n_out)
# Weights and biases
print("weights :", norm.weights[0])
print("biases :", norm.weights[1], "\n")

# Inputs
print("-- Inputs --")
print("x :", x)

# Outputs
y = norm(x)
print("-- Outputs --")
print("y :", y)
```

    Normal shape: (4,) Data Type: <class 'tuple'>
    Shapes Trax: ShapeDtype{shape:(4,), dtype:float64} Data Type: <class 'trax.shapes.ShapeDtype'>
    -- Properties --
    name : LayerNorm
    expected inputs : 1
    promised outputs : 1
    weights : [1. 1. 1. 1.]
    biases : [0. 0. 0. 0.]

    -- Inputs --
    x : [0. 1. 2. 3.]
    -- Outputs --
    y : [-1.3416404  -0.44721344  0.44721344  1.3416404 ]

## Custom Layers
This is where things start getting more interesting!
You can create your own custom layers too and define custom functions for computations by using `tl.Fn`. Let me show you how.

```python
help(tl.Fn)
```

```python
# Define a custom layer
# In this example you will create a layer to calculate the input times 2

def TimesTwo():
    layer_name = "TimesTwo" #don't forget to give your custom layer a name to identify

    # Custom function for the custom layer
    def func(x):
        return x * 2

    return tl.Fn(layer_name, func)

# Test it
times_two = TimesTwo()

# Inspect properties
print("-- Properties --")
print("name :", times_two.name)
print("expected inputs :", times_two.n_in)
print("promised outputs :", times_two.n_out, "\n")

# Inputs
x = np.array([1, 2, 3])
print("-- Inputs --")
print("x :", x, "\n")

# Outputs
y = times_two(x)
print("-- Outputs --")
print("y :", y)
```

    -- Properties --
    name : TimesTwo
    expected inputs : 1
    promised outputs : 1

    -- Inputs --
    x : [1 2 3]

    -- Outputs --
    y : [2 4 6]

## Combinators
You can combine layers to build more complex layers. Trax provides a set of objects named combinator layers to make this happen. Combinators are themselves layers, so behavior commutes.

### Serial Combinator
This is the most common and easiest to use. For example could build a simple neural network by combining layers into a single layer using the `Serial` combinator. This new layer then acts just like a single layer, so you can inspect intputs, outputs and weights. Or even combine it into another layer! Combinators can then be used as trainable models. _Try adding more layers_

**Note:As you must have guessed, if there is serial combinator, there must be a parallel combinator as well. Do try to explore about combinators and other layers from the trax documentation and look at the repo to understand how these layers are written.**

```python
# help(tl.Serial)
# help(tl.Parallel)
```

```python
# Serial combinator
serial = tl.Serial(
    tl.LayerNorm(),         # normalize input
    tl.Relu(),              # convert negative values to zero
    times_two,              # the custom layer you created above, multiplies the input recieved from above by 2

    ### START CODE HERE
#     tl.Dense(n_units=2),  # try adding more layers. eg uncomment these lines
#     tl.Dense(n_units=1),  # Binary classification, maybe? uncomment at your own peril
#     tl.LogSoftmax()       # Yes, LogSoftmax is also a layer
    ### END CODE HERE
)

# Initialization
x = np.array([-2, -1, 0, 1, 2]) #input
serial.init(shapes.signature(x)) #initialising serial instance

print("-- Serial Model --")
print(serial,"\n")
print("-- Properties --")
print("name :", serial.name)
print("sublayers :", serial.sublayers)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out)
print("weights & biases:", serial.weights, "\n")

# Inputs
print("-- Inputs --")
print("x :", x, "\n")

# Outputs
y = serial(x)
print("-- Outputs --")
print("y :", y)
```

    -- Serial Model --
    Serial[
      LayerNorm
      Relu
      TimesTwo
    ]

    -- Properties --
    name : Serial
    sublayers : [LayerNorm, Relu, TimesTwo]
    expected inputs : 1
    promised outputs : 1
    weights & biases: [(DeviceArray([1, 1, 1, 1, 1], dtype=int32), DeviceArray([0, 0, 0, 0, 0], dtype=int32)), (), ()]

    -- Inputs --
    x : [-2 -1  0  1  2]

    -- Outputs --
    y : [0.        0.        0.        1.4142132 2.8284264]

## JAX
Just remember to lookout for which numpy you are using, the regular ol' numpy or Trax's JAX compatible numpy. Both tend to use the alias np so watch those import blocks.

**Note:There are certain things which are still not possible in fastmath.numpy which can be done in numpy so you will see in assignments we will switch between them to get our work done.**

```python
# Numpy vs fastmath.numpy have different data types
# Regular ol' numpy
x_numpy = np.array([1, 2, 3])
print("good old numpy : ", type(x_numpy), "\n")

# Fastmath and jax numpy
x_jax = fastmath.numpy.array([1, 2, 3])
print("jax trax numpy : ", type(x_jax))
```

    good old numpy :  <class 'numpy.ndarray'>

    jax trax numpy :  <class 'jax.interpreters.xla.DeviceArray'>

## Summary
Trax is a concise framework, built on TensorFlow, for end to end machine learning. The key building blocks are layers and combinators. This notebook is just a taste, but sets you up with some key inuitions to take forward into the rest of the course and assignments where you will build end to end models.