Models

README.md

@@ -98,15 +98,13 @@ Summed processes keep track of all subprocesses, e.g. the total net heating rate
 
 
 ```python
-from jaco.processes import FreeFreeEmission
-system += FreeFreeEmission("H+")
-FreeFreeEmission("H+").bibliography
+system.heat
 ```
 
 
 
 
-    ['1978ppim.book.....S']
+$\displaystyle - \frac{1.55 \cdot 10^{-26} C_{2} n_{He+} n_{e-}}{T^{0.3647}} - \frac{1.46719838641439 \cdot 10^{-26} C_{2} \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{5.86879354565754 \cdot 10^{-26} C_{2} \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}} - \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{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}$
 
 
 
@@ -140,21 +138,21 @@ sol = system.solve(knowns, guesses,tol=1e-3, verbose=True,careful_steps=30)
 print(sol)
 ```
 
-    Undetermined symbols: {C_2, y, T, x_He++, n_Htot, x_H+, x_He+}
-    C_2 not specified; assuming C_2=1.0.
+    Undetermined symbols: {x_He+, y, n_Htot, x_H+, T, C_2, x_He++}
     y not specified; assuming y=0.09254634923706946.
-    Free symbols: {x_He+, y, T, x_He++, n_Htot, x_H+, C_2}
+    C_2 not specified; assuming C_2=1.0.
+    Free symbols: {x_He+, y, n_Htot, x_H+, T, C_2, x_He++}
     Known values: ['T', 'n_Htot']
-    Assumed values: ['C_2', 'y']
-    Equations solved: ['He+', 'He++', 'H+']
-    It's solvin time. Solving for {'He+', 'H+', 'He++'} based on input {'n_Htot', 'T'} and assumptions about {'y', 'C_2'}
-    num_iter average=35.819149017333984 min=18 max=68
-    {'He+': Array([2.8003510e-16, 2.8003862e-16, 2.8003910e-16, ..., 7.6925389e-06,
-           7.6924471e-06, 7.6923561e-06], dtype=float32), 'He++': Array([4.1297267e-17, 4.1297541e-17, 4.1297644e-17, ..., 9.2538655e-02,
-           9.2538655e-02, 9.2538655e-02], dtype=float32), 'H+': Array([2.5647242e-16, 2.5647083e-16, 2.5647025e-16, ..., 9.9999940e-01,
-           9.9999940e-01, 9.9999940e-01], dtype=float32), 'H': Array([1.0000000e+00, 1.0000000e+00, 1.0000000e+00, ..., 5.9604645e-07,
-           5.9604645e-07, 5.9604645e-07], dtype=float32), 'He': Array([9.254635e-02, 9.254635e-02, 9.254635e-02, ..., 7.450581e-09,
-           7.450581e-09, 7.450581e-09], dtype=float32), 'e-': Array([6.1910205e-16, 6.1910453e-16, 6.1910464e-16, ..., 1.1850843e+00,
+    Assumed values: ['y', 'C_2']
+    Equations solved: ['He++', 'H+', 'He+']
+    It's solvin time. Solving for {'He++', 'He+', 'H+'} based on input {'T', 'n_Htot'} and assumptions about {'y', 'C_2'}
+    num_iter average=35.86250305175781 min=18 max=68
+    {'He+': Array([8.259601e-17, 8.259580e-17, 8.259559e-17, ..., 7.692539e-06,
+           7.692447e-06, 7.692356e-06], dtype=float32), 'H+': Array([4.1381906e-16, 4.1381869e-16, 4.1381975e-16, ..., 9.9999940e-01,
+           9.9999940e-01, 9.9999940e-01], dtype=float32), 'He++': Array([1.4600168e-18, 1.4600135e-18, 1.4600135e-18, ..., 9.2538655e-02,
+           9.2538655e-02, 9.2538655e-02], dtype=float32), 'He': Array([9.254635e-02, 9.254635e-02, 9.254635e-02, ..., 7.450581e-09,
+           7.450581e-09, 7.450581e-09], dtype=float32), 'H': Array([1.0000000e+00, 1.0000000e+00, 1.0000000e+00, ..., 5.9604645e-07,
+           5.9604645e-07, 5.9604645e-07], dtype=float32), 'e-': Array([4.9933508e-16, 4.9933450e-16, 4.9933535e-16, ..., 1.1850843e+00,
            1.1850843e+00, 1.1850843e+00], dtype=float32)}
 
 
@@ -190,11 +188,11 @@ Suppose you just want the RHS of the system you're solving, or its Jacobian, bec
 print(system.generate_code(('H+','He+','He++'),language='c'))
 ```
 
