API Documentation

class jaco.Equation(lhs, rhs, **options)

Sympy equation where we overload addition/subtraction to apply those operations to the RHS, for summing rate equations

get_summand(other): Value-check the operand and return the quantity to be summed in the operation: the expression itself if an expression, or the RHS

class jaco.EquationSystem

Dict of symbolic expressions with certain superpowers for manipulating sets of conservation equations.

property chemical_species: Returns a tuple of all chemical species detected within the network

copy(): Return a shallow copy of the dict.

do_conservation_reductions(time_dependent_vars): Eliminate equations from the system using known conservation laws.

generate_code(solve_vars, time_dependent=[], language='Fortran', jac=True, do_cse=True): Generates numerical code that implements the system RHS and/or Jacobian in the specified language.

property jacobian: 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.

package_solution(sol, guessvals, guesses, paramvals, subsystem, symbolic_keys): Takes the output of newton_rootsolve and packages it into a dict containing the system solution for all variables, including substitutions for eliminated variables.

property rhs: Return as dict of rhs-lhs instead of equations

property rhs_scaled: Returns a scaled version of the the RHS pulling out the usual factors affecting collision rates

set_time_dependence(time_dependent_vars): Insert backward-difference formulae or set to steady state

solve(knowns, guesses, time_dependent=[], dt=None, verbose=False, tol=0.001, careful_steps=20, 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 of species and their equilibrium abundances relative to H or raw number densities (depending on value of normalize_to_H)

Return type:

dict

solver_functions(solve_vars, time_dependent=[], return_jac=False, return_dict=False): Returns the RHS of the system to solve and its Jacobian, applying simplifications

subs(expr, replacement): Substitute symbolic expressions throughout the whole network.

property symbols: Returns the set of all symbols in the equations

class jaco.Process(name='', bibliography=[])

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.

Construct an empty Process instance

Parameters:: name (str, optional) – Name of the process

generate_code(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.

property heat: Energy lost from gas per unit volume and time

solve(known_quantities, guess, time_dependent=[], dt=None, verbose=False, tol=0.001, 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 of solved quantities

Return type:

dict

solver_functions(solve_vars, time_dependent=[], return_jac=False, return_dict=False): Returns the RHS of the system to solve and its Jacobian, applying simplifications

jaco.processes.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 instance describing the cooling process
Return type:: NBodyProcess

jaco.processes.GasPhaseRecombination(ion=None) → Recombination

Return a recombination process representing gas-phase (e.g. radiative) recombination

Parameters:: ion (str, optional) – Ionic species getting recombined. If None, function will return a composite process of all gas-phase recombination processes with known rates.
Returns:: process – Process instance describing the gas-phase recombination process
Return type:: Recombination

class jaco.processes.Ionization(species: str, rate_per_volume=0)

Class describing an ionization process. Could be collisional, photo, or cosmic ray-induced.

Implements method for setting the chemistry network terms

Construct an empty Process instance

Parameters:: name (str, optional) – Name of the process

update_network(): Sets up rate terms in the associated chemistry network for each species involved

jaco.processes.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

Return type:

An NBodyProcess instance whose heat attribute is the line cooling process’s cooling rate in erg cm^-3

class jaco.processes.NBodyProcess(colliding_species, rate_coefficient=0, heat_rate_coefficient=0, name: str = '', bibliography=[])

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. Species should be repeated if more than 1 is involved.
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

Construct an empty Process instance

Parameters:: name (str, optional) – Name of the process

property heat_rate_coefficient: Returns the heat rate coefficient of the N-body process

property nprod: Returns the product of species number densities

property num_colliding_species: Number of colliding species

property rate: Returns the number of events per unit time and volume

class jaco.processes.Recombination(ion: str)

Class describing a recombination process: ion + e- -> recombined species + hν

Implements method for setting the chemistry network terms

Parameters:: ion (str) – Ionic species being recombined

Construct an empty Process instance

Parameters:: name (str, optional) – Name of the process

property rate_coefficient: Returns the rate coefficient of the recombination process

update_network(): Sets up rate terms in the associated chemistry network for each ion involved

class jaco.processes.ThermalProcess(heating_rate, name='', bibliography=[])

Generic heating/cooling process

Construct an empty Process instance

Parameters:: name (str, optional) – Name of the process

jaco.numerics.newton_rootsolve(func, guesses, params=[], jacfunc=None, tolfunc=None, rtol=1e-06, 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 – Shape (N,n) array of solutions

Return type:

array_like