Galaxy PSM Dataset
This class organizes power spectrum multipoles into a format expected by PyTorch.
- class mentat_lss.dataset.pk_galaxy_dataset(data_dir: str, type: str, frac=1.0)[source]
Custom dataset storing large sets of galaxy power spectrum multipoles
- get_normalized_galaxy_power_spectra(idx)[source]
Returns the normalized power spectrum multipoles corresponding to idx
- Parameters:
idx (int or torch.Tensor) – index (or set of indexes) to access
- Returns:
normalized power spectrum to access.
- Return type:
galaxy_ps[idx] (torch.Tensor)
- get_true_galaxy_power_spectra(idx, ps_fid: Tensor, sqrt_eigvals: Tensor, Q: Tensor, Q_inv: Tensor)[source]
Returns the galaxy power spectrum multipoles in units of (Mpc/h)^3 corresponding to idx
- Parameters:
idx (int or torch.Tensor) – index (or set of indexes) to access
ps_fid (torch.Tensor) – fiducial power spectrum used to reverse normalization. Expected shape is [nps*nz, nk*nl]
sqrt_eigvals (torch.Tensor) – square root eigenvalues of the inverse covariance matrix. Expected shape is [nps*nz, nk*nl]
Q (torch.Tensor) – eigenvectors of the inverse covariance matrix. Expected shape is [nps*nz, nk*nl, nk*nl]
Q_inv (torch.Tensor) – inverse eigenvectors of the inverse covariance matrix. Expected shape is [nps*nz, nk*nl, nk*nl]
- Returns:
galaxy power spectrum in units of (Mpc/h)^3 to access. has shape [b, nps, nz, nk, nl] or [nps, nz, nk, nl]
- Return type:
galaxy_ps[idx] (torch.Tensor)
- normalize_data(ps_fid: Tensor, sqrt_eigvals: Tensor, Q: Tensor)[source]
Normalizes the reshapes the data
- Parameters:
ps_fid (torch.Tensor) – fiducial power spectrum multipoles in units of (Mpc/h)^3 used for normalization. Should have shape [nps, z, nk*nl]
ps_nw_fid (torch.Tensor) – NOTE: currently not used.
sqrt_eigvals (torch.Tensor) – set of sqrt eigenvalues used for normalization. Should have shape [ps, z, nk*nl]
Q (torch.Tensor) – set of eigenvectors used for normalization. Should have shape [ps, z, nk*nl, nk*nl]