-    /* Computes the RHS function and Jacobian to solve for [x_He+, x_He++, x_H+]
+    /* Computes the RHS function and Jacobian to solve for [x_He+, x_H+, x_He++]
 
     This code was auto-generated by jaco v0.1.1 and is not intended to be modified or maintained by human beings.
 
-    INDEX CONVENTION: (0: x_He+) (1: x_He++) (2: x_H+)
+    INDEX CONVENTION: (0: x_He+) (1: x_H+) (2: x_He++)
     */
 
     x0 = sqrt(T); 
@@ -241,18 +239,18 @@ print(system.generate_code(('H+','He+','He++'),language='c'))
     x41 = -5.8500000000000005e-11*x0*x15*x28*x29*x4 + x26*x40;
 
     rhs_result[0] = 2.3800000000000001e-11*x0*x15*x20*x21*x4*x5 + x11 - x19 - x23*x5;
-    rhs_result[1] = -x11 + x19;
-    rhs_result[2] = 5.8500000000000005e-11*x0*x15*x28*x29*x4*x5 - x27*x_Hplus;
+    rhs_result[1] = 5.8500000000000005e-11*x0*x15*x28*x29*x4*x5 - x27*x_Hplus;
+    rhs_result[2] = -x11 + x19;
 
     jac_result[0] = -x22*x6 - x31 - x36 - x37;
-    jac_result[1] = 4.7600000000000002e-11*x0*x15*x20*x21*x4 + x10 - 2*x23 - x31 + x38 - x39;
-    jac_result[2] = -x35 - x37;
-    jac_result[3] = x36;
-    jac_result[4] = -x10 - x38 + x39;
-    jac_result[5] = x18 - x34;
-    jac_result[6] = -x41;
-    jac_result[7] = 1.1700000000000001e-10*x0*x15*x28*x29*x4 - 2.8324293174022895e-10*x24*x25*x40;
-    jac_result[8] = -x27 - 5.8500000000000005e-11*x29*x30 - x41;
+    jac_result[1] = -x35 - x37;
+    jac_result[2] = 4.7600000000000002e-11*x0*x15*x20*x21*x4 + x10 - 2*x23 - x31 + x38 - x39;
+    jac_result[3] = -x41;
+    jac_result[4] = -x27 - 5.8500000000000005e-11*x29*x30 - x41;
+    jac_result[5] = 1.1700000000000001e-10*x0*x15*x28*x29*x4 - 2.8324293174022895e-10*x24*x25*x40;
+    jac_result[6] = x36;
+    jac_result[7] = x18 - x34;
+    jac_result[8] = -x10 - x38 + x39;
 
 
 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:

docs/requirements.txt

@@ -7,3 +7,4 @@ astropy
 jax
 sympy
 matplotlib
+periodictable

docs/source/Quickstart.rst

@@ -78,16 +78,14 @@ heating rate is:
 
 .. code:: ipython3
 
-    from jaco.processes import FreeFreeEmission
-    system += FreeFreeEmission("H+")
-    FreeFreeEmission("H+").bibliography
+    system.heat
 
 
 
 
-.. parsed-literal::
+.. math::
 
-    ['1978ppim.book.....S']
+    \displaystyle - \frac{1.55 \cdot 10^{-26} C_{2} n_{He+} n_{e-}}{T^{0.3647}} - \frac{1.46719838641439 \cdot 10^{-26} C_{2} \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{5.86879354565754 \cdot 10^{-26} C_{2} \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}} - \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{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}
 
 
 
@@ -132,21 +130,21 @@ The ``solve`` method calls the JAX solver and computes the solution:
 
 .. parsed-literal::
 
-    Undetermined symbols: {C_2, y, T, x_He++, n_Htot, x_H+, x_He+}
-    C_2 not specified; assuming C_2=1.0.
+    Undetermined symbols: {x_He+, y, n_Htot, x_H+, T, C_2, x_He++}
     y not specified; assuming y=0.09254634923706946.
-    Free symbols: {x_He+, y, T, x_He++, n_Htot, x_H+, C_2}
+    C_2 not specified; assuming C_2=1.0.
+    Free symbols: {x_He+, y, n_Htot, x_H+, T, C_2, x_He++}
     Known values: ['T', 'n_Htot']
