reaxion 0.1.1__py3-none-any.whl

reaxion/__init__.py

@@ -0,0 +1,5 @@
+from .misc import *
+from .numerics import *
+from .process import Process
+from .equation_system import EquationSystem
+from .equation import Equation

reaxion/data/__init__.py

@@ -0,0 +1 @@
+from .solar_abundances import SolarAbundances

reaxion/data/atomic_weights.py

@@ -0,0 +1,123 @@
+atomic_weights = {
+    "Z": float("nan"),
+    "H": 1.008,
+    "He": 4.002602,
+    "Li": 6.94,
+    "Be": 9.0121831,
+    "B": 10.81,
+    "C": 12.011,
+    "N": 14.007,
+    "O": 15.999,
+    "F": 18.998403163,
+    "Ne": 20.1797,
+    "Na": 22.98976928,
+    "Mg": 24.305,
+    "Al": 26.9815384,
+    "Si": 28.085,
+    "P": 30.973761998,
+    "S": 32.06,
+    "Cl": 35.45,
+    "Ar": 39.95,
+    "K": 39.0983,
+    "Ca": 40.078,
+    "Sc": 44.955907,
+    "Ti": 47.867,
+    "V": 50.9415,
+    "Cr": 51.9961,
+    "Mn": 54.938043,
+    "Fe": 55.845,
+    "Co": 58.933194,
+    "Ni": 58.6934,
+    "Cu": 63.546,
+    "Zn": 65.38,
+    "Ga": 69.723,
+    "Ge": 72.630,
+    "As": 74.921595,
+    "Se": 78.971,
+    "Br": 79.904,
+    "Kr": 83.798,
+    "Rb": 85.4678,
+    "Sr": 87.62,
+    "Y": 88.905838,
+    "Zr": 91.224,
+    "Nb": 92.90637,
+    "Mo": 95.95,
+    "Tc": 97,
+    "Ru": 101.07,
+    "Rh": 102.90549,
+    "Pd": 106.42,
+    "Ag": 107.8682,
+    "Cd": 112.414,
+    "In": 114.818,
+    "Sn": 118.710,
+    "Sb": 121.760,
+    "Te": 127.60,
+    "I": 126.90447,
+    "Xe": 131.293,
+    "Cs": 132.90545196,
+    "Ba": 137.327,
+    "La": 138.90547,
+    "Ce": 140.116,
+    "Pr": 140.90766,
+    "Nd": 144.242,
+    "Pm": 145,
+    "Sm": 150.36,
+    "Eu": 151.964,
+    "Gd": 157.25,
+    "Tb": 158.925354,
+    "Dy": 162.500,
+    "Ho": 164.930329,
+    "Er": 167.259,
+    "Tm": 168.934219,
+    "Yb": 173.045,
+    "Lu": 174.9668,
+    "Hf": 178.486,
+    "Ta": 180.94788,
+    "W": 183.84,
+    "Re": 186.207,
+    "Os": 190.23,
+    "Ir": 192.217,
+    "Pt": 195.084,
+    "Au": 196.966570,
+    "Hg": 200.592,
+    "Tl": 204.38,
+    "Pb": 207.2,
+    "Bi": 208.98040,
+    "Po": 209,
+    "At": 210,
+    "Rn": 222,
+    "Fr": 223,
+    "Ra": 226,
+    "Ac": 227,
+    "Th": 232.0377,
+    "Pa": 231.03588,
+    "U": 238.02891,
+    "Np": 237,
+    "Pu": 244,
+    "Am": 243,
+    "Cm": 247,
+    "Bk": 247,
+    "Cf": 251,
+    "Es": 252,
+    "Fm": 257,
+    "Md": 258,
+    "No": 259,
+    "Lr": 262,
+    "Rf": 267,
+    "Db": 270,
+    "Sg": 269,
+    "Bh": 270,
+    "Hs": 270,
+    "Mt": 278,
+    "Ds": 281,
+    "Rg": 281,
+    "Cn": 285,
+    "Nh": 286,
+    "Fl": 289,
+    "Mc": 289,
+    "Lv": 293,
+    "Ts": 293,
+    "Og": 294,
+}
+
+bibliography = "https://iupac.qmul.ac.uk/AtWt/"

reaxion/data/solar_abundances.py

