# jaco

[![Python package](https://github.com/mikegrudic/jaco/actions/workflows/test.yml/badge.svg)](https://github.com/mikegrudic/jaco/actions/workflows/test.yml)
[![Readthedocs Status][docs-badge]][docs-link]
[![codecov](https://codecov.io/github/mikegrudic/jaco/graph/badge.svg?token=OWJQMWGABZ)](https://codecov.io/github/mikegrudic/jaco)

[docs-link]:           https://jaco.readthedocs.io
[docs-badge]:          https://readthedocs.org/projects/jaco/badge

`jaco` is a flexible, object-oriented implementation for systems of ISM microphysics and chemistry equations, with numerical solvers implemented in JAX, and interfaces for embedding the equations and their Jacobians into other codes.

## Do we really need yet another ISM code?

`jaco` might be interesting because it combines two powerful concepts:
1. **Object-oriented implementation of microphysics and chemistry via the `Process` class**, which implements methods for representing physical processes, composing them into a network in a fully-symbolic `sympy` representation. OOP is nice here because if you want to add a new process to `jaco`, you typically only have to do it in one file. Rate expressions never have to be repeated in-code. Most processes one would want to implement follow very common patterns (e.g. 2-body processes), so class inheritance is also used to minimize new lines of code. 
Once you've constructed your system, `jaco` can give you the symbolic equations to manipulate and analyze as you please. If you want to solve the equations numerically, `Process` has methods for substituting known values into numerical solvers. It can also automatically generate compilable implementations of the RHS of the system to embed in your choice of simulation code and plug into your choice of solver.
2. **Fast, differentiable implementation of nonlinear algebraic and differential-algebraic equation solvers with JAX**, implemented in its functional programming paradigm (e.g. `jaco.numerics.newton_rootsolve`). These can achieve excellent numerical throughput running natively on GPUs - in fact, crunching iterates in-place is essentially the best-case application of numerics on GPUs. Differentiability enables sensitivity analysis with respect to all parameters in a single pass, instead of constructing a grid of `N` parameter variations for `N` parameters. This makes it easier in principle to directly answer questions like "How sensitive is this temperature to the abundance of C or the ionization energy of H?", etc.

## Roadmap

`jaco` is in an early prototyping phase right now. Here are some things I would eventually like to add:
* Models: packages of physical processes and their implementations conveniently grouped together, including commonly-used or historically-important networks such as KWH 1996, Nelson & Langer 1997, Glover 2010, Grackle, Gong+2017, FIRE-3/STARFORGE, etc.
* Flexible implementation of a reduced network suitable for RHD simulations in GIZMO and potentially other codes.
* Dust and radiation physics: add the dust energy equation and evolution of photon number densities to the network.
* Interfaces to convert from other existing chemistry network formats to the `Process` representation.
* Solver robustness upgrades: thermochemical networks can be quite challenging numerically, due to how steeply terms switch on with increasing `T`. In can be hard to get a solution without good initial guesses.
* If possible, glue interface allowing an existing compiled hydro code to call the JAX solvers on-the-fly.

## Installation

Clone the repo and run `pip install .` from the directory, or install the latest release from pypi via `pip install jaco`.

# Quickstart: Collisional Ionization Equilibrium

Example of using `jaco` to solve for collisional ionization equilibrium (CIE) for a hydrogen-helium mixture and plot the ionization states as a function of temperature.


```python
%matplotlib inline
%config InlineBackend.figure_format='retina'
import numpy as np
from matplotlib import pyplot as plt
import sympy as sp
```

## Simple processes
A simple process is defined by a single reaction, with a specified rate.

Let's inspect the structure of a single process, the gas-phase recombination of H+: `H+ + e- -> H + hν` 


```python
from jaco.processes import CollisionalIonization, GasPhaseRecombination

process = GasPhaseRecombination("H+")
print(f"Name: {process.name}")
print(f"Heating rate coefficient: {process.heat_rate_coefficient}")
print(f"Heating rate per cm^-3: {process.heat}"),
print(f"Rate coefficient: {process.rate_coefficient}")
print(f"Recombination rate per cm^-3: {process.rate}")
print(f"RHS of e- number density equation: {process.network['e-']}")
```

    Name: Gas-phase recombination of H+
    Heating rate coefficient: -1.46719838641439e-26*sqrt(T)/((0.00119216696847702*sqrt(T) + 1.0)**1.748*(0.563615123664978*sqrt(T) + 1.0)**0.252)
    Heating rate per cm^-3: -1.46719838641439e-26*sqrt(T)*n_H+*n_e-/((0.00119216696847702*sqrt(T) + 1.0)**1.748*(0.563615123664978*sqrt(T) + 1.0)**0.252)
    Rate coefficient: 1.41621465870114e-10/(sqrt(T)*(0.00119216696847702*sqrt(T) + 1.0)**1.748*(0.563615123664978*sqrt(T) + 1.0)**0.252)
    Recombination rate per cm^-3: 1.41621465870114e-10*n_H+*n_e-/(sqrt(T)*(0.00119216696847702*sqrt(T) + 1.0)**1.748*(0.563615123664978*sqrt(T) + 1.0)**0.252)
    RHS of e- number density equation: Eq(Derivative(n_e-(t), t), -1.41621465870114e-10*n_H+*n_e-/(sqrt(T)*(0.00119216696847702*sqrt(T) + 1.0)**1.748*(0.563615123664978*sqrt(T) + 1.0)**0.252))


Note that all symbolic representations assume CGS units as is standard in ISM physics.

## Composing processes
Now let's define our full network as a sum of simple processes


```python
processes = [CollisionalIonization(s) for s in ("H", "He", "He+")] + [GasPhaseRecombination(i) for i in ("H+", "He+", "He++")]
system = sum(processes)

system.subprocesses
```




    [Collisional Ionization of H,
     Collisional Ionization of He,
     Collisional Ionization of He+,
     Gas-phase recombination of H+,
     Gas-phase recombination of He+,
     Gas-phase recombination of He++]



Summed processes keep track of all subprocesses, e.g. the total net heating rate is:


```python
system.heat
```




$\displaystyle - \frac{1.55 \cdot 10^{-26} n_{He+} n_{e-}}{T^{0.3647}} - \frac{1.2746917300104 \cdot 10^{-21} \sqrt{T} n_{H} n_{e-} e^{- \frac{157809.1}{T}}}{\frac{\sqrt{10} \sqrt{T}}{1000} + 1} - \frac{1.46719838641439 \cdot 10^{-26} \sqrt{T} n_{H+} n_{e-}}{\left(0.00119216696847702 \sqrt{T} + 1.0\right)^{1.748} \left(0.563615123664978 \sqrt{T} + 1.0\right)^{0.252}} - \frac{9.37661057635428 \cdot 10^{-22} \sqrt{T} n_{He} n_{e-} e^{- \frac{285335.4}{T}}}{\frac{\sqrt{10} \sqrt{T}}{1000} + 1} - \frac{4.9524176975855 \cdot 10^{-22} \sqrt{T} n_{He+} n_{e-} e^{- \frac{631515}{T}}}{\frac{\sqrt{10} \sqrt{T}}{1000} + 1} - \frac{5.86879354565754 \cdot 10^{-26} \sqrt{T} n_{He++} n_{e-}}{\left(0.00119216696847702 \sqrt{T} + 1.0\right)^{1.748} \left(0.563615123664978 \sqrt{T} + 1.0\right)^{0.252}}$



Summing processes also sums all chemical and gas/dust cooling/heating rates. 

## Solving ionization equilibrium

We would like to solve for ionization equilibrium given a temperature $T$, overall H number density $n_{\rm H,tot}$.  We define a dictionary of those input quantities and also one for the initial guesses of the number densities of the species in the reduced network.


```python
Tgrid = np.logspace(3,6,10**6)
ngrid = np.ones_like(Tgrid) * 100

knowns = {"T": Tgrid, "n_Htot": ngrid}

guesses = {
    "H": 0.5*np.ones_like(Tgrid),
    "He": 1e-5*np.ones_like(Tgrid),
    "He+": 1e-5*np.ones_like(Tgrid)
}
```

Note that by default, the solver only directly solves for $n_{\rm H}$, $n_{\rm He}$ and $n_{\rm He+}$ because $n_{\rm H+}$, $n_{\rm He++}$, and $n_{\rm e-}$ are eliminated by conservation equations. So we only need initial guesses for those 3 quantities. By default the solver takes abundances $x_i = n_i / n_{\rm H,tot}$ as inputs and outputs.

The `solve` method calls the JAX solver and computes the solution:


```python
sol = system.solve(knowns, guesses,tol=1e-3)
print(sol)
```

    {'He': Array([9.2546351e-02, 9.2546351e-02, 9.2546351e-02, ..., 2.7493625e-09,
           2.7493037e-09, 2.7492442e-09], dtype=float32), 'H': Array([9.9999994e-01, 9.9999994e-01, 9.9999994e-01, ..., 6.0612075e-07,
           6.0611501e-07, 6.0610921e-07], dtype=float32), 'He+': Array([3.1222404e-13, 3.1222396e-13, 3.1222374e-13, ..., 7.6922206e-06,
           7.6921306e-06, 7.6920396e-06], dtype=float32), 'He++': Array([0.        , 0.        , 0.        , ..., 0.09253865, 0.09253865,
           0.09253865], dtype=float32), 'H+': Array([5.9604645e-08, 5.9604645e-08, 5.9604645e-08, ..., 9.9999940e-01,
           9.9999940e-01, 9.9999940e-01], dtype=float32), 'e-': Array([5.9604957e-08, 5.9604957e-08, 5.9604957e-08, ..., 1.1850843e+00,
           1.1850843e+00, 1.1850843e+00], dtype=float32)}



```python
for i, xi in sorted(sol.items()):
    plt.loglog(Tgrid, xi, label=i)
plt.legend(labelspacing=0)
plt.ylabel("$x_i$")
plt.xlabel("T (K)")
plt.ylim(1e-4,3)
```




    (0.0001, 3)




    
![png](CIE_files/CIE_15_1.png)
    


## Generating code

Suppose you just want the RHS of the system you're solving, or its Jacobian, because you have a better solver and/or want to embed these equations in some old C or Fortran code without any dependencies. You can do that too with `generate_code`.


```python
print(system.generate_code(('H','He','He+'),language='c'))
```

    # Computes the RHS function and Jacobianto solve for [x_He, x_H, x_Heplus]
    
    # INDEX CONVENTION: (0: x_He) (1: x_H) (2: x_Heplus)
    
    x0 = 1.0/T; 
    x1 = sqrt(T); 
    x2 = pow(n_Htot, 2); 
    x3 = 1.0/((1.0/1000.0)*sqrt(10)*x1 + 1); 
    x4 = x1*x2*x3; 
    x5 = x4*exp(-285335.40000000002*x0); 
    x6 = x5*x_He; 
    x7 = x_H - 1; 
    x8 = -x7 - 2*x_He - x_Heplus + 2*y; 
    x9 = 2.3800000000000001e-11*x8; 
    x10 = 1.0/x1; 
    x11 = x2*(0.0019*pow(T, -1.5)*(1 + 0.29999999999999999*exp(-94000.0*x0))*exp(-470000.0*x0) + 1.9324160622805846e-10*x10*pow(0.00016493478118851054*x1 + 1.0, -1.7891999999999999)*pow(4.8416074481177231*x1 + 1.0, -0.21079999999999999)); 
    x12 = x11*x_Heplus; 
    x13 = -x12*x8 + x6*x9; 
    x14 = exp(-157809.10000000001*x0); 
    x15 = x14*x4; 
    x16 = x15*x_H; 
    x17 = 5.8500000000000005e-11*x16; 
    x18 = -x7; 
    x19 = pow(0.0011921669684770192*x1 + 1.0, -1.748); 
    x20 = pow(0.56361512366497779*x1 + 1.0, -0.252); 
    x21 = -x_He - x_Heplus + y; 
    x22 = x10*x2; 
    x23 = x22*pow(0.00059608348423850961*x1 + 1.0, -1.748)*pow(0.2818075618324889*x1 + 1.0, -0.252); 
    x24 = 5.664858634804579e-10*x23; 
    x25 = x24*x8; 
    x26 = exp(-631515*x0); 
    x27 = x26*x4; 
    x28 = 4.7600000000000002e-11*x6; 
    x29 = x5*x9; 
    x30 = 2*x12; 
    x31 = -x12 + 2.3800000000000001e-11*x6; 
    x32 = x11*x8 + x31; 
    x33 = x19*x20*x22; 
    x34 = x18*x33; 
    x35 = 1.4162146587011448e-10*x34; 
    x36 = -5.68e-12*x1*x2*x26*x3*x_Heplus + x21*x24;
    
    rhs_result[0] = -x13;
    rhs_result[1] = 1.4162146587011448e-10*x10*x18*x19*x2*x20*x8 - x17*x8;
    rhs_result[2] = x13 + x21*x25 - 5.68e-12*x27*x8*x_Heplus;
    
    jac_result[0] = x28 - x29 - x30;
    jac_result[1] = x31;
    jac_result[2] = x32;
    jac_result[3] = 1.1700000000000001e-10*x16 - 2.8324293174022895e-10*x34;
    jac_result[4] = 5.8500000000000005e-11*x1*x14*x2*x3*x_H - 5.8500000000000005e-11*x15*x8 - 1.4162146587011448e-10*x33*x8 - x35;
    jac_result[5] = x17 - x35;
    jac_result[6] = 1.136e-11*x1*x2*x26*x3*x_Heplus - 1.1329717269609158e-9*x21*x23 - x25 - x28 + x29 + x30;
    jac_result[7] = -x31 - x36;
    jac_result[8] = -x25 - 5.68e-12*x27*x8 - x32 - x36;


Let's break down what happened there. First, jaco is generating the symbolic functions needed to solve the system, as it needs to do before it solves the system with its own solver:


```python
func, jac, _ = system.network.solver_functions(('H','He','He+'),return_jac=True)
```

Here `func` represents the set of functions $f_i$ such that $f_i = 0$ solves the system. `jac` encodes the Jacbian of f $J_{ij} = \frac{\partial f_i}{\partial x_j}$ of derivatives with respect to the solved variables. Note that the two have many common expressions - before being implemented, one should employ common expression elimination to simplify the code and evaluate the functions more efficiently:


```python
cse, (cse_func, cse_jac) = sp.cse((sp.Matrix(func),sp.Matrix(jac)))

cse
```




    [(x0, 1/T),
     (x1, sqrt(T)),
     (x2, n_Htot**2),
     (x3, 1/(sqrt(10)*x1/1000 + 1)),
     (x4, x1*x2*x3),
     (x5, x4*exp(-285335.4*x0)),
     (x6, x5*x_He),
     (x7, x_H - 1),
     (x8, -x7 - 2*x_He - x_He+ + 2*y),
     (x9, 2.38e-11*x8),
     (x10, 1/x1),
     (x11,
      x2*(0.0019*(1 + 0.3*exp(-94000.0*x0))*exp(-470000.0*x0)/T**1.5 + 1.93241606228058e-10*x10/((0.000164934781188511*x1 + 1.0)**1.7892*(4.84160744811772*x1 + 1.0)**0.2108))),
     (x12, x11*x_He+),
     (x13, -x12*x8 + x6*x9),
     (x14, exp(-157809.1*x0)),
     (x15, x14*x4),
     (x16, x15*x_H),
     (x17, 5.85e-11*x16),
     (x18, -x7),
     (x19, (0.00119216696847702*x1 + 1.0)**(-1.748)),
     (x20, (0.563615123664978*x1 + 1.0)**(-0.252)),
     (x21, -x_He - x_He+ + y),
     (x22, x10*x2),
     (x23,
      x22/((0.00059608348423851*x1 + 1.0)**1.748*(0.281807561832489*x1 + 1.0)**0.252)),
     (x24, 5.66485863480458e-10*x23),
     (x25, x24*x8),
     (x26, exp(-631515*x0)),
     (x27, x26*x4),
     (x28, 4.76e-11*x6),
     (x29, x5*x9),
     (x30, 2*x12),
     (x31, -x12 + 2.38e-11*x6),
     (x32, x11*x8 + x31),
     (x33, x19*x20*x22),
     (x34, x18*x33),
     (x35, 1.41621465870114e-10*x34),
     (x36, -5.68e-12*x1*x2*x26*x3*x_He+ + x21*x24)]




```python
cse_func
```




$\displaystyle \left[\begin{matrix}- x_{13}\\1.41621465870114 \cdot 10^{-10} x_{10} x_{18} x_{19} x_{2} x_{20} x_{8} - x_{17} x_{8}\\x_{13} + x_{21} x_{25} - 5.68 \cdot 10^{-12} x_{27} x_{8} x_{He+}\end{matrix}\right]$




```python
cse_jac
```




$\displaystyle \left[\begin{matrix}x_{28} - x_{29} - x_{30} & x_{31} & x_{32}\\1.17 \cdot 10^{-10} x_{16} - 2.83242931740229 \cdot 10^{-10} x_{34} & 5.85 \cdot 10^{-11} x_{1} x_{14} x_{2} x_{3} x_{H} - 5.85 \cdot 10^{-11} x_{15} x_{8} - 1.41621465870114 \cdot 10^{-10} x_{33} x_{8} - x_{35} & x_{17} - x_{35}\\1.136 \cdot 10^{-11} x_{1} x_{2} x_{26} x_{3} x_{He+} - 1.13297172696092 \cdot 10^{-9} x_{21} x_{23} - x_{25} - x_{28} + x_{29} + x_{30} & - x_{31} - x_{36} & - x_{25} - 5.68 \cdot 10^{-12} x_{27} x_{8} - x_{32} - x_{36}\end{matrix}\right]$



One can then take these expressions and convert them to the syntax of the code you wish to embed them in: 


```python
from sympy.codegen.ast import Assignment
for expr in cse:
    print(sp.ccode(Assignment(*expr),standard='c99'))

rhs_result = sp.MatrixSymbol('rhs_result', len(func), 1)
jac_result = sp.MatrixSymbol('jac_result', len(func),len(func))
print()
print(sp.ccode(Assignment(rhs_result, cse_func),standard='c99'))
print()
print(sp.ccode(Assignment(jac_result, cse_jac),standard='c99'))
```

    x0 = 1.0/T;
    x1 = sqrt(T);
    x2 = pow(n_Htot, 2);
    x3 = 1.0/((1.0/1000.0)*sqrt(10)*x1 + 1);
    x4 = x1*x2*x3;
    x5 = x4*exp(-285335.40000000002*x0);
    x6 = x5*x_He;
    x7 = x_H - 1;
    x8 = -x7 - 2*x_He - x_He+ + 2*y;
    x9 = 2.3800000000000001e-11*x8;
    x10 = 1.0/x1;
    x11 = x2*(0.0019*pow(T, -1.5)*(1 + 0.29999999999999999*exp(-94000.0*x0))*exp(-470000.0*x0) + 1.9324160622805846e-10*x10*pow(0.00016493478118851054*x1 + 1.0, -1.7891999999999999)*pow(4.8416074481177231*x1 + 1.0, -0.21079999999999999));
    x12 = x11*x_He+;
    x13 = -x12*x8 + x6*x9;
    x14 = exp(-157809.10000000001*x0);
    x15 = x14*x4;
    x16 = x15*x_H;
    x17 = 5.8500000000000005e-11*x16;
    x18 = -x7;
    x19 = pow(0.0011921669684770192*x1 + 1.0, -1.748);
    x20 = pow(0.56361512366497779*x1 + 1.0, -0.252);
    x21 = -x_He - x_He+ + y;
    x22 = x10*x2;
    x23 = x22*pow(0.00059608348423850961*x1 + 1.0, -1.748)*pow(0.2818075618324889*x1 + 1.0, -0.252);
    x24 = 5.664858634804579e-10*x23;
    x25 = x24*x8;
    x26 = exp(-631515*x0);
    x27 = x26*x4;
    x28 = 4.7600000000000002e-11*x6;
    x29 = x5*x9;
    x30 = 2*x12;
    x31 = -x12 + 2.3800000000000001e-11*x6;
    x32 = x11*x8 + x31;
    x33 = x19*x20*x22;
    x34 = x18*x33;
    x35 = 1.4162146587011448e-10*x34;
    x36 = -5.68e-12*x1*x2*x26*x3*x_He+ + x21*x24;
    
    rhs_result[0] = -x13;
    rhs_result[1] = 1.4162146587011448e-10*x10*x18*x19*x2*x20*x8 - x17*x8;
    rhs_result[2] = x13 + x21*x25 - 5.68e-12*x27*x8*x_He+;
    
    jac_result[0] = x28 - x29 - x30;
    jac_result[1] = x31;
    jac_result[2] = x32;
    jac_result[3] = 1.1700000000000001e-10*x16 - 2.8324293174022895e-10*x34;
    jac_result[4] = 5.8500000000000005e-11*x1*x14*x2*x3*x_H - 5.8500000000000005e-11*x15*x8 - 1.4162146587011448e-10*x33*x8 - x35;
    jac_result[5] = x17 - x35;
    jac_result[6] = 1.136e-11*x1*x2*x26*x3*x_He+ - 1.1329717269609158e-9*x21*x23 - x25 - x28 + x29 + x30;
    jac_result[7] = -x31 - x36;
    jac_result[8] = -x25 - 5.68e-12*x27*x8 - x32 - x36;