-    Assumed values: ['C_2', 'y']
-    Equations solved: ['He+', 'He++', 'H+']
-    It's solvin time. Solving for {'He+', 'H+', 'He++'} based on input {'n_Htot', 'T'} and assumptions about {'y', 'C_2'}
-    num_iter average=35.819149017333984 min=18 max=68
-    {'He+': Array([2.8003510e-16, 2.8003862e-16, 2.8003910e-16, ..., 7.6925389e-06,
-           7.6924471e-06, 7.6923561e-06], dtype=float32), 'He++': Array([4.1297267e-17, 4.1297541e-17, 4.1297644e-17, ..., 9.2538655e-02,
-           9.2538655e-02, 9.2538655e-02], dtype=float32), 'H+': Array([2.5647242e-16, 2.5647083e-16, 2.5647025e-16, ..., 9.9999940e-01,
-           9.9999940e-01, 9.9999940e-01], dtype=float32), 'H': Array([1.0000000e+00, 1.0000000e+00, 1.0000000e+00, ..., 5.9604645e-07,
-           5.9604645e-07, 5.9604645e-07], dtype=float32), 'He': Array([9.254635e-02, 9.254635e-02, 9.254635e-02, ..., 7.450581e-09,
-           7.450581e-09, 7.450581e-09], dtype=float32), 'e-': Array([6.1910205e-16, 6.1910453e-16, 6.1910464e-16, ..., 1.1850843e+00,
+    Assumed values: ['y', 'C_2']
+    Equations solved: ['He++', 'H+', 'He+']
+    It's solvin time. Solving for {'He++', 'He+', 'H+'} based on input {'T', 'n_Htot'} and assumptions about {'y', 'C_2'}
+    num_iter average=35.86250305175781 min=18 max=68
+    {'He+': Array([8.259601e-17, 8.259580e-17, 8.259559e-17, ..., 7.692539e-06,
+           7.692447e-06, 7.692356e-06], dtype=float32), 'H+': Array([4.1381906e-16, 4.1381869e-16, 4.1381975e-16, ..., 9.9999940e-01,
+           9.9999940e-01, 9.9999940e-01], dtype=float32), 'He++': Array([1.4600168e-18, 1.4600135e-18, 1.4600135e-18, ..., 9.2538655e-02,
+           9.2538655e-02, 9.2538655e-02], dtype=float32), 'He': Array([9.254635e-02, 9.254635e-02, 9.254635e-02, ..., 7.450581e-09,
+           7.450581e-09, 7.450581e-09], dtype=float32), 'H': Array([1.0000000e+00, 1.0000000e+00, 1.0000000e+00, ..., 5.9604645e-07,
+           5.9604645e-07, 5.9604645e-07], dtype=float32), 'e-': Array([4.9933508e-16, 4.9933450e-16, 4.9933535e-16, ..., 1.1850843e+00,
            1.1850843e+00, 1.1850843e+00], dtype=float32)}
 
 
@@ -189,11 +187,11 @@ can do that too with ``generate_code``.
 
 .. parsed-literal::
 
-    /* Computes the RHS function and Jacobian to solve for [x_He+, x_He++, x_H+]
+    /* Computes the RHS function and Jacobian to solve for [x_He+, x_H+, x_He++]
 
     This code was auto-generated by jaco v0.1.1 and is not intended to be modified or maintained by human beings.
 
-    INDEX CONVENTION: (0: x_He+) (1: x_He++) (2: x_H+)
+    INDEX CONVENTION: (0: x_He+) (1: x_H+) (2: x_He++)
     */
 
     x0 = sqrt(T); 
@@ -240,18 +238,18 @@ can do that too with ``generate_code``.
     x41 = -5.8500000000000005e-11*x0*x15*x28*x29*x4 + x26*x40;
 
     rhs_result[0] = 2.3800000000000001e-11*x0*x15*x20*x21*x4*x5 + x11 - x19 - x23*x5;
-    rhs_result[1] = -x11 + x19;
-    rhs_result[2] = 5.8500000000000005e-11*x0*x15*x28*x29*x4*x5 - x27*x_Hplus;
+    rhs_result[1] = 5.8500000000000005e-11*x0*x15*x28*x29*x4*x5 - x27*x_Hplus;
+    rhs_result[2] = -x11 + x19;
 
     jac_result[0] = -x22*x6 - x31 - x36 - x37;
-    jac_result[1] = 4.7600000000000002e-11*x0*x15*x20*x21*x4 + x10 - 2*x23 - x31 + x38 - x39;
-    jac_result[2] = -x35 - x37;
-    jac_result[3] = x36;
-    jac_result[4] = -x10 - x38 + x39;
-    jac_result[5] = x18 - x34;
-    jac_result[6] = -x41;
-    jac_result[7] = 1.1700000000000001e-10*x0*x15*x28*x29*x4 - 2.8324293174022895e-10*x24*x25*x40;
-    jac_result[8] = -x27 - 5.8500000000000005e-11*x29*x30 - x41;
+    jac_result[1] = -x35 - x37;
+    jac_result[2] = 4.7600000000000002e-11*x0*x15*x20*x21*x4 + x10 - 2*x23 - x31 + x38 - x39;
+    jac_result[3] = -x41;
+    jac_result[4] = -x27 - 5.8500000000000005e-11*x29*x30 - x41;
+    jac_result[5] = 1.1700000000000001e-10*x0*x15*x28*x29*x4 - 2.8324293174022895e-10*x24*x25*x40;
+    jac_result[6] = x36;
+    jac_result[7] = x18 - x34;
+    jac_result[8] = -x10 - x38 + x39;
 
 
 Let’s break down what happened there. First, jaco is generating the

experiments/ccodegen.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 1,
    "id": "561bc739",
    "metadata": {},
    "outputs": [],

src/jaco/codegen/gizmo/gizmo.py

@@ -29,7 +29,7 @@ def generate_funcjac_code(system, cse=True):
     for i, p in enumerate(paramsvars):
         funcjac = funcjac.subs(p, P[i])
         assignments.append(Assignment(P[i], p))
-        index_defs.append(f"#define IDX_{p} {i}")
+        index_defs.append(f"#define INDEX_{p} {i}")
     cg = C99CodeGen(cse=cse)
     routine = cg.routine("microphysics_func_jac", funcjac)  # cg.routine("func",sp.Matrix(func + sp.flatten(jac)))
     cg.write([routine], "microphysics_func_jac", to_files=True)

src/jaco/eos/H2_partition_function.py

@@ -1,4 +1,5 @@
 import sympy as sp
+from jaco.math import logistic
 
 
 def H2_energy_over_kbT(T):
@@ -9,4 +10,4 @@ def H2_energy_over_kbT(T):
     p2 = ((0.102728060235312 * sp.log(T) - 2.42078687298443) * sp.log(T) + 19.4951634944834) * sp.log(
         T
     ) - 51.5605088374882
-    return 1.5 + 1 / (1 + sp.exp(-p1)) + 1 / (1 + sp.exp(-p2))
+    return 1.5 + logistic(p1) + logistic(p2)

src/jaco/math.py

@@ -0,0 +1,13 @@
+"""Special mathematical functions"""
+
+import sympy as sp
+
+
+def logistic(x):
+    """Symbolic implementation of the logistic function.
+
+    f(x) = 1/(1+exp(-x))
+
+    Uses tanh which is better behaved numerically in extreme limits
+    """
+    return 0.5 * (1 + sp.tanh(x / 2))

src/jaco/models/wind_comparison/cooling.py

@@ -2,6 +2,7 @@
 from jaco.symbols import piecewise_powerlaw, T, n_
 from jaco.processes import ThermalProcess
 import sympy as sp
+from jaco.math import logistic
 
 T_cooling_curve = np.array(
     [0.99999999e1, 1.0e02, 6.0e03, 1.75e04, 4.0e04, 8.7e04, 2.30e05, 3.6e05, 1.5e06, 3.50e06, 2.6e07, 1.0e12]
@@ -44,4 +45,4 @@
 
 lambda_cooling = piecewise_powerlaw(T_cooling_curve, lambda_cooling_curve, T, extrapolate=True)
 cooling = ThermalProcess(-lambda_cooling * n_("H") ** 2, name="Cooling")
-heating = ThermalProcess(2e-26 * n_("H") / (1 + sp.exp(sp.Min((T - 15000) / 1000, 100))), name="Heating")
+heating = ThermalProcess(2e-26 * n_("H") * logistic(-(T - 15000) / 1000), name="Heating")

src/jaco/symbols.py

@@ -104,7 +104,8 @@ def piecewise_linear(X, Y, x, extrapolate=False):
 
     for i in range(len(X) - 1):
         cases.append((Y[i] + slopes[i] * (x - X[i]), (x >= X[i]) & (x < X[i + 1])))
-    return sp.piecewise_exclusive(sp.Piecewise(*cases))
+    #    cases.append((sp.nan, True))
+    return sp.Piecewise(*cases)
 
 
 def piecewise_powerlaw(X, Y, x, extrapolate=False):