# Using Numba for GPUs

*Vasileios Karakasis, **Theofilos Manitaras**, CSCS*

You may use Numba to generate GPU code from high-level Python functions. Both CUDA and ROCm GPUs (NVIDIA and AMD, respectively) are supported, but we are going to focus only on the CUDA capabilities of Numba.

There is an almost one-to-one mapping between the Numba high-level Python abstractions and the different CUDA constructs. This practically means that, as a programmer, you need to take care explicitly of the host/device memory management and you need to be aware of the CUDA programming model.

This demo will teach you to write your first CUDA kernels using Numba. It will cover the basic CUDA programming principles, but it should be enough to kick start you in GPU programming. More specifically, we will cover the following topics:

- Writing a GPU kernel
- Moving data to/from the GPU.
- Spawning a GPU kernel
- Profiling a GPU kernel
- Optimizing memory accesses
- Making use of the shared memory

We will not cover CUDA streams.

## Verify that Numba sees the GPU and understands CUDA

First thing is to check if Numba can detect the GPU. You can achieve this by running the `numba` executable that comes with Numba's installation:

```python
!numba -s
```

If this command does not work out of the box for your Numba installation, you may have to set the `CUDA_HOME` environment variable to point to your CUDA installation.

## Writing the first kernel

Here is a vector addition kernel that takes two vectors as input and writes the sum in a third vector:

```python
import numba.cuda as cuda

@cuda.jit
def _vecadd_cuda(z, x, y):
    i = cuda.grid(1)
    N = x.shape[0]
    if i >= N:
        return

    z[i] = x[i] + y[i]

```

Pretty simple, right? Let's explain each line of this code.

The `@cuda.jit` decorator will compile the following function into a CUDA kernel at runtime. We will see the cost of it later on.

A CUDA kernel in Numba is a Python function that does *not* return a value. This is in accordance with CUDA, where kernel functions are declared `void`. The arguments of the function can be either Numpy arrays or scalars of a Numba recognized type.

CUDA kernels specify the work to be done by a single GPU thread. Due to the massive hardware parallelism available on the GPU, a single GPU thread gets only a tiny portion of work, in this case, the sumation of a single element in the vectors. Since each element is independent from each other, we can safely create as many threads as the elements of the target vectors and spawn them (we will see how later).

Threads on the GPU as organized in groups, called CUDA blocks. There are constraints on the maximum number of threads a block can contain. For the P100 GPUs on Daint, this is 1024. This means that you will need multiple blocks to calculate the sum of large vectors. The blocks that comprise a kernel form the *grid*. Threads inside a block are numbered sequentially and the blocks inside the grid are numbered, too. The figure below shows the arrangement of threads for the vector addition example.

In order to obtain the i-th element of the vector given a block size $B$, you would have to calculate the following:

\begin{equation}
i = i_{b}B + i_{t}
\end{equation}

where $i_{b}$ is the block index and $i_{t}$ is the thread index. CUDA provides this information and Numba makes it available, so that the above statement would be written as follows:

```python
i = cuda.blockIdx.x*cuda.blockDim.x + threadIdx.x
```

Blocks and grids can be three dimensional, thus the `.x` attribute in all of these variables (the other dimensions are accessed through the `.y` and `.z` attributes).

Since obtaining the absolute position of a thread in a CUDA is quite common operation, Numba offers the convenience function `grid()`, which essentially does automatically the above calculation:

```python
i = cuda.grid(1)
```

The argument passed to the `grid()` function is the number of dimensions of the grid.

The next three statements in the code simply check that we don't overrun the arrays in case that their dimension is not a multiple of the block size. In this case, some threads of the last block will remain idle:

```python
    N = x.shape[0]
    if i >= N:
        return
```

> The threads inside a block are run in batches of 32 at once, called *warps*. All the threads of the warp execute the same instruction. If the program control flow diverges for some of the threads of the warp, e.g., due to an `if` condition, the branches will be executed sequentially by disabling the non participating threads. This condition is called *warp divergence* and may incur a performance penalty.

Finally, the `z[i] = x[i] + y[i]` computes the actual sum.

## Preparing the data for the GPU

All Numba CUDA kernels operate on data residing on the GPU. That means that the `x`, `y` and `z` arrays must be  transferred from the host (CPU) to the device (accelerator) before calling the CUDA kernel.

Let's create first two vectors on the host:

```python
import numpy as np

N = 1024
x = np.random.rand(N)
y = np.random.rand(N)
```

Let's also compute our reference result:

```python
z_ref = x + y
```

Now it's time to transfer the vectors to the GPU:

```python
d_x = cuda.to_device(x)
d_y = cuda.to_device(y)
```

The `d_x` and `d_y` are Numpy *array-like* objects that have their data mapped in the GPU memory:

```python
d_x
```

For the `z` array, we don't need to create it on the host and copy it, since it is essentially an output only array. We can simply allocate it directly on the GPU and copy it out when the kernel finishes.

```python
d_z = cuda.device_array_like(x)
```

This will simply allocate an array like `x`, i.e., with the same element type, shape and order, on the GPU.

Data is ready, now it's time to call the kernel!

## Invoking a CUDA kernel

When invoking a CUDA kernel you have to specify the block size to use and the corresponding grid size. Picking the right size of block is not always straightforward, but usually values between 64 and 256 are good enough.

