# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
from astropy.io import fits
import matplotlib.pyplot as plt
import healpy as hp
from scipy.interpolate import interp1d
import time

nside = 1024
lmax = 1024*3-1

def iqu2teb(IQU, nside, lmax=None):
    print('a')
    alms = hp.map2alm(IQU, lmax=lmax, pol=True)
    print('b')
    return hp.alm2map(alms, nside=nside, lmax=lmax, pol=False)

def teb2iqu(TEB, nside, lmax=None):
    alms = hp.map2alm(TEB, lmax=lmax, pol=False)
    return hp.alm2map(alms, nside=nside, lmax=lmax, pol=True)


# -

# # Make a field with T_QU covariance, make a copy upgraded to 2048, compare the power spectra of both 

noise_diag = hp.read_map('map_noise_cov_des1024_QU.fits', field=(0,1))
noise_diag = np.concatenate([noise_diag[0],noise_diag[1]])

plt.plot(noise_diag[:])

tau = np.min(noise_diag)
print(tau)
T_pixel = np.ones(noise_diag.shape)*tau
noise_bar_diag = noise_diag - tau

random_T_field = np.hstack([np.random.normal(0., np.sqrt(tau), hp.nside2npix(nside)),
                            np.random.normal(0., np.sqrt(tau), hp.nside2npix(nside))])

random_T_field.shape

hp.mollview(random_T_field[:hp.nside2npix(nside)], sub=(1,2,1), title='Q')
hp.mollview(random_T_field[hp.nside2npix(nside):], sub=(1,2,2), title='U')

power_pol = hp.anafast([random_T_field[:hp.nside2npix(nside)]*0.,
                        random_T_field[:hp.nside2npix(nside)],
                        random_T_field[hp.nside2npix(nside):]],
                       lmax = nside*3-1, pol = True)

tau_ell = np.mean([power_pol[1][50:nside*2], power_pol[2][50:nside*2]]) 

tau_ell

plt.plot(power_pol[1])
plt.plot(power_pol[2])
plt.plot([tau_ell]*nside*2, ':', c='black')

# # Now with ud grade

random_T_field_2048 = hp.ud_grade(np.array([random_T_field[:hp.nside2npix(nside)]*0.,
                                            random_T_field[:hp.nside2npix(nside)],
                                            random_T_field[hp.nside2npix(nside):]]),
                                  2048)

random_T_field_2048

power_pol_2048 = hp.anafast(random_T_field_2048,
                       lmax = 2048*3-1, pol = True)

power_pol_2048_reduced = hp.anafast(random_T_field_2048,
                       lmax = 1024*3-1, pol = True)

power_pol_2048_reduced2 = hp.anafast(random_T_field_2048,
                       lmax = 512*3-1, pol = True)

tau_ell_2048 = np.mean([power_pol_2048[1][50:2048*2], power_pol_2048[2][50:2048*2]]) 

plt.plot(power_pol_2048[1])
plt.plot(power_pol_2048[2])
plt.plot([tau_ell_2048]*nside*2, c='black')
plt.plot([tau_ell]*nside*2,':', c='black')

plt.plot(power_pol_2048_reduced[1])
plt.plot(power_pol_2048_reduced[2])
plt.plot([tau_ell]*nside*2,':', c='black')

plt.plot(power_pol_2048[1])
plt.plot(power_pol_2048_reduced[2])
plt.plot(power_pol_2048_reduced2[2])
plt.plot([tau_ell]*nside*2,':', c='black')

_ = plt.loglog(power_pol[1], label = '1024 original', alpha=0.7)
_ = plt.plot(power_pol_2048[1], label = '2048 ud_grade', alpha=0.4)
_ = plt.plot([tau_ell]*nside*2, ':', c='black')
_ = plt.ylabel(r'$C_\ell$'), plt.xlabel(r'$\ell$'), plt.legend()

# # Should check that tau is transforming correctly
#
# # Can also test that this framework gives approx the same result for a single iteration when signal power and tau ell are of similar order

random_T_field_2048_variance = hp.ud_grade(np.array([random_T_field[:hp.nside2npix(nside)]*0.,
                                            random_T_field[:hp.nside2npix(nside)],
                                            random_T_field[hp.nside2npix(nside):]]),
                                  2048, power=2)



StarcoderdataJupyterScriptsDedupFiltered