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 aspect_ratio#

Return the aspect ratio of the simulation domain in the x direction. nx = aspect_ratio * resolution; ny = nz = resolution. Cell spacing dx = dy = dz is preserved.

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 domain_size#

Return the physical domain size in each dimension (kpc).

property dx#

Return the (uniform) cell size of the simulation (kpc). dx = dy = dz = box_size / resolution; Lx = aspect_ratio * box_size.

property grid#

Return the simulation grid (xx, yy, zz) in kpc. Shape is (nx, ny, nz); Lx = aspect_ratio * box_size, Ly = Lz = box_size.

property kgrid#

Return the simulation spectral grid (kx, ky, kz) in kpc^-1. For a non-cubic domain the fundamental mode differs along x vs y/z. After jaxdecomp’s pfft3d the array axis order is (Y, Z, X), so the returned k-arrays are already in that transposed layout.

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 nx#

Return the number of cells in the x direction (elongated axis).

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. Passes dx directly so the result is correct for non-cubic domains.

property resolution#

Return the (linear) resolution of the simulation (ny == nz)

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 shape#

Return the grid shape (nx, ny, nz).

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