reaxion 0.1.1__py3-none-any.whl

reaxion/misc.py

@@ -0,0 +1,66 @@
+"""Various convenience routines used throughout the package"""
+
+from string import digits
+
+
+def species_charge(species: str) -> int:
+    """Returns the charge number of a species from its name"""
+    if species[-1] not in ("-", "+"):
+        return 0
+    if "++" in species:
+        return 2
+    elif "--" in species:
+        return -2
+
+    base = base_species(species)
+    suffix = species.split(base)[-1]
+    if suffix == "+":
+        return 1
+    if suffix == "-":
+        return -1
+    elif "+" in suffix:
+        return int(suffix.rstrip("+"))
+    else:
+        return -int(suffix.rstrip("-"))
+
+
+def is_an_ion(species: str) -> bool:
+    return species_charge(species) != 0 and species != "e-"
+
+
+def base_species(species: str) -> str:
+    """Removes the charge suffix from a species"""
+    base = species.rstrip(digits + "-+")
+    return base
+
+
+def charge_suffix(charge: int) -> str:
+    """Returns the suffix for an input charge number"""
+    match charge:
+        case 0:
+            return ""
+        case 1:
+            return "+"
+        case 2:
+            return "++"
+        case -1:
+            return "-"
+        case -2:
+            return "--"
+        case _:
+            if charge < -2:
+                return str(abs(charge)) + "-"
+            else:
+                return str(abs(charge)) + "+"
+
+
+def ionize(species: str) -> str:
+    """Returns the symbol of the species produced by removing an electron from the input species"""
+    charge = species_charge(species)
+    return base_species(species) + charge_suffix(charge + 1)
+
+
+def recombine(species: str) -> str:
+    """Returns the symbol of the species produced by adding an electron to the input species"""
+    charge = species_charge(species)
+    return base_species(species) + charge_suffix(charge - 1)

reaxion/numerics/__init__.py

@@ -0,0 +1 @@
+from .solvers import *

reaxion/numerics/solvers.py

