<jupyter_start><jupyter_text>Compute theoretical power spectra<jupyter_code>import os

os.environ['XLA_FLAGS']='--xla_gpu_cuda_data_dir=/gpfslocalsys/cuda/10.1.2'

import jax_cosmo as jc
import numpy as np
%load_ext autoreload
%autoreload 2 
%pylab inline
cosmo = jc.Planck15()
z_array = np.load('../data/ktng/z_array.npy')
pz_array = np.load('../data/ktng/pz_array.npy')
nzs_s = jc.redshift.kde_nz(z_array, pz_array, bw=0.011)
nzs = [nzs_s]
probes = [ jc.probes.WeakLensing(nzs, sigma_e=0.26) ]
elle = np.logspace(1,5)
cls = jc.angular_cl.angular_cl(cosmo, elle, probes)
loglog(elle, cls[0])
ylabel(r'$C_\ell$')
xlabel(r'$\ell$');
title(r'Angular $C_\ell$');<jupyter_output><empty_output><jupyter_text>Compute kappaTNG maps power spectrum and compare<jupyter_code>from jax_lensing.utils import load_dataset
from jax_lensing.spectral import measure_power_spectrum
resolution = 0.29
map_size = 360

train_set = load_dataset("kappatng",
                         100,
                         map_size,
                         0.2,
                         "90%")
batch = next(train_set)
maps = batch['x']
maps.shape
resolution = 0.29
pixel_size = np.pi * resolution / 180. / 60. #rad/pixel


figure(figsize=(10,5))
subplot(121)
ell_map, ps_sum = measure_power_spectrum(maps[0, ..., 0], pixel_size)
loglog(ell_map, ps_sum, alpha=0.25, c="lightgrey")
for i in range(1, maps.shape[0]):
    ell_map, ps_map = measure_power_spectrum(maps[i, ..., 0], pixel_size)
    ps_sum += ps_map
    loglog(ell_map, ps_map, alpha=.5, c="lightgrey")
    
ps_mean = ps_sum/maps.shape[0]
loglog(ell_map, ps_mean, label='mean over {} maps'.format(maps.shape[0]))
loglog(elle, cls[0], label='theory')
ylabel(r'$C_\ell$')
xlabel(r'$\ell$');
title(r'Angular $C_\ell$');
legend()

subplot(122)
cls_interp = np.interp(ell_map, elle, cls[0])
loglog(ell_map, np.abs(cls_interp-ps_mean))
title('Difference btw theory and map PS')
xlabel(r'$\ell$')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15827338129496402, 'fraction_numerical': 0.030832476875642344, 'mean_word_length': 4.43854748603352, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '178651', 'license': 'MIT', 'n_tokens_mistral': 838, 'n_tokens_neox': 792, 'n_words': 139, 'path': 'notebooks/Check_power_spectrum.ipynb', 'repo_name': 'CosmoStat/jax-lensing', 'size': '1946'}