Jaxion API

Jaxion API#

simulation#

class jaxion.simulation.Simulation(params, sharding=None, checkpoint_number=None)#

Simulation: The base class for an astrophysics simulation.

Parameters:

  • params (dict) – The Python dictionary that contains the simulation parameters. Params can also be a string path to a checkpoint directory to load a saved simulation.

  • sharding (jax.sharding.NamedSharding, optional) – jax sharding used for distributed (multi-GPU) simulations

property axion_mass#

Return the axion particle mass in the simulation (M_sun)

property box_size#

Return the box size of the simulation (kpc)

property dx#

Return the cell size size of the simulation (kpc)

property grid#

Return the simulation grid

property kgrid#

Return the simulation spectral grid

property m_per_hbar#

Return the mass per hbar in the simulation (M_sun / hbar)

property nt#

Return the number of timesteps

property num_particles#

Return the number of particles in the simulation

property params#

Return the parameters of the simulation

property potential#

Return the gravitational potential

property quantum_velocity#

Return the dark matter velocity field from the wavefunction

property resolution#

Return the (linear) resolution of the simulation

property rho_bar#

Return the mean density of the simulation

run()#

Run the simulation

property scattering_length#

Return the axion self-interaction scattering length in the simulation (kpc)

property sound_speed#

Return the isothermal gas sound speed in the simulation (km/s)

constants#

Physical constants in units of: [L] = kpc, [V] = km/s, [M] = Msun

note: other units are derived from these base units, e.g., [T] = [L]/[V] = kpc / (km/s) ~= 0.978 Gyr

jaxion.constants#

alias of {‘electron_volt’: 8.05478173e-56, ‘gravitational_constant’: 4.30241002e-06, ‘reduced_planck_constant’: 1.71818134e-87, ‘speed_of_light’: 299792.458}

analysis#

jaxion.radial_power_spectrum(data_cube, kx, ky, kz, box_size)#

Computes the radially averaged power spectral density of a 3D datacube.

Parameters:

  • data_cube (jnp.ndarray) – 3D data cube, must be cubic

  • kx (jnp.ndarray) – wavenumber grids in each dimension

  • ky (jnp.ndarray) – wavenumber grids in each dimension

  • kz (jnp.ndarray) – wavenumber grids in each dimension

  • box_size (float) – physical size of box

Returns:

  • Pf (jnp.ndarray) – radial power spectrum

  • k (jnp.ndarray) – wavenumbers

  • total_power (float) – total power