@@ -0,0 +1,98 @@
+import jax, jax.numpy as jnp
+
+
+def newton_rootsolve(
+    func,
+    guesses,
+    params=[],
+    jacfunc=None,
+    tolfunc=None,
+    rtol=1e-6,
+    max_iter=100,
+    careful_steps=1,
+    nonnegative=False,
+    return_num_iter=False,
+):
+    """
+    Solve the system f(X,p) = 0 for X, where both f and X can be vectors of arbitrary length and p is a set of fixed
+    parameters passed to f. Broadcasts and parallelizes over an arbitrary number of initial guesses and parameter
+    choices.
+
+    Parameters
+    ----------
+    func: callable
+        A JAX function of signature f(X,params) that implements the function we wish to rootfind, where X and params
+        are arrays of shape (n,) and (n_p,) for dimension n and parameter number n_p. In general can return an array of
+        shape (m,)
+    guesses: array_like
+        Shape (n,) or (N,n) array_like where N is the number of guesses + corresponding parameter choices
+    params: array_like
+        Shape (n,) or (N,n_p) array_like where N is the number of guesses + corresponding parameter choices
+    jacfunc: callable, optional
+        Function with the same signature as f that returns the Jacobian of f - will be computed with autodiff from f if
+        not specified.
+    rtol: float, optional
+        Relative tolerance - iteration will terminate if relative change in all quantities is less than this value.
+    atol: float, optional
+        Absolute tolerance: iteration will terminate if the value computed by tolfunc goes below this value.
+    careful_steps: int, optional
+        Number of "careful" initial steps to take, gradually ramping up the step size in the Newton iteration
+
+    Returns
+    -------
+    X: array_like
+        Shape (N,n) array of solutions
+    """
+    BIG, SMALL = 1e37, 1e-37
+
+    guesses = jnp.array(guesses)
+    guesses = jnp.where(nonnegative, guesses.clip(SMALL), guesses)
+    params = jnp.array(params)
+    if len(guesses.shape) < 2:
+        guesses = jnp.atleast_2d(guesses).T
+    if len(params.shape) < 2:
+        params = jnp.atleast_2d(params).T
+
+    if jacfunc is None:
+        jac = jax.jacfwd(func)
+
+    if tolfunc is None:
+
+        def tolfunc(X, *params):
+            return X
+
+    def solve(guess, params):
+        """Function to be called in parallel that solves the root problem for one guess and set of parameters"""
+
+        def iter_condition(arg):
+            """Iteration condition for the while loop: check if we are within desired tolerance."""
+            X, dx, num_iter = arg
+            fac = jnp.min(jnp.array([(num_iter + 1.0) / careful_steps, 1.0]))
+            tol2, tol1 = tolfunc(X, *params), tolfunc(X - dx, *params)
+            tolcheck = jnp.any(jnp.abs(tol1 - tol2) > rtol * jnp.abs(tol1) * fac)
+            return jnp.any(jnp.abs(dx) > fac * rtol * jnp.abs(X)) & (num_iter < max_iter) & tolcheck
+
+        def X_new(arg):
+            """Returns the next Newton iterate and the difference from previous guess."""
+            X, _, num_iter = arg
+            fac = jnp.min(jnp.array([(num_iter + 1.0) / careful_steps, 1.0]))
+            J = jac(X, *params)
+            #  condition number is nice but possibly very slow due to batching, e.g. https://github.com/jax-ml/jax/issues/11321
+            # there is no reason for this from a pure FLOPS standpoint!
+            #            cond = jsp.linalg.cond(J)  # , p=2)
+            #            dx = jnp.where(cond < 1e30, -jnp.linalg.solve(J, func(X, *params)) * fac, jnp.zeros_like(X))
+            dx = -jnp.linalg.solve(J, func(X, *params)) * fac
+            dx_finite = jnp.all(jnp.isfinite(dx))
+            Xnew = jnp.where(dx_finite, jnp.where(nonnegative, (X + dx).clip(SMALL), X + dx), X)
+            return Xnew, dx, num_iter + 1
+
+        init_val = guess, 100 * guess, 0
+        X, _, num_iter = jax.lax.while_loop(iter_condition, X_new, init_val)
+
+        return X, num_iter
+
+    X, num_iter = jax.vmap(solve)(guesses, params)
+    if return_num_iter:
+        return X, num_iter
+    else:
+        return X

reaxion/numerics/tests/test_linear.py

@@ -0,0 +1,33 @@
+import numpy as np
+import jax, jax.numpy as jnp
+from ..solvers import newton_rootsolve
+
+
+def test_newton_rootsolve_linear(N=10**5, dim=10):
+    """Test: solve a linear system of dim variables with the Newton solver and check the solution.
+    Solver should get the right answer to machine precision.
+    """
+
+    A = np.random.rand(N, dim, dim)
+    cond = np.linalg.cond(A)
+    A, cond = A[cond < 1e3], cond[cond < 1e3]  # just take the easy ones, we don't need to be fancy here
+    N = len(A)
+    b = np.random.rand(N, dim, 1)
+    sol = np.linalg.solve(A, b)[:, :, 0]
+    b = b[:, :, 0]
+
+    @jax.jit
+    def func(x, *params):
+        """horrid function for mapping X, *params signature to linear residual Ax-b"""
+        dim = x.shape[0]
+        A = jnp.array(params)[: dim * dim].reshape((dim, dim))
+        b = jnp.array(params[dim * dim : dim * (dim + 1)])
+        return jnp.matmul(A, x) - b
+
+    p = np.c_[A.reshape(N, dim * dim), b.reshape(N, dim)]
+    guesses = np.copy(b)
+    sol_newton = newton_rootsolve(func, guesses, p)  # [:,:,None]
+
+    error = np.sum((sol_newton - sol) ** 2, axis=1) ** 0.5
+    norm = np.sum(sol**2, axis=1)
+    assert np.all(error < 1e-2 * norm)

reaxion/numerics/tests/test_newton_rootsolve.py

