diff --git a/cupix/px_data/data_healpix.py b/cupix/px_data/data_healpix.py new file mode 100644 index 0000000..6306721 --- /dev/null +++ b/cupix/px_data/data_healpix.py @@ -0,0 +1,160 @@ +import numpy as np +import h5py + +from cupix.px_data import read_healpix_px, px_window + + +class HealpixPxArchive(object): + '''Collection of PX measurements from healpixels''' + + def __init__(self, fname, list_hp=None, list_px=None, + compute_window=False): + '''Setup collection, from file or by hand''' + + if fname is not None: + assert (list_hp is None) and (list_px is None) + self.fname = fname + print('read healpix info from', fname) + # read information from file (Sindhu's format) + reader = read_healpix_px.HealpixPxReader(fname) + # use information to construct list_px + self.list_hp, self.list_px = reader.get_list_px(compute_window) + else: + self.fname = None + assert len(list_hp) == len(list_px) + self.list_hp = list_hp + self.list_px = list_px + + print(f'Got Px for {len(self.list_hp)} healpixels') + + # get binning information + self.z_bins = self.list_px[0].z_bins + self.t_bins = self.list_px[0].t_bins + self.k_bins = self.list_px[0].k_bins + + # length in Angstroms + self.L_A = self.list_px[0].L_A + # mean pixel resolution in Angstroms (could be z dependent eventually) + self.sig_A = self.list_px[0].sig_A + + for px in self.list_px: + assert len(self.z_bins) == len(px.z_bins) + assert len(self.t_bins) == len(px.t_bins) + assert len(self.k_bins) == len(px.k_bins) + assert self.L_A == px.L_A + assert self.sig_A == px.sig_A + + return + + + def get_mean_px(self): + list_px_z = [] + for iz, z_bin in enumerate(self.z_bins): + list_px_zt = [] + for it, t_bin in enumerate(self.t_bins): + # this should be properly computed, not hard-coded + Nk = len(self.k_bins) + # add the contribution from each healpixel + F_m = np.zeros(Nk) + W_m = np.zeros(Nk) + for px in self.list_px: + px_zt = px.list_px_z[iz].list_px_zt[it] + F_m += px_zt.F_m + W_m += px_zt.W_m + px_zt = px_window.Px_zt_w.from_unnormalized( + z_bin, t_bin, self.k_bins, + F_m=F_m, W_m=W_m, + L_A=self.L_A, sig_A=self.sig_A, + compute_window=False) + list_px_zt.append(px_zt) + px_z = px_window.Px_z_w(self.t_bins, list_px_zt) + list_px_z.append(px_z) + mean_px = px_window.Px_w(self.z_bins, list_px_z) + return mean_px + + + def get_mean_and_cov(self, compute_window=False): + list_px_z = [] + for iz, z_bin in enumerate(self.z_bins): + list_px_zt = [] + for it, t_bin in enumerate(self.t_bins): + # collect PX from each healpixel into (N_hp, Nk) array + Nk = len(self.k_bins) + N_hp = len(self.list_hp) + all_P_m = np.empty((N_hp,Nk)) + all_V_m = np.empty((N_hp,Nk)) + for ipix in range(N_hp): + px_zt = self.list_px[ipix].list_px_z[iz].list_px_zt[it] + all_P_m[ipix] = px_zt.P_m + all_V_m[ipix] = px_zt.V_m + # compute weighted mean and covariance (from Picca) + mean_P_m = (all_P_m * all_V_m).sum(axis=0) + sum_V_m = all_V_m.sum(axis=0) + w = sum_V_m > 0. + mean_P_m[w] /= sum_V_m[w] + meanless_P_V_m = all_V_m * (all_P_m - mean_P_m) + C_mn = meanless_P_V_m.T.dot(meanless_P_V_m) + sum_V2 = sum_V_m * sum_V_m[:, None] + w = sum_V2 > 0. + C_mn[w] /= sum_V2[w] + # check if we need to compute the mean window matrix + if compute_window: + mean_U_mn = self.get_mean_window_zt(iz, it, all_V_m) + else: + mean_U_mn = None + # setup object (normalized, with covariance) + px_zt = px_window.Px_zt_w(z_bin, t_bin, self.k_bins, + P_m=mean_P_m, V_m=sum_V_m, C_mn=C_mn, + F_m=None, W_m=None, U_mn=mean_U_mn) + list_px_zt.append(px_zt) + px_z = px_window.Px_z_w(self.t_bins, list_px_zt) + list_px_z.append(px_z) + mean_px = px_window.Px_w(self.z_bins, list_px_z) + return mean_px + + + def get_mean_window_zt(self, iz, it): + '''Compute mean window matrix, for a given (z,t) bin''' + + # collect PX from each healpixel into (N_hp, Nk) array + Nk = len(self.k_bins) + N_hp = len(self.list_hp) + all_P_m = np.empty((N_hp, Nk)) + all_V_m = np.empty((N_hp, Nk)) + # collect individual window matrices (per healpixel) + all_U_mn = np.empty((N_hp, Nk, Nk)) + for ipix in range(N_hp): + px_zt = self.list_px[ipix].list_px_z[iz].list_px_zt[it] + all_P_m[ipix] = px_zt.P_m + all_V_m[ipix] = px_zt.V_m + if px_zt.U_mn is None: + # compute U_mn for each healpixel + R2_m = px_window.compute_R2_m(px_zt.k_bins, px_zt.sig_A) + U_mn = px_window.compute_U_mn(px_zt.W_m, R2_m, px_zt.L_A) + all_U_mn[ipix] = U_mn + else: + all_U_mn[ipix] = px_zt.U_mn + # compute weighted mean + mean_P_m = (all_P_m * all_V_m).sum(axis=0) + sum_V_m = all_V_m.sum(axis=0) + w = sum_V_m > 0. + mean_P_m[w] /= sum_V_m[w] + # weighted sum of all the windows + sum_U_mn = np.einsum('im,imn->mn', all_V_m, all_U_mn) + mean_U_mn = sum_U_mn / sum_V_m[:,None] + + return sum_V_m, mean_P_m, mean_U_mn + + + def rebin(self, rebin_t_factor, rebin_k_factor, include_k_0=True): + '''Return a new Px archive, after rebinning in theta (t) and k''' + + new_list_px = [] + for px in self.list_px: + new_px = px.rebin(rebin_t_factor, rebin_k_factor, include_k_0) + new_list_px.append(new_px) + + new_archive = HealpixPxArchive(fname=None, + list_hp=self.list_hp, list_px=new_list_px) + + return new_archive diff --git a/cupix/px_data/px_binning.py b/cupix/px_data/px_binning.py new file mode 100644 index 0000000..837ef80 --- /dev/null +++ b/cupix/px_data/px_binning.py @@ -0,0 +1,161 @@ +import numpy as np + + +class Bin_z(object): + '''Description of a particular redshift bin''' + + def __init__(self, min_z, max_z, label=None): + '''Construct by providing edges, and plotting label''' + + self.min_z = min_z + self.max_z = max_z + + if label is None: + self.label = f'{min_z} < z < {max_z}' + else: + self.label = label + + + def mean(self): + return 0.5*(self.min_z+self.max_z) + + + +class Bin_t(object): + '''Description of a particular theta bin''' + + def __init__(self, min_t, max_t, label=None): + '''Construct by providing edges, and plotting label''' + + self.min_t = min_t + self.max_t = max_t + if label is None: + self.label = f'{min_t} < theta [arcmin] < {max_t}' + else: + self.label = label + + + def mean(self): + return 0.5*(self.min_t+self.max_t) + + + def B_t(self, t): + '''Contribution of theta value to this bin''' + + # this could be made more efficient by allowing arrays, if needed + if (t >= self.min_t) and (t < self.max_t): + return 1.0 + else: + return 0.0 + + + def B_t_bins(self, t_bins): + '''Same than B_t for a list of narrow theta bins''' + + t_a = np.array( [ t_bin.mean() for t_bin in t_bins ]) + B_a = np.zeros_like(t_a) + in_bin = [ (t >= self.min_t) and (t < self.max_t) for t in t_a ] + B_a[in_bin] = 1.0 + return B_a + + +class Discrete_k(object): + '''Description of a particular wavenumber k (no binning)''' + + def __init__(self, k, label=None): + '''Construct by providing value and plotting label''' + + self.k = k + if label is None: + self.label = f'k = {k} 1/A' + else: + self.label = label + + + +class Bin_k(object): + '''Description of a particular k bin''' + + def __init__(self, min_k, max_k, label=None): + '''Construct by providing edges, and plotting label''' + + self.min_k = min_k + self.max_k = max_k + if label is None: + self.label = f'{min_k} < k [1/A] < {max_k}' + else: + self.label = label + + + def mean(self): + return 0.5*(self.min_k+self.max_k) + + + def B_k(self, k): + '''Contribution of k value to this k-bin''' + + # this could be made more efficient by allowing arrays, if needed + if (k >= self.min_k) and (k < self.max_k): + return 1.0 + else: + return 0.0 + + + def B_k_bins(self, k_bins): + '''Same than B_k, for list of narrow bins''' + + k_m = np.array( [ k_bin.k for k_bin in k_bins ]) + B_m = np.zeros_like(k_m) + in_bin = [ (k >= self.min_k) and (k < self.max_k) for k in k_m ] + B_m[in_bin] = 1.0 + return B_m + + +def get_coarser_k_bins(in_k_bins, rebin_factor, include_k_0=True): + '''Return coarser bins, rebin_factor coarser than in_k_bins. + If include_k_0=True, include k=0 instead of k_Ny.''' + + # this function for now should be called on Discrete_k bins + assert isinstance(in_k_bins[0], Discrete_k) + in_k = [ k_bin.k for k_bin in in_k_bins ] + in_dk = in_k[1]-in_k[0] + in_Nk = len(in_k) + in_k_max = np.max(np.abs(in_k)) + #print(f'input Nk = {in_Nk}, dk = {in_dk:.4f}, k_max = {in_k_max:.4f}') + + # extra factor of 2 here since in_k_bins include negative frequencies + out_Nk = int(in_Nk / 2 / rebin_factor) + if include_k_0: + # include k=0 in first bin + k_shift = -0.5*in_dk + else: + # include k_Ny in last bin + k_shift = +0.5*in_dk + out_k_edges = np.linspace(k_shift, in_k_max+k_shift, num = out_Nk+1, endpoint=True) + out_dk = out_k_edges[1] - out_k_edges[0] + out_k_max = np.max(np.abs(out_k_edges)) + #print(f'out Nk = {out_Nk}, dk = {out_dk:.4f}, k_max = {out_k_max:.4f}') + out_k_bins = [] + for ik in range(out_Nk): + k_bin = Bin_k(out_k_edges[ik], out_k_edges[ik+1]) + out_k_bins.append(k_bin) + return out_k_bins + + +def get_coarser_t_bins(in_t_bins, rebin_factor): + '''Return coarser bins, rebin_factor coarser than in_t_bins''' + + in_Nt = len(in_t_bins) + assert in_Nt % rebin_factor == 0, 'size mismatch' + in_min_t = [ t_bin.min_t for t_bin in in_t_bins ] + in_max_t = [ t_bin.max_t for t_bin in in_t_bins ] + out_min_t = in_min_t[::rebin_factor] + out_max_t = in_max_t[rebin_factor-1::rebin_factor] + out_Nt = len(out_min_t) + out_t_bins = [] + for it in range(out_Nt): + #print(f'{it}, {out_min_t[it]} < t < {out_max_t[it]}') + t_bin = Bin_t(out_min_t[it], out_max_t[it]) + out_t_bins.append(t_bin) + return out_t_bins + diff --git a/cupix/px_data/px_window.py b/cupix/px_data/px_window.py new file mode 100644 index 0000000..74a41db --- /dev/null +++ b/cupix/px_data/px_window.py @@ -0,0 +1,336 @@ +import numpy as np +from numba import njit + +from cupix.px_data import px_binning, px_ztk + + +class Px_w(px_ztk.BasePx): + '''Derived BasePx object, with information related to window matrix.''' + + def __init__(self, z_bins, list_px_z): + super().__init__(z_bins, list_px_z) + + # check consistency of FFT length and pixel resolution + self.L_A = list_px_z[0].L_A + self.sig_A = list_px_z[0].sig_A + for px_z in list_px_z: + assert self.L_A == px_z.L_A + assert self.sig_A == px_z.sig_A + + return + + + def rebin_t(self, rebin_factor): + '''Return a new Px_w object, rebinned in theta''' + + # for each z bin, rebin in theta + new_list_px_z = [] + for px_z in self.list_px_z: + new_px_z = px_z.rebin_t(rebin_factor) + new_list_px_z.append(new_px_z) + + new_px = Px_w(self.z_bins, new_list_px_z) + + return new_px + + + def rebin_k(self, rebin_factor, include_k_0=True): + '''Return a new Px_w object, rebinned in k''' + + # for each z bin, rebin in k + new_list_px_z = [] + for px_z in self.list_px_z: + new_px_z = px_z.rebin_k(rebin_factor, include_k_0) + new_list_px_z.append(new_px_z) + + new_px = Px_w(self.z_bins, new_list_px_z) + + return new_px + + + def rebin(self, rebin_t_factor, rebin_k_factor, include_k_0=True): + '''Return a new Px_w object, rebinned in theta and k''' + + # for each z bin, rebin in theta + new_list_px_z = [] + for px_z in self.list_px_z: + new_px_z = px_z.rebin(rebin_t_factor, rebin_k_factor, include_k_0) + new_list_px_z.append(new_px_z) + + new_px = Px_w(self.z_bins, new_list_px_z) + + return new_px + + + +class Px_z_w(px_ztk.Px_z): + '''Derived Px_z object, with information related to window matrix.''' + + def __init__(self, t_bins, list_px_zt): + super().__init__(t_bins, list_px_zt) + + # check consistency of FFT length and pixel resolution + self.L_A = list_px_zt[0].L_A + self.sig_A = list_px_zt[0].sig_A + for px_zt in list_px_zt: + assert self.L_A == px_zt.L_A + assert self.sig_A == px_zt.sig_A + + return + + + def rebin_t(self, rebin_factor): + '''Return a new Px_z_w object, rebinned in theta''' + + new_t_bins = px_binning.get_coarser_t_bins(self.t_bins, rebin_factor) + new_list_px_zt = [] + for new_t_bin in new_t_bins: + #print('new_t', new_t_bin.min_t, new_t_bin.max_t) + new_F_am = np.zeros(len(self.k_bins)) + new_V_am = np.zeros_like(new_F_am) + for in_t_bin, in_px_zt in zip(self.t_bins, self.list_px_zt): + t_a = in_t_bin.mean() + B_a = new_t_bin.B_t(t_a) + if B_a > 0: + #print(t_a, B_a) + new_F_am += B_a * in_px_zt.F_m + new_V_am += B_a * in_px_zt.V_m + + # normalize Px (for bins measured) + mask = new_V_am>0 + new_P_am = np.zeros_like(new_F_am) + new_P_am[mask] = new_F_am[mask] / new_V_am[mask] + new_px_zt = Px_zt_w.from_rebinning(self.z_bin, new_t_bin, self.k_bins, + P_m=new_P_am, V_m=new_V_am, U_mn=None) + new_list_px_zt.append(new_px_zt) + + new_px_z = Px_z_w(new_t_bins, new_list_px_zt) + + return new_px_z + + + def rebin_k(self, rebin_factor, include_k_0=True): + '''Return a new Px_z_w object, rebinned in k''' + + # for each theta bin, rebin in k + new_list_px_zt = [] + for px_zt in self.list_px_zt: + new_px_zt = px_zt.rebin_k(rebin_factor, include_k_0) + new_list_px_zt.append(new_px_zt) + + new_px_z = Px_z_w(self.t_bins, new_list_px_zt) + + return new_px_z + + + def rebin(self, rebin_t_factor, rebin_k_factor, include_k_0=True): + '''Return a new Px_w_z object, rebinned in theta and k''' + + # get first a new object, rebinned in k + new_px_z = self.rebin_k(rebin_k_factor, include_k_0) + + # ask the new object to rebin in theta + return new_px_z.rebin_t(rebin_t_factor) + + + def rebin_model(self, raw_model, raw_px_z, convolve=True): + '''Rebin model, and convolve with window matrix''' + + # theoretical prediction, in raw bins (2D array) + raw_Nt, raw_Nk = raw_model.shape + print(f'input raw model, shape = {raw_Nt}, {raw_Nk}') + assert raw_Nt == len(raw_px_z.t_bins), 'size mismatch' + assert raw_Nk == len(raw_px_z.k_bins), 'size mismatch' + + # get window matrix (and weights) for each (original) theta bin + list_V_m = [ px_zt.V_m for px_zt in raw_px_z.list_px_zt] + list_U_mn = [ px_zt.U_mn for px_zt in raw_px_z.list_px_zt] + print(f'got {len(list_U_mn)} window matrices') + + # for each (original) theta bin, construct rectangular window + rb_Nk = len(self.k_bins) + print(f'will rebin windows to shape {rb_Nk} x {raw_Nk}') + + # for each rebinned k bin, figure out contributions + B_M_m = [] + for iM in range(rb_Nk): + k_bin = self.k_bins[iM] + B_m = k_bin.B_k_bins(raw_px_z.k_bins) + print(iM, np.nonzero(B_m)[0]) + B_M_m.append(B_m) + + # compute rectangular window matrix (one per original theta bin) + list_U_Mn = [] + list_V_M = [] + list_P_M = [] + for raw_it in range(raw_Nt): + V_m = list_V_m[raw_it] + U_mn = list_U_mn[raw_it] + # V_M = sum_m (V_m B_M_m) + V_M = np.zeros(rb_Nk) + # U_Mn = sum_m (V_m B_M_m U_mn) / V_M + U_Mn = np.zeros([rb_Nk, raw_Nk]) + # P_M = sum_n U_Mn P_n + P_M = np.zeros(rb_Nk) + list_P_M.append(P_M) + list_V_M.append(V_M) + list_U_Mn.append(U_Mn) + + + # rebin model into coarser theta bins + rb_Nt = len(self.t_bins) + B_A_a = [] + for iA in range(rb_Nt): + t_bin = self.t_bins[iA] + B_a = t_bin.B_t_bins(raw_px_z.t_bins) + print(iA, np.nonzero(B_a)[0]) + B_A_a.append(B_a) + + + # convolve here with the window, rebin, etc + rb_model = np.zeros([len(self.t_bins), len(self.k_bins)]) + + return rb_model + + + +class Px_zt_w(px_ztk.Px_zt): + '''Derived Px_zt object, with information related to window matrix''' + + def __init__(self, z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, U_mn, + C_mn=None, L_A=None, sig_A=None): + '''Provide extra information related to window / weights''' + + super().__init__( + z_bin=z_bin, + t_bin=t_bin, + k_bins=k_bins, + P_m=P_m, + V_m=V_m, + C_mn=C_mn) + + # sum of the product of FFT of delta * weights + self.F_m = F_m + # sum of the product of FFT of weights + self.W_m = W_m + # window matrix + self.U_mn = U_mn + # length of FFT grid (in Angstroms) + self.L_A = L_A + # mean pixel resolution (in Angstroms) + self.sig_A = sig_A + + return + + + @classmethod + def from_unnormalized(cls, z_bin, t_bin, k_bins, + F_m, W_m, L_A, sig_A, compute_window=False): + '''Construct object from unnormalized quantities''' + + R2_m = compute_R2_m(k_bins, sig_A) + P_m, V_m = normalize_Px(F_m, W_m, R2_m, L_A) + C_mn = None + if compute_window: + R2_m = compute_R2_m(k_bins, sig_A) + U_mn = compute_U_mn(W_m, R2_m, L_A) + else: + U_mn = None + + return cls(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, U_mn, C_mn, L_A, sig_A) + + + @classmethod + def from_rebinning(cls, z_bin, t_bin, k_bins, P_m, V_m, U_mn): + '''Construct object from rebinning thinner bins''' + + # these might be used in further rebinning + F_m = P_m * V_m + # for now, these are not needed + W_m = None + C_mn = None + L_A = None + sig_A = None + + return cls(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, U_mn, C_mn, L_A, sig_A) + + + def rebin_k(self, rebin_factor, include_k_0=True): + '''Return a new Px_zt_w object, rebinned in k''' + + in_k_bins = self.k_bins + in_Nk = len(in_k_bins) + new_k_bins = px_binning.get_coarser_k_bins(in_k_bins, rebin_factor, include_k_0) + new_Nk = len(new_k_bins) + new_F_m = np.zeros(new_Nk) + new_V_m = np.zeros(new_Nk) + # should rebin window matrix U_mn here as well + for new_ik in range(new_Nk): + new_k_bin = new_k_bins[new_ik] + #print('new_ik', new_ik, new_k_bin.min_k, new_k_bin.max_k) + for in_ik in range(in_Nk): + in_k = in_k_bins[in_ik].k + B_m = new_k_bin.B_k(in_k) + if B_m > 0: + #print('in_ik', in_ik, in_k, B_m) + new_F_m[new_ik] += B_m * self.F_m[in_ik] + new_V_m[new_ik] += B_m * self.V_m[in_ik] + + # normalize Px (for bins measured) + mask = new_V_m>0 + new_P_m = np.zeros_like(new_F_m) + new_P_m[mask] = new_F_m[mask] / new_V_m[mask] + new_px_tz = Px_zt_w.from_rebinning(self.z_bin, self.t_bin, new_k_bins, + P_m=new_P_m, V_m=new_V_m, U_mn=None) + return new_px_tz + + + +def normalize_Px(F_m, W_m, R2_m, L_A): + '''Compute P_m and V_m given unnormalized measurements''' + # compute normalization factors, used + V_m = compute_V_m(W_m, R2_m, L_A) + P_m = np.zeros_like(V_m) + P_m[V_m>0] = (F_m[V_m>0] / V_m[V_m>0]).real + return P_m, V_m + + +@njit +def compute_U_mn_fast(W_m, R2_m, V_m_L): + '''numba function to speed up the matrix operation''' + Nk = len(W_m) + U_mn = np.zeros((Nk, Nk)) + for m in range(Nk): + U_mn[m, m] = W_m[0] * R2_m[m] / V_m_L[m] + for n in range(m): + diff = m - n + U_mn[m, n] = W_m[diff] * R2_m[n] / V_m_L[m] + for n in range(m+1, Nk): + diff = Nk + m - n + U_mn[m, n] = W_m[diff] * R2_m[n] / V_m_L[m] + return U_mn + + +def compute_R2_m(k_bins, sig_A): + '''Resolution kernel''' + k_m = np.array( [ k_bin.k for k_bin in k_bins ] ) + R2_m = np.exp(-(k_m * sig_A)**2) + return R2_m + + +def compute_U_mn(W_m, R2_m, L_A): + '''Compute window matrix''' + # normalization + V_m = compute_V_m(W_m, R2_m, L_A) + Nk = len(V_m) + if np.all(V_m == 0): + return np.zeros((Nk, Nk)) + V_m_L = V_m * L_A + return compute_U_mn_fast(W_m, R2_m, V_m_L) + + +def compute_V_m(W_m, R2_m, L_A): + '''Compute normalization factor for Px''' + # convolve W and R2 arrays + W_R2 = np.fft.ifft(np.fft.fft(W_m)* np.fft.fft(R2_m)) + return np.abs(W_R2) / L_A diff --git a/cupix/px_data/px_ztk.py b/cupix/px_data/px_ztk.py new file mode 100644 index 0000000..d112f03 --- /dev/null +++ b/cupix/px_data/px_ztk.py @@ -0,0 +1,76 @@ +import numpy as np + + +class Px_zt(object): + '''Collection of Px measurements at many k bins, same (z,t)''' + + def __init__(self, z_bin, t_bin, k_bins, P_m, V_m=None, C_mn=None): + '''Construct Px measurement at bin (z,t), for different k-bins''' + + # P_m, V_m and C_mn should be numpy arrays + self.Nk = len(k_bins) + assert self.Nk == P_m.size, 'size mismatch' + self.k_bins = k_bins + self.P_m = P_m + + # before combining different healpixels, we do not have a covariance + if C_mn is not None: + assert C_mn.size == self.Nk**2 + self.C_mn = C_mn + else: + self.C_mn = None + + # this is used to weight different bins when rebinning + if V_m is not None: + assert V_m.size == self.Nk + self.V_m = V_m + else: + self.V_m = None + + # store information about this particular bin (z, t) + self.z_bin = z_bin + self.t_bin = t_bin + + +class Px_z(object): + '''Collection of Px measurements at many (t,k) bins, same z''' + + def __init__(self, t_bins, list_px_zt): + '''Construct from list of theta, Px(z, theta)''' + + # store information about theta bins, and their Px + assert len(t_bins) == len(list_px_zt), 'size mismatch' + self.t_bins = t_bins + self.list_px_zt = list_px_zt + + # store information about this z-bin + self.z_bin = list_px_zt[0].z_bin + for px_zt in list_px_zt: + assert px_zt.z_bin.label == self.z_bin.label, 'inconsistent binning' + + # finally get k bins, and check consistency + self.k_bins = list_px_zt[0].k_bins + for px_z in list_px_zt: + assert len(px_z.k_bins) == len(self.k_bins), 'inconsistent binning' + + +class BasePx(object): + '''Base class to store measurements of the cross power spectrum''' + + def __init__(self, z_bins, list_px_z): + '''Construct from lists of z bins, Px(z)''' + + # store information about z bins, and their Px + assert len(z_bins) == len(list_px_z), 'size mismatch' + self.z_bins = z_bins + self.list_px_z = list_px_z + + # get also theta bins, and check consistency + self.t_bins = list_px_z[0].t_bins + for px_z in list_px_z: + assert len(px_z.t_bins) == len(self.t_bins), 'inconsistent binning' + + # finally get k bins, and check consistency + self.k_bins = list_px_z[0].k_bins + for px_z in list_px_z: + assert len(px_z.k_bins) == len(self.k_bins), 'inconsistent binning' diff --git a/cupix/px_data/read_healpix_px.py b/cupix/px_data/read_healpix_px.py new file mode 100644 index 0000000..44d8ab3 --- /dev/null +++ b/cupix/px_data/read_healpix_px.py @@ -0,0 +1,158 @@ +import numpy as np +import h5py + +from cupix.px_data import px_binning, px_ztk, px_window + + +class HealpixPxReader(object): + '''Class to read current data model by Sindhu, with healpixels, before rebinning''' + + def __init__(self, fname, verbose=False): + '''Read PX file and store information''' + + print(f'will read healpixel PX from {fname}') + + with h5py.File(fname, 'r') as f: + if verbose: print(f.keys()) + + # Load shared datasets + self.k_arr = f['k_arr'][:] + + # Load attributes + self.N_fft = f.attrs['N_fft'] + self.pw_A = f.attrs['pixel_width_A'] + self.L_A = self.N_fft * self.pw_A + # mean spectral resolution (made up number for now) + self.sig_A = self.pw_A + if verbose: print(f'N_fft = {self.N_fft}, pw_A = {self.pw_A}') + + group_names = [ + (group_name, float(group_name.split('_')[3])) # (name, theta_min) + for group_name in f.keys() + if group_name.startswith('z_') + ] + + # Sort by theta_min + group_names.sort(key=lambda x: x[1]) # sorts by the second item (theta_min) + + self.px = {} + self.px_weights = {} + self.theta_bins = [] + self.z_bins = [] + + # Now loop over the sorted group names + for group_name, theta_min in group_names: + if verbose: print('loop over', group_name) + + parts = group_name.split('_') + z_bin = float(parts[1]) + theta_max = float(parts[4]) + + group = f[group_name] + + key = (z_bin, theta_min, theta_max) + + self.px[key] = group['px'][:] + # this is what we would call W_m (not V_m) + self.px_weights[key] = group['px_weights'][:] + + self.theta_bins.append((theta_min, theta_max)) + self.z_bins.append(z_bin) + + print('finished reading PX file') + # get first PX to get number of healpixels + test_px = next(iter(self.px.values())) + self.N_hp, N_fft = test_px.shape + assert self.N_fft == N_fft + + return + + + def get_z_bins(self): + '''Get list of redshift bins (Bin_z objects)''' + + zs = np.array(sorted(set(self.z_bins))) + #print(zs) + dz = zs[1] - zs[0] + #print(dz) + z_bins = [] + for z in zs: + z_bin = px_binning.Bin_z(min_z=z-0.5*dz,max_z=z+0.5*dz) + z_bins.append(z_bin) + return z_bins + + + def get_t_bins(self): + '''Get list of theta bins (Bin_t objects)''' + + theta_min = [theta_bin[0] for theta_bin in self.theta_bins] + min_t = np.array(sorted(set(theta_min))) + N_t = len(min_t) + theta_max = [theta_bin[1] for theta_bin in self.theta_bins] + max_t = np.array(sorted(set(theta_max))) + assert N_t == len(max_t) + t_bins = [] + for it in range(N_t): + t_bin = px_binning.Bin_t(min_t=min_t[it], max_t=max_t[it]) + t_bins.append(t_bin) + return t_bins + + + def get_k_bins(self): + '''Get list of wavenumbers (Discrete_k objects)''' + k_bins = [] + for k in self.k_arr: + k_bin = px_binning.Discrete_k(k) + k_bins.append(k_bin) + return k_bins + + + def get_list_px(self, compute_window=False, verbose=False): + '''Get a list of PX (Px_w objects), one per healpixel''' + + # get list of bins + z_bins = self.get_z_bins() + t_bins = self.get_t_bins() + k_bins = self.get_k_bins() + + # collect list of healpixels and their px + list_hp = range(self.N_hp) + list_px = [] + for pix in list_hp: + if verbose: print('setting PX for healpixel', pix) + list_px_z = [] + for z_bin in z_bins: + if verbose: print('setting PX for z', z_bin.label) + list_px_zt = [] + for t_bin in t_bins: + if verbose: print('setting PX for theta', t_bin.label) + # try to find PX for this bin and healpixel + mean_z = 0.5*(z_bin.min_z+z_bin.max_z) + key = (mean_z, t_bin.min_t, t_bin.max_t) + if verbose: print(key) + F_m = self.px[key][pix] + W_m = self.px_weights[key][pix] + R2_m = np.ones_like(self.k_arr) + # set NaNs to zero + F_m[np.isnan(F_m)] = 0 + W_m[np.isnan(W_m)] = 0 + # create Px_zt_w object from these + px_zt = px_window.Px_zt_w.from_unnormalized( + z_bin, t_bin, k_bins, + F_m=F_m, W_m=W_m, + L_A=self.L_A, sig_A=self.sig_A, + compute_window=compute_window) + list_px_zt.append(px_zt) + # create Px_z object from these (if not empty) + px_z = px_window.Px_z_w(t_bins, list_px_zt) + list_px_z.append(px_z) + # create Px_w object from these (if not empty) + px = px_window.Px_w(z_bins, list_px_z) + list_px.append(px) + + # count empty healpixels + N_empty = list_px.count(None) + print(f'{N_empty} out of {len(list_px)} healpixels are empty') + + return list_hp, list_px + diff --git a/notebooks/development/likelihood.py b/notebooks/development/likelihood.py new file mode 100644 index 0000000..6d190d7 --- /dev/null +++ b/notebooks/development/likelihood.py @@ -0,0 +1,280 @@ +# %% [markdown] +# # Interfaces between likelihood, data and theory + +# %% +import numpy as np + +from lace.cosmo import camb_cosmo +from cupix.px_data import data_healpix, px_window, px_ztk + + +# %% +class Emulator(object): + '''Interface between ForestFlow and the rest of the code''' + + def __init__(self, config): + '''Setup the emulator given input arguments''' + print('setting up emulator') + self.config = config + + def emulate_px_Mpc(self, z, rt, kp, emu_params): + '''Emulate Px at one z, for several (rt, kp), given emu_params''' + print('asking ForestFlow to emulate Px') + # dummy + px = np.zeros([len(rt), len(kp)]) + return px + + +# %% +class Theory(object): + '''Lya theory object to model observed px (with contaminants)''' + + def __init__(self, fid_cosmo, emulator, contaminants): + '''Setup the theory object given input arguments''' + + print('setting up theory') + self.fid_cosmo = fid_cosmo + self.emulator = emulator + self.contaminants = contaminants + + + @classmethod + def from_config(cls, config): + '''Setup theory from config file (dict)''' + + fid_cosmo = camb_cosmo.get_cosmology_from_dictionary(config) + emulator = Emulator(config) + contaminants = None + return cls(fid_cosmo, emulator, contaminants) + + + def model_px(self, px_z, like_params): + '''Given likelihood parameters, and data object, model Px''' + + # px_z containts Px at a given z + z = px_z.z_bin.mean() + print(f'modeling px for z = {z}') + + # figure out the relevant emulator parameters (if needed) + emu_params = self.get_emu_params(z, like_params) + + # get values of (rt, kp) in Mpc + rt_Mpc, kp_Mpc = self.get_rt_kp_Mpc(px_z) + + # ask emulator for prediction in Mpc + px_Mpc = self.emulator.emulate_px_Mpc(z, rt_Mpc, kp_Mpc, emu_params) + + # you would here convert to px (Angstroms, arcmin) + px_arcmin_A = px_Mpc + + # then you would add contaminants + px_arcmin_A *= 1.0 + + return px_arcmin_A + + + def get_emu_params(self, z, like_params): + '''Figure out emulator params given likelihood params''' + + # similar to cup1d.likelihood.get_emulator_calls, but at one z? + # dummy for now + emu_params = {'mF': 0.7} + return emu_params + + + def get_rt_kp_Mpc(self, px_z): + '''Convert coordinates to Mpc, for a given (not binned) Px_z object''' + + # redshift of the measurement + z = px_z.z_bin.mean() + print(f'in get_rt_kp_Mpc, z = {z}') + # theta bins (in arcmin), could use t_min, t_max here + theta_arcmin = [t_bin.mean() for t_bin in px_z.t_bins] + #print('theta_arcmin =', theta_arcmin) + # kp bins (in inverse Angstroms) + kp_A = [k_bin.k for k_bin in px_z.k_bins] + + # convert to comoving separations with angular comoving distance (dummy) + rt_Mpc = 1.0 * np.array(theta_arcmin) + # convert to comoving wavenumbers with Hubble rate (dummy) + kp_Mpc = 1.0 * np.array(kp_A) + return rt_Mpc, kp_Mpc + + +# %% +class Likelihood(object): + '''Lya theory object to model observed px (with contaminants)''' + + def __init__(self, data, theory): + '''Setup the likelihood object given input arguments''' + print('setting up likelihood') + self.data = data + self.theory = theory + self.z_bins = data.z_bins + + + @classmethod + def from_config(cls, config): + '''Setup likelihood from config file (dict)''' + + data = Data(config) + theory = Theory.from_config(config) + return cls(data, theory) + + + def get_chi2(self, like_params): + '''Given input parameters, return chi2''' + + # sum over redshift bins + sum_chi2=0 + for iz in range(len(self.z_bins)): + chi2 = self.get_chi2_z(iz, like_params) + sum_chi2 += chi2 + + return sum_chi2 + + + def get_chi2_z(self, iz, like_params): + '''Compute chi2 for a single z bin''' + + z = self.z_bins[iz].mean() + print(f'add chi2 for z = {z}') + + # get convolved theoretical prediction, after rebinning + model_z = self.convolved_model_z(iz, like_params) + print(f'colvolved model prediction, shape = {model_z.shape}') + + # get rebinned data + px_z = self.data.rebinned_px.list_px_z[iz] + Nt = len(px_z.t_bins) + assert (Nt == model_z.shape[0]), 'size mismatch' + + # sum over (rebinned) theta bins + sum_chi2=0 + for it in range(Nt): + data = px_z.list_px_zt[it].P_m + model = model_z[it] + cov = px_z.list_px_zt[it].C_mn + # (data-theory) C^{-1} (data-theory) + chi2 = 1 + np.sum(data-model) * 0.0 + sum_chi2 += chi2 + + return sum_chi2 + + + def convolved_model_z(self, iz, like_params): + '''Compute convolved model for rebinned data''' + + # get theoretical prediction, in raw bins (2D array) + raw_px_z = self.data.raw_px.list_px_z[iz] + raw_model = self.theory.model_px(raw_px_z, like_params) + Nt, Nk = raw_model.shape + print(f'model prediction, shape = {Nt}, {Nk}') + + # rebin and convolve model, using window matrix + rb_px_z = self.data.rebinned_px.list_px_z[iz] + rb_model = rb_px_z.rebin_model(raw_model, raw_px_z, convolve=True) + + # convolve here with the window, rebin, etc + rb_px_z = self.data.rebinned_px.list_px_z[iz] + rb_model = np.zeros([len(rb_px_z.t_bins), len(rb_px_z.k_bins)]) + + return rb_model + + +# %% +class Data(object): + '''Contains Px measurement, including rebinning if needed''' + + def __init__(self, config): + '''Provide raw (unbinned) data and rebinned data (with cov)''' + + if 'fname' in config: + fname = config['fname'] + else: + basedir = '/Users/afont/Codes/cupix/data/px_measurements/Lyacolore/' + fname = basedir + '/px-nhp_41-zbins_4-thetabins_40.hdf5' + #fname = basedir + '/px-nhp_150-zbins_4-thetabins_40.hdf5' + + # collection of Px from multiple healpixels + archive = data_healpix.HealpixPxArchive(fname) + print('read Px archive') + + # combine healpixels to get raw data (no binning) + self.raw_px = archive.get_mean_and_cov() + print('got raw Px data (averaged)') + + # rebinned archive, will be used to compute covariance + rebin_t_factor = config['rebin_t_factor'] + rebin_k_factor = config['rebin_k_factor'] + rebinned_archive = archive.rebin(rebin_t_factor, rebin_k_factor) + print('got rebinned archive') + + # mean of rebinned archive (with cov) + self.rebinned_px = rebinned_archive.get_mean_and_cov() + print('got rebinned Px data (averaged)') + + # both datasets should have the same z bins + self.z_bins = self.rebinned_px.z_bins + assert len(self.z_bins) == len(self.raw_px.z_bins) + + +# %% +config = {} +emulator = Emulator(config) + +# %% +fid_cosmo = camb_cosmo.get_cosmology_from_dictionary(config) +#print(f'H_0 = {fid_cosmo.H0}') +camb_cosmo.print_info(fid_cosmo) +camb_cosmo.dkms_dhMpc(fid_cosmo, z=3) + +# %% +theory = Theory(fid_cosmo, emulator, contaminants=None) + +# %% +# theory_2 = Theory.from_config(config) + +# %% +basedir = '/Users/afont/Codes/cupix/data/px_measurements/Lyacolore/' +config = { + 'fname': basedir + '/px-nhp_41-zbins_4-thetabins_40.hdf5', + 'rebin_t_factor': 4, + 'rebin_k_factor': 4 +} +data = Data(config) +print(f'Raw Px has Nt = {len(data.raw_px.t_bins)}, Nk = {len(data.raw_px.k_bins)}') +print(f'Rebinned Px has Nt = {len(data.rebinned_px.t_bins)}, Nk = {len(data.rebinned_px.k_bins)}') + +# %% + +# %% + +# %% +likelihood = Likelihood(data, theory) +# like_2 = Likelihood.from_config(config) + +# %% +like_params = {'mF': 0.8} +iz=0 +px_z = likelihood.data.raw_px.list_px_z[iz] +#rt_Mpc, kp_Mpc = likelihood.theory.get_rt_kp_Mpc(px_z) +#model_px = likelihood.theory.model_px(px_z, like_params) +likelihood.get_chi2_z(iz, like_params) +#likelihood.get_chi2(like_params) + +# %% +test = [0, 1, 2] +test.append(None) + +# %% +test + +# %% +test.count(None) + +# %% +test.count(None) == len(test) + +# %% +np.all( diff --git a/notebooks/development/rebinned_covariance.py b/notebooks/development/rebinned_covariance.py new file mode 100644 index 0000000..22137f9 --- /dev/null +++ b/notebooks/development/rebinned_covariance.py @@ -0,0 +1,102 @@ +# %% [markdown] +# # Compute covariance of rebinned PX measurement + +# %% +import numpy as np +import matplotlib.pyplot as plt + +# %% +from cupix.px_data import data_healpix, px_window, px_ztk + +# %% +# hardcoded for now +basedir = '/Users/afont/Codes/cupix/data/px_measurements/Lyacolore/' +#fname = basedir + '/px-nhp_41-zbins_3-thetabins_7.hdf5' +#fname = basedir + '/px-nhp_41-zbins_4-thetabins_40.hdf5' +fname = basedir + '/px-nhp_150-zbins_4-thetabins_40.hdf5' +print(fname) + +# %% +# load PX measurement from different healpixels (before rebinning) +archive = data_healpix.HealpixPxArchive(fname) + +# %% +# combine measurement from all healpixels, and covariance (before rebinning) +mean_px = archive.get_mean_and_cov() + +# %% +# rebin all measurements in the original archive, create a new one +rebin_t_factor=4 +rebin_k_factor=4 +rebinned_archive = archive.rebin( + rebin_t_factor=rebin_t_factor, + rebin_k_factor=rebin_k_factor +) + +# %% +# get a new measurement combining all healpixels, after rebinning +mean_rebinned_px = rebinned_archive.get_mean_and_cov() + + +# %% +def compare_px(iz, rebinned_it, kmax=0.5): + '''Plot the mean Px, before and after rebinning''' + + # plot mean px, before rebinning + k_m = np.array( [ k_bin.k for k_bin in mean_px.k_bins ] ) + # range of original theta bins to plot + it_min = rebinned_it * rebin_t_factor + it_max = (rebinned_it+1) * rebin_t_factor + for px_zt in mean_px.list_px_z[iz].list_px_zt[it_min:it_max]: + min_t = px_zt.t_bin.min_t + max_t = px_zt.t_bin.max_t + label=f'{px_zt.t_bin.min_t} < $\\theta$ < {px_zt.t_bin.max_t}' + mask = np.abs(k_m) < kmax + plt.plot(k_m[mask], px_zt.P_m[mask], alpha=0.5, label=label) + + # plot mean px, after rebinning + k_m = np.array( [ k_bin.mean() for k_bin in mean_rebinned_px.k_bins ] ) + px_zt = mean_rebinned_px.list_px_z[iz].list_px_zt[rebinned_it] + min_t = px_zt.t_bin.min_t + max_t = px_zt.t_bin.max_t + label=f'{px_zt.t_bin.min_t} < $\\theta$ < {px_zt.t_bin.max_t}' + yerr = np.sqrt(np.diagonal(px_zt.C_mn)) + mask = k_m < kmax + plt.errorbar(k_m[mask], px_zt.P_m[mask], yerr=yerr[mask], label=label) + plt.legend() + plt.xlim(0, kmax) + + +# %% +compare_px(iz=0, rebinned_it=4) + +# %% +compare_px(iz=1, rebinned_it=2) + + +# %% +def plot_px(px, iz, theta_min, theta_max, kmax=0.5): + '''Plot Px at given z, for theta bins within range''' + k_m = np.array( [ k_bin.mean() for k_bin in px.k_bins ] ) + for px_zt in px.list_px_z[iz].list_px_zt: + min_t = px_zt.t_bin.min_t + max_t = px_zt.t_bin.max_t + if min_t > theta_min and max_t < theta_max: + label=f'{px_zt.t_bin.min_t} < $\\theta$ < {px_zt.t_bin.max_t}' + yerr = np.sqrt(np.diagonal(px_zt.C_mn)) + mask = k_m < kmax + plt.errorbar(k_m[mask], px_zt.P_m[mask], yerr=yerr[mask], label=label) + plt.legend() + plt.axhline(y=0, ls=':', color='gray') + plt.xlim(0, 0.5) + + +# %% +plot_px(mean_rebinned_px, iz=0, theta_min=0, theta_max=1000) + +# %% +plot_px(mean_rebinned_px, iz=1, theta_min=5, theta_max=1000) + +# %% + +# %% diff --git a/notebooks/development/test_healpixel_data.py b/notebooks/development/test_healpixel_data.py new file mode 100644 index 0000000..8471694 --- /dev/null +++ b/notebooks/development/test_healpixel_data.py @@ -0,0 +1,192 @@ +# %% [markdown] +# # PX measurements from healpixels +# +# ### Read archive, compute mean and covariance and make plots + +# %% +import time +import numpy as np +import matplotlib.pyplot as plt + +# %% +from cupix.px_data import data_healpix, px_window, px_ztk + +# %% +# hardcoded for now +basedir = '/Users/afont/Codes/cupix/data/px_measurements/Lyacolore/' +#fname = basedir + '/px-nhp_41-zbins_3-thetabins_7.hdf5' +fname = basedir + '/px-nhp_41-zbins_4-thetabins_40.hdf5' +#fname = basedir + '/px-nhp_150-zbins_4-thetabins_40.hdf5' +print(fname) + +# %% +# load PX measurement from different healpixels +start = time.time() +archive = data_healpix.HealpixPxArchive(fname) +#archive = data_healpix.HealpixPxArchive(fname, compute_window=True) +end = time.time() +print('time spent =', end - start) + +# %% [markdown] +# ### What fraction of bins in the archive are empty? + +# %% +total_bins = len(archive.t_bins) * len(archive.z_bins) +print(total_bins) +empty_bins = np.zeros(len(archive.list_px)) +for pix, px in enumerate(archive.list_px): + for iz, px_z in enumerate(px.list_px_z): + for it, px_zt in enumerate(px_z.list_px_zt): + sum_W = np.sum(px_zt.W_m) + if sum_W==0: + empty_bins[pix] += 1 +print(np.mean(empty_bins)/total_bins) +plt.hist(empty_bins/total_bins,bins=10); + +# %% +# number of Fourier modes (should be equal to number of FFT points, at least for now) +Nk = len(archive.k_bins) +print(f'N_k = {Nk}') +# number of healpixels +N_hp = len(archive.list_hp) +print(f'Got {N_hp} healpixels') + +# %% +# test code to compute mean PX (summing F_m, W_m from healpixels) +mean_px_v1 = archive.get_mean_px() + +# %% +# test code to compute mean PX and covariance (using P_m, V_m from healpixels) +mean_px_v2 = archive.get_mean_and_cov() + + +# %% +def compare_means(iz, it, kmax=0.5): + px_zt_v1 = mean_px_v1.list_px_z[iz].list_px_zt[it] + px_zt_v2 = mean_px_v2.list_px_z[iz].list_px_zt[it] + k = np.array( [k_bin.k for k_bin in archive.k_bins] ) + mask = np.abs(k) < kmax + plt.plot(k[mask], px_zt_v1.P_m[mask], label='mean v1') + plt.plot(k[mask], px_zt_v2.P_m[mask], ':', label='mean v2') + plt.xlim(0,kmax) + plt.legend() +# plt.ylim(-0.1, 0.2) + + +# %% +compare_means(1,8) + + +# %% +def plot_weights(iz, it, kmax=0.5): + px_zt_v1 = mean_px_v1.list_px_z[iz].list_px_zt[it] + px_zt_v2 = mean_px_v2.list_px_z[iz].list_px_zt[it] + k = np.array( [k_bin.k for k_bin in archive.k_bins] ) + mask = np.abs(k) < kmax + for px in archive.list_px: + px_zt = px.list_px_z[iz].list_px_zt[it] +# plt.semilogy(k[mask], px_zt.V_m[mask], alpha=0.1) + plt.semilogy(k[mask], px_zt_v1.V_m[mask], label='total weights v1') + plt.semilogy(k[mask], px_zt_v2.V_m[mask], ':', label='total weights v2') + plt.xlim(0,kmax) + plt.legend() + + +# %% +plot_weights(1,5) + +# %% [markdown] +# ## Rebin the mean measurement + +# %% +rebin_8_4 = mean_px_v1.rebin(rebin_t_factor=8, rebin_k_factor=4, include_k_0=True) +rebin_4_4 = mean_px_v1.rebin(rebin_t_factor=4, rebin_k_factor=4, include_k_0=True) +rebin_2_4 = mean_px_v1.rebin(rebin_t_factor=2, rebin_k_factor=4, include_k_0=True) +rebin_1_4 = mean_px_v1.rebin(rebin_t_factor=1, rebin_k_factor=4, include_k_0=True) +# k bin +k_m = [ k_bin.mean() for k_bin in rebin_4_4.k_bins ] +print(len(k_m)) + + +# %% +def compare_theta(iz=0, theta_min=1.0, theta_max=2.0, kmax=0.5): + k_m = np.array( [k_bin.mean() for k_bin in rebin_8_4.k_bins] ) + mask = np.abs(k_m) < kmax + for px in [rebin_8_4, rebin_4_4, rebin_2_4, rebin_1_4]: + for px_zt in px.list_px_z[iz].list_px_zt: + min_t = px_zt.t_bin.min_t + max_t = px_zt.t_bin.max_t + mean_t = px_zt.t_bin.mean() + if mean_t < theta_max and mean_t > theta_min: + label=f'{px_zt.t_bin.min_t} < $\\theta$ < {px_zt.t_bin.max_t}' + plt.plot(k_m[mask], px_zt.P_m[mask], label=label) + plt.legend() + plt.xlim(0, kmax) +# plt.ylim(-0.1,0.2) + + +# %% +compare_theta(iz=0, theta_min=1.0, theta_max=2.0) + +# %% +compare_theta(iz=1, theta_min=5.0, theta_max=10.0) + +# %% [markdown] +# ## Rebin the entire archive, then compute the mean + +# %% +archive_4_4 = archive.rebin(rebin_t_factor=4, rebin_k_factor=4, include_k_0=True) + +# %% +mean_4_4 = archive_4_4.get_mean_and_cov() + + +# %% +def compare_means(iz, it, kmax=0.5): + k_m = np.array( [ k_bin.mean() for k_bin in rebin_4_4.k_bins ] ) + mask = k_m < kmax + px_zt_1 = rebin_4_4.list_px_z[iz].list_px_zt[it] + plt.plot(k_m[mask], px_zt_1.P_m[mask], label='first average, then rebin') + px_zt_2 = mean_4_4.list_px_z[iz].list_px_zt[it] + plt.plot(k_m[mask], px_zt_2.P_m[mask], ':', label='first rebin, then average') + z_bin = rebin_4_4.z_bins[iz] + t_bin = rebin_4_4.t_bins[it] + plt.title(f"{z_bin.min_z} < z < {z_bin.max_z} , {t_bin.min_t}' < theta < {t_bin.max_t}' ") + plt.legend() + + +# %% +compare_means(iz=2,it=6) + +# %% +compare_means(iz=1,it=0) + + +# %% [markdown] +# ### Plot mean PX and errorbars + +# %% +def plot_px(mean_px, iz, it, kmax=0.5): + k_m = np.array( [ k_bin.mean() for k_bin in mean_px.k_bins ] ) + px_zt = mean_px.list_px_z[iz].list_px_zt[it] + yerr = np.sqrt(np.diagonal(px_zt.C_mn)) + mask = k_m < kmax + plt.errorbar(k_m[mask], px_zt.P_m[mask], yerr=yerr[mask]) + #plt.xlim(0,kmax) + z_bin = mean_px.z_bins[iz] + zlabel = f"{z_bin.min_z:.2f} < z < {z_bin.max_z:.2f}" + t_bin = mean_px.t_bins[it] + tlabel = f"{t_bin.min_t:.3f}' < theta < {t_bin.max_t:.3f}'" + plt.title(zlabel + ' , ' + tlabel) + plt.axhline(y=0, color='gray', ls=':') + + +# %% +plot_px(mean_4_4, iz=0, it=2) + +# %% +plot_px(mean_4_4, iz=1, it=5) + +# %% + +# %% diff --git a/notebooks/development/test_rebin.py b/notebooks/development/test_rebin.py new file mode 100644 index 0000000..c54f242 --- /dev/null +++ b/notebooks/development/test_rebin.py @@ -0,0 +1,121 @@ +# %% [markdown] +# # Rebin a PX measurement + +# %% +import numpy as np +import matplotlib.pyplot as plt + +# %% +from cupix.px_data import data_healpix, px_window, px_ztk, px_binning + +# %% +# hardcoded for now +basedir = '/Users/afont/Codes/cupix/data/px_measurements/Lyacolore/' +#fname = basedir + '/px-nhp_41-zbins_3-thetabins_7.hdf5' +fname = basedir + '/px-nhp_41-zbins_4-thetabins_40.hdf5' +print(fname) + +# %% +# load PX measurement from different healpixels +compute_window=False +archive = data_healpix.HealpixPxArchive(fname, compute_window=compute_window) + +# %% +# define new, coarser k bins +new_k_bins = px_binning.get_coarser_k_bins(archive.k_bins, rebin_factor=4, include_k_0=True) + + +# %% [markdown] +# ### Rebin in k + +# %% +def sum_weights(pix, iz, it): + test_px_zt = archive.list_px[pix].list_px_z[iz].list_px_zt[it] + return np.sum(test_px_zt.V_m) + + +# %% +# find a good PX measurement to work with +iz=1 +it=4 +for pix in archive.list_hp: + if sum_weights(pix, iz, it)>0: + print(pix) + ipix = pix + break + +# %% +# rebin only a redshift bin and theta bin +test_px_zt = archive.list_px[ipix].list_px_z[iz].list_px_zt[it] +new_px_zt = test_px_zt.rebin_k(rebin_factor=4, include_k_0=True) + + +# %% +def compare(px_zt_1, px_zt_2, kmax=0.2): + k1 = np.array( [k_bin.k for k_bin in px_zt_1.k_bins] ) + p1 = px_zt_1.P_m + k2 = np.array( [k_bin.mean() for k_bin in px_zt_2.k_bins] ) + p2 = px_zt_2.P_m + plt.plot(k1[np.abs(k1)