> The block size has a direct effect on the *occupancy* of each GPU SM, i.e., how much the actual hardware threads of the SM are utilized, and as a result to the occupancy of the whole GPU. Performance-wise, though, it does not have such a big impact. Nvidia provides a nice tool for calculating the occupancy of the GPU, that you can find [here](https://docs.nvidia.com/cuda/cuda-occupancy-calculator/CUDA_Occupancy_Calculator.xls).

For our kernel, we will select a block size of 128 threads.

```python
block_size = 128
```

Having decided the block size we need to set up the grid:

```python
num_blocks = N // block_size
if N % block_size:
    num_blocks += 1
```

Now we are ready to call the kernel!

```python
_vecadd_cuda[num_blocks, block_size](d_z, d_x, d_y)
```

Both the block and the grid could have been two-dimensional (not in this example), in which case you would just define them as tuples.

## Copying back the results

As mentioned before, kernels operate on GPU data only. We need a way to transfer back to the host the result, which is the `d_z` array. Here's how:

```python
res = d_z.copy_to_host()
```

Let's validate the result though to make sure that everything has worked fine:

```python
assert np.allclose(z_ref, res)
```

## Putting it all together

For completeness and easy reference, here is the whole vector addition example:

```python
import numba.cuda as cuda
import numpy as np

@cuda.jit
def _vecadd_cuda(z, x, y):
    '''The CUDA kernel'''
    i = cuda.grid(1)
    N = x.shape[0]
    if i >= N:
        return

    z[i] = x[i] + y[i]

# Set up the host vectors
N = 1000
x = np.random.rand(N)
y = np.random.rand(N)

# Copy and allocate data on the device
d_x = cuda.to_device(x)
d_y = cuda.to_device(y)
d_z = cuda.device_array_like(x)

# Set up the kernel invocation
block_size = 128
num_blocks = N // block_size
if N % block_size:
    num_blocks += 1

# Call the kernel
_vecadd_cuda[num_blocks, block_size](d_z, d_x, d_y)

# Copy back the result to the host
res = d_z.copy_to_host()

# Validate the result
assert np.allclose(x + y, res)
```

### Exercise

> 1. Make the array sufficiently large and time the Numpy version of the sum `x + y`.
> 2. Now time the call to the CUDA kernel with `%timeit -n1 -r1`. What do you see?
> 3. Try timing the CUDA kernels with `%timeit -n1 -r2`. What is happening?


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

## Authors: Zach Etienne, Phil Chang, and Leo Werneck
### Formatting improvements courtesy Brandon Clark

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

**Notebook Status:**  Validated

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

### NRPy+ Source Code for this module:

* [TOV/TOV_Solver.py](../edit/TOV/TOV_Solver.py); ([**NRPy+ Tutorial module reviewing mathematical formulation and equations solved**](Tutorial-ADM_Initial_Data-TOV.ipynb)); ([**start-to-finish NRPy+ Tutorial module demonstrating that initial data satisfy Hamiltonian constraint**](Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data.ipynb)): Tolman-Oppenheimer-Volkoff (TOV) initial data; defines all ADM variables and nonzero $T^{\mu\nu}$ components in Spherical basis.
* [BSSN/ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py](../edit/BSSN/ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py); [\[**tutorial**\]](Tutorial-ADM_Initial_Data-Converting_Numerical_ADM_Spherical_or_Cartesian_to_BSSNCurvilinear.ipynb): *Numerical* Spherical ADM$\to$Curvilinear BSSN converter function
* [BSSN/BSSN_constraints.py](../edit/BSSN/BSSN_constraints.py); [\[**tutorial**\]](Tutorial-BSSN_constraints.ipynb): Hamiltonian constraint in BSSN curvilinear basis/coordinates

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

The entire algorithm is outlined as follows, with links to the relevant NRPy+ tutorial notebooks listed at each step:

1. Allocate memory for gridfunctions, including temporary storage for the Method of Lines time integration [(**NRPy+ tutorial on NRPy+ Method of Lines algorithm**)](Tutorial-Method_of_Lines-C_Code_Generation.ipynb).
1. Set gridfunction values to initial data
    * [**NRPy+ tutorial on TOV initial data**](Tutorial-ADM_Initial_Data-TOV.ipynb)
    * [**NRPy+ tutorial on validating TOV initial data**](Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data.ipynb).
1. Evaluate the Hamiltonian constraint violation
    * [**NRPy+ tutorial on BSSN constraints**](Tutorial-BSSN_constraints.ipynb)
1. Repeat above steps at two numerical resolutions to confirm convergence of Hamiltonian constraint violation to zero.

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

This notebook is organized as follows

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

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

```python
# Step P1: Import needed NRPy+ core modules:
from outputC import lhrh,outCfunction,outputC  # NRPy+: Core C code output module
import NRPy_param_funcs as par   # NRPy+: Parameter interface
import sympy as sp               # SymPy: The Python computer algebra package upon which NRPy+ depends
import finite_difference as fin  # NRPy+: Finite difference C code generation module
import grid as gri               # NRPy+: Functions having to do with numerical grids
import indexedexp as ixp         # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import reference_metric as rfm   # NRPy+: Reference metric support
import cmdline_helper as cmd     # NRPy+: Multi-platform Python command-line interface
import shutil, os, sys           # Standard Python modules for multiplatform OS-level functions

# Step P2: Create C code output directory:
Ccodesdir = os.path.join("TOVID_Ccodes/")
# First remove C code output directory if it exists
# Courtesy https://stackoverflow.com/questions/303200/how-do-i-remove-delete-a-folder-that-is-not-empty
# !rm -r ScalarWaveCurvilinear_Playground_Ccodes
shutil.rmtree(Ccodesdir, ignore_errors=True)
# Then create a fresh directory
cmd.mkdir(Ccodesdir)

# Step P3: Create executable output directory:
outdir = os.path.join(Ccodesdir,"output/")
cmd.mkdir(outdir)

# Step 1: Set the spatial dimension parameter
#         to three this time, and then read
#         the parameter as DIM.
par.set_parval_from_str("grid::DIM",3)
DIM = par.parval_from_str("grid::DIM")

# Step 2: Set some core parameters, including CoordSystem MoL timestepping algorithm,
#                                 FD order, floating point precision, and CFL factor:
# Choices are: Spherical, SinhSpherical, SinhSphericalv2, Cylindrical, SinhCylindrical,
#              SymTP, SinhSymTP
CoordSystem     = "Spherical"

# Step 2.a: Set defaults for Coordinate system parameters.
#           These are perhaps the most commonly adjusted parameters,
#           so we enable modifications at this high level.

# domain_size     = 7.5 # SET BELOW BASED ON TOV STELLAR RADIUS

# sinh_width sets the default value for:
#   * SinhSpherical's params.SINHW
#   * SinhCylindrical's params.SINHW{RHO,Z}
#   * SinhSymTP's params.SINHWAA
sinh_width      = 0.4 # If Sinh* coordinates chosen

# sinhv2_const_dr sets the default value for:
#   * SinhSphericalv2's params.const_dr
#   * SinhCylindricalv2's params.const_d{rho,z}
sinhv2_const_dr = 0.05# If Sinh*v2 coordinates chosen

# SymTP_bScale sets the default value for:
#   * SinhSymTP's params.bScale
SymTP_bScale    = 0.5 # If SymTP chosen

# Step 2.b: Set the order of spatial finite difference derivatives;
#           and the core data type.
FD_order  = 4        # Finite difference order: even numbers only, starting with 2. 12 is generally unstable
REAL      = "double" # Best to use double here.

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

# Step 4: Set the finite differencing order to FD_order (set above).
par.set_parval_from_str("finite_difference::FD_CENTDERIVS_ORDER", FD_order)

# Step 5: Set the direction=2 (phi) axis to be the symmetry axis; i.e.,
#         axis "2", corresponding to the i2 direction.
#         This sets all spatial derivatives in the phi direction to zero.
par.set_parval_from_str("indexedexp::symmetry_axes","2")

# Step 6: The MoLtimestepping interface is only used for memory allocation/deallocation
import MoLtimestepping.C_Code_Generation as MoL
from MoLtimestepping.RK_Butcher_Table_Dictionary import Butcher_dict
RK_method = "Euler" # DOES NOT MATTER; Again MoL interface is only used for memory alloc/dealloc.
RK_order  = Butcher_dict[RK_method][1]
cmd.mkdir(os.path.join(Ccodesdir,"MoLtimestepping/"))
MoL.MoL_C_Code_Generation(RK_method, RHS_string      = "", post_RHS_string = "",
    outdir = os.path.join(Ccodesdir,"MoLtimestepping/"))

# Step 7: Polytropic EOS setup
# For EOS_type, choose either "SimplePolytrope" or "PiecewisePolytrope"
EOS_type = "SimplePolytrope"
# If "PiecewisePolytrope" is chosen as EOS_type, you
#   must also choose the name of the EOS, which can
#   be any of the following:
# 'PAL6', 'SLy', 'APR1', 'APR2', 'APR3', 'APR4',
# 'FPS', 'WFF1', 'WFF2', 'WFF3', 'BBB2', 'BPAL12',
# 'ENG', 'MPA1', 'MS1', 'MS2', 'MS1b', 'PS', 'GS1',
# 'GS2', 'BGN1H1', 'GNH3', 'H1', 'H2', 'H3', 'H4',
# 'H5', 'H6', 'H7', 'PCL2', 'ALF1', 'ALF2', 'ALF3',
# 'ALF4'
EOS_name = 'SLy' # <-- IGNORED IF EOS_type is not PiecewisePolytrope.
```

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

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

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

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

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

```python
##########################
# Polytropic EOS example #
##########################
import TOV.Polytropic_EOSs as ppeos

if EOS_type == "SimplePolytrope":
    # Set neos = 1 (single polytrope)
    neos = 1

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

    # Set Gamma_poly_tab
    Gamma_poly_tab = [2.0]

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

    # Set the eos quantities
    eos = ppeos.set_up_EOS_parameters__complete_set_of_input_variables(neos,rho_poly_tab,Gamma_poly_tab,K_poly_tab0)
    rho_baryon_central = 0.129285
elif EOS_type == "PiecewisePolytrope":
    eos = ppeos.set_up_EOS_parameters__Read_et_al_input_variables(EOS_name)
    rho_baryon_central=2.0
else:
    print("""Error: unknown EOS_type. Valid types are 'SimplePolytrope' and 'PiecewisePolytrope' """)
    sys.exit(1)

import TOV.TOV_Solver as TOV
M_TOV, R_Schw_TOV, R_iso_TOV = TOV.TOV_Solver(eos,
                                              outfile="outputTOVpolytrope.txt",
                                              rho_baryon_central=rho_baryon_central,
                                              return_M_RSchw_and_Riso = True,
                                              verbose = True)

# domain_size sets the default value for:
#   * Spherical's params.RMAX
#   * SinhSpherical*'s params.AMAX
#   * Cartesians*'s -params.{x,y,z}min & .{x,y,z}max
#   * Cylindrical's -params.ZMIN & .{Z,RHO}MAX
#   * SinhCylindrical's params.AMPL{RHO,Z}
#   * *SymTP's params.AMAX
domain_size = 2.0 * R_iso_TOV
```

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

## Step 2.a: Interpolate the TOV data file as needed to set up ADM spacetime quantities in spherical basis (for input into the `Converting_Numerical_ADM_Spherical_or_Cartesian_to_BSSNCurvilinear` module) and $T^{\mu\nu}$ in the chosen reference metric basis \[Back to [top](#toc)\]
$$\label{tov_interp}$$

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

**METRIC DATA IN TERMS OF ADM QUANTITIES**

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

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

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

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

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

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

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

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

and since the 3-metric is diagonal we get

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

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

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

```python
thismodule = "TOVID"
rbar,theta,rho,P,expnu,exp4phi = par.Cparameters("REAL",thismodule,
                                                 ["rbar","theta","rho","P","expnu","exp4phi"],1e300)
IDalpha = sp.sqrt(expnu)
gammaSphDD = ixp.zerorank2(DIM=3)
gammaSphDD[0][0] = exp4phi
gammaSphDD[1][1] = exp4phi*rbar**2
gammaSphDD[2][2] = exp4phi*rbar**2*sp.sin(theta)**2

T4SphUU = ixp.zerorank2(DIM=4)
T4SphUU[0][0] = rho/expnu
T4SphUU[1][1] = P/exp4phi
T4SphUU[2][2] = P/(exp4phi*rbar**2)
T4SphUU[3][3] = P/(exp4phi*rbar**2*sp.sin(theta)**2)
```

```python
expr_list = [IDalpha]
name_list = ["*alpha"]
for i in range(3):
    for j in range(i,3):
        expr_list.append(gammaSphDD[i][j])
        name_list.append("*gammaDD"+str(i)+str(j))

desc = """This function takes as input either (x,y,z) or (r,th,ph) and outputs
all ADM quantities in the Cartesian or Spherical basis, respectively."""
name = "ID_TOV_ADM_quantities"
outCparams = "preindent=1,outCverbose=False,includebraces=False"

outCfunction(
    outfile=os.path.join(Ccodesdir, name + ".h"), desc=desc, name=name,
    params="""  const REAL xyz_or_rthph[3],

                const ID_inputs other_inputs,

                REAL *gammaDD00,REAL *gammaDD01,REAL *gammaDD02,REAL *gammaDD11,REAL *gammaDD12,REAL *gammaDD22,
                REAL *KDD00,REAL *KDD01,REAL *KDD02,REAL *KDD11,REAL *KDD12,REAL *KDD22,
                REAL *alpha,
                REAL *betaU0,REAL *betaU1,REAL *betaU2,
                REAL *BU0,REAL *BU1,REAL *BU2""",
    body="""
      // Set trivial metric quantities:
      *KDD00 = *KDD01 = *KDD02 = 0.0;
      /**/     *KDD11 = *KDD12 = 0.0;
      /**/              *KDD22 = 0.0;
      *betaU0 = *betaU1 = *betaU2 = 0.0;
      *BU0 = *BU1 = *BU2 = 0.0;

      // Next set gamma_{ij} in spherical basis
      const REAL rbar  = xyz_or_rthph[0];
      const REAL theta = xyz_or_rthph[1];
      const REAL phi   = xyz_or_rthph[2];

      REAL rho,rho_baryon,P,M,expnu,exp4phi;
      TOV_interpolate_1D(rbar,other_inputs.Rbar,other_inputs.Rbar_idx,other_inputs.interp_stencil_size,
                         other_inputs.numlines_in_file,
                         other_inputs.r_Schw_arr,other_inputs.rho_arr,other_inputs.rho_baryon_arr,other_inputs.P_arr,other_inputs.M_arr,
                         other_inputs.expnu_arr,other_inputs.exp4phi_arr,other_inputs.rbar_arr,
                         &rho,&rho_baryon,&P,&M,&expnu,&exp4phi);\n"""+
    outputC(expr_list,name_list, "returnstring",outCparams),
    enableCparameters=False)
```

    Output C function ID_TOV_ADM_quantities() to file TOVID_Ccodes/ID_TOV_ADM_quantities.h

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

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

$$
{\rm Jac\_dUSph\_dDrfmUD[mu][nu]} = \frac{\partial x^\mu_{\rm Sph}}{\partial x^\nu_{\rm rfm}},
$$

via exact differentiation (courtesy SymPy), and the inverse Jacobian
$$
{\rm Jac\_dUrfm\_dDSphUD[mu][nu]} = \frac{\partial x^\mu_{\rm rfm}}{\partial x^\nu_{\rm Sph}},
$$

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

$$
T^{\mu\nu}_{\rm rfm} =
\frac{\partial x^\mu_{\rm rfm}}{\partial x^\delta_{\rm Sph}}
\frac{\partial x^\nu_{\rm rfm}}{\partial x^\sigma_{\rm Sph}} T^{\delta\sigma}_{\rm Sph}
$$

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

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

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

Jac4_dUSphorCart_dDrfmUD = ixp.zerorank2(DIM=4)
Jac4_dUSphorCart_dDrfmUD[0][0] = sp.sympify(1)
for i in range(DIM):
    for j in range(DIM):
        Jac4_dUSphorCart_dDrfmUD[i+1][j+1] = sp.diff(r_th_ph_or_Cart_xyz_oID_xx[i],rfm.xx[j])

Jac4_dUrfm_dDSphorCartUD, dummyDET = ixp.generic_matrix_inverter4x4(Jac4_dUSphorCart_dDrfmUD)

# Perform Jacobian operations on T^{mu nu} and gamma_{ij}
T4UU = ixp.register_gridfunctions_for_single_rank2("AUXEVOL","T4UU","sym01",DIM=4)

IDT4UU = ixp.zerorank2(DIM=4)
for mu in range(4):
    for nu in range(4):
        for delta in range(4):
            for sigma in range(4):
                IDT4UU[mu][nu] += \
                     Jac4_dUrfm_dDSphorCartUD[mu][delta]*Jac4_dUrfm_dDSphorCartUD[nu][sigma]*T4SphUU[delta][sigma]

lhrh_list = []
for mu in range(4):
    for nu in range(mu,4):
        lhrh_list.append(lhrh(lhs=gri.gfaccess("auxevol_gfs","T4UU"+str(mu)+str(nu)),rhs=IDT4UU[mu][nu]))

desc = """This function takes as input either (x,y,z) or (r,th,ph) and outputs
all ADM quantities in the Cartesian or Spherical basis, respectively."""
name = "ID_TOV_TUPMUNU_xx0xx1xx2"
outCparams = "preindent=1,outCverbose=False,includebraces=False"
outCfunction(
    outfile=os.path.join(Ccodesdir, name + ".h"), desc=desc, name=name,
    params="""const paramstruct *restrict params,REAL *restrict xx[3],
              const ID_inputs other_inputs,REAL *restrict auxevol_gfs""",
    body=outputC([rfm.xxSph[0],rfm.xxSph[1],rfm.xxSph[2]],
                 ["const REAL rbar","const REAL theta","const REAL ph"],"returnstring",
                 "CSE_enable=False,includebraces=False")+"""
      REAL rho,rho_baryon,P,M,expnu,exp4phi;
      TOV_interpolate_1D(rbar,other_inputs.Rbar,other_inputs.Rbar_idx,other_inputs.interp_stencil_size,
                         other_inputs.numlines_in_file,
                         other_inputs.r_Schw_arr,other_inputs.rho_arr,other_inputs.rho_baryon_arr,other_inputs.P_arr,other_inputs.M_arr,
                         other_inputs.expnu_arr,other_inputs.exp4phi_arr,other_inputs.rbar_arr,
                         &rho,&rho_baryon,&P,&M,&expnu,&exp4phi);\n"""+
    fin.FD_outputC("returnstring",lhrh_list,params="outCverbose=False,includebraces=False"),
    loopopts="AllPoints,Read_xxs")
```

    Output C function ID_TOV_TUPMUNU_xx0xx1xx2() to file TOVID_Ccodes/ID_TOV_TUPMUNU_xx0xx1xx2.h

# Step 3: Convert ADM initial data to BSSN-in-curvilinear coordinates \[Back to [top](#toc)\]
$$\label{adm_id_spacetime}$$

This is an automated process, taken care of by [`BSSN.ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear`](../edit/BSSN.ADM_Numerical_Spherical_or_Cartesian_to_BSSNCurvilinear.py), and documented [in this tutorial notebook](Tutorial-ADM_Initial_Data-Converting_Numerical_ADM_Spherical_or_Cartesian_to_BSSNCurvilinear.ipynb).

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

    Output C function ID_BSSN_lambdas() to file TOVID_Ccodes/ID_BSSN_lambdas.h
    Output C function ID_ADM_xx0xx1xx2_to_BSSN_xx0xx1xx2__ALL_BUT_LAMBDAs() to file TOVID_Ccodes/ID_ADM_xx0xx1xx2_to_BSSN_xx0xx1xx2__ALL_BUT_LAMBDAs.h
    Output C function ID_BSSN__ALL_BUT_LAMBDAs() to file TOVID_Ccodes/ID_BSSN__ALL_BUT_LAMBDAs.h

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

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

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

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

```python
# Enable rfm_precompute infrastructure, which results in
#   BSSN RHSs that are free of transcendental functions,
#   even in curvilinear coordinates, so long as
#   ConformalFactor is set to "W" (default).
cmd.mkdir(os.path.join(Ccodesdir,"rfm_files/"))
par.set_parval_from_str("reference_metric::enable_rfm_precompute","True")
par.set_parval_from_str("reference_metric::rfm_precompute_Ccode_outdir",os.path.join(Ccodesdir,"rfm_files/"))

import BSSN.Enforce_Detgammahat_Constraint as EGC
enforce_detg_constraint_symb_expressions = EGC.Enforce_Detgammahat_Constraint_symb_expressions()

# Now register the Hamiltonian as a gridfunction.
H = gri.register_gridfunctions("AUX","H")
# Then define the Hamiltonian constraint and output the optimized C code.
import BSSN.BSSN_constraints as bssncon
import BSSN.BSSN_stress_energy_source_terms as Bsest
bssncon.BSSN_constraints(add_T4UUmunu_source_terms=False)
Bsest.BSSN_source_terms_for_BSSN_constraints(T4UU)
bssncon.H += Bsest.sourceterm_H

# Now that we are finished with all the rfm hatted
#           quantities in generic precomputed functional
#           form, let's restore them to their closed-
#           form expressions.
par.set_parval_from_str("reference_metric::enable_rfm_precompute","False") # Reset to False to disable rfm_precompute.
rfm.ref_metric__hatted_quantities()

desc="Evaluate the Hamiltonian constraint"
name="Hamiltonian_constraint"
outCfunction(
    outfile  = os.path.join(Ccodesdir,name+".h"), desc=desc, name=name,
    params   = """rfm_struct *restrict rfmstruct,const paramstruct *restrict params,
                  REAL *restrict in_gfs, REAL *restrict auxevol_gfs, REAL *restrict aux_gfs""",
    body     = fin.FD_outputC("returnstring",lhrh(lhs=gri.gfaccess("aux_gfs", "H"), rhs=bssncon.H),
                              params="outCverbose=False"),
    loopopts = "InteriorPoints,enable_rfm_precompute")
```

    Output C function Hamiltonian_constraint() to file TOVID_Ccodes/Hamiltonian_constraint.h

## Step 4.b: Set up boundary condition functions for chosen singular, curvilinear coordinate system \[Back to [top](#toc)\]
$$\label{bc_functs}$$

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

```python
import CurviBoundaryConditions.CurviBoundaryConditions as cbcs
cbcs.Set_up_CurviBoundaryConditions(os.path.join(Ccodesdir,"boundary_conditions/"),Cparamspath=os.path.join("../"))
```

    Wrote to file "TOVID_Ccodes/boundary_conditions/parity_conditions_symbolic_dot_products.h"
    Evolved parity: ( aDD00:4, aDD01:5, aDD02:6, aDD11:7, aDD12:8, aDD22:9,
        alpha:0, betU0:1, betU1:2, betU2:3, cf:0, hDD00:4, hDD01:5, hDD02:6,
        hDD11:7, hDD12:8, hDD22:9, lambdaU0:1, lambdaU1:2, lambdaU2:3, trK:0,
        vetU0:1, vetU1:2, vetU2:3 )
    Auxiliary parity: ( H:0 )
    AuxEvol parity: ( T4UU00:0, T4UU01:1, T4UU02:2, T4UU03:3, T4UU11:4,
        T4UU12:5, T4UU13:6, T4UU22:7, T4UU23:8, T4UU33:9 )
    Wrote to file "TOVID_Ccodes/boundary_conditions/EigenCoord_Cart_to_xx.h"

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

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

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

```python
# Set up the C function for the det(gammahat) = det(gammabar)
EGC.output_Enforce_Detgammahat_Constraint_Ccode(Ccodesdir,
                                                exprs=enforce_detg_constraint_symb_expressions)
```

    Output C function enforce_detgammahat_constraint() to file TOVID_Ccodes/enforce_detgammahat_constraint.h

## Step 4.d: Output C codes needed for declaring and setting Cparameters; also set `free_parameters.h` \[Back to [top](#toc)\]
$$\label{cparams_rfm_and_domainsize}$$

Based on declared NRPy+ Cparameters, first we generate `declare_Cparameters_struct.h`, `set_Cparameters_default.h`, and `set_Cparameters[-SIMD].h`.

Then we output `free_parameters.h`, which sets initial data parameters, as well as grid domain & reference metric parameters, applying `domain_size` and `sinh_width`/`SymTP_bScale` (if applicable) as set above

```python
# Step 3.d.i: Generate declare_Cparameters_struct.h, set_Cparameters_default.h, and set_Cparameters[-SIMD].h
par.generate_Cparameters_Ccodes(os.path.join(Ccodesdir))

# Step 3.d.ii: Set free_parameters.h
# Output to $Ccodesdir/free_parameters.h reference metric parameters based on generic
#    domain_size,sinh_width,sinhv2_const_dr,SymTP_bScale,
#    parameters set above.
rfm.out_default_free_parameters_for_rfm(os.path.join(Ccodesdir,"free_parameters.h"),
                                        domain_size,sinh_width,sinhv2_const_dr,SymTP_bScale)

# Step 3.d.iii: Generate set_Nxx_dxx_invdx_params__and__xx.h:
rfm.set_Nxx_dxx_invdx_params__and__xx_h(Ccodesdir)

# Step 3.d.iv: Generate xx_to_Cart.h, which contains xx_to_Cart() for
#               (the mapping from xx->Cartesian) for the chosen
#               CoordSystem:
rfm.xx_to_Cart_h("xx_to_Cart","./set_Cparameters.h",os.path.join(Ccodesdir,"xx_to_Cart.h"))

# Step 3.d.v: Generate declare_Cparameters_struct.h, set_Cparameters_default.h, and set_Cparameters[-SIMD].h
par.generate_Cparameters_Ccodes(os.path.join(Ccodesdir))
```

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

```python
# Part P0: Define REAL, set the number of ghost cells NGHOSTS (from NRPy+'s FD_CENTDERIVS_ORDER)

with open(os.path.join(Ccodesdir,"TOV_Playground_REAL__NGHOSTS.h"), "w") as file:
    file.write("""
// Part P0.a: Set the number of ghost cells, from NRPy+'s FD_CENTDERIVS_ORDER
#define NGHOSTS """+str(int(FD_order/2)+1)+"""
// Part P0.b: Set the numerical precision (REAL) to double, ensuring all floating point
//            numbers are stored to at least ~16 significant digits
#define REAL """+REAL+"""
// Part P0.c: Set TOV stellar parameters
#define TOV_Mass  """+str(M_TOV)+"""
#define TOV_Riso  """+str(R_iso_TOV)+"\n")
```

```python
%%writefile $Ccodesdir/TOV_Playground.c

// Step P0: Define REAL and NGHOSTS. This header is generated by NRPy+.
#include "TOV_Playground_REAL__NGHOSTS.h"

#include "rfm_files/rfm_struct__declare.h"

#include "declare_Cparameters_struct.h"

// Step P1: Import needed header files
#include "stdio.h"
#include "stdlib.h"
#include "math.h"
#ifndef M_PI
#define M_PI 3.141592653589793238462643383279502884L
#endif
#ifndef M_SQRT1_2
#define M_SQRT1_2 0.707106781186547524400844362104849039L
#endif

// Step P2: Declare the IDX4S(gf,i,j,k) macro, which enables us to store 4-dimensions of
//           data in a 1D array. In this case, consecutive values of "i"
//           (all other indices held to a fixed value) are consecutive in memory, where
//           consecutive values of "j" (fixing all other indices) are separated by
//           Nxx_plus_2NGHOSTS0 elements in memory. Similarly, consecutive values of
//           "k" are separated by Nxx_plus_2NGHOSTS0*Nxx_plus_2NGHOSTS1 in memory, etc.
#define IDX4S(g,i,j,k) \
( (i) + Nxx_plus_2NGHOSTS0 * ( (j) + Nxx_plus_2NGHOSTS1 * ( (k) + Nxx_plus_2NGHOSTS2 * (g) ) ) )
#define IDX4ptS(g,idx) ( (idx) + (Nxx_plus_2NGHOSTS0*Nxx_plus_2NGHOSTS1*Nxx_plus_2NGHOSTS2) * (g) )
#define IDX3S(i,j,k) ( (i) + Nxx_plus_2NGHOSTS0 * ( (j) + Nxx_plus_2NGHOSTS1 * ( (k) ) ) )
#define LOOP_REGION(i0min,i0max, i1min,i1max, i2min,i2max) \
  for(int i2=i2min;i2<i2max;i2++) for(int i1=i1min;i1<i1max;i1++) for(int i0=i0min;i0<i0max;i0++)
#define LOOP_ALL_GFS_GPS(ii) _Pragma("omp parallel for") \
  for(int (ii)=0;(ii)<Nxx_plus_2NGHOSTS_tot*NUM_EVOL_GFS;(ii)++)

// Step P3: Set UUGF and VVGF macros, as well as xx_to_Cart()
#include "boundary_conditions/gridfunction_defines.h"

// Step P4: Set xx_to_Cart(const paramstruct *restrict params,
//                     REAL *restrict xx[3],
//                     const int i0,const int i1,const int i2,
//                     REAL xCart[3]),
//           which maps xx->Cartesian via
//    {xx[0][i0],xx[1][i1],xx[2][i2]}->{xCart[0],xCart[1],xCart[2]}
#include "xx_to_Cart.h"

// Step P5: Defines set_Nxx_dxx_invdx_params__and__xx(const int EigenCoord, const int Nxx[3],
//                                       paramstruct *restrict params, REAL *restrict xx[3]),
//          which sets params Nxx,Nxx_plus_2NGHOSTS,dxx,invdx, and xx[] for
//          the chosen Eigen-CoordSystem if EigenCoord==1, or
//          CoordSystem if EigenCoord==0.
#include "set_Nxx_dxx_invdx_params__and__xx.h"

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

// Step P8: Include function for enforcing detgammabar constraint.
#include "enforce_detgammahat_constraint.h"

// Step P4: Declare initial data input struct:
//          stores data from initial data solver,
//          so they can be put on the numerical grid.
typedef struct __ID_inputs {
    REAL Rbar;
    int Rbar_idx;
    int interp_stencil_size;
    int numlines_in_file;
    REAL *r_Schw_arr,*rho_arr,*rho_baryon_arr,*P_arr,*M_arr,*expnu_arr,*exp4phi_arr,*rbar_arr;
} ID_inputs;

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

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

// Step P10: Declare function necessary for setting up the initial data.
// Step P10.a: Define BSSN_ID() for BrillLindquist initial data

// Step P10.b: Set the generic driver function for setting up BSSN initial data
void initial_data(const paramstruct *restrict params,const bc_struct *restrict bcstruct,
                  const rfm_struct *restrict rfmstruct,
                  REAL *restrict xx[3], REAL *restrict auxevol_gfs, REAL *restrict in_gfs) {
#include "set_Cparameters.h"
        // Step 1: Set up TOV initial data
    // Step 1.a: Read TOV initial data from data file
    // Open the data file:
    char filename[100];
    sprintf(filename,"./outputTOVpolytrope.txt");
    FILE *in1Dpolytrope = fopen(filename, "r");
    if (in1Dpolytrope == NULL) {
        fprintf(stderr,"ERROR: could not open file %s\n",filename);
        exit(1);
    }
    // Count the number of lines in the data file:
    int numlines_in_file = count_num_lines_in_file(in1Dpolytrope);
    // Allocate space for all data arrays:
    REAL *r_Schw_arr     = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *rho_arr        = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *rho_baryon_arr = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *P_arr          = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *M_arr          = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *expnu_arr      = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *exp4phi_arr    = (REAL *)malloc(sizeof(REAL)*numlines_in_file);
    REAL *rbar_arr       = (REAL *)malloc(sizeof(REAL)*numlines_in_file);

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

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

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

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

    // Step 1.b: Interpolate data from data file to set BSSN gridfunctions
    ID_BSSN__ALL_BUT_LAMBDAs(params,xx,TOV_in, in_gfs);
    apply_bcs_curvilinear(params, bcstruct, NUM_EVOL_GFS, evol_gf_parity, in_gfs);
    enforce_detgammahat_constraint(rfmstruct, params,                   in_gfs);
    ID_BSSN_lambdas(params, xx, in_gfs);
    apply_bcs_curvilinear(params, bcstruct, NUM_EVOL_GFS, evol_gf_parity, in_gfs);
    enforce_detgammahat_constraint(rfmstruct, params,                   in_gfs);

    ID_TOV_TUPMUNU_xx0xx1xx2(params,xx,TOV_in,auxevol_gfs);

    free(rbar_arr);
    free(rho_arr);
    free(rho_baryon_arr);
    free(P_arr);
    free(M_arr);
    free(expnu_arr);
}

// Step P11: Declare function for evaluating Hamiltonian constraint (diagnostic)
#include "Hamiltonian_constraint.h"

// main() function:
// Step 0: Read command-line input, set up grid structure, allocate memory for gridfunctions, set up coordinates
// Step 1: Set up initial data to an exact solution
// Step 2: Start the timer, for keeping track of how fast the simulation is progressing.
// Step 3: Integrate the initial data forward in time using the chosen RK-like Method of
//         Lines timestepping algorithm, and output periodic simulation diagnostics
// Step 3.a: Output 2D data file periodically, for visualization
// Step 3.b: Step forward one timestep (t -> t+dt) in time using
//           chosen RK-like MoL timestepping algorithm
// Step 3.c: If t=t_final, output conformal factor & Hamiltonian
//           constraint violation to 2D data file
// Step 3.d: Progress indicator printing to stderr
// Step 4: Free all allocated memory
int main(int argc, const char *argv[]) {
    paramstruct params;
#include "set_Cparameters_default.h"

    // Step 0a: Read command-line input, error out if nonconformant
    if((argc != 4) || atoi(argv[1]) < NGHOSTS || atoi(argv[2]) < NGHOSTS || atoi(argv[3]) < 2 /* FIXME; allow for axisymmetric sims */) {
        fprintf(stderr,"Error: Expected three command-line arguments: ./BrillLindquist_Playground Nx0 Nx1 Nx2,\n");
        fprintf(stderr,"where Nx[0,1,2] is the number of grid points in the 0, 1, and 2 directions.\n");
        fprintf(stderr,"Nx[] MUST BE larger than NGHOSTS (= %d)\n",NGHOSTS);
        exit(1);
    }
    // Step 0b: Set up numerical grid structure, first in space...
    const int Nxx[3] = { atoi(argv[1]), atoi(argv[2]), atoi(argv[3]) };
    if(Nxx[0]%2 != 0 || Nxx[1]%2 != 0 || Nxx[2]%2 != 0) {
        fprintf(stderr,"Error: Cannot guarantee a proper cell-centered grid if number of grid cells not set to even number.\n");
        fprintf(stderr,"       For example, in case of angular directions, proper symmetry zones will not exist.\n");
        exit(1);
    }

    // Step 0c: Set free parameters, overwriting Cparameters defaults
    //          by hand or with command-line input, as desired.
#include "free_parameters.h"

   // Step 0d: Uniform coordinate grids are stored to *xx[3]
    REAL *xx[3];
    // Step 0d.i: Set bcstruct
    bc_struct bcstruct;
    {
        int EigenCoord = 1;
        // Step 0d.ii: Call set_Nxx_dxx_invdx_params__and__xx(), which sets
        //             params Nxx,Nxx_plus_2NGHOSTS,dxx,invdx, and xx[] for the
        //             chosen Eigen-CoordSystem.
        set_Nxx_dxx_invdx_params__and__xx(EigenCoord, Nxx, &params, xx);
        // Step 0d.iii: Set Nxx_plus_2NGHOSTS_tot
#include "set_Cparameters-nopointer.h"
        const int Nxx_plus_2NGHOSTS_tot = Nxx_plus_2NGHOSTS0*Nxx_plus_2NGHOSTS1*Nxx_plus_2NGHOSTS2;
        // Step 0e: Find ghostzone mappings; set up bcstruct
#include "boundary_conditions/driver_bcstruct.h"
        // Step 0e.i: Free allocated space for xx[][] array
        for(int i=0;i<3;i++) free(xx[i]);
    }

    // Step 0f: Call set_Nxx_dxx_invdx_params__and__xx(), which sets
    //          params Nxx,Nxx_plus_2NGHOSTS,dxx,invdx, and xx[] for the
    //          chosen (non-Eigen) CoordSystem.
    int EigenCoord = 0;
    set_Nxx_dxx_invdx_params__and__xx(EigenCoord, Nxx, &params, xx);

    // Step 0g: Set all C parameters "blah" for params.blah, including
    //          Nxx_plus_2NGHOSTS0 = params.Nxx_plus_2NGHOSTS0, etc.
#include "set_Cparameters-nopointer.h"
    const int Nxx_plus_2NGHOSTS_tot = Nxx_plus_2NGHOSTS0*Nxx_plus_2NGHOSTS1*Nxx_plus_2NGHOSTS2;

    // Step 0j: Error out if the number of auxiliary gridfunctions outnumber evolved gridfunctions.
    //              This is a limitation of the RK method. You are always welcome to declare & allocate
    //              additional gridfunctions by hand.
    if(NUM_AUX_GFS > NUM_EVOL_GFS) {
        fprintf(stderr,"Error: NUM_AUX_GFS > NUM_EVOL_GFS. Either reduce the number of auxiliary gridfunctions,\n");
        fprintf(stderr,"       or allocate (malloc) by hand storage for *diagnostic_output_gfs. \n");
        exit(1);
    }

    // Step 0k: Allocate memory for gridfunctions
#include "MoLtimestepping/RK_Allocate_Memory.h"
    REAL *restrict auxevol_gfs = (REAL *)malloc(sizeof(REAL) * NUM_AUXEVOL_GFS * Nxx_plus_2NGHOSTS_tot);

    // Step 0l: Set up precomputed reference metric arrays
    // Step 0l.i: Allocate space for precomputed reference metric arrays.
#include "rfm_files/rfm_struct__malloc.h"

    // Step 0l.ii: Define precomputed reference metric arrays.
    {
    #include "set_Cparameters-nopointer.h"
    #include "rfm_files/rfm_struct__define.h"
    }

    // Step 1: Set up initial data to an exact solution
    initial_data(&params,&bcstruct, &rfmstruct, xx, auxevol_gfs, y_n_gfs);

    // Step 1b: Apply boundary conditions, as initial data
    //          are sometimes ill-defined in ghost zones.
    //          E.g., spherical initial data might not be
    //          properly defined at points where r=-1.
    apply_bcs_curvilinear(&params, &bcstruct, NUM_EVOL_GFS,evol_gf_parity, y_n_gfs);
    enforce_detgammahat_constraint(&rfmstruct, &params, y_n_gfs);

    // Evaluate Hamiltonian constraint violation
    Hamiltonian_constraint(&rfmstruct, &params, y_n_gfs,auxevol_gfs, diagnostic_output_gfs);
    char filename[100];
    sprintf(filename,"out%d.txt",Nxx[0]);
    FILE *out2D = fopen(filename, "w");
    LOOP_REGION(NGHOSTS,Nxx_plus_2NGHOSTS0-NGHOSTS,
                NGHOSTS,Nxx_plus_2NGHOSTS1-NGHOSTS,
                NGHOSTS,Nxx_plus_2NGHOSTS2-NGHOSTS) {
        REAL xx0 = xx[0][i0];
        REAL xx1 = xx[1][i1];
        REAL xx2 = xx[2][i2];
        REAL xCart[3];
        xx_to_Cart(&params,xx,i0,i1,i2,xCart);
        int idx = IDX3S(i0,i1,i2);
        fprintf(out2D,"%e %e %e %e\n",xCart[1]/TOV_Mass,xCart[2]/TOV_Mass, y_n_gfs[IDX4ptS(CFGF,idx)],
                log10(fabs(diagnostic_output_gfs[IDX4ptS(HGF,idx)])));
    }
    fclose(out2D);

    // Step 4: Free all allocated memory
#include "rfm_files/rfm_struct__freemem.h"
#include "boundary_conditions/bcstruct_freemem.h"
#include "MoLtimestepping/RK_Free_Memory.h"
    free(auxevol_gfs);
    for(int i=0;i<3;i++) free(xx[i]);

    return 0;
}
```

    Writing TOVID_Ccodes//TOV_Playground.c

```python
import cmdline_helper as cmd
cmd.C_compile(os.path.join(Ccodesdir,"TOV_Playground.c"), "TOV_Playground")
cmd.delete_existing_files("out96.txt")
cmd.Execute("TOV_Playground", "96 96 2", "out96.txt")
```

    Compiling executable...
    (EXEC): Executing `gcc -std=gnu99 -Ofast -fopenmp -march=native -funroll-loops TOVID_Ccodes/TOV_Playground.c -o TOV_Playground -lm`...
    (BENCH): Finished executing in 1.6104352474212646 seconds.
    Finished compilation.
    (EXEC): Executing `taskset -c 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 ./TOV_Playground 96 96 2`...
    (BENCH): Finished executing in 0.20788025856018066 seconds.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

# fig, ax = plt.subplots()

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

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

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

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

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

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

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

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

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

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

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

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

```python
import cmdline_helper as cmd    # NRPy+: Multi-platform Python command-line interface
cmd.output_Jupyter_notebook_to_LaTeXed_PDF("Tutorial-Start_to_Finish-BSSNCurvilinear-Setting_up_TOV_initial_data")
```

    Created Tutorial-Start_to_Finish-BSSNCurvilinear-
        Setting_up_TOV_initial_data.tex, and compiled LaTeX file to PDF file
        Tutorial-Start_to_Finish-BSSNCurvilinear-
        Setting_up_TOV_initial_data.pdf


# Coherent dark states and polarization switching
Studying the effect of polarization switching on coherent dark states in a 9-level system. The system is made of two ground states, one excited state but all with J = 1 for a total of nine levels. Basically just 3x the 3-level system studied in "Coherent dark states in a 3-level system.ipynb". Note that the system actually has two polarization dark states which significantly complicates things.

## Imports
Start by importing the necessary packages

```python
%load_ext autoreload
%autoreload 2

import joblib
import matplotlib.pyplot as plt
plt.style.use("ggplot")
import numpy as np
import scipy
import qutip
from sympy import Symbol

from toy_systems.couplings import FirstRankCouplingJ, ToyEnergy
from toy_systems.dark_states import get_dark_states
from toy_systems.decays import CouplingDecay, ToyDecay
from toy_systems.hamiltonian import Hamiltonian
from toy_systems.operators import JRotation
from toy_systems.quantum_system import QuantumSystem
from toy_systems.states import Basis, BasisState, JQuantumNumbers, ToyQuantumNumbers
from toy_systems.utils import generate_P_op
```

    The autoreload extension is already loaded. To reload it, use:
      %reload_ext autoreload

## Set up states and basis
We start by defining the three states of the system: the ground states $|g0\rangle$ and $|g1\rangle$, and the excited state $|e\rangle$>

```python
g0s = [BasisState(qn=JQuantumNumbers(label="g0", J = 1, mJ = mJ)) for mJ in np.arange(-1,2)]
g1s = [BasisState(qn=JQuantumNumbers(label="g1", J = 1, mJ = mJ)) for mJ in np.arange(-1,2)]
es = [BasisState(qn=JQuantumNumbers(label="e", J = 1, mJ = mJ)) for mJ in np.arange(-1,2)]
dump = [BasisState(qn = ToyQuantumNumbers(label="dump"))] # A state where the excited state can decay if desired

# Define basis
basis = Basis(g0s+g1s+es+dump)
basis.print()
```

    |0> = JQuantumNumbers(J=1, mJ=-1, label='g0')
    |1> = JQuantumNumbers(J=1, mJ=0, label='g0')
    |2> = JQuantumNumbers(J=1, mJ=1, label='g0')
    |3> = JQuantumNumbers(J=1, mJ=-1, label='g1')
    |4> = JQuantumNumbers(J=1, mJ=0, label='g1')
    |5> = JQuantumNumbers(J=1, mJ=1, label='g1')
    |6> = JQuantumNumbers(J=1, mJ=-1, label='e')
    |7> = JQuantumNumbers(J=1, mJ=0, label='e')
    |8> = JQuantumNumbers(J=1, mJ=1, label='e')
    |9> = |dump>

## Define energies, couplings, decays and quantum system
I'm going to define the system in the rotating frame as usual.

### Energies

```python
δ = Symbol('delta') # Energy splitting between |g0> and |g1
Δ = Symbol('Delta') # Detuning of drive field from 0
E0 = ToyEnergy(g0s, -δ/2)
E1 = ToyEnergy(g1s, +δ/2)
Ee = ToyEnergy(es, Δ)
```

### Couplings
Will treat the problem as if it has two time dependent laser fields: one polarized along z and the other one along x. The polarization will rotate back and forth between the two directions

```python
Ωz = Symbol('Omega_z') # Drive field Rabi rate for z-polarization
Ωx = Symbol('Omega_x') # Drive field Rabi rate for x-polarization

# A condition to make sure the ground states don't get coupled to each other
def both_not_ground(state1, state2):
    return (state1.qn.label == 'e') or (state2.qn.label == 'e')

def both_not_excited(state1, state2):
    return not ((state1.qn.label == 'e') and (state2.qn.label == 'e'))

def not_dump(state1, state2):
    return (state1.qn.label != "dump") and (state2.qn.label != "dump")

coupling_z = FirstRankCouplingJ(Ωz, p_car=np.array((0,0,1)), other_conds = [both_not_ground, both_not_excited, not_dump], time_dep = "(t<0)")
coupling_x = FirstRankCouplingJ(Ωx, p_car=np.array((1,0,0)), other_conds = [both_not_ground, both_not_excited, not_dump], time_dep = "(t>0)")
```

### Decays
Defining a decay from all $|e\rangle$ to all $|g0\rangle$ and $|g1\rangle$ as permitted by angular momentum:

```python
# Define dipole couplings that connect excited state to ground states
decay_couplings = [
    FirstRankCouplingJ(1, p_car=np.array((1,0,0)), other_conds=[both_not_ground, both_not_excited, not_dump]),
    FirstRankCouplingJ(1, p_car=np.array((0,1,0)), other_conds=[both_not_ground, both_not_excited, not_dump]),
    FirstRankCouplingJ(1, p_car=np.array((0,0,1)), other_conds=[both_not_ground, both_not_excited, not_dump]),
]

decays = [CouplingDecay(e, Symbol("Gamma")/2, decay_couplings) for e in es] + [ToyDecay(e, ground = dump[0], gamma = Symbol("Gamma_d")) for e in es]
```

### Define a QuantumSystem
The QuantumSystem object combines the basis, Hamiltonian and decays to make setting parameters for time evolution using QuTiP more convenient.

```python
# Define the system
system = QuantumSystem(
    basis=basis,
    couplings=[E0, E1, Ee, coupling_z, coupling_x],
    decays=decays,
)

# Get representations of the Hamiltonian and the decays that will be accepted by qutip
Hqobj, c_qobj = system.get_qobjs()
```

## Time evolution
No matter what state the system starts in, it should always end up in the dark state, from which it will slowly evolve out since the dark state is not an eigenstate of the Hamiltonian.

```python
# Get a pointer to the time-evolution arguments
args = Hqobj.args
print("Keys for setting arguments:")
print(f"args = {args}")
```

    Keys for setting arguments:
    args = {'delta': 1, 'Delta': 1, 'Omega_z': 1, 'Omega_x': 1, 'Gamma': 1, 'Gamma_d': 1}

```python
test_coupling_z = FirstRankCouplingJ(Ωz, p_car=np.array((0,0,1)), other_conds = [both_not_ground, both_not_excited])

bright_state, dark_states, pol_dark_states = get_dark_states([g0s[0], g1s[0]], es[0], [test_coupling_z])

print(f"|B_z> =\n{bright_state}")

if len(dark_states) != 0:
    print(f"\n|D_z> =\n{dark_states[0]}")

if len(pol_dark_states) != 0:
    print(f"\n|D_z> =\n{pol_dark_states[0]}")
```

    |B_z> =
    [-0.71+0.00j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    -0.71+0.00j x JQuantumNumbers(J=1, mJ=-1, label='g1')]

    |D_z> =
    -0.71+0.00j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    0.71-0.00j x JQuantumNumbers(J=1, mJ=-1, label='g1')

```python
# Set the parameters for the system
args['delta'] = 0.1
args['Delta'] = 0
args['Omega_z'] = 1
args['Omega_x'] = 1
args['Gamma'] = 1
args['Gamma_d'] = 0

# Find the dark and bright states for each mJ and each polarization
bright_states_z = []
dark_states_z = []
pol_dark_states_z = []
bright_states_x = []
dark_states_x = []
pol_dark_states_x = []
bright_states_y = []
dark_states_y = []
pol_dark_states_y = []

for g0, g1, e in zip(g0s, g1s, es):
    test_coupling_z = FirstRankCouplingJ(Ωz, p_car=np.array((0,0,1)), other_conds = [both_not_ground, both_not_excited])

    # Find the bright and dark states for z-polarization
    bright_states, dark_states, pol_dark_states = get_dark_states([g0, g1], e, [test_coupling_z])
    bright_states_z += bright_states
    dark_states_z += dark_states
    pol_dark_states_z += pol_dark_states

# Find the dark and bright states for x- and y-polarization by rotating the basis
# X-polarized
rotation_x = JRotation(np.pi/2, np.array((0,1,0)))
for bs in bright_states_z:
    bright_states_x.append(bs.apply_operator(basis, rotation_x))

for ds in dark_states_z:
    dark_states_x.append(ds.apply_operator(basis, rotation_x))

for ds in pol_dark_states_z:
    pol_dark_states_x.append(ds.apply_operator(basis, rotation_x))

# Y-polarized
rotation_y = JRotation(np.pi/2, np.array((1,0,0)))
for bs in bright_states_z:
    bright_states_y.append(bs.apply_operator(basis, rotation_y))

for ds in dark_states_z:
    dark_states_y.append(ds.apply_operator(basis, rotation_y))

for ds in pol_dark_states_z:
    pol_dark_states_y.append(ds.apply_operator(basis, rotation_y))

# Print the bright and dark states for each polarization:
print("X-polarization:\n")
for i, state in enumerate(bright_states_x):
    print(f"\n|Bx{i}> = ")
    state.print_state()

for i, state in enumerate(dark_states_x):
    print(f"\n|Dx{i}> = ")
    state.print_state()

print("\n\nY-polarization:\n")
for i, state in enumerate(bright_states_y):
    print(f"\n|By{i}> = ")
    state.print_state()

for i, state in enumerate(dark_states_y):
    print(f"\n|Dy{i}> = ")
    state.print_state()

print("\n\nZ-polarization:\n")
for i, state in enumerate(bright_states_z):
    print(f"\n|Bz{i}> = ")
    state.print_state()

for i, state in enumerate(dark_states_z):
    print(f"\n|Dz{i}> = ")
    state.print_state()
```

    X-polarization:

    |Bx0> =
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    +0.5000+0.0000j x JQuantumNumbers(J=1, mJ=0, label='g0')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')
    +0.5000+0.0000j x JQuantumNumbers(J=1, mJ=0, label='g1')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    |Bx1> =
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    +0.5000+0.0000j x JQuantumNumbers(J=1, mJ=0, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')
    +0.5000+0.0000j x JQuantumNumbers(J=1, mJ=0, label='g1')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    |Dx0> =
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    +0.5000+0.0000j x JQuantumNumbers(J=1, mJ=0, label='g0')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')
    -0.5000+0.0000j x JQuantumNumbers(J=1, mJ=0, label='g1')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    |Dx1> =
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    +0.5000+0.0000j x JQuantumNumbers(J=1, mJ=0, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')
    -0.5000+0.0000j x JQuantumNumbers(J=1, mJ=0, label='g1')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    Y-polarization:

    |By0> =
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    0.0000+0.5000j x JQuantumNumbers(J=1, mJ=0, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')
    0.0000+0.5000j x JQuantumNumbers(J=1, mJ=0, label='g1')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    |By1> =
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    0.0000-0.5000j x JQuantumNumbers(J=1, mJ=0, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')
    0.0000-0.5000j x JQuantumNumbers(J=1, mJ=0, label='g1')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    |Dy0> =
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    0.0000+0.5000j x JQuantumNumbers(J=1, mJ=0, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')
    0.0000-0.5000j x JQuantumNumbers(J=1, mJ=0, label='g1')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    |Dy1> =
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    0.0000-0.5000j x JQuantumNumbers(J=1, mJ=0, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    +0.3536+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')
    0.0000+0.5000j x JQuantumNumbers(J=1, mJ=0, label='g1')
    -0.3536+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    Z-polarization:

    |Bz0> =
    -0.7071+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    -0.7071+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')

    |Bz1> =
    +0.7071+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    +0.7071+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

    |Dz0> =
    -0.7071+0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g0')
    +0.7071-0.0000j x JQuantumNumbers(J=1, mJ=-1, label='g1')

    |Dz1> =
    +0.7071+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g0')
    -0.7071+0.0000j x JQuantumNumbers(J=1, mJ=1, label='g1')

```python
# Generate a Qobj representing the initial state
# psi0 = (1*g0s[0]).qobj(basis)
psi0 = (1*dark_states_z[0]).qobj(basis)

# Operators for getting probability of being in each state as a function of time
P_g0 = generate_P_op(g0s, basis)
P_g1 = generate_P_op(g1s, basis)
P_e = generate_P_op(es, basis)
P_B_z = generate_P_op(bright_states_z, basis)
P_D_z = generate_P_op(dark_states_z, basis)
P_D_pol_z = generate_P_op(pol_dark_states_z, basis)
P_B_x = generate_P_op(bright_states_x, basis)
P_D_x = generate_P_op(dark_states_x, basis)
P_D_pol_x = generate_P_op(pol_dark_states_x, basis)
P_B_y = generate_P_op(bright_states_y, basis)
P_D_y = generate_P_op(dark_states_y, basis)
P_dump = generate_P_op(dump, basis)
P_mJm1 = generate_P_op([g0s[0], g1s[0]], basis)
P_mJ0 = generate_P_op([g0s[1], g1s[1]], basis)
P_mJp1 = generate_P_op([g0s[2], g1s[2]], basis)
P_ops = [P_g0, P_g1, P_e, P_B_z, P_D_z, P_B_x, P_D_x, P_B_y, P_D_y, P_mJm1, P_mJ0, P_mJp1, P_D_pol_z, P_D_pol_x]

# Times at which result is requested
times = np.linspace(-10,10,1001)/args["delta"]

# Setting the max_step is sometimes necessary
options = qutip.solver.Options(method = 'adams', nsteps=10000, max_step=1e0, rhs_reuse=False)

# Setup a progress bar
pb = qutip.ui.progressbar.EnhancedTextProgressBar()

# Run the time-evolution
result = qutip.mesolve(Hqobj, psi0, times, c_ops = c_qobj, e_ops = P_ops, progress_bar=pb, options = options)
```

     Total run time:   0.74s*] Elapsed 0.74s / Remaining 00:00:00:00

Plot the results

```python
fig, ax = plt.subplots(3, 1, figsize = (16,18))
ax[0].plot(times, result.expect[0], label = "P_g0")
ax[0].plot(times, result.expect[1], label = "P_g1", ls = '--')
ax[0].plot(times, result.expect[2], label = "P_e")
ax[0].legend()
ax[0].set_ylabel("Population in each state")
ax[0].set_title("Energy eigenstate basis")

ln = []
ln +=ax[1].plot(times, result.expect[3], label = "P_B_z")
ln +=ax[1].plot(times, result.expect[4], label = "P_D_z")
ln +=ax[1].plot(times, result.expect[12], label = "P_D_pol_z")
ln +=ax[1].plot(times, result.expect[2], label = "P_e")
ax[1].set_ylabel("Population in each state")
ax[1].set_title("Dark and bright state basis for z-polarization")
ax1c = ax[1].twinx()
ax1c.grid(False)
ln +=coupling_z.plot_time_dep(times, args, ax = ax1c, c = 'k', ls = '--', label = "z-pol mag")
ax1c.set_ylabel("Magnitude of z-polarization")
ax[1].legend(ln, [l.get_label() for l in ln])

ln2 = []
ln2+=ax[2].plot(times, result.expect[5], label = "P_B_x")
ln2+=ax[2].plot(times, result.expect[6], label = "P_D_x")
ln2+=ax[2].plot(times, result.expect[13], label = "P_D_pol_x")
ln2+=ax[2].plot(times, result.expect[2], label = "P_e")
ax[2].legend()
ax[2].set_ylabel("Population in each state")
ax[2].set_title("Dark and bright state basis for x-polarization")
ax2c = ax[2].twinx()
ax2c.grid(False)
ln2 +=coupling_x.plot_time_dep(times, args, ax = ax2c, c = 'k', ls = '--', label = "x-pol mag")
ax2c.set_ylabel("Magnitude of x-polarization")
ax[2].legend(ln2, [l.get_label() for l in ln2])

print(f"\nPopulation in excited state at the end: {result.expect[2][-1]*100:.1e} %")
print(f"Photons per unit time: {scipy.integrate.trapezoid(result.expect[2], x = times)/(times[-1]-times[0]):.2e}")
```

So what is happening here?:
- I'm starting the system in the z-polarization coherent dark state with mJ = -1.
- For x-polarization the initial state is an even mixture of a coherent and polarization dark state
- Population slowly gets transferred to the z-polarization polarization dark state with mJ = 0, which is an even mixture of bright and dark for x-polarized light
- Once the polarization gets flipped from z to x at t = 0, the bright state component for x-polarization allows some transitions to proceed quickly until the bright state is depleted.
- After the bright state is depleted, the coherent dark state again starts to deplete at a rate that is proportional to $\delta$ and accumulates in the polarization dark state.
- Would be more clear in a system that doesn't have any polarization dark states, so try J = 2 for the excited state (separate notebook for that).

```python

```

examples/Coherent dark states and polarization switching.ipynb


# Interactive notebook for using `helmpy` for parameter inference from data

The mean field or 'deterministic' models which can be run with the `helmpy.run_meanfield` method can be used to provide an efficient way to perform approximate posterior parameter inference with respect to some dataset(s) using the `helmpy.fit_data` method, but before we illustrate how to do this, it will be informative to describe the theoretical background necessary to perform this inference. Note that this notebook will assume prior knowledge with the interface of `helmpy`, so to make sure that this is familiar it is suggested that one reads through helmpy_examples.ipynb before continuing here.

---

## Theoretical background

Here we outline the basic theoretical background for performing Bayesian parameter inference using mean field helminth models under the assumption that the system close to a state of endemic equilibrium. Bear in mind that inference with the mean field model will typically underestimate the variance of the posterior over parameters in comparison to the fully individual-based stochastic model and so it should be used when population sizes are large so that this additional variance is minimised.

The formalism we will outline here assumes that the diagnostic data are either Kato-Katz counts (for _Ascaris lumbricoides, Trichuris trichiura, Schistosoma mansoni_ and hookworm diagnostic testing) or urine filtration counts (for _Schistosoma haematobium_ diagnostic testing). As ever in any canonical Bayesian problem, specification of the likelihood function ${\cal L}$ is not the end of the story. To infer a full joint posterior distribution ${\cal P}$ given a dataset ${\cal D}$ over the collection of transmission and diagnostic parameters $\{ R_0(a), M (a,t_0),k,z,\lambda_{\rm d},k_{\rm d} \}$, Bayes' rule here reads

$$
\begin{equation}
{\cal P}[ R_0(a), M (a,t_0),k,z,\lambda_{\rm d},k_{\rm d}  \vert {\cal D}] = \frac{1}{{\cal E}} \pi [R_0(a)]\, \pi [M (a,t_0)] \, \pi (k) \, \pi (z) \, \pi (\lambda_{\rm d})\, \pi (k_{\rm d}) \,{\cal L}[{\cal D}\vert R_0(a), M (a,t_0),k,z,\lambda_{\rm d},k_{\rm d} ] \,,
\end{equation}
$$

where ${\cal E}$ is a normalisation constant and $\pi [R_0(a)]$, $\pi [M (a,t_0)]$, $\pi (k)$, $\pi (z)$, $\pi (\lambda_{\rm d})$ and $\pi (k_{\rm d})$ are the prior distributions over $R_0(a)$, $M (a,t_0)$ $k$, $z$, $\lambda_{\rm d}$ and $k_{\rm d}$ which are the age-dependent contributions to the basic reproduction number, the age-dependent initial condition to the mean total worm burden per individual (which, when combined with $R_0(a)$ and the other parameters, specify the dynamics $M(a,t)$), worm aggregation, density dependent fecundity factor, the number of diagnostically-detected eggs per female worm and measured diagnostic aggregation parameter, respectively (all assumed to be independent of each other _a priori_).

Kato-katz and urine filtration counts typically follow a distribution which appears to be negative binomial in shape. By computing the mean diagnostically-detected egg count $\hat{{\sf e}}_{\rm d}=\lambda_{\rm d}\hat{{\sf e}}$ from the transmission parameters $\{ R_0(a), M (a,t_0),k,z,\lambda_{\rm d},k_{\rm d}\}$, the likelihood distribution which will be used for the inference of these parameters is therefore likely to be well-approximated by

$$
\begin{equation}
{\cal L}[{\cal D}\vert  R_0(a), M (a,t_0),k,z,\lambda_{\rm d},k_{\rm d} ] = \prod_{\forall {\sf e}_i\in {\cal D}}{\rm NB}\bigg\{ {\sf e}_i; \frac{\lambda_{\rm d}}{2}\hat{{\sf e}}[M(a,t),k,z],k_{\rm d}\bigg\} \,,
\end{equation}
$$

where $M(a,t)$ is the age and time-dependent total mean worm burden which can be fully specified from $R_0(a)$, the other parameters, and the initial conditions $M (a,t_0)$.

Note that one should refer to, e.g., [Anderson & May (1991)](https://global.oup.com/academic/product/infectious-diseases-of-humans-9780198540403?cc=gb&lang=en&) or [Hardwick et al. (2020)](https://www.sciencedirect.com/science/article/pii/S002251931930445X) for the motivations behind the calculation of the proportionality factor for the first moment of the egg count distribution $\hat{{\sf e}}$, which, for example, in the case of STH (fully polygamous male worms) is analytic (it is not in the case of monogamous schistosomes)

$$
\begin{align}
\hat{{\sf e}}_{\rm STH}[M(a,t),k,t] &= \phi [M(a,t);k,z]\, f [M(a,t);k,z] M(a,t) \\
f[M(a,t);k,z] &\equiv \left[ 1+(1-z)\frac{M(a,t)}{k}\right]^{-(k+1)} \\
\phi [M(a,t);k,z] &\equiv 1-\left[ \frac{1+(1-z)M(a,t)/k}{1+(2-z)M(a,t)/(2k)}\right]^{k+1} \,.
\end{align}
$$

From [Anderson & May (1991)](https://global.oup.com/academic/product/infectious-diseases-of-humans-9780198540403?cc=gb&lang=en&), the mean field (or 'deterministic') transmission dynamics of the helminth infections considered here (STH and schistosomes) with age structure can be described by the following system

$$
\begin{equation}
\frac{\partial M}{\partial t} + \frac{\partial M}{\partial a} = \Lambda (a,t) - \mu_1 M(a,t) \,.
\end{equation}
$$

This equation may be converted to a differential equation with respect to time only, while discretising the mean worm burden into age bins $\{ a_i\}$, by integrating over $a$ using a survival rate kernel $S(a)$ like so

$$
\begin{equation}
M_i(t) \equiv M(a_i,t) = \frac{\int^{a_{i+1/2}}_{a_{i-1/2}}{\rm d}a \, M(a,t) S(a)}{\int^{\infty}_{0}{\rm d}a S(a)} \,.
\end{equation}
$$

Choosing the $S(a) = e^{-\mu a}$ (where $\mu$ is the human death rate) and assuming an age-constant $\Lambda (a_i,t)$ within the bin (as we have assumed before in the fitting procedure), one may obtain the following first-order differential equation corresponding to the dynamics in the $i$-th age bin

$$
\begin{equation}
\frac{{\rm d} M_i}{{\rm d} t} = \Lambda (a_i,t) - (\mu + \mu_1)M_i(t) \,.
\end{equation}
$$

In order to obtain this equation above, we have assumed that the boundary flux between age bins must vanish

$$
\begin{equation}
\left.\frac{\partial M_i}{\partial a} \right\vert_{a_{i+1/2}}=0 \,,
\end{equation}
$$

due to an approximated instananeous change in the worm burden for the individual (as a consequence of individuals changing force of infection $\Lambda (a_i,t) \rightarrow \Lambda (a_{i+1},t)$) - note that this also sets the other boundary flux

$$
\begin{equation}
\left.\frac{\partial M_i}{\partial a} \right\vert_{a_{i-1/2}}=\left.\frac{\partial M_{i-1}}{\partial a} \right\vert_{a_{i-1/2}}=\left.\frac{\partial M_{(i-1)}}{\partial a} \right\vert_{a_{(i-1)+1/2}}=0 \,.
\end{equation}
$$

As a side note: in a fully stochastic individual-based model, the change in the expected worm burden will occur over a timescale of $1/\mu_1$, so ensuring that the age bin widths are wider than this timescale is a necessity for this approximation remain accurate. Note also that the birth rate into the first age bin should be set to $\mu$ to match the simulation.

Similarly, one may obtain an equation for the age-binned dynamics (see: [Hardwick et al. (2021)](https://www.sciencedirect.com/science/article/pii/S1755436521000013)) of the force of infection $\Lambda_i \equiv \Lambda (a_i,t)$

$$
\begin{equation}
\frac{{\rm d}\Lambda_i}{{\rm d}t} = \mu_2(\mu + \mu_1)R_{0,i}\bigg\{ \sum_{j=1}^{N_a}\frac{N_j}{N_{\rm tot}}\hat{{\sf e}}[M_j(t),k,z] \bigg\} - \mu_2 \Lambda_i\,.
\end{equation}
$$

The solution to the system $[M_i(t),\Lambda_i(t)]$ can hence be inserted into the negative binomial likelihood ${\rm NB}\{ {\sf e}_i; \lambda_{\rm d}\hat{{\sf e}}[M(a,t),k,z],k_{\rm d}\}$ to perform the inference. Note also that given the rapid equilibration of the infectious reservoir ${\rm d}\Lambda_i /{\rm d}t\rightarrow 0$, we need not specify $\Lambda_i(t_0)$ independently in the inference, but instead may identify

$$
\begin{equation}
\Lambda_i(t) = (\mu + \mu_1)R_{0,i}\sum_{j=1}^{N_a}\frac{N_j}{N_{\rm tot}} \, \hat{{\sf e}}[M_j(t),k,z] \,,
\end{equation}
$$

where $N_i$ is the number of people within an age group (and $N_{\rm tot}$ in total) and $R_{0,i}=R_0(a_i)$ is an age-dependent coefficient which contributes to the basic reproduction number in this bin. By inserting $\Lambda (a_i,t)$ into the equation for ${\rm d} M_i/{\rm d} t$ above, the nonlinear dynamical system of equations that this generates is

$$
\begin{equation}
\frac{{\rm d} M_{i}}{{\rm d} t} = (\mu +\mu_1) R_{0,i}\sum_{j=1}^{N_a}\left\{ \frac{N_j}{N_{\rm tot}}\, \hat{{\sf e}}[M_j(t),k,z] \right\} - (\mu +\mu_1)M_i(t) \,,
\end{equation}
$$

where value of the overall $R_0$ may be obtained through the relation

$$
\begin{equation}
R_{0}=\frac{1}{N_{\rm tot}}\sum^{N_a}_{i=1} N_iR_{0,i} \,.
\end{equation}
$$

Note also that, at equilibrium $M(a_i,t)\rightarrow M(a_i)\,\, \forall i$, the value of $R_0$ is constrained to

$$
\begin{equation}
R_{0} = \frac{\sum_{i=1}^{N_a} N_iM_i}{\sum_{j=1}^{N_a} N_j \,\hat{{\sf e}}(M_j,k,z)} \,.
\end{equation}
$$

To include migration between clusters in the inference, the ${\rm d}\Lambda_i/{\rm d}t$ equation above would be modified by the expectations of compound Poisson processes modelling the net ingoing and outgoing eggs/larvae (see: [Hardwick et al. (2021)](https://www.sciencedirect.com/science/article/pii/S1755436521000013)). We will not handle this case here though.

---

## Setup with mock data

Before illustrating how to use `helmpy` for inference, first we must import it...

```python
import sys
path_to_helmpy = '/Users/Rob/work/helmpy' # Give your path to helmpy here
sys.path.append(path_to_helmpy + '/source/')
from helmpy import helmpy

import time
import numpy as np
import matplotlib.pyplot as plt
```

...and make up some mock (here we shall assume full Kato-Katz intensity counts) data in 2 age categories to use...

```python
# Mean egg counts from data
meanegg_age1 = 10.0
meanegg_age2 = 20.0
# Variance of egg count data
varegg = 3000.0
# Kato-Katz samples drawn for each age group
kksamps_age1 = np.random.negative_binomial(meanegg_age1**2.0/np.abs(varegg-meanegg_age1),meanegg_age1/varegg,size=150)
kksamps_age2 = np.random.negative_binomial(meanegg_age2**2.0/np.abs(varegg-meanegg_age2),meanegg_age2/varegg,size=150)
# Combine Kato-Katz samples into list
kksamps = [kksamps_age1,kksamps_age2]
```

## Parameter inference and visualisation

We will now illustrate how to use `helmpy.fit_data` with an STH instance. This will generate the parameter samples from the posterior distribution (see ${\cal P}[ R_{0,i}, M_i (t_0),k,z,\lambda_{\rm d},k_{\rm d}  \vert {\cal D}]$ above) using the dataset we have made up above. Of course, it will not be possible to infer all of these parameters from the Kato-Katz data alone, so we shall make some prior assumptions (i.e., Dirac delta priors) for their values in some cases. In particular (see: [Truscott et al. (2019)](https://parasitesandvectors.biomedcentral.com/articles/10.1186/s13071-019-3686-2)) here for STH we will assume $\lambda_{\rm d} = 3.05$ and $z=e^{-\gamma}=e^{-0.005}\simeq 0.995$. We will also keep things simple running the simulation for 10 years so that it has long enough to get close to the dynamical endemic equilibrium.

Assuming that one has already used the `helmpy` interface, the following code should make sense now...

```python
hp = helmpy('STH',path_to_helmpy,suppress_terminal_output=True)     # New helmpy instance
hp.parameter_dictionary['mu'] = [0.014,0.014]                       # Human death rate (per year)
hp.parameter_dictionary['mu1'] = [0.5,0.5]                          # Adult worm death rate (per year)
hp.parameter_dictionary['mu2'] = [26.0,26.0]                        # Reservoir (eggs and larvae) death rate
hp.parameter_dictionary['gam'] = [0.005,0.005]                      # Density dependent fecundity: z = exp(-gam)
hp.parameter_dictionary['Np'] = [250,250]                           # Number of people within grouping
hp.parameter_dictionary['spi'] = [1,1]                              # Spatial index number of grouping

# Setting the same arbitrary initial conditions for the force of infection since in the cases of either STH
# or SCH, the equilibration of the force of infection is very rapid anyway
hp.initial_conditions['FOI'] = [5.0,5.0]

# Set the parameter lambda_d
hp.data_specific_parameters['KatoKatz'] = [3.05]

# Initialise number of Monte Carlo walkers and iterations - these are less than you need but it
# is quicker to see how the code works for just this example
nwalkers = 20
niterations = 100

# Set the deterministic model runtime and timestep in years
runtime = 10.0
timestep = 0.02

# Initialise ensemble of walkers and their initial conditions
walker_initconds = []

# Add lnM0s to ensemble (normal distribution initial ensemble with [mean,std])
walker_initconds.append([0.5,0.5])
walker_initconds.append([0.5,0.5])
# Add lnR0s to ensemble (normal distribution initial ensemble with [mean,std])
walker_initconds.append([0.5,0.5])
walker_initconds.append([0.5,0.5])
# Add lnk to ensemble (normal distribution initial ensemble with [mean,std])
walker_initconds.append([-2.5,0.5])
# Add lnkd to ensemble (normal distribution initial ensemble with [mean,std])
walker_initconds.append([-2.0,0.5])

# Set filename
output_filename = 'default_example'

# Parameter labels for the corner plot
names = ['lnM01','lnM02','lnR01','lnR02','lnk','lnkd']

# Run with mock data
hp.fit_data(kksamps,
            walker_initconds,
            'fit_' + output_filename,
            num_walkers=nwalkers,
            num_iterations=niterations,
            runtime=runtime,
            timestep=timestep,
            plot_labels=names)
```

Note that this procedure may need to be performed many times and the chain diagnostics should be examined to make sure the sampling is converging! We've only done a few iterations above for speed but there's definitely more to be done. We can see the posterior samples stored in the `helmpy.data_samples` variable...

```python
hp.data_samples
```

    array([[ 1.22986137,  0.34354855,  2.63850277,  2.87589218, -2.47471124,
            -2.78339066],
           [ 1.50692794,  0.30907869,  2.82559968,  2.96398481, -2.50665834,
            -2.79762252],
           [ 1.50692794,  0.30907869,  2.82559968,  2.96398481, -2.50665834,
            -2.79762252],
           ...,
           [ 0.50895313, -2.43135162, -0.93751167,  2.95455729,  0.21552532,
            -2.4981832 ],
           [ 0.50895313, -2.43135162, -0.93751167,  2.95455729,  0.21552532,
            -2.4981832 ],
           [ 0.50895313, -2.43135162, -0.93751167,  2.95455729,  0.21552532,
            -2.4981832 ]])

## Posterior predictive forecasting with the samples

Once the samples have been obtained, forecasting with a full stochastic individual-based simulation can be achieved by using the `helmpy.run_full_stoch` method immediately. This will use a uniform random subset of the posterior samples and run a fully individual-based stochastic simulation with them initialised with a realisation for each sample. By multiple reruns of this method, we can therefore generate good forecasting statistics for whatever purpose. We do this like so...

```python
# Set the number of realisations
realisations = 100

# Run the stochastic simulation for the results
runtime = 20.0
do_nothing_timescale = 0.01
hp.run_full_stoch(runtime,realisations,do_nothing_timescale,'fit_sim_' + output_filename)
```

The output can also be found in the usual place...

```python
# Load plot data
forecast_runs = np.loadtxt(path_to_helmpy + '/data/fit_sim_' + output_filename + '.txt')

# Generate plot for mean and 68% credible intervals
plt.plot(forecast_runs.T[0],forecast_runs.T[1],color='r')
plt.fill_between(forecast_runs.T[0],forecast_runs.T[4],forecast_runs.T[3],color='r',alpha=0.2)
plt.show()
```

Looks like $z$ and/or $\lambda_{\rm d}$ need to be changed as there are far too many worms in this toy example!

Alternatively to running with the same `helmpy` instance, we can simply set the `helmpy.data_samples` variable of a new instance with the above obtained samples (as long as its parameter configurations are the same) to generate the forecasts with this new instance. This is useful when running on HPC clusters since the data analysis can be performed seperately to the forecasting runs.

Using the deterministic fitting code in conjunction with stochastic simulations is a convenient way to generate predictions for a whole range of scenarios, especially when used in conjunction with the other features within `helmpy`. Note, however, there are important limitations to this method that were discussed at the beginning of this notebook.

```python

```


# Multi Label Binary Classification with CatBoost

CatBoost since the 1.0.0 version supports multilabel binary classification. A target for this mode is a matrix of size $NxK$ where $N$ is the number of objects and $K$ is the number of classes (a.k.a. labels).

There are two loss functions for multilabel binary classification - $MultiLogloss$ and $MultiCrossEntropy$.

For $MultiLogloss$ loss function $y_{ik}$ should be in $\{0, 1\}$:
$$
Y = \begin{pmatrix}
1 & 0 & 0 & ... & 1 & 1\\
0 & 1 & 0 & ... & 1 & 0\\
... & ... & ... & y_{ik} & ... & ...\\
1 & 0 & 1 & ... & 0 & 1\\
\end{pmatrix}_{NxK}
$$

For $MultiCrossEntropy$ loss function $y_{ik}$ should be in $[0, 1]$:
$$
Y = \begin{pmatrix}
0.2 & 0.8 & 0.0 & ... & 0.7 & 0.3\\
0 & 0.6 & 0.3 & ... & 0.9 & 0.0\\
... & ... & ... & y_{ik} & ... & ...\\
1.0 & 0.0 & 0.3 & ... & 0.3 & 0.8\\
\end{pmatrix}_{NxK}
$$

The formula is the same for both loss functions:
$$
MultiLogloss = MultiCrossEntropy = \displaystyle\frac{-\sum\limits_{k=1}^{K} \sum\limits_{i=1}^{N} w_{i} (y_{ik} \log p_{ik} + (1 - t_{ik}) \log (1 - p_{ik}) )}{K\sum\limits_{i=1}^{N}w_{i}}{ ,}
$$

$$
  \mbox{where }p_{ik} = \sigma(a_{ik}) = \frac{e^{a_{ik}}}{1 + e^{a_{ik}}}{ ,}
$$

$$
 a_{ik} \mbox{- raw value of model for } i\mbox{-th object and } j\mbox{-th class}
$$

Let's try multilabel classification mode on synthetic dataset.

```python
# !pip install catboost
# !pip install sklearn
```

```python
from catboost import CatBoostClassifier, Pool
from catboost.utils import eval_metric
from sklearn.datasets import make_multilabel_classification
from sklearn.model_selection import train_test_split
```

#### Generate synthetic dataset

```python
X, Y = make_multilabel_classification(n_samples=500, n_features=20, n_classes=5, random_state=0)
X_train, X_test, Y_train, Y_test = train_test_split(X, Y)
train_pool = Pool(X_train, Y_train)
test_pool = Pool(X_test, Y_test)
```

#### Train model

```python
clf = CatBoostClassifier(
    loss_function='MultiLogloss',
    eval_metric='HammingLoss',
    iterations=500,
    class_names=['A', 'B', 'C', 'D', 'E']
)
clf.fit(train_pool, eval_set=test_pool, metric_period=10, plot=True, verbose=50)
```

    MetricVisualizer(layout=Layout(align_self='stretch', height='500px'))

    Learning rate set to 0.033623
    0:	learn: 0.2778667	test: 0.3328000	best: 0.3328000 (0)	total: 51ms	remaining: 25.4s
    50:	learn: 0.0917333	test: 0.2496000	best: 0.2496000 (50)	total: 135ms	remaining: 1.19s
    100:	learn: 0.0432000	test: 0.2304000	best: 0.2304000 (100)	total: 217ms	remaining: 857ms
    150:	learn: 0.0160000	test: 0.2352000	best: 0.2304000 (100)	total: 301ms	remaining: 695ms
    200:	learn: 0.0058667	test: 0.2352000	best: 0.2304000 (100)	total: 390ms	remaining: 580ms
    250:	learn: 0.0016000	test: 0.2320000	best: 0.2304000 (100)	total: 480ms	remaining: 476ms
    300:	learn: 0.0000000	test: 0.2384000	best: 0.2304000 (100)	total: 572ms	remaining: 378ms
    350:	learn: 0.0000000	test: 0.2320000	best: 0.2304000 (100)	total: 663ms	remaining: 281ms
    400:	learn: 0.0000000	test: 0.2368000	best: 0.2304000 (100)	total: 756ms	remaining: 187ms
    450:	learn: 0.0000000	test: 0.2336000	best: 0.2304000 (100)	total: 851ms	remaining: 92.5ms
    499:	learn: 0.0000000	test: 0.2384000	best: 0.2304000 (100)	total: 947ms	remaining: 0us

    bestTest = 0.2304
    bestIteration = 100

    Shrink model to first 101 iterations.

    <catboost.core.CatBoostClassifier at 0x7f14ab052b50>

#### Predict for test data

```python
test_predict = clf.predict(X_test)
```

## Metrics

**Common deifinitions:**

$N$ - the number of objects

$K$ - the number of labels (classes)

$y_{ik} \in {0, 1}$ - the target value of k-th label for i-th object

$c_{ik} \in {0, 1}$ - the predicted value of k-th label for i-th object

$p_{ik} \in [0, 1]$ - the predicted probability of k-th label for i-th object

### Accuracy

Classic accuracy for multilabel classifications for each object compares whole target vector with whole predicted vector:
$$Accuracy_i =
\begin{equation}
    \begin{cases}
      1, \mbox{if }\forall k: c_{ik} = y_{ik} \\
      0, \mbox{otherwise}
    \end{cases}\,
\end{equation}
$$

And total accuracy is an average accuracy for all objects:
$$Accuracy = \frac{\sum\limits_{i=1}^{N} w_i Accuracy_i}{\sum\limits_{i=1}^{N} w_i} $$

```python
accuracy = eval_metric(Y_test, test_predict, 'Accuracy')[0]
print(f'Accuracy: {accuracy}')
```

    Accuracy: 0.368

### Accuracy per class

You may want to calculate accuracy for each class individually, then you should specify `type` parameter for Accuracy metric like this:

```python
accuracy_per_class = eval_metric(Y_test, test_predict, 'Accuracy:type=PerClass')
for cls, value in zip(clf.classes_, accuracy_per_class):
    print(f'Accuracy for class {cls}: {value}')
```

    Accuracy for class A: 0.776
    Accuracy for class B: 0.712
    Accuracy for class C: 0.744
    Accuracy for class D: 0.776
    Accuracy for class E: 0.84

Accuracy per class is usually better because it more sensible for model changes. But we can't use it as `eval_metric` for early stopping and best model selection, because it returns several values.
Instead of accuracy per class we can use $HammingLoss$.

### HammingLoss

Essentially $HammingLoss$ is a mean accuracy per class subtracted from $1$. So the lower value of HammingLoss is better.

$$
HammingLoss = \frac{\sum\limits_{i=1}^{N} \sum\limits_{j=1}^{K} I(c_{ik}, y_{ik})}{NK}
$$

$$
\mbox{where } I(c, y) =
\begin{equation}
    \begin{cases}
      1, \mbox{if } c = y \\
      0, \mbox{otherwise}
    \end{cases}\,
\end{equation}
$$

```python
hamming = eval_metric(Y_test, test_predict, 'HammingLoss')[0]
print(f'HammingLoss: {hamming:.4f}')
mean_accuracy_per_class = sum(accuracy_per_class) / len(accuracy_per_class)
print(f'MeanAccuracyPerClass: {mean_accuracy_per_class:.4f}')
print(f'HammingLoss + MeanAccuracyPerClass = {hamming + mean_accuracy_per_class}')
```

    HammingLoss: 0.2304
    MeanAccuracyPerClass: 0.7696
    HammingLoss + MeanAccuracyPerClass = 1.0

### Precision, Recall, F1

These metrics are calculated for each class individually:

```python
for metric in ('Precision', 'Recall', 'F1'):
    print(metric)
    values = eval_metric(Y_test, test_predict, metric)
    for cls, value in zip(clf.classes_, values):
        print(f'class={cls}: {value:.4f}')
    print()
```

    Precision
    class=A: 0.8333
    class=B: 0.6066
    class=C: 0.7826
    class=D: 0.8235
    class=E: 0.9655

    Recall
    class=A: 0.3750
    class=B: 0.7551
    class=C: 0.4000
    class=D: 0.5600
    class=E: 0.5957

    F1
    class=A: 0.5172
    class=B: 0.6727
    class=C: 0.5294
    class=D: 0.6667
    class=E: 0.7368

```python

```