@@ -0,0 +1,34 @@
+import numpy as np
+import jax, jax.numpy as jnp
+import reaxion
+from reaxion.numerics.solvers import newton_rootsolve
+
+
+def test_newton_rootsolve(N=10**5):
+    """Test: solve the system x^p = a for various choices of p and a and check solution
+
+    NOTE: this newton iteration does not converge in general, so we will not catch all possible bugs that might return
+    nonfinite values...
+    """
+    import numpy as np
+
+    p = 0.1 + np.random.rand(N) * 10
+    a = 0.1 + np.random.rand(N)
+    params = jnp.c_[p, a]
+    guess = jnp.ones(N)
+
+    exact = np.atleast_2d(a ** (1.0 / p)).T
+
+    def func(x, *params):
+        return x ** params[0] - params[1]
+
+    func = jax.jit(func)
+
+    sol = newton_rootsolve(func, guess, params, nonnegative=True)
+    converged = jnp.all(jnp.isfinite(sol), axis=1)
+    assert converged.sum() > 0.9 * N
+    assert jnp.all(jnp.isclose(sol[converged], exact[converged], rtol=1e-3, atol=0))
+
+
+if __name__ == "__main__":
+    test_newton_rootsolve()

reaxion/process.py

@@ -0,0 +1,126 @@
+"""Implementation of base Process class with methods for managing and solving systems of equations"""
+
+from .symbols import n_, d_dt
+from .equation import Equation
+from .equation_system import EquationSystem
+
+
+class Process:
+    """
+    Top-level class containing a description of a microscopic process
+
+    Most importantly, this implements the procedure for combining processes to build up a network for chemistry
+    + conservation equations.
+    """
+
+    def __init__(self, name="", bibliography={}):
+        """Construct an empty Process instance
+
+        Parameters
+        ----------
+        name: str, optional
+            Name of the process
+        """
+        self.name = name
+        self.initialize_network()
+        self.rate = 0
+        self.heat = 0
+        self.bibliography = bibliography
+        self.subprocesses = [self]
+
+    def __repr__(self):
+        """Print the name in print()"""
+        return self.name
+
+    def __add__(self, other):
+        """Sum 2 processes together: define a new process whose rates are the sum of the input process"""
+        if other == 0:  # necessary for native sum() routine to work
+            return self
+
+        attrs_to_sum = "heat", "subprocesses", "network"  # all rates
+
+        sum_process = Process()
+        sum_process.rate = None  # "rate" ceases to be meaningful for composite processes
+        for summed_attr in attrs_to_sum:
+            attr1, attr2 = getattr(self, summed_attr), getattr(other, summed_attr)
+            if attr1 is None or attr2 is None:
+                setattr(sum_process, summed_attr, None)
+            else:
+                setattr(sum_process, summed_attr, attr1 + attr2)
+
+        sum_process.name = f"{self.name} + {other.name}"
+        return sum_process
+
+    def __radd__(self, other):
+        return self.__add__(other)
+
+    def initialize_network(self):
+        self.network = EquationSystem()  # this is a dict for which unknown keys are initialized to 0 by default
+
+    @property
+    def heat(self):
+        """Energy lost from gas per unit volume and time"""
+        return self.__heat
+
+    @heat.setter
+    def heat(self, value):
+        """Ensures that the network is always updated when we update the heat"""
+        self.__heat = value
+        self.network["heat"] = Equation(d_dt(n_("heat")), value)
+
+    def solve(
+        self,
+        known_quantities,
+        guess,
+        time_dependent=[],
+        dt=None,
+        verbose=False,
+        tol=1e-3,
+        careful_steps=10,
+    ):
+        """
+        Solves the equations for a set of desired quantities given a set of known quantities
+
+        Parameters
+        ----------
+        known_quantities: dict
+            Dict of symbolic quantities and their values that will be plugged into the network solve as known quantities.
+            Can be arrays if you want to substitute multiple values. If T is included here, we solve for chemical
+            equilibrium. If T is not included, solve for thermochemical equilibrium.
+        guess: dict, optional
+            Dict of symbolic quantities and their values that will be plugged into the network solve as guesses for the
+            unknown quantities. Can be arrays if you want to substitute multiple values. Will default to trying sensible
+            guesses for recognized quantities.
+        normalize_to_H: bool, optional
+            Whether to return abundances normalized by the number density of H nucleons (default: True)
+        reduce_network: bool, optional
+            Whether to solve the reduced version of the network substituting conservation laws (default: True)
+        tol: float, optional
+            Desired relative error in chemical abundances (default: 1e-3)
+        careful_steps: int, optional
+            Number of careful initial steps in the Newton solve before full step size is used - try increasing this if
+            your solve has trouble converging.
+
+        Returns
+        -------
+        soldict: dict
+            Dict of solved quantities
+        """
+
+        return self.network.solve(
+            known_quantities,
+            guess,
+            time_dependent=time_dependent,
+            tol=tol,
+            careful_steps=careful_steps,
+            dt=dt,
+            verbose=verbose,
+        )
+
+    def solver_functions(self, solve_vars, time_dependent=[], return_jac=False, return_dict=False):
+        """Returns the RHS of the system to solve and its Jacobian, applying simplifications"""
+        return self.network.solver_functions(solve_vars, time_dependent, return_jac, return_dict)
+
+    def generate_code(self, solve_vars, time_dependent=[], language="c", jac=True, cse=True):
+        """Generates numerical code that implements the system RHS and/or Jacobian in the specified language."""
+        return self.network.generate_code(solve_vars, time_dependent, language, jac, cse)