@@ -0,0 +1,49 @@
+from .atomic_weights import atomic_weights
+
+
+class SolarAbundancesClass:
+    """Container for solar abundances with methods to convert between mass fraction and abundance per H"""
+
+    bibliography = ["2009ARA&A..47..481A"]
+
+    @property
+    def mass_fraction(self):
+        """Returns a hard-coded dict of Solar abundance mass fractions"""
+        return {
+            "Z": 0.0142,
+            "He": 0.27030,
+            "C": 2.53e-3,
+            "N": 7.41e-4,
+            "O": 6.13e-3,
+            "Ne": 1.34e-3,
+            "Mg": 7.57e-4,
+            "Si": 7.12e-4,
+            "S": 3.31e-4,
+            "Ca": 6.87e-5,
+            "Fe": 1.38e-3,
+        }
+
+    @property
+    def abundance_per_H(self):
+        """Returns dictionary of abundances per H nucleon"""
+        return {species: f / (1 - f) / atomic_weights[species] for species, f in self.mass_fraction.items()}
+
+    def x(self, species):
+        return self.get_abundance(species)
+
+    def get_mass_fraction(self, species: str) -> float:
+        """Returns the mass fraction of a given species"""
+        if species in self.mass_fraction:
+            return self.mass_fraction[species]
+        else:
+            raise NotImplementedError(f"Solar abundance of {species} not available.")
+
+    def get_abundance(self, species: str) -> float:
+        """Returns the abundance per H nuclear of an input species"""
+        if species in self.mass_fraction:
+            return self.abundance_per_H[species]
+        else:
+            raise NotImplementedError(f"Solar abundance of {species} not available.")
+
+
+SolarAbundances = SolarAbundancesClass()

reaxion/eos/__init__.py

@@ -0,0 +1 @@
+# from .eos import u, rho

reaxion/eos/eos.py

@@ -0,0 +1,3 @@
+"""Should eventually contain routines to generate the symbolic expression for internal energy and pressure, given
+the abundances available in a network.
+"""

reaxion/equation.py

@@ -0,0 +1,41 @@
+"""Implementation of Equation class for representing conservation laws"""
+
+import sympy as sp
+
+
+class Equation(sp.core.relational.Equality):
+    """Sympy equation where we overload addition/subtraction to apply those operations to the RHS, for summing rate
+    equations"""
+
+    def get_summand(self, other):
+        """Value-check the operand and return the quantity to be summed in the operation: the expression itself if an expression, or the RHS"""
+        if isinstance(other, sp.core.relational.Equality):
+            if self.lhs != other.lhs:
+                raise ValueError(
+                    "Tried to sum incompatible equations. Equation summation only defined for differential equations with the same LHS."
+                )
+            else:
+                return other.rhs
+        elif isinstance(other, sp.logic.boolalg.BooleanAtom):
+            return 0
+        else:
+            return other
+
+    def __add__(self, other):
+        summand = self.get_summand(other)
+        return Equation(self.lhs, self.rhs + summand)
+
+    def __sub__(self, other):
+        summand = self.get_summand(other)
+        return Equation(self.lhs, self.rhs - summand)
+
+    def __radd__(self, other):
+        return self.__add__(other)
+
+    def __iadd__(self, other):
+        self = self + other
+        return self
+
+    def __isub__(self, other):
+        self = self - other
+        return self

reaxion/equation_system.py

