cosmax.spectral_ops package

Submodules

cosmax.spectral_ops.generator module

class cosmax.spectral_ops.generator.Generator(elements: int, size: float = 1.0)

Bases: SpectralOperation

Generates initial conditions given a power spectrum

Parameters:

  • elements – number of grid points in each dimension

  • size – size of the box in real space

elements

number of grid points in each dimension

Type:

int

size

size of the box in real space

Type:

float

k

wavenumber of each bin

index_grid

index of the bin for each wavenumber

Type:

jax.Array

index_grid: Array
n_bins: int

cosmax.spectral_ops.potential module

class cosmax.spectral_ops.potential.Potential(elements: int)

Bases: SpectralOperation

Gravitational potential from a 3D density field.

Parameters:

elements – number of grid points in each dimension

cosmax.spectral_ops.power_spectrum module

class cosmax.spectral_ops.power_spectrum.PowerSpectrum(elements: int, bins: int, size: float = 1.0)

Bases: SpectralOperation

Power spectrum from a 3D density field.

Parameters:

  • elements – number of grid points in each dimension

  • bins – number of bins for the power spectrum

  • size – size of the box in real space

elements

number of grid points in each dimension

Type:

int

size

size of the box in real space

Type:

float

bins

number of bins for the power spectrum

Type:

int

k

wavenumber of each bin

index_grid

index of the bin for each wavenumber

Type:

jax.Array

n_modes

number of modes in each bin

Type:

jax.Array

bins: int
index_grid: Array
n_modes: Array

cosmax.spectral_ops.spectral_loss module

class cosmax.spectral_ops.spectral_loss.SpectralLoss(elements: int, bins: int, size: float = 1.0)

Bases: SpectralOperation

Compute the MSE in spectral space for each wavenumber

Parameters:

  • elements – number of grid points in each dimension

  • bins – number of bins for the power spectrum

  • size – size of the box in real space

elements

number of grid points in each dimension

Type:

int

size

jax.Array

Type:

float

bins

number of bins for the power spectrum

Type:

int

index_grid

index of the bin for each wavenumber

Type:

jax.Array

n_modes

number of modes in each bin

Type:

jax.Array

bins: int
index_grid: Array
n_modes: Array

cosmax.spectral_ops.spectral_op module

class cosmax.spectral_ops.spectral_op.SpectralOperation(elements: int, size: float = 1.0)

Bases: object

Base class for spectral operations

Parameters:

  • elements – number of grid points in each dimension

  • size – size of the box in real space

elements

number of grid points in each dimension

Type:

int

size

size of the box in real space

Type:

float

frequencies

frequencies in fourier space

Type:

jax.Array

real_frequencies

real frequencies in fourier space

k_mag

magnitude of the wavenumber

elements: int
frequencies: Array
k_mag

alias of Array

nyquist: int
size: float

Module contents