reaxion/processes/__init__.py

@@ -0,0 +1,7 @@
+from ..process import *
+from .nbody_process import *
+from .thermal_process import *
+from .freefree_emission import *
+from .line_cooling import *
+from .ionization import *
+from .recombination import *

reaxion/processes/freefree_emission.py

@@ -0,0 +1,32 @@
+"""Implementation of free-free emission as a 2-body process"""
+
+from .nbody_process import NBodyProcess
+from ..misc import species_charge
+from ..symbols import T
+import sympy as sp
+
+
+def gaunt_factor(T):
+    """Fit to data in table 3.3 of Spitzer (1978)"""
+    return 1.1 + 0.34 * sp.exp(-((5.5 - sp.log(T) / sp.log(10)) ** 2) / 3.0)  # 1996ApJS..105...19K
+
+
+def FreeFreeEmission(ion: str) -> NBodyProcess:
+    """Returns a free-free emisison process (i.e. bremmsstrahlung) for the input ion
+
+    Parameters
+    ----------
+    ion: str
+        Ion species
+
+    Returns
+    -------
+    process: NBodyProcess
+        `NBodyProcess` instance describing the cooling process
+    """
+    process = NBodyProcess({ion, "e-"})
+    charge = species_charge(ion)
+    if charge <= 0:
+        raise ValueError(f"{ion} does not appear to be a cation - cannot do bremmstrahlung.")
+    process.heat_rate_coefficient = -1.42e-27 * gaunt_factor(T) * charge**2 * sp.sqrt(T)  # 1996ApJS..105...19K
+    return process

reaxion/processes/ionization.py