@@ -0,0 +1,380 @@
+"""Implementation of EquationSystem for representing, manipulating, and constructing systems of conservation laws"""
+
+import sympy as sp
+from .symbols import d_dt, n_, x_, t, BDF, n_Htot, internal_energy
+from .data import SolarAbundances
+from jax import numpy as jnp
+import numpy as np
+from .numerics import newton_rootsolve
+from astropy import units
+from .equation import Equation
+from sympy.codegen.ast import Assignment
+
+
+class EquationSystem(dict):
+    """Dict of symbolic expressions with certain superpowers for manipulating sets of conservation equations."""
+
+    def copy(self):
+        new = EquationSystem()
+        for k in self:
+            new[k] = self[k]
+        return new
+
+    def __getitem__(self, __key: str):
+        """Dict getitem method where we initialize a differential equation for the conservation of a species if the key
+        does not exist"""
+        if __key not in self:
+            self.__setitem__(__key, Equation(d_dt(n_(__key)), 0))  # technically should only be n_ if this is a species
+            # need to make sure that d/dt's don't add up when composing equations
+        return super().__getitem__(__key)
+
+    def __add__(self, other):
+        """Return a dict whose values are the sum of the values of the operands"""
+        keys = self.keys() | other.keys()
+        new = EquationSystem()
+        for k in keys:
+            new[k] = self[k] + other[k]
+        return new
+
+    @property
+    def symbols(self):
+        """Returns the set of all symbols in the equations"""
+        all = set()
+        for e in self.values():
+            all.update(e.free_symbols)
+        if t in all:  # leave time out
+            all.remove(t)
+        return all
+
+    @property
+    def jacobian(self):
+        """Returns a dict of dicts representing the Jacobian of the RHS of the system. Keys are the names of the
+        conserved quantities and subkeys are the variable of differentiation.
+        """
+        return {k: {s: sp.diff(e.rhs, s) for s in self.symbols} for k, e in self.items()}
+
+    def subs(self, expr, replacement):
+        """Substitute symbolic expressions throughout the whole network."""
+        for k, e in self.items():
+            self[k] = e.subs(expr, replacement)
+
+    def reduced(self, knowns, time_dependent=[]):
+        subsystem = self.copy()
+        subsystem.set_time_dependence(time_dependent)
+        subsystem.do_conservation_reductions(time_dependent)
+        if "T" in (str(k) for k in knowns) and "T" not in time_dependent:
+            del subsystem["heat"]
+        return subsystem
+
+    def set_time_dependence(self, time_dependent_vars):
+        """Insert backward-difference formulae or set to steady state"""
+        # put in backward differences
+        for q in self:
+            if q in time_dependent_vars:  # insert backward-difference formula
+                self[q] = Equation(BDF(q), self[q].rhs)
+            else:
+                self[q] = Equation(0, self[q].rhs)
+        if "T" in time_dependent_vars:  # special behaviour
+            self["heat"] = Equation(BDF("T"), self["heat"].rhs)
+            if "u" not in self:
+                self["u"] = Equation(0, sp.Symbol("u") - internal_energy)
+
+    def do_conservation_reductions(self, time_dependent_vars):
+        """Eliminate equations from the system using known conservation laws."""
+        self.substitutions = []
+
+        # since we have n_Htot let's convert all other n's to x's
+        for s in self.symbols:
+            if "n_" in str(s) and "Htot" not in str(s):
+                species = str(s).split("_")[1]
+                self.substitutions.append((s, n_Htot * x_(species)))
+
+        # charge neutrality
+        if "e-" not in time_dependent_vars:
+            self.substitutions.append((x_("e-"), x_("H+") + x_("He+") + 2 * x_("He++")))
+            del self["e-"]
+
+        #  general: sum(n_(species containing H) / (number of H in species))  - n_("H_2") / 2 #
+        if "H+" not in time_dependent_vars:
+            self.substitutions.append((x_("H+"), 1 - x_("H")))
+            if "H+" in self:
+                del self["H+"]
+
+        if "He++" not in time_dependent_vars:
+            y = sp.Symbol("y")
+            self.substitutions.append((x_("He++"), y - x_("He") - x_("He+")))
+            if "He++" in self:
+                del self["He++"]
+
+        for expr, sub in self.substitutions:
+            self.subs(expr, sub)
+
+            # general: substitute highest ionization state with n_Htot * x_element - sum of lower ionization states
+
+    @property
+    def rhs(self):
+        """Return as dict of rhs-lhs instead of equations"""
+        return {k: e.rhs - e.lhs for k, e in self.items()}
+
+    @property
+    def rhs_scaled(self):
+        """Returns a scaled version of the the RHS pulling out the usual factors affecting collision rates"""
+        return [r for r in self.rhs.values()]  # / (T**0.5 * n_Htot * n_Htot * 1e-12)
+
+    def solve(
+        self,
+        knowns,
+        guesses,
+        time_dependent=[],
+        dt=None,
+        verbose=False,
+        tol=1e-3,
+        careful_steps=10,
+        symbolic_keys=False,
+    ):
+        """
+        Solves for equilibrium after substituting a set of known quantities, e.g. temperature, metallicity,
+        etc.
+
+        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.
+        guesses: dict
+            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 (NOT IMPLEMENTED YET)
+        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 species and their equilibrium abundances relative to H or raw number densities (depending on
+            value of normalize_to_H)
+        """
+
+        def printv(*a, **k):
+            """Print only if locally verbose=True"""
+            if verbose:
+                print(*a, **k)
+
+        # first: check knowns and guesses are all same size
+        num_params = np.array([len(np.array(guesses[g])) for g in guesses] + [len(np.array(knowns[g])) for g in knowns])
+        if not np.all(num_params == num_params[0]):
+            raise ValueError("Input parameters and initial guesses must all have the same shape.")
+        num_params = num_params[0]
+
+        if dt is not None:
+            knowns["Δt"] = np.repeat(dt.to(units.s), num_params)
+
+        if "u" in guesses or "T" in time_dependent:
+            self["u"] = Equation(0, internal_energy - sp.Symbol("u"))
+        subsystem = self.reduced(knowns, time_dependent)
+        symbols = subsystem.symbols
+        num_equations = len(subsystem)
+
+        # are there any symbols for which we can make a reasonable assumption or directly solve the steady-state approximation?
+        prescriptions = {"y": SolarAbundances.x("He"), "Y": SolarAbundances.mass_fraction["He"], "Z": 1.0}
+        assumed_values = {}
+        if len(symbols) > num_equations + len(knowns):
+            undetermined_symbols = symbols.difference(set(sp.Symbol(g) for g in guesses))
+            printv(f"Undetermined symbols: {undetermined_symbols}")
+            for s in undetermined_symbols:
+                # if we have a prescription for this quantity, plug it in here. This should eventually be specified at the model level.
+                if str(s) in prescriptions:
+                    # case 1: we have given a value, which we should add to the list of knowns
+                    assumed_values[str(s)] = np.repeat(prescriptions[str(s)], num_params)
+                    printv(f"{s} not specified; assuming {s}={prescriptions[str(s)]}.")
+                    symbols = subsystem.symbols
+                    # case 2: we have given an expression in terms of the other available quantities: we need to subs it
+
+        # ok now we should have number of symbols unknowns + knowns
+        printv(
+            f"Free symbols: {symbols}\nKnown values: {list(knowns)}\nAssumed values: {list(assumed_values)}\nEquations solved: {list(subsystem.rhs)}"
+        )
+        if len(symbols) != len(knowns | assumed_values) + len(subsystem):
+            raise ValueError(
+                f"Number of free symbols is {len(symbols)} != number of knowns {len(knowns)} + number of assumptions {len(assumed_values)} + number of equations {len(subsystem)}\n"
+            )
+        else:
+            printv(
+                f"It's solvin time. Solving for {set(guesses)} based on input {set(knowns)} and assumptions about {set(assumed_values)}"
+            )
+
+        guessvals = {}
+        paramvals = {}
+        for s in subsystem.symbols:
+            for g in guesses:
+                if g == str(s) or f"x_{g}" == str(s):
+                    guessvals[s] = guesses[g]
+            for k in knowns | assumed_values:
+                if k == str(s) or f"x_{k}" == str(s):
+                    paramvals[s] = (knowns | assumed_values)[k]
+
+        lambda_args = [list(guessvals.keys()), list(paramvals.keys())]
+        func = sp.lambdify(lambda_args, subsystem.rhs_scaled, modules="jax", cse=True)
+
+        tolerance_vars = [x_("H"), x_("He+") + x_("He"), 1 - x_("H")]
+        if "T" in guesses:
+            tolerance_vars += [sp.Symbol("T")]
+        if "u" in guesses:
+            tolerance_vars += [sp.Symbol("u"), subsystem["heat"].rhs]
+            # , subsystem["heat"]]  # converge on the internal energy and  cooling rate
+        tolfunc = sp.lambdify(lambda_args, tolerance_vars, modules="jax", cse=True)
+
+        def f_numerical(X, *params):
+            """JAX function to rootfind"""
+            return jnp.array(func(X, params))
+
+        def tolerance_func(X, *params):
+            """Solution will terminate if the relative change in this quantity is < tol"""
+            return jnp.array(tolfunc(X, params))
+
+        # option to bail here and just provide the RHS
+
+        # jacfunc = sp.lambdify(
+        #     lambda_args, [[sp.diff(a, g) for g in guessvals] for a in subsystem.rhs_scaled]
+        # )  # , modules="jax", cse=True
+
+        sol, num_iter = newton_rootsolve(
+            f_numerical,
+            jnp.array([g for g in guessvals.values()]).T,
+            jnp.array([p for p in paramvals.values()]).T,
+            tolfunc=tolerance_func,
+            rtol=tol,
+            careful_steps=careful_steps,
+            nonnegative=True,
+            return_num_iter=True,
+        )
+
+        soldict = self.package_solution(sol, guessvals, guesses, paramvals, subsystem, symbolic_keys)
+
+        return soldict
+
+    def package_solution(self, sol, guessvals, guesses, paramvals, subsystem, symbolic_keys):
+        # now repack the solution
+        soldict = {}
+        for i, g in enumerate(guessvals):
+            soldict[g] = sol[:, i]
+        # do a reverse-pass on the substitutions we made to get all quantities
+        values_to_subs = soldict | paramvals
+        for expr, sub in reversed(subsystem.substitutions):
+            if expr in soldict:
+                continue
+            if "n_" in str(expr):
+                continue
+            soldict[expr] = sp.lambdify(list(sub.free_symbols), sub)(
+                *[values_to_subs[s] for s in list(sub.free_symbols)]
+            )  # should probably make a function of this
+            values_to_subs |= soldict
+        if not symbolic_keys:
+            soldict = {str(k): v for k, v in soldict.items()}
+            # if we have a bunch of x_'s, should also link up keys in the original input format, e.g. H->x_H
+            if np.any(["x_" in k for k in guesses]):  # if we specified abundances with x_ notation, return same
+                return soldict
+            soldict2 = {}  # otherwise return with input format where keys are simple species strings
+            for k in soldict:
+                if "x_" in str(k):
+                    soldict2[str(k).replace("x_", "")] = soldict[k]
+                else:
+                    soldict2[k] = soldict[k]
+            soldict = soldict2
+
+        return soldict
+
+    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"""
+
+        solve_vars = list(solve_vars)
+        if "u" in solve_vars or "T" in time_dependent:
+            self["u"] = Equation(0, internal_energy - sp.Symbol("u"))
+            solve_vars.append("u")
+
+        knowns = self.symbols.difference(solve_vars)
+        subsystem = self.reduced(knowns, time_dependent)
+
+        rhs = {}
+        for s in subsystem.symbols:
+            for g in solve_vars:
+                if str(s) == "T" and "T" in solve_vars:
+                    rhs[s] = subsystem.rhs["heat"]
+                elif str(g) == str(s) or f"x_{g}" == str(s):
+                    rhs[s] = subsystem.rhs[g]
+
+        if return_jac:
+            jac = {}
+            for s, expr in rhs.items():
+                jac[s] = {s2: sp.diff(expr, s2) for s2 in rhs}
+
+            if return_dict:
+                return rhs, jac
+            else:
+                return (
+                    list(rhs.values()),
+                    [[jac[s1][s2] for s2 in rhs] for s1 in jac],
+                    {s: i for i, s in enumerate(rhs)},
+                )
+
+        if return_dict:
+            return rhs
+        else:
+            return list(rhs.values()), {s: i for i, s in enumerate(rhs)}
+
+    def generate_code(self, solve_vars, time_dependent=[], language="Fortran", jac=True, cse=True, sanitize=True):
+        """Generates numerical code that implements the system RHS and/or Jacobian in the specified language."""
+        func, jac, indices = self.solver_functions(solve_vars, time_dependent, return_jac=jac)
+
+        def printer(x, language="c"):
+            match language.lower():
+                case "fortran":
+                    return sp.fcode(x, standard=2008)
+                case "c":
+                    return sp.ccode(x, standard="c99")
+                case "python":
+                    return sp.pycode(x)
+                case "c++":
+                    return sp.cxxcode(x, standard="c++11")
+
+        codeblocks = []
+
+        header = "# Computes the RHS function "
+        if jac:
+            header += "and Jacobian "
+        header += f"to solve for {list(indices.keys())}\n\n"
+
+        header += "# INDEX CONVENTION: " + " ".join(f"({i}: {s})" for s, i in indices.items())
+
+        codeblocks.append(header)
+
+        if cse:
+            cse, (func, jac) = sp.cse((sp.Matrix(func), sp.Matrix(jac)))
+            block = []
+            for expr in cse:
+                block.append(printer(Assignment(*expr), language))
+            codeblocks.append(" \n".join(block))
+
+        rhs_result = sp.MatrixSymbol("rhs_result", len(func), 1)
+        codeblocks.append(printer(Assignment(rhs_result, func), language))
+
+        if jac:
+            jac_result = sp.MatrixSymbol("jac_result", len(func), len(func))
+            codeblocks.append(printer(Assignment(jac_result, jac), language))
+
+        code = "\n\n".join(codeblocks)
+        if sanitize:
+            sanitized_code = ""
+            replacements = {"+": "plus"}
+            for i, char in enumerate(code):
+                if char in replacements:
+                    if code[i - 1].split():  # if preceding character is not whitespace
+                        sanitized_code += replacements[char]
+                        continue
+                sanitized_code += char
+            code = sanitized_code
+        return code

reaxion/localstate.py

@@ -0,0 +1,21 @@
+"""Specifies class for the local thermal and chemical state"""
+
+
+class State:
+    """Internal energy, abundances, dust mass fraction, radiation field, cosmic ray ionization rate,
+    ortho-para ratio,"""
+
+    f_ortho: float = 0.75
+    hydrogen_massfrac: float = 0.7381
+    metallicity: float = 0.0134
+
+    def __init__(self):
+        pass
+
+    @property
+    def density(self, species):
+        return 0
+
+
+#   @property
+#    def