@@ -0,0 +1,95 @@
+"""Implementation of ionization process"""
+
+from ..process import Process
+from ..misc import ionize
+from ..symbols import T, T5, T3, T6, n_e, n_
+import sympy as sp
+from astropy import units as u
+
+
+class Ionization(Process):
+    """
+    Class describing an ionization process. Could be collisional, photo, or cosmic ray-induced.
+
+    Implements method for setting the chemistry network terms
+    """
+
+    def __init__(self, species: str, rate_per_volume=0):
+        self.species = species
+        self.ionized_species = ionize(species)
+        self.__ionization_energy = None
+        super().__init__()
+        self.__rate_per_volume = rate_per_volume
+        self.update_network()
+
+    @property
+    def rate(self):
+        return self.__rate_per_volume
+
+    @rate.setter
+    def rate(self, value):
+        self.__rate_per_volume = value
+        self.update_network()
+
+    @property
+    def ionization_energy(self):
+        if self.__ionization_energy is None:
+            self.__ionization_energy = ionization_energy(self.species)
+        return self.__ionization_energy
+
+    def update_network(self):
+        """Sets up rate terms in the associated chemistry network for each species involved"""
+        if self.rate is None:
+            return
+        self.network[self.species] -= self.rate
+        self.network[self.ionized_species] += self.rate
+        self.network["e-"] += self.rate
+
+
+def ionization_energy(species, unit=u.erg):
+    """Return the energy in erg required to ionize a species"""
+    # NOTE: come back and get this from a proper datafile
+    energies_eV = {"H": 13.6, "He": 24.59, "He+": 54.42}
+    return energies_eV[species] * u.eV.to(unit)
+
+
+collisional_ionization_cooling_rates = {
+    "H": 1.27e-21 * sp.sqrt(T) * sp.exp(-157809.1 / T) / (1 + sp.sqrt(T5)),  # 1996ApJS..105...19K
+    "He": 9.38e-22 * sp.sqrt(T) * sp.exp(-285335.4 / T) / (1 + sp.sqrt(T5)),  # 1996ApJS..105...19K
+    "He+": 4.95e-22 * sp.sqrt(T) * sp.exp(-631515 / T) / (1 + sp.sqrt(T5)),  # 1996ApJS..105...19K
+}
+
+collisional_ionization_rates = {
+    "H": 5.85e-11 * sp.sqrt(T) * sp.exp(-157809.1 / T) / (1 + sp.sqrt(T5)),  # 1996ApJS..105...19K
+    "He": 2.38e-11 * sp.sqrt(T) * sp.exp(-285335.4 / T) / (1 + sp.sqrt(T5)),  # 1996ApJS..105...19K
+    "He+": 5.68e-12 * sp.sqrt(T) * sp.exp(-631515 / T) / (1 + sp.sqrt(T5)),  # 1996ApJS..105...19K
+}
+
+
+def CollisionalIonization(species=None) -> Ionization:
+    """Return an ionization process representing collisional ionization of the input species.
+
+    Parameters
+    ----------
+    species: str, optional
+        Species being collisionally ionized. If None, we compose all collisional ionization processes for all ions rates are known.
+
+    Returns
+    -------
+    process: Ionization
+        `Ionization` instance describing the collisional ionization process
+    """
+
+    if species is None:
+        return sum([CollisionalIonization(s) for s in collisional_ionization_rates], Process())
+
+    process = Ionization(species)
+    process.name = f"Collisional Ionization of {species}"
+    nprod = n_(species) * n_e
+
+    if species not in collisional_ionization_rates:
+        raise NotImplementedError(f"{species} does not have an available collisional ionization coefficient.")
+    process.rate = collisional_ionization_rates[species] * nprod
+    process.heat = -process.ionization_energy * process.rate
+
+    return process

reaxion/processes/line_cooling.py

@@ -0,0 +1,56 @@
+import sympy as sp
+from ..process import Process
+from .nbody_process import NBodyProcess
+from ..symbols import T, T5, n_
+from ..data import SolarAbundances
+
+# put analytic fits for cooling efficiencies
+line_cooling_coeffs = {
+    "H": {"e-": 7.5e-19 * sp.exp(-118348 / T) / (1 + sp.sqrt(T5))},  # 1996ApJS..105...19K
+    "He+": {"e-": 5.54e-17 * T**-0.397 * sp.exp(-473638 / T) / (1 + sp.sqrt(T5))},  # 1996ApJS..105...19K
+    "C+": {  # 2023MNRAS.519.3154H
+        "e-": 1e-27 * 4890 / sp.sqrt(T) * sp.exp(-91.211 / T) / SolarAbundances.x("C"),
+        "H": 1e-27 * 0.47 * T**0.15 * sp.exp(-91.211 / T) / SolarAbundances.x("C"),
+    },
+}
+
+
+def LineCoolingSimple(emitter: str, collider=None) -> NBodyProcess:
+    """Returns a 2-body process representing cooling via excitations from collisions of given pair of species
+
+    This is the simple approximation where everything is well below critical density and no ambient radiation field.
+    eventually would like to have a class that considers collisions from all available colliders, given just the
+    energies, deexcitation coefficients, temperature, and statistical weights...
+
+    Parameters
+    ----------
+    emitter: str
+        Emitting excited species
+    collider: str, optional
+        Exciting colliding species. If None, will look up all known
+
+    Returns
+    -------
+    An NBodyProcess instance whose heat attribute is the line cooling process's cooling rate in erg cm^-3
+    """
+
+    if emitter not in line_cooling_coeffs:
+        raise NotImplementedError(f"Line cooling not implemented for {emitter}")
+
+    coeffs = line_cooling_coeffs[emitter]
+
+    if collider is None:  # if we haven't specified a collider, just take all of them and return the sum
+        p = [LineCoolingSimple(emitter, c) for c in coeffs]
+        return sum(p)  # type: ignore # have to put a 0-process in here as start variable or it will try to add 0 + process
+
+    process = NBodyProcess({emitter, collider})
+    if collider not in line_cooling_coeffs[emitter]:
+        raise NotImplementedError(f"Excitation by collisions with {collider} not implemented for {emitter}")
+
+    process.heat_rate_coefficient = -line_cooling_coeffs[emitter][collider]
+    process.name = f"{emitter}-{collider} Line Cooling"
+
+    return process
+
+
+# def LineCooling(emitter: str)

reaxion/processes/nbody_process.py

@@ -0,0 +1,71 @@
+"""Class specifying an N-body collisional process with generic methods"""
+
+from ..process import Process
+from ..symbols import n_
+import sympy as sp
+
+
+class NBodyProcess(Process):
+    """Process implementing special methods specific to n-body processes, whose rates
+    all follow the pattern
+
+    rate per volume = k * prod_i(n_i) for i in species
+
+    rate and heat are promoted from attributes to properties implemented to compute this pattern.
+
+    Parameters
+    ----------
+    colliding_species:
+        iterable of strings representing the colliding species.
+    rate_coefficient: sympy.core.symbol.Symbol, optional
+        Symbol symbol expression for k
+    heat_rate_coefficient: sympy.core.symbol.Symbol, optional
+        Symbol symbol expression for the heat rate coefficient = average radiated energy * k
+    name: str, optional
+        Name of the process
+    """
+
+    def __init__(self, colliding_species, rate_coefficient=0, heat_rate_coefficient=0, name: str = ""):
+        self.name = name
+        self.initialize_network()
+        self.colliding_species = colliding_species
+        self.rate_coefficient = rate_coefficient
+        self.heat_rate_coefficient = heat_rate_coefficient
+        self.subprocesses = [self]
+
+    @property
+    def heat_rate_coefficient(self):
+        """Returns the heat rate coefficient of the N-body process"""
+        return self.__heat_rate_coefficient
+
+    @heat_rate_coefficient.setter
+    def heat_rate_coefficient(self, value):
+        """Ensures that the network is always updated when we update the rate coefficient"""
+        self.__heat_rate_coefficient = value
+        self.heat = value * sp.prod([n_(c) for c in self.colliding_species])
+
+    @property
+    def rate(self):
+        """Returns the number of events per unit time and volume"""
+        if self.rate_coefficient is None:
+            return None
+        return self.rate_coefficient * sp.prod([n_(c) for c in self.colliding_species])
+
+    # @property
+    # def heat(self):
+    #     """Returns the energy radiated per unit time and volume"""
+    #     return super().heat
+
+    # #        self.__heat = self.heat_rate_coefficient * sp.prod([sp.Symbol("n_" + c) for c in self.colliding_species])
+    # #        return self.__heat
+
+    # @heat.setter
+    # def heat(self, value):
+    #     raise NotImplementedError(
+    #         "Cannot directly set the heat value of an N-body process - set the rate coefficient instead."
+    #     )
+
+    @property
+    def num_colliding_species(self):
+        """Number of colliding species"""
+        return len(self.colliding_species)