From b49791720a8fea6db3d21a0f7632a4f448610039 Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Thu, 24 Jul 2025 10:20:22 +0200
Subject: [PATCH 01/12] work towards alternative data storage

---
 cupix/px_data/healpixel_data.py | 109 ++++++++++++++++++++++++++++++++
 cupix/px_data/px_binning.py     |  99 +++++++++++++++++++++++++++++
 cupix/px_data/px_ztk.py         |  80 +++++++++++++++++++++++
 3 files changed, 288 insertions(+)
 create mode 100644 cupix/px_data/healpixel_data.py
 create mode 100644 cupix/px_data/px_binning.py
 create mode 100644 cupix/px_data/px_ztk.py

diff --git a/cupix/px_data/healpixel_data.py b/cupix/px_data/healpixel_data.py
new file mode 100644
index 0000000..c9c6dff
--- /dev/null
+++ b/cupix/px_data/healpixel_data.py
@@ -0,0 +1,109 @@
+import numpy as np
+import h5py
+
+import px_ztk
+import px_binning
+
+
+
+class Px_reader(object):
+    '''Class to read current data model by Sindhu, with healpixels, before rebinning'''
+
+    def __init__(self, fname):
+        '''Read PX file and store informatio'''
+
+        print(f'will read healpixel PX from {fname}')
+    
+        with h5py.File(fname, 'r') as f:
+            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']
+            print(f'N_fft = {N_fft}, pw_A = {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)
+
+            # Now loop over the sorted group names
+            for group_name, theta_min in group_names:
+                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'][:]
+                # is this W_m or V_m? I'm guessing V_m?
+                self.px_weights[key] = group['px_weights'][:]
+                self.p1d[key[0]] = group['p1d'][:]
+                self.num_pairs[key] = group['no_of_pairs'][:]
+
+                self.theta_bins.append((theta_min, theta_max))
+                self.z_bins.append(z_bin)
+
+
+    def get_list_px(self):
+        '''Return list of PX objects, one per healpixel'''
+
+        # I need to check with Sindhu how to do this
+        healpixels = list_of_healpixels
+        list_px = []
+
+        for pix in healpixels:
+            print('setting PX for healpixel', pix)
+            
+            z_bins = []
+            list_px_z = []
+
+            for z in self.z_bins:
+                print('setting PX for z', z)
+                z_bin = px_binning.Bin_z(min_z=z-0.5*dz,max_z=z+0.5*dz)
+                z_bins.append(z_bin)
+
+                t_bins = []
+                list_px_zt = []
+
+                for (min_t, max_t) in self.theta_bins:
+                    print('setting PX for theta', min_t, max_t)
+                    t_bin = px_binning.Bin_t(min_t=min_t, max_t=max_t)
+                    t_bins.append(t_bin)
+
+                    # try to find PX for this bin and healpixel
+                    key = (z, min_t, max_t)
+                    P_m = self.px[key][pix]
+                    V_m = self.px_weights[key][pix]
+
+                    # define k bins (discrete values here)
+                    k_bins = []
+                    for k in self.k_arr:
+                        print('setting PX for k', k)
+                        k_bin = px_binning.Discrete_k(k)
+                        k_bins.append(k_bin)
+
+                    # create Px_zt object from these
+                    px_zt = px_ztk.Px_zt(z_bin, t_bin, k_bins, px_ztk=P_m, V_ztk=V_m)
+                    list_px_zt.append(px_zt)
+
+                # create Px_z object from these
+                px_z = px_ztk.Px_z(t_bins, list_px_ztk)
+                list_px_z.append(px_z)
+
+            # create BasePx object from these
+            px = px_ztk.BaseDataPx(z_bins, list_px_z)
+            list_px.append(px)
+                
+        return healpixels, list_px
diff --git a/cupix/px_data/px_binning.py b/cupix/px_data/px_binning.py
new file mode 100644
index 0000000..a504b8d
--- /dev/null
+++ b/cupix/px_data/px_binning.py
@@ -0,0 +1,99 @@
+import numpy as np
+
+
+class Bin_z(object):
+    '''Description of a particular redshift bin'''
+
+    def __init__(self, min_z, max_z, mean_z=None, 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
+
+        # compute mean redshift in bin if needed
+        if mean_z is None:
+            self.mean_z = 0.5*(self.min_z + self.max_z)
+        else:
+            self.mean_z = mean_z
+
+
+
+class Bin_t(object):
+    '''Description of a particular theta bin'''
+
+    def __init__(self, min_t, max_t, mean_t=None, 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'{tmin} < theta [arcmin] < {tmax}'
+        else:
+            self.label = label
+
+        # compute mean separation in bin if needed
+        if mean_t is None:
+            self.mean_t = 0.5*(self.min_t + self.max_t)
+        else:
+            self.mean_t = mean_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
+
+
+
+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.mean_k = k
+        if label is None:
+            self.label = f'k = {mean_k} 1/A'
+        else:
+            self.label = label
+
+
+
+class Bin_k(object):
+    '''Description of a particular k bin'''
+
+    def __init__(self, min_k, max_k, mean_k=None, 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
+
+        # compute mean wavenumber in bin if needed
+        if mean_k is None:
+            self.mean_k = 0.5*(self.min_k + self.max_k)
+        else:
+            self.mean_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
+
diff --git a/cupix/px_data/px_ztk.py b/cupix/px_data/px_ztk.py
new file mode 100644
index 0000000..7101e85
--- /dev/null
+++ b/cupix/px_data/px_ztk.py
@@ -0,0 +1,80 @@
+import numpy as np
+
+
+class Px_zt(object):
+    '''Px measurement at one (z,t), including many k bins'''
+
+    def __init__(self, z_bin, t_bin, k_bins, px_ztk, cov_px_ztk=None, V_ztk=None):
+        '''Construct Px measurement at bin (z,t), for different k-bins'''
+
+        # px_ztk and cov_px_ztk should numpy (nd)arrays
+        self.Nk = len(k_bins)
+        assert self.Nk == px_ztk.size, 'size mismatch'
+        self.k_bins = k_bins
+        self.px = px_ztk
+
+        # before combining different healpixels, we do not have a covariance
+        if cov_px_ztk is not None:
+            assert cov_px_ztk.size == self.Nk**2
+            self.cov = cov_px_ztk
+        else:
+            self.cov = None
+
+        # this is used to weight differnet bins when rebinning
+        if V_ztk is not None:
+            assert V_ztk.size == self.Nk
+            self.V_ztk = V_ztk
+        else:
+            self.V_ztk = None
+
+        # eventually, we should also keep track of W_ztk for the window matrix
+        # self.W_ztk = None
+
+        # store information about this particular bin (z, t)
+        self.z_bin = z_bin
+        self.t_bin = t_bin
+
+
+
+class Px_z(object):
+    '''Px measurement at one z, including many (t,k) bins'''
+
+    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 == self.z_bin, '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 px_z.k_bins == self.k_bins, 'inconsistent binning'
+
+
+class BaseDataPx(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 px_z.t_bins == 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 px_z.k_bins == self.k_bins, 'inconsistent binning'

From afa645c455c6bb790384a6e957455d2a7932beed Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Fri, 25 Jul 2025 15:46:27 +0200
Subject: [PATCH 02/12] working on healpix coaddition

---
 cupix/px_data/data_healpix.py    |  69 ++++++++++++++
 cupix/px_data/healpixel_data.py  | 109 ----------------------
 cupix/px_data/px_binning.py      |   7 +-
 cupix/px_data/px_window.py       |  80 +++++++++++++++++
 cupix/px_data/px_ztk.py          |  36 ++++----
 cupix/px_data/read_healpix_px.py | 150 +++++++++++++++++++++++++++++++
 6 files changed, 319 insertions(+), 132 deletions(-)
 create mode 100644 cupix/px_data/data_healpix.py
 delete mode 100644 cupix/px_data/healpixel_data.py
 create mode 100644 cupix/px_data/px_window.py
 create mode 100644 cupix/px_data/read_healpix_px.py

diff --git a/cupix/px_data/data_healpix.py b/cupix/px_data/data_healpix.py
new file mode 100644
index 0000000..7c2a21b
--- /dev/null
+++ b/cupix/px_data/data_healpix.py
@@ -0,0 +1,69 @@
+import numpy as np
+import h5py
+
+from cupix.px_data import read_healpix_px, px_ztk, px_window
+
+
+class HealpixPxArchive(object):
+    '''Collection of PX measurements from healpixels'''
+
+    def __init__(self, fname, list_hp=None, list_px=None):
+        '''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()
+        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
+        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)
+
+        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
+                pw_A = 0.8
+                Nk = len(self.k_bins)
+                N_fft = Nk
+                L_A = pw_A * N_fft
+                # add the contribution from each healpixel
+                F_m = np.zeros(Nk, dtype='complex')
+                W_m = np.zeros(Nk, dtype='complex')
+                T_m = np.zeros(Nk, dtype='complex')
+                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
+                    T_m += px_zt.T_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, T_m=T_m, L=L_A)
+                list_px_zt.append(px_zt)
+            px_z = px_ztk.Px_z(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
+
+
diff --git a/cupix/px_data/healpixel_data.py b/cupix/px_data/healpixel_data.py
deleted file mode 100644
index c9c6dff..0000000
--- a/cupix/px_data/healpixel_data.py
+++ /dev/null
@@ -1,109 +0,0 @@
-import numpy as np
-import h5py
-
-import px_ztk
-import px_binning
-
-
-
-class Px_reader(object):
-    '''Class to read current data model by Sindhu, with healpixels, before rebinning'''
-
-    def __init__(self, fname):
-        '''Read PX file and store informatio'''
-
-        print(f'will read healpixel PX from {fname}')
-    
-        with h5py.File(fname, 'r') as f:
-            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']
-            print(f'N_fft = {N_fft}, pw_A = {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)
-
-            # Now loop over the sorted group names
-            for group_name, theta_min in group_names:
-                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'][:]
-                # is this W_m or V_m? I'm guessing V_m?
-                self.px_weights[key] = group['px_weights'][:]
-                self.p1d[key[0]] = group['p1d'][:]
-                self.num_pairs[key] = group['no_of_pairs'][:]
-
-                self.theta_bins.append((theta_min, theta_max))
-                self.z_bins.append(z_bin)
-
-
-    def get_list_px(self):
-        '''Return list of PX objects, one per healpixel'''
-
-        # I need to check with Sindhu how to do this
-        healpixels = list_of_healpixels
-        list_px = []
-
-        for pix in healpixels:
-            print('setting PX for healpixel', pix)
-            
-            z_bins = []
-            list_px_z = []
-
-            for z in self.z_bins:
-                print('setting PX for z', z)
-                z_bin = px_binning.Bin_z(min_z=z-0.5*dz,max_z=z+0.5*dz)
-                z_bins.append(z_bin)
-
-                t_bins = []
-                list_px_zt = []
-
-                for (min_t, max_t) in self.theta_bins:
-                    print('setting PX for theta', min_t, max_t)
-                    t_bin = px_binning.Bin_t(min_t=min_t, max_t=max_t)
-                    t_bins.append(t_bin)
-
-                    # try to find PX for this bin and healpixel
-                    key = (z, min_t, max_t)
-                    P_m = self.px[key][pix]
-                    V_m = self.px_weights[key][pix]
-
-                    # define k bins (discrete values here)
-                    k_bins = []
-                    for k in self.k_arr:
-                        print('setting PX for k', k)
-                        k_bin = px_binning.Discrete_k(k)
-                        k_bins.append(k_bin)
-
-                    # create Px_zt object from these
-                    px_zt = px_ztk.Px_zt(z_bin, t_bin, k_bins, px_ztk=P_m, V_ztk=V_m)
-                    list_px_zt.append(px_zt)
-
-                # create Px_z object from these
-                px_z = px_ztk.Px_z(t_bins, list_px_ztk)
-                list_px_z.append(px_z)
-
-            # create BasePx object from these
-            px = px_ztk.BaseDataPx(z_bins, list_px_z)
-            list_px.append(px)
-                
-        return healpixels, list_px
diff --git a/cupix/px_data/px_binning.py b/cupix/px_data/px_binning.py
index a504b8d..ee1b346 100644
--- a/cupix/px_data/px_binning.py
+++ b/cupix/px_data/px_binning.py
@@ -32,7 +32,7 @@ def __init__(self, min_t, max_t, mean_t=None, label=None):
         self.min_t = min_t
         self.max_t = max_t
         if label is None:
-            self.label = f'{tmin} < theta [arcmin] < {tmax}'
+            self.label = f'{min_t} < theta [arcmin] < {max_t}'
         else:
             self.label = label
 
@@ -60,9 +60,10 @@ class Discrete_k(object):
     def __init__(self, k, label=None):
         '''Construct by providing value and plotting label'''
 
-        self.mean_k = k
+        #self.mean_k = k
+        self.k = k
         if label is None:
-            self.label = f'k = {mean_k} 1/A'
+            self.label = f'k = {k} 1/A'
         else:
             self.label = label
 
diff --git a/cupix/px_data/px_window.py b/cupix/px_data/px_window.py
new file mode 100644
index 0000000..27aec84
--- /dev/null
+++ b/cupix/px_data/px_window.py
@@ -0,0 +1,80 @@
+import numpy as np
+from cupix.px_data import px_binning, px_ztk
+
+class Px_w(px_ztk.BasePx):
+    '''Derived BasePx object, with information related to window matrix.
+       This object should only be used without rebinning. '''
+
+    def __init__(self, z_bins, list_px_z):
+        super().__init__(z_bins, list_px_z)
+
+        # if doing rebinning, do not use this object (ill-defined)
+        assert isinstance(self.k_bins[0], px_binning.Discrete_k)
+
+        return
+
+
+class Px_zt_w(px_ztk.Px_zt):
+    '''Derived Px_zt object, with information related to window matrix.
+       This object should only be used without rebinning. '''
+
+    def __init__(self, z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn=None):
+        '''Provide extra information related to window / weights'''
+
+        # if doing rebinning, do not use this object (ill-defined)
+        assert isinstance(k_bins[0], px_binning.Discrete_k)
+
+        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)
+
+        # (complex) sum of the product of FFT of delta * weights
+        self.F_m = F_m
+        # (complex) sum of the product of FFT of weights
+        self.W_m = W_m
+        # (complex) sum of the product of FFT of resolution * weights
+        self.T_m = T_m
+        # window matrix
+        self.U_mn = U_mn
+
+        return
+
+
+    @classmethod
+    def from_unnormalized(cls, z_bin, t_bin, k_bins, F_m, W_m, T_m, L):
+        '''Construct object from unnormalized quantities'''
+
+        P_m, V_m = normalize_Px(F_m, W_m, T_m, L)
+        C_mn = None
+        # do not compute the window matrix for now
+        U_mn = None
+
+        return Px_zt_w(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn)
+
+
+def normalize_Px(F_m, W_m, T_m, L):
+    '''Compute P_m and V_m given unnormalized measurements'''
+
+    # compute normalization factors, used 
+    V_m = compute_V_m(W_m, T_m, L)
+    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
+
+
+def compute_V_m(W_m, T_m, L):
+    '''Compute normalization factor for Px'''
+
+    # effective resolution kernel (squared)
+    R2_m = np.zeros_like(W_m)
+    R2_m[W_m>0] = T_m[W_m>0] / W_m[W_m>0]
+ 
+    # 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
diff --git a/cupix/px_data/px_ztk.py b/cupix/px_data/px_ztk.py
index 7101e85..ca135da 100644
--- a/cupix/px_data/px_ztk.py
+++ b/cupix/px_data/px_ztk.py
@@ -2,42 +2,38 @@
 
 
 class Px_zt(object):
-    '''Px measurement at one (z,t), including many k bins'''
+    '''Collection of Px measurements at many k bins, same (z,t)'''
 
-    def __init__(self, z_bin, t_bin, k_bins, px_ztk, cov_px_ztk=None, V_ztk=None):
+    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'''
 
-        # px_ztk and cov_px_ztk should numpy (nd)arrays
+        # P_m, V_m and C_mn should be numpy arrays
         self.Nk = len(k_bins)
-        assert self.Nk == px_ztk.size, 'size mismatch'
+        assert self.Nk == P_m.size, 'size mismatch'
         self.k_bins = k_bins
-        self.px = px_ztk
+        self.P_m = P_m
 
         # before combining different healpixels, we do not have a covariance
-        if cov_px_ztk is not None:
-            assert cov_px_ztk.size == self.Nk**2
-            self.cov = cov_px_ztk
+        if C_mn is not None:
+            assert C_mn.size == self.Nk**2
+            self.C_mn = C_mn
         else:
-            self.cov = None
+            self.C_mn = None
 
-        # this is used to weight differnet bins when rebinning
-        if V_ztk is not None:
-            assert V_ztk.size == self.Nk
-            self.V_ztk = V_ztk
+        # 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_ztk = None
-
-        # eventually, we should also keep track of W_ztk for the window matrix
-        # self.W_ztk = None
+            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):
-    '''Px measurement at one z, including many (t,k) bins'''
+    '''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)'''
@@ -58,7 +54,7 @@ def __init__(self, t_bins, list_px_zt):
             assert px_z.k_bins == self.k_bins, 'inconsistent binning'
 
 
-class BaseDataPx(object):
+class BasePx(object):
     '''Base class to store measurements of the cross power spectrum'''
 
     def __init__(self, z_bins, list_px_z):
diff --git a/cupix/px_data/read_healpix_px.py b/cupix/px_data/read_healpix_px.py
new file mode 100644
index 0000000..efbad5b
--- /dev/null
+++ b/cupix/px_data/read_healpix_px.py
@@ -0,0 +1,150 @@
+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
+            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, 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
+                    key = (z_bin.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
+                    T_m = W_m * R2_m
+                    # 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, T_m=T_m, L=self.L_A)
+                    list_px_zt.append(px_zt)
+                # create Px_z object from these
+                px_z = px_ztk.Px_z(t_bins, list_px_zt)
+                list_px_z.append(px_z)
+            # create Px_w object from these
+            px = px_window.Px_w(z_bins, list_px_z)
+            list_px.append(px)
+
+        return list_hp, list_px
+

From 4f3ebd39e6ef411024fcee0ae2ce050015169f31 Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Mon, 28 Jul 2025 09:25:45 +0200
Subject: [PATCH 03/12] added example notebook

---
 cupix/px_data/data_healpix.py                | 34 ++++++++
 notebooks/development/test_healpixel_data.py | 84 ++++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 notebooks/development/test_healpixel_data.py

diff --git a/cupix/px_data/data_healpix.py b/cupix/px_data/data_healpix.py
index 7c2a21b..ea95e1f 100644
--- a/cupix/px_data/data_healpix.py
+++ b/cupix/px_data/data_healpix.py
@@ -67,3 +67,37 @@ def get_mean_px(self):
         return mean_px
 
 
+    def get_mean_and_cov(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):
+                # 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]
+                # 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, T_m=None, U_mn=None)
+                list_px_zt.append(px_zt)
+            px_z = px_ztk.Px_z(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
+
diff --git a/notebooks/development/test_healpixel_data.py b/notebooks/development/test_healpixel_data.py
new file mode 100644
index 0000000..9e124cb
--- /dev/null
+++ b/notebooks/development/test_healpixel_data.py
@@ -0,0 +1,84 @@
+# %% [markdown]
+# # PX measurements from healpixels
+#
+# ### Read archive, compute mean and covariance and make plots
+
+# %%
+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_38.hdf5'
+print(fname) 
+
+# %%
+# load PX measurement from different healpixels
+archive = data_healpix.HealpixPxArchive(fname)
+
+# %%
+# 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}')
+N_fft = Nk
+# pixel width, in Angstroms
+pw_A = 0.8
+# length of FFT grid, in Angstroms
+L_A = N_fft * pw_A
+print(f'L = {L_A} A')
+# 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 plot_px(iz, it):
+    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 = [k_bin.k for k_bin in archive.k_bins]
+    for px in archive.list_px:
+        px_zt = px.list_px_z[iz].list_px_zt[it]
+        plt.plot(k, px_zt.P_m, alpha=0.1)
+    plt.plot(k, px_zt_v1.P_m, label='mean v1')
+    plt.plot(k, px_zt_v2.P_m, label='mean v2')
+    plt.xlim(0,0.5)
+    plt.legend()
+
+
+# %%
+plot_px(1,2)
+
+
+# %%
+def plot_weights(iz, it):
+    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 = [k_bin.k for k_bin in archive.k_bins]
+    for px in archive.list_px:
+        px_zt = px.list_px_z[iz].list_px_zt[it]
+        plt.semilogy(k, px_zt.V_m, alpha=0.1)
+    plt.semilogy(k, px_zt_v1.V_m, label='total weights v1')
+    plt.semilogy(k, px_zt_v2.V_m, label='total weights v2')
+    plt.xlim(0,0.5)
+    plt.legend()
+
+
+# %%
+plot_weights(1,5)
+
+# %%
+
+# %%

From 60a73cec9f3f22310ac0af9ef8154ce46016af25 Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Tue, 29 Jul 2025 08:32:06 +0200
Subject: [PATCH 04/12] more on rebinning

---
 cupix/px_data/data_healpix.py    |  6 +--
 cupix/px_data/px_binning.py      | 84 ++++++++++++++++++++++++--------
 cupix/px_data/px_window.py       | 63 ++++++++++++++++++------
 cupix/px_data/read_healpix_px.py |  3 +-
 4 files changed, 117 insertions(+), 39 deletions(-)

diff --git a/cupix/px_data/data_healpix.py b/cupix/px_data/data_healpix.py
index ea95e1f..c61e095 100644
--- a/cupix/px_data/data_healpix.py
+++ b/cupix/px_data/data_healpix.py
@@ -49,9 +49,9 @@ def get_mean_px(self):
                 N_fft = Nk
                 L_A = pw_A * N_fft
                 # add the contribution from each healpixel
-                F_m = np.zeros(Nk, dtype='complex')
-                W_m = np.zeros(Nk, dtype='complex')
-                T_m = np.zeros(Nk, dtype='complex')
+                F_m = np.zeros(Nk) #, dtype='complex')
+                W_m = np.zeros(Nk) #, dtype='complex')
+                T_m = np.zeros(Nk) #, dtype='complex')
                 for px in self.list_px:
                     px_zt = px.list_px_z[iz].list_px_zt[it]
                     F_m += px_zt.F_m
diff --git a/cupix/px_data/px_binning.py b/cupix/px_data/px_binning.py
index ee1b346..f12aafb 100644
--- a/cupix/px_data/px_binning.py
+++ b/cupix/px_data/px_binning.py
@@ -4,7 +4,7 @@
 class Bin_z(object):
     '''Description of a particular redshift bin'''
 
-    def __init__(self, min_z, max_z, mean_z=None, label=None):
+    def __init__(self, min_z, max_z, label=None):
         '''Construct by providing edges, and plotting label'''
 
         self.min_z = min_z
@@ -15,18 +15,16 @@ def __init__(self, min_z, max_z, mean_z=None, label=None):
         else:
             self.label = label
 
-        # compute mean redshift in bin if needed
-        if mean_z is None:
-            self.mean_z = 0.5*(self.min_z + self.max_z)
-        else:
-            self.mean_z = mean_z
+
+    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, mean_t=None, label=None):
+    def __init__(self, min_t, max_t, label=None):
         '''Construct by providing edges, and plotting label'''
 
         self.min_t = min_t
@@ -36,18 +34,16 @@ def __init__(self, min_t, max_t, mean_t=None, label=None):
         else:
             self.label = label
 
-        # compute mean separation in bin if needed
-        if mean_t is None:
-            self.mean_t = 0.5*(self.min_t + self.max_t)
-        else:
-            self.mean_t = mean_t
+
+    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):
+        if (t >= self.min_t) and (t < self.max_t):
             return 1.0
         else:
             return 0.0
@@ -60,7 +56,6 @@ class Discrete_k(object):
     def __init__(self, k, label=None):
         '''Construct by providing value and plotting label'''
 
-        #self.mean_k = k
         self.k = k
         if label is None:
             self.label = f'k = {k} 1/A'
@@ -72,7 +67,7 @@ def __init__(self, k, label=None):
 class Bin_k(object):
     '''Description of a particular k bin'''
 
-    def __init__(self, min_k, max_k, mean_k=None, label=None):
+    def __init__(self, min_k, max_k, label=None):
         '''Construct by providing edges, and plotting label'''
 
         self.min_k = min_k
@@ -82,19 +77,66 @@ def __init__(self, min_k, max_k, mean_k=None, label=None):
         else:
             self.label = label
 
-        # compute mean wavenumber in bin if needed
-        if mean_k is None:
-            self.mean_k = 0.5*(self.min_k + self.max_k)
-        else:
-            self.mean_k
+
+    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):
+        if (k >= self.min_k) and (k < self.max_k):
             return 1.0
         else:
             return 0.0
 
+
+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
index 27aec84..c5b8df6 100644
--- a/cupix/px_data/px_window.py
+++ b/cupix/px_data/px_window.py
@@ -2,28 +2,20 @@
 from cupix.px_data import px_binning, px_ztk
 
 class Px_w(px_ztk.BasePx):
-    '''Derived BasePx object, with information related to window matrix.
-       This object should only be used without rebinning. '''
+    '''Derived BasePx object, with information related to window matrix.'''
 
     def __init__(self, z_bins, list_px_z):
         super().__init__(z_bins, list_px_z)
 
-        # if doing rebinning, do not use this object (ill-defined)
-        assert isinstance(self.k_bins[0], px_binning.Discrete_k)
-
         return
 
 
 class Px_zt_w(px_ztk.Px_zt):
-    '''Derived Px_zt object, with information related to window matrix.
-       This object should only be used without rebinning. '''
+    '''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, T_m, U_mn, C_mn=None):
         '''Provide extra information related to window / weights'''
 
-        # if doing rebinning, do not use this object (ill-defined)
-        assert isinstance(k_bins[0], px_binning.Discrete_k)
-
         super().__init__(
                 z_bin=z_bin, 
                 t_bin=t_bin, 
@@ -32,11 +24,11 @@ def __init__(self, z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn=Non
                 V_m=V_m,
                 C_mn=C_mn)
 
-        # (complex) sum of the product of FFT of delta * weights
+        # sum of the product of FFT of delta * weights
         self.F_m = F_m
-        # (complex) sum of the product of FFT of weights
+        # sum of the product of FFT of weights
         self.W_m = W_m
-        # (complex) sum of the product of FFT of resolution * weights
+        # sum of the product of FFT of resolution * weights
         self.T_m = T_m
         # window matrix
         self.U_mn = U_mn
@@ -50,12 +42,55 @@ def from_unnormalized(cls, z_bin, t_bin, k_bins, F_m, W_m, T_m, L):
 
         P_m, V_m = normalize_Px(F_m, W_m, T_m, L)
         C_mn = None
-        # do not compute the window matrix for now
+        # should compute the window matrix here
         U_mn = None
 
         return Px_zt_w(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn)
 
 
+    @classmethod
+    def from_rebinning(cls, z_bin, t_bin, k_bins, P_m, V_m, U_mn):
+        '''Construct object from rebinning thinner bins'''
+
+        # for now, these are not needed
+        F_m = None
+        W_m = None
+        T_m = None
+        C_mn = None
+
+        return Px_zt_w(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn)
+
+
+    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, T_m, L):
     '''Compute P_m and V_m given unnormalized measurements'''
 
diff --git a/cupix/px_data/read_healpix_px.py b/cupix/px_data/read_healpix_px.py
index efbad5b..268924f 100644
--- a/cupix/px_data/read_healpix_px.py
+++ b/cupix/px_data/read_healpix_px.py
@@ -125,7 +125,8 @@ def get_list_px(self, verbose=False):
                 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
-                    key = (z_bin.mean_z, t_bin.min_t, t_bin.max_t)
+                    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]

From 64959e1b168b3d04d15a9e1b558783a99bdc6cf0 Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Tue, 29 Jul 2025 16:18:32 +0200
Subject: [PATCH 05/12] first working version of rebinning

---
 cupix/px_data/px_binning.py      |   6 +-
 cupix/px_data/px_window.py       | 112 ++++++++++++++++++++++++++++++-
 cupix/px_data/px_ztk.py          |   8 +--
 cupix/px_data/read_healpix_px.py |   2 +-
 4 files changed, 119 insertions(+), 9 deletions(-)

diff --git a/cupix/px_data/px_binning.py b/cupix/px_data/px_binning.py
index f12aafb..d5d2d50 100644
--- a/cupix/px_data/px_binning.py
+++ b/cupix/px_data/px_binning.py
@@ -102,7 +102,7 @@ def get_coarser_k_bins(in_k_bins, rebin_factor, include_k_0=True):
     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}')
+    #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)
@@ -115,7 +115,7 @@ def get_coarser_k_bins(in_k_bins, rebin_factor, include_k_0=True):
     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}')
+    #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])
@@ -135,7 +135,7 @@ def get_coarser_t_bins(in_t_bins, 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]}')
+        #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
index c5b8df6..42c292b 100644
--- a/cupix/px_data/px_window.py
+++ b/cupix/px_data/px_window.py
@@ -10,6 +10,114 @@ def __init__(self, z_bins, list_px_z):
         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)
+
+        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)
+            # should rebin window matrix U_mn here as well
+            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)
+
+
+
 class Px_zt_w(px_ztk.Px_zt):
     '''Derived Px_zt object, with information related to window matrix'''
 
@@ -52,8 +160,9 @@ def from_unnormalized(cls, z_bin, t_bin, k_bins, F_m, W_m, T_m, L):
     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
-        F_m = None
         W_m = None
         T_m = None
         C_mn = None
@@ -81,6 +190,7 @@ def rebin_k(self, rebin_factor, include_k_0=True):
                     #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)
diff --git a/cupix/px_data/px_ztk.py b/cupix/px_data/px_ztk.py
index ca135da..d112f03 100644
--- a/cupix/px_data/px_ztk.py
+++ b/cupix/px_data/px_ztk.py
@@ -46,12 +46,12 @@ def __init__(self, t_bins, 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 == self.z_bin, 'inconsistent binning'
+            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 px_z.k_bins == self.k_bins, 'inconsistent binning'
+            assert len(px_z.k_bins) == len(self.k_bins), 'inconsistent binning'
 
 
 class BasePx(object):
@@ -68,9 +68,9 @@ def __init__(self, z_bins, 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 px_z.t_bins == self.t_bins, 'inconsistent binning'
+            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 px_z.k_bins == self.k_bins, 'inconsistent binning'
+            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
index 268924f..2d77678 100644
--- a/cupix/px_data/read_healpix_px.py
+++ b/cupix/px_data/read_healpix_px.py
@@ -141,7 +141,7 @@ def get_list_px(self, verbose=False):
                             F_m=F_m, W_m=W_m, T_m=T_m, L=self.L_A)
                     list_px_zt.append(px_zt)
                 # create Px_z object from these
-                px_z = px_ztk.Px_z(t_bins, list_px_zt)
+                px_z = px_window.Px_z_w(t_bins, list_px_zt)
                 list_px_z.append(px_z)
             # create Px_w object from these
             px = px_window.Px_w(z_bins, list_px_z)

From dd7f802200c31fe145e8707629fccad52b462dd6 Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Tue, 29 Jul 2025 16:45:47 +0200
Subject: [PATCH 06/12] updated notebooks

---
 notebooks/development/test_healpixel_data.py |   4 +-
 notebooks/development/test_rebin.py          | 119 +++++++++++++++++++
 2 files changed, 121 insertions(+), 2 deletions(-)
 create mode 100644 notebooks/development/test_rebin.py

diff --git a/notebooks/development/test_healpixel_data.py b/notebooks/development/test_healpixel_data.py
index 9e124cb..2fe7aee 100644
--- a/notebooks/development/test_healpixel_data.py
+++ b/notebooks/development/test_healpixel_data.py
@@ -13,8 +13,8 @@
 # %%
 # 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_38.hdf5'
+#fname = basedir + '/px-nhp_41-zbins_3-thetabins_7.hdf5'
+fname = basedir + '/px-nhp_41-zbins_4-thetabins_40.hdf5'
 print(fname) 
 
 # %%
diff --git a/notebooks/development/test_rebin.py b/notebooks/development/test_rebin.py
new file mode 100644
index 0000000..f9d62dc
--- /dev/null
+++ b/notebooks/development/test_rebin.py
@@ -0,0 +1,119 @@
+# %% [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
+archive = data_healpix.HealpixPxArchive(fname)
+
+# %%
+# 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)
+
+
+# %%
+iz=1
+it=4
+for pix in range(10):
+    print(pix, sum_weights(pix, iz, it))
+
+# %%
+# rebin only a redshift bin and theta bin
+ipix = 5
+iz = 1
+it = 4
+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=2, include_k_0=True)
+
+
+# %%
+def compare(px_zt_1, px_zt_2):
+    k1 = [k_bin.k for k_bin in px_zt_1.k_bins]
+    p1 = px_zt_1.P_m
+    k2 = [k_bin.mean() for k_bin in px_zt_2.k_bins]
+    p2 = px_zt_2.P_m    
+    plt.plot(k1, p1, 'o')
+    plt.plot(k2, p2, 'x')
+    plt.xlim(0,0.1)
+
+
+# %%
+compare(test_px_zt, new_px_zt)
+
+# %% [markdown]
+# ## Rebin in theta
+
+# %%
+new_t_bins = px_binning.get_coarser_t_bins(archive.t_bins, rebin_factor=4)
+
+# %%
+test_px_z = archive.list_px[ipix].list_px_z[iz]
+
+# %%
+new_px_z = test_px_z.rebin_t(rebin_factor=4)
+
+
+# %%
+def compare(test_px_z, new_px_z, new_it, plot_V_m=False):
+    rebin_factor=int(len(test_px_z.t_bins)/len(new_px_z.t_bins))
+    print(rebin_factor)
+    k_m = [k_bin.k for k_bin in test_px_zt.k_bins]
+    for px_zt in test_px_z.list_px_zt[new_it*rebin_factor:(new_it+1)*rebin_factor]:
+        if plot_V_m:
+            plt.plot(k_m, px_zt.V_m, alpha=0.2)
+        else:
+            plt.plot(k_m, px_zt.P_m, alpha=0.2)
+    if plot_V_m:
+        plt.plot(k_m, new_px_z.list_px_zt[new_it].V_m)
+    else:
+        plt.plot(k_m, new_px_z.list_px_zt[new_it].P_m)
+    plt.xlim([0,0.5])
+    #plt.ylim([-0.01, 0.02])
+
+
+# %%
+compare(test_px_z, new_px_z, 4)
+
+# %%
+compare(test_px_z, new_px_z, 4, plot_V_m=True)
+
+# %% [markdown]
+# ### Rebin the entire PX measurement
+
+# %%
+# rebin the entire Px object
+test_px = archive.list_px[ipix]
+test_px.rebin_k(rebin_factor=4)
+test_px.rebin_t(rebin_factor=4)
+new_px = test_px.rebin(rebin_t_factor=4, rebin_k_factor=4, include_k_0=True)
+
+# %%
+# rebin only a redshift bin
+test_px_z = archive.list_px[ipix].list_px_z[iz]
+test_px_z.rebin_k(rebin_factor=4)
+test_px_z.rebin_t(rebin_factor=4)
+new_px_z = test_px_z.rebin(rebin_t_factor=4, rebin_k_factor=4, include_k_0=True)
+
+# %%

From 4935e288c5e09ed9d9963dffcf795508c1bd7699 Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Tue, 29 Jul 2025 19:06:43 +0200
Subject: [PATCH 07/12] rebin healpix archive

---
 cupix/px_data/data_healpix.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/cupix/px_data/data_healpix.py b/cupix/px_data/data_healpix.py
index c61e095..60f9ebd 100644
--- a/cupix/px_data/data_healpix.py
+++ b/cupix/px_data/data_healpix.py
@@ -1,7 +1,7 @@
 import numpy as np
 import h5py
 
-from cupix.px_data import read_healpix_px, px_ztk, px_window
+from cupix.px_data import read_healpix_px, px_window
 
 
 class HealpixPxArchive(object):
@@ -61,7 +61,7 @@ def get_mean_px(self):
                                 z_bin, t_bin, self.k_bins,
                                 F_m=F_m, W_m=W_m, T_m=T_m, L=L_A)
                 list_px_zt.append(px_zt)
-            px_z = px_ztk.Px_z(self.t_bins, list_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
@@ -96,8 +96,21 @@ def get_mean_and_cov(self):
                                 P_m=mean_P_m, V_m=sum_V_m, C_mn=C_mn,
                                 F_m=None, W_m=None, T_m=None, U_mn=None)
                 list_px_zt.append(px_zt)
-            px_z = px_ztk.Px_z(self.t_bins, list_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 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

From bfb72daa95a99fcfafdba0ee64be40ba6f1072fe Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Tue, 29 Jul 2025 22:18:53 +0200
Subject: [PATCH 08/12] rebinned covariance

---
 notebooks/development/rebinned_covariance.py | 73 +++++++++++++++
 notebooks/development/test_healpixel_data.py | 94 ++++++++++++++++++--
 notebooks/development/test_rebin.py          | 11 +--
 3 files changed, 168 insertions(+), 10 deletions(-)
 create mode 100644 notebooks/development/rebinned_covariance.py

diff --git a/notebooks/development/rebinned_covariance.py b/notebooks/development/rebinned_covariance.py
new file mode 100644
index 0000000..dd9ba59
--- /dev/null
+++ b/notebooks/development/rebinned_covariance.py
@@ -0,0 +1,73 @@
+# %% [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'
+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):
+    '''Plot the mean Px, before and after rebinning'''    
+    
+    # plot mean px, before rebinning
+    k_m = [ 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}'
+        plt.plot(k_m, px_zt.P_m, alpha=0.5, label=label)
+
+    # plot mean px, after rebinning
+    k_m = [ 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))
+    plt.errorbar(k_m, px_zt.P_m, yerr=yerr, label=label)
+    plt.legend()
+    plt.xlim(0, 0.5)
+
+
+# %%
+compare_px(iz=0, rebinned_it=4)
+
+# %%
+compare_px(iz=1, rebinned_it=7)
+
+# %%
diff --git a/notebooks/development/test_healpixel_data.py b/notebooks/development/test_healpixel_data.py
index 2fe7aee..5883adb 100644
--- a/notebooks/development/test_healpixel_data.py
+++ b/notebooks/development/test_healpixel_data.py
@@ -45,21 +45,19 @@
 
 
 # %%
-def plot_px(iz, it):
+def compare_means(iz, it):
     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 = [k_bin.k for k_bin in archive.k_bins]
-    for px in archive.list_px:
-        px_zt = px.list_px_z[iz].list_px_zt[it]
-        plt.plot(k, px_zt.P_m, alpha=0.1)
     plt.plot(k, px_zt_v1.P_m, label='mean v1')
     plt.plot(k, px_zt_v2.P_m, label='mean v2')
     plt.xlim(0,0.5)
     plt.legend()
+#    plt.ylim(-0.1, 0.2)
 
 
 # %%
-plot_px(1,2)
+compare_means(0,4)
 
 
 # %%
@@ -79,6 +77,92 @@ def plot_weights(iz, it):
 # %%
 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):
+    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, px_zt.P_m, label=label)
+    plt.legend()
+    plt.xlim(0, 0.5)
+
+
+# %%
+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):
+    k_m = [ k_bin.mean() for k_bin in rebin_4_4.k_bins ]
+    px_zt_1 = rebin_4_4.list_px_z[iz].list_px_zt[it]
+    plt.plot(k_m, px_zt_1.P_m, label='first average, then rebin')
+    px_zt_2 = mean_4_4.list_px_z[iz].list_px_zt[it]
+    plt.plot(k_m, px_zt_2.P_m, label='first rebin, then average')
+    plt.xlim(0,0.8)
+    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):
+    k_m = [ 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))
+    plt.errorbar(k_m, px_zt.P_m, yerr=yerr)
+    plt.xlim(0,0.8)
+    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
index f9d62dc..83050cb 100644
--- a/notebooks/development/test_rebin.py
+++ b/notebooks/development/test_rebin.py
@@ -34,16 +34,17 @@ def sum_weights(pix, iz, it):
 
 
 # %%
+# find a good PX measurement to work with
 iz=1
 it=4
-for pix in range(10):
-    print(pix, sum_weights(pix, iz, it))
+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
-ipix = 5
-iz = 1
-it = 4
 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=2, include_k_0=True)
 

From be645fa3dc03c475ad890496638cdd0c2c0c9987 Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Wed, 30 Jul 2025 11:17:12 +0200
Subject: [PATCH 09/12] added notebook with interfaces

---
 cupix/px_data/px_window.py                   |   4 +-
 notebooks/development/likelihood.py          | 197 +++++++++++++++++++
 notebooks/development/rebinned_covariance.py |  23 ++-
 3 files changed, 221 insertions(+), 3 deletions(-)
 create mode 100644 notebooks/development/likelihood.py

diff --git a/cupix/px_data/px_window.py b/cupix/px_data/px_window.py
index 42c292b..ea6005a 100644
--- a/cupix/px_data/px_window.py
+++ b/cupix/px_data/px_window.py
@@ -153,7 +153,7 @@ def from_unnormalized(cls, z_bin, t_bin, k_bins, F_m, W_m, T_m, L):
         # should compute the window matrix here
         U_mn = None
 
-        return Px_zt_w(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn)
+        return cls(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn)
 
 
     @classmethod
@@ -167,7 +167,7 @@ def from_rebinning(cls, z_bin, t_bin, k_bins, P_m, V_m, U_mn):
         T_m = None
         C_mn = None
 
-        return Px_zt_w(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn)
+        return cls(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn)
 
 
     def rebin_k(self, rebin_factor, include_k_0=True):
diff --git a/notebooks/development/likelihood.py b/notebooks/development/likelihood.py
new file mode 100644
index 0000000..4b46d2f
--- /dev/null
+++ b/notebooks/development/likelihood.py
@@ -0,0 +1,197 @@
+# %% [markdown]
+# # Interfaces between likelihood, data and theory
+
+# %%
+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, kp, 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()
+        # theta bins (in arcmin), could use t_min, t_max here
+        theta_arcmin = [t_bin.mean() for t_bin in px_z.t_bins()]
+        # 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
+
+    
+    @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'''
+
+        # (data-theory) C^{-1} (data-theory)
+        return 0.0
+    
+
+# %%
+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)')
+
+
+# %%
+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}
+likelihood.get_chi2(like_params)
+
+# %%
diff --git a/notebooks/development/rebinned_covariance.py b/notebooks/development/rebinned_covariance.py
index dd9ba59..03210ef 100644
--- a/notebooks/development/rebinned_covariance.py
+++ b/notebooks/development/rebinned_covariance.py
@@ -12,7 +12,8 @@
 # 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_41-zbins_4-thetabins_40.hdf5'
+fname = basedir + '/px-nhp_150-zbins_4-thetabins_40.hdf5'
 print(fname) 
 
 # %%
@@ -70,4 +71,24 @@ def compare_px(iz, rebinned_it):
 # %%
 compare_px(iz=1, rebinned_it=7)
 
+
+# %%
+def plot_px(px, iz, theta_min, theta_max):
+    '''Plot Px at given z, for theta bins within range'''
+    k_m = [ 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))
+            plt.errorbar(k_m, px_zt.P_m, yerr=yerr, 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=3, theta_max=100)
+
 # %%

From 50a02792b847165223d4dea91b5a37ccfd718c5c Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Wed, 30 Jul 2025 16:02:01 +0200
Subject: [PATCH 10/12] working towards likelihood interface

---
 cupix/px_data/px_binning.py                  | 19 +++++
 cupix/px_data/px_window.py                   | 60 +++++++++++++
 notebooks/development/likelihood.py          | 90 ++++++++++++++++++--
 notebooks/development/test_healpixel_data.py |  7 +-
 notebooks/development/test_rebin.py          | 20 +++--
 5 files changed, 177 insertions(+), 19 deletions(-)

diff --git a/cupix/px_data/px_binning.py b/cupix/px_data/px_binning.py
index d5d2d50..837ef80 100644
--- a/cupix/px_data/px_binning.py
+++ b/cupix/px_data/px_binning.py
@@ -49,6 +49,15 @@ def B_t(self, t):
             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)'''
@@ -92,6 +101,16 @@ def B_k(self, k):
             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.'''
diff --git a/cupix/px_data/px_window.py b/cupix/px_data/px_window.py
index ea6005a..6fed44f 100644
--- a/cupix/px_data/px_window.py
+++ b/cupix/px_data/px_window.py
@@ -117,6 +117,66 @@ def rebin(self, rebin_t_factor, rebin_k_factor, include_k_0=True):
         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'''
diff --git a/notebooks/development/likelihood.py b/notebooks/development/likelihood.py
index 4b46d2f..52df60f 100644
--- a/notebooks/development/likelihood.py
+++ b/notebooks/development/likelihood.py
@@ -2,6 +2,8 @@
 # # 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
 
@@ -60,7 +62,7 @@ def model_px(self, px_z, like_params):
         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, kp, emu_params)
+        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
@@ -85,10 +87,12 @@ def get_rt_kp_Mpc(self, px_z):
 
         # 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()]
+        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()]
+        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)
@@ -106,6 +110,7 @@ def __init__(self, data, theory):
         print('setting up likelihood')
         self.data = data
         self.theory = theory
+        self.z_bins = data.z_bins
 
     
     @classmethod
@@ -120,9 +125,62 @@ def from_config(cls, config):
     def get_chi2(self, like_params):
         '''Given input parameters, return chi2'''
 
-        # (data-theory) C^{-1} (data-theory)
-        return 0.0
+        # 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):
@@ -156,6 +214,10 @@ def __init__(self, config):
         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 = {}
@@ -171,7 +233,7 @@ def __init__(self, config):
 theory = Theory(fid_cosmo, emulator, contaminants=None)
 
 # %%
-theory_2 = Theory.from_config(config)
+# theory_2 = Theory.from_config(config)
 
 # %%
 basedir = '/Users/afont/Codes/cupix/data/px_measurements/Lyacolore/'
@@ -185,13 +247,23 @@ def __init__(self, config):
 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)
+
+# %%
+likelihood = Likelihood(data, theory)
+# like_2 = Likelihood.from_config(config)
 
 # %%
 like_params = {'mF': 0.8}
-likelihood.get_chi2(like_params)
+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)
+
+# %%
+np.exp(-(0.8*0.5)**2)
 
 # %%
diff --git a/notebooks/development/test_healpixel_data.py b/notebooks/development/test_healpixel_data.py
index 5883adb..c3350c4 100644
--- a/notebooks/development/test_healpixel_data.py
+++ b/notebooks/development/test_healpixel_data.py
@@ -14,7 +14,8 @@
 # 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_41-zbins_4-thetabins_40.hdf5'
+fname = basedir + '/px-nhp_150-zbins_4-thetabins_40.hdf5'
 print(fname) 
 
 # %%
@@ -57,7 +58,7 @@ def compare_means(iz, it):
 
 
 # %%
-compare_means(0,4)
+compare_means(1,8)
 
 
 # %%
@@ -67,7 +68,7 @@ def plot_weights(iz, it):
     k = [k_bin.k for k_bin in archive.k_bins]
     for px in archive.list_px:
         px_zt = px.list_px_z[iz].list_px_zt[it]
-        plt.semilogy(k, px_zt.V_m, alpha=0.1)
+#        plt.semilogy(k, px_zt.V_m, alpha=0.1)
     plt.semilogy(k, px_zt_v1.V_m, label='total weights v1')
     plt.semilogy(k, px_zt_v2.V_m, label='total weights v2')
     plt.xlim(0,0.5)
diff --git a/notebooks/development/test_rebin.py b/notebooks/development/test_rebin.py
index 83050cb..695be65 100644
--- a/notebooks/development/test_rebin.py
+++ b/notebooks/development/test_rebin.py
@@ -46,7 +46,7 @@ def sum_weights(pix, iz, it):
 # %%
 # 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=2, include_k_0=True)
+new_px_zt = test_px_zt.rebin_k(rebin_factor=4, include_k_0=True)
 
 
 # %%
@@ -63,6 +63,12 @@ def compare(px_zt_1, px_zt_2):
 # %%
 compare(test_px_zt, new_px_zt)
 
+# %%
+new_px_zt.k_bins[-1].mean()
+
+# %%
+test_px_zt.k_bins[512].k
+
 # %% [markdown]
 # ## Rebin in theta
 
@@ -95,10 +101,10 @@ def compare(test_px_z, new_px_z, new_it, plot_V_m=False):
 
 
 # %%
-compare(test_px_z, new_px_z, 4)
+compare(test_px_z, new_px_z, 7)
 
 # %%
-compare(test_px_z, new_px_z, 4, plot_V_m=True)
+compare(test_px_z, new_px_z, 3, plot_V_m=True)
 
 # %% [markdown]
 # ### Rebin the entire PX measurement
@@ -106,15 +112,15 @@ def compare(test_px_z, new_px_z, new_it, plot_V_m=False):
 # %%
 # rebin the entire Px object
 test_px = archive.list_px[ipix]
-test_px.rebin_k(rebin_factor=4)
-test_px.rebin_t(rebin_factor=4)
+_ = test_px.rebin_k(rebin_factor=4)
+_ = test_px.rebin_t(rebin_factor=4)
 new_px = test_px.rebin(rebin_t_factor=4, rebin_k_factor=4, include_k_0=True)
 
 # %%
 # rebin only a redshift bin
 test_px_z = archive.list_px[ipix].list_px_z[iz]
-test_px_z.rebin_k(rebin_factor=4)
-test_px_z.rebin_t(rebin_factor=4)
+_ = test_px_z.rebin_k(rebin_factor=4)
+_ = test_px_z.rebin_t(rebin_factor=4)
 new_px_z = test_px_z.rebin(rebin_t_factor=4, rebin_k_factor=4, include_k_0=True)
 
 # %%

From 865d7c67a514e86c0bf88df264cfa1dadc9aa277 Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Wed, 30 Jul 2025 19:49:39 +0200
Subject: [PATCH 11/12] compute window

---
 cupix/px_data/data_healpix.py    |  8 +++---
 cupix/px_data/px_window.py       | 44 +++++++++++++++++++++++++++++---
 cupix/px_data/read_healpix_px.py | 13 +++++++---
 3 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/cupix/px_data/data_healpix.py b/cupix/px_data/data_healpix.py
index 60f9ebd..19c442e 100644
--- a/cupix/px_data/data_healpix.py
+++ b/cupix/px_data/data_healpix.py
@@ -7,7 +7,8 @@
 class HealpixPxArchive(object):
     '''Collection of PX measurements from healpixels'''
 
-    def __init__(self, fname, list_hp=None, list_px=None):
+    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:
@@ -17,7 +18,7 @@ def __init__(self, fname, list_hp=None, list_px=None):
             # 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()
+            self.list_hp, self.list_px = reader.get_list_px(compute_window)
         else:
             self.fname = None
             assert len(list_hp) == len(list_px)
@@ -59,7 +60,8 @@ def get_mean_px(self):
                     T_m += px_zt.T_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, T_m=T_m, L=L_A)
+                                F_m=F_m, W_m=W_m, T_m=T_m, 
+                                L=L_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)
diff --git a/cupix/px_data/px_window.py b/cupix/px_data/px_window.py
index 6fed44f..458213c 100644
--- a/cupix/px_data/px_window.py
+++ b/cupix/px_data/px_window.py
@@ -1,6 +1,7 @@
 import numpy as np
 from cupix.px_data import px_binning, px_ztk
 
+
 class Px_w(px_ztk.BasePx):
     '''Derived BasePx object, with information related to window matrix.'''
 
@@ -71,7 +72,6 @@ def rebin_t(self, rebin_factor):
             #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)
-            # should rebin window matrix U_mn here as well
             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)
@@ -205,13 +205,16 @@ def __init__(self, z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn=Non
 
 
     @classmethod
-    def from_unnormalized(cls, z_bin, t_bin, k_bins, F_m, W_m, T_m, L):
+    def from_unnormalized(cls, z_bin, t_bin, k_bins, 
+                F_m, W_m, T_m, L, compute_window=False):
         '''Construct object from unnormalized quantities'''
 
         P_m, V_m = normalize_Px(F_m, W_m, T_m, L)
         C_mn = None
-        # should compute the window matrix here
-        U_mn = None
+        if compute_window:
+            U_mn = compute_U_mn(W_m, T_m, L)
+        else:
+            U_mn = None
 
         return cls(z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn)
 
@@ -272,6 +275,39 @@ def normalize_Px(F_m, W_m, T_m, L):
     return P_m, V_m
 
 
+def compute_U_mn(W_m, T_m, L):
+    '''Compute window matrix'''
+
+    # effective resolution kernel (squared)
+    Nk = len(W_m)
+    R2_m = np.zeros(Nk)
+    R2_m[W_m>0] = T_m[W_m>0] / W_m[W_m>0]
+ 
+    # normalization
+    V_m = compute_V_m(W_m, T_m, L)
+
+    # setup window matrix
+    U_mn = np.zeros([Nk, Nk])
+
+    # avoid unnecessary computations
+    if np.sum(V_m) == 0:
+        #print('empty bin')
+        return U_mn
+
+    for m in range(Nk):
+        if V_m[m] >0:
+            V_m_L = V_m[m]*L
+            U_mn[m,m] = W_m[0] * R2_m[0] / V_m_L # m = n
+            for n in range(m):
+                U_mn[m,n] = W_m[m-n] * R2_m[n] / V_m_L # m > n
+            for n in range(m+1,Nk):
+                U_mn[n,m] = W_m[Nk+m-n] *  R2_m[n] / V_m_L # m < n
+        else:
+            print('ignore bin', m)
+
+    return U_mn
+
+
 def compute_V_m(W_m, T_m, L):
     '''Compute normalization factor for Px'''
 
diff --git a/cupix/px_data/read_healpix_px.py b/cupix/px_data/read_healpix_px.py
index 2d77678..4af7faf 100644
--- a/cupix/px_data/read_healpix_px.py
+++ b/cupix/px_data/read_healpix_px.py
@@ -105,7 +105,7 @@ def get_k_bins(self):
         return k_bins
 
 
-    def get_list_px(self, verbose=False):
+    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 
@@ -138,14 +138,19 @@ def get_list_px(self, verbose=False):
                     # 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, T_m=T_m, L=self.L_A)
+                            F_m=F_m, W_m=W_m, T_m=T_m, 
+                            L=self.L_A, compute_window=compute_window)
                     list_px_zt.append(px_zt)
-                # create Px_z object from these
+                # 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
+            # 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
 

From 36b95169deb45445155e1dd42869c8ca3b41e6da Mon Sep 17 00:00:00 2001
From: andreufont <font.andreu@gmail.com>
Date: Thu, 31 Jul 2025 19:57:14 +0200
Subject: [PATCH 12/12] last updates

---
 cupix/px_data/data_healpix.py                |  64 +++++++++--
 cupix/px_data/px_window.py                   | 109 +++++++++++--------
 cupix/px_data/read_healpix_px.py             |   8 +-
 notebooks/development/likelihood.py          |  15 ++-
 notebooks/development/rebinned_covariance.py |  32 ++++--
 notebooks/development/test_healpixel_data.py |  87 +++++++++------
 notebooks/development/test_rebin.py          |  29 ++---
 notebooks/development/test_window.py         |  93 ++++++++++++++++
 8 files changed, 313 insertions(+), 124 deletions(-)
 create mode 100644 notebooks/development/test_window.py

diff --git a/cupix/px_data/data_healpix.py b/cupix/px_data/data_healpix.py
index 19c442e..6306721 100644
--- a/cupix/px_data/data_healpix.py
+++ b/cupix/px_data/data_healpix.py
@@ -31,10 +31,18 @@ def __init__(self, fname, list_hp=None, list_px=None,
         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
 
@@ -45,23 +53,19 @@ def get_mean_px(self):
             list_px_zt = []
             for it, t_bin in enumerate(self.t_bins):
                 # this should be properly computed, not hard-coded
-                pw_A = 0.8
                 Nk = len(self.k_bins)
-                N_fft = Nk
-                L_A = pw_A * N_fft
                 # add the contribution from each healpixel
-                F_m = np.zeros(Nk) #, dtype='complex')
-                W_m = np.zeros(Nk) #, dtype='complex')
-                T_m = np.zeros(Nk) #, dtype='complex')
+                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
-                    T_m += px_zt.T_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, T_m=T_m, 
-                                L=L_A, compute_window=False)
+                                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)
@@ -69,7 +73,7 @@ def get_mean_px(self):
         return mean_px
 
 
-    def get_mean_and_cov(self):
+    def get_mean_and_cov(self, compute_window=False):
         list_px_z = []
         for iz, z_bin in enumerate(self.z_bins):
             list_px_zt = []
@@ -93,10 +97,15 @@ def get_mean_and_cov(self):
                 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, T_m=None, U_mn=None)
+                                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)
@@ -104,6 +113,39 @@ def get_mean_and_cov(self):
         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'''
 
diff --git a/cupix/px_data/px_window.py b/cupix/px_data/px_window.py
index 458213c..74a41db 100644
--- a/cupix/px_data/px_window.py
+++ b/cupix/px_data/px_window.py
@@ -1,4 +1,6 @@
 import numpy as np
+from numba import njit
+
 from cupix.px_data import px_binning, px_ztk
 
 
@@ -8,6 +10,13 @@ class Px_w(px_ztk.BasePx):
     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
 
 
@@ -60,6 +69,13 @@ class Px_z_w(px_ztk.Px_z):
     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
 
 
@@ -181,7 +197,8 @@ def rebin_model(self, raw_model, raw_px_z, convolve=True):
 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, T_m, U_mn, C_mn=None):
+    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__(
@@ -196,27 +213,31 @@ def __init__(self, z_bin, t_bin, k_bins, P_m, V_m, F_m, W_m, T_m, U_mn, C_mn=Non
         self.F_m = F_m
         # sum of the product of FFT of weights
         self.W_m = W_m
-        # sum of the product of FFT of resolution * weights
-        self.T_m = T_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, T_m, L, compute_window=False):
+                F_m, W_m, L_A, sig_A, compute_window=False):
         '''Construct object from unnormalized quantities'''
 
-        P_m, V_m = normalize_Px(F_m, W_m, T_m, L)
+        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:
-            U_mn = compute_U_mn(W_m, T_m, L)
+            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, T_m, U_mn, C_mn)
+        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
@@ -227,10 +248,11 @@ def from_rebinning(cls, z_bin, t_bin, k_bins, P_m, V_m, U_mn):
         F_m = P_m * V_m
         # for now, these are not needed
         W_m = None
-        T_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, T_m, U_mn, C_mn)
+        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):
@@ -264,58 +286,51 @@ def rebin_k(self, rebin_factor, include_k_0=True):
 
 
 
-def normalize_Px(F_m, W_m, T_m, L):
+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, T_m, L)
+    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
 
 
-def compute_U_mn(W_m, T_m, L):
-    '''Compute window matrix'''
-
-    # effective resolution kernel (squared)
+@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)
-    R2_m = np.zeros(Nk)
-    R2_m[W_m>0] = T_m[W_m>0] / W_m[W_m>0]
- 
-    # normalization
-    V_m = compute_V_m(W_m, T_m, L)
+    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
 
-    # setup window matrix
-    U_mn = np.zeros([Nk, Nk])
 
-    # avoid unnecessary computations
-    if np.sum(V_m) == 0:
-        #print('empty bin')
-        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
 
-    for m in range(Nk):
-        if V_m[m] >0:
-            V_m_L = V_m[m]*L
-            U_mn[m,m] = W_m[0] * R2_m[0] / V_m_L # m = n
-            for n in range(m):
-                U_mn[m,n] = W_m[m-n] * R2_m[n] / V_m_L # m > n
-            for n in range(m+1,Nk):
-                U_mn[n,m] = W_m[Nk+m-n] *  R2_m[n] / V_m_L # m < n
-        else:
-            print('ignore bin', m)
 
-    return U_mn
+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, T_m, L):
+def compute_V_m(W_m, R2_m, L_A):
     '''Compute normalization factor for Px'''
-
-    # effective resolution kernel (squared)
-    R2_m = np.zeros_like(W_m)
-    R2_m[W_m>0] = T_m[W_m>0] / W_m[W_m>0]
- 
     # 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
+    return np.abs(W_R2) / L_A
diff --git a/cupix/px_data/read_healpix_px.py b/cupix/px_data/read_healpix_px.py
index 4af7faf..44d8ab3 100644
--- a/cupix/px_data/read_healpix_px.py
+++ b/cupix/px_data/read_healpix_px.py
@@ -22,6 +22,8 @@ def __init__(self, fname, verbose=False):
             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 = [
@@ -134,12 +136,12 @@ def get_list_px(self, compute_window=False, verbose=False):
                     # set NaNs to zero
                     F_m[np.isnan(F_m)] = 0
                     W_m[np.isnan(W_m)] = 0
-                    T_m = W_m * R2_m
                     # 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, T_m=T_m, 
-                            L=self.L_A, compute_window=compute_window)
+                            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)
diff --git a/notebooks/development/likelihood.py b/notebooks/development/likelihood.py
index 52df60f..6d190d7 100644
--- a/notebooks/development/likelihood.py
+++ b/notebooks/development/likelihood.py
@@ -71,7 +71,7 @@ def model_px(self, px_z, like_params):
         px_arcmin_A *= 1.0
         
         return px_arcmin_A
-        
+
 
     def get_emu_params(self, z, like_params):
         '''Figure out emulator params given likelihood params'''
@@ -264,6 +264,17 @@ def __init__(self, config):
 #likelihood.get_chi2(like_params)
 
 # %%
-np.exp(-(0.8*0.5)**2)
+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
index 03210ef..22137f9 100644
--- a/notebooks/development/rebinned_covariance.py
+++ b/notebooks/development/rebinned_covariance.py
@@ -14,7 +14,7 @@
 #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) 
+print(fname)
 
 # %%
 # load PX measurement from different healpixels (before rebinning)
@@ -39,11 +39,11 @@
 
 
 # %%
-def compare_px(iz, rebinned_it):
+def compare_px(iz, rebinned_it, kmax=0.5):
     '''Plot the mean Px, before and after rebinning'''    
     
     # plot mean px, before rebinning
-    k_m = [ k_bin.k for k_bin in mean_px.k_bins ]
+    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
@@ -51,44 +51,52 @@ def compare_px(iz, 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}'
-        plt.plot(k_m, px_zt.P_m, alpha=0.5, label=label)
+        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 = [ k_bin.mean() for k_bin in mean_rebinned_px.k_bins ]
+    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))
-    plt.errorbar(k_m, px_zt.P_m, yerr=yerr, label=label)
+    mask = k_m < kmax
+    plt.errorbar(k_m[mask], px_zt.P_m[mask], yerr=yerr[mask], label=label)
     plt.legend()
-    plt.xlim(0, 0.5)
+    plt.xlim(0, kmax)
 
 
 # %%
 compare_px(iz=0, rebinned_it=4)
 
 # %%
-compare_px(iz=1, rebinned_it=7)
+compare_px(iz=1, rebinned_it=2)
 
 
 # %%
-def plot_px(px, iz, theta_min, theta_max):
+def plot_px(px, iz, theta_min, theta_max, kmax=0.5):
     '''Plot Px at given z, for theta bins within range'''
-    k_m = [ k_bin.mean() for k_bin in px.k_bins ]
+    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))
-            plt.errorbar(k_m, px_zt.P_m, yerr=yerr, label=label)
+            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=3, theta_max=100)
+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
index c3350c4..8471694 100644
--- a/notebooks/development/test_healpixel_data.py
+++ b/notebooks/development/test_healpixel_data.py
@@ -4,6 +4,7 @@
 # ### Read archive, compute mean and covariance and make plots
 
 # %%
+import time
 import numpy as np
 import matplotlib.pyplot as plt
 
@@ -14,24 +15,38 @@
 # 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'
+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}')
-N_fft = Nk
-# pixel width, in Angstroms
-pw_A = 0.8
-# length of FFT grid, in Angstroms
-L_A = N_fft * pw_A
-print(f'L = {L_A} A')
 # number of healpixels
 N_hp = len(archive.list_hp)
 print(f'Got {N_hp} healpixels')
@@ -46,13 +61,14 @@
 
 
 # %%
-def compare_means(iz, it):
+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 = [k_bin.k for k_bin in archive.k_bins]
-    plt.plot(k, px_zt_v1.P_m, label='mean v1')
-    plt.plot(k, px_zt_v2.P_m, label='mean v2')
-    plt.xlim(0,0.5)
+    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)
 
@@ -62,16 +78,17 @@ def compare_means(iz, it):
 
 
 # %%
-def plot_weights(iz, it):
+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 = [k_bin.k for k_bin in archive.k_bins]
+    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, px_zt.V_m, alpha=0.1)
-    plt.semilogy(k, px_zt_v1.V_m, label='total weights v1')
-    plt.semilogy(k, px_zt_v2.V_m, label='total weights v2')
-    plt.xlim(0,0.5)
+#        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()
 
 
@@ -92,7 +109,9 @@ def plot_weights(iz, it):
 
 
 # %%
-def compare_theta(iz=0, theta_min=1.0, theta_max=2.0):
+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
@@ -100,9 +119,10 @@ def compare_theta(iz=0, theta_min=1.0, theta_max=2.0):
             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, px_zt.P_m, label=label)
+                plt.plot(k_m[mask], px_zt.P_m[mask], label=label)
     plt.legend()
-    plt.xlim(0, 0.5)
+    plt.xlim(0, kmax)
+#    plt.ylim(-0.1,0.2)
 
 
 # %%
@@ -122,13 +142,13 @@ def compare_theta(iz=0, theta_min=1.0, theta_max=2.0):
 
 
 # %%
-def compare_means(iz, it):
-    k_m = [ k_bin.mean() for k_bin in rebin_4_4.k_bins ]
+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, px_zt_1.P_m, label='first average, then rebin')
+    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, px_zt_2.P_m, label='first rebin, then average')
-    plt.xlim(0,0.8)
+    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}' ")
@@ -146,12 +166,13 @@ def compare_means(iz, it):
 # ### Plot mean PX and errorbars
 
 # %%
-def plot_px(mean_px, iz, it):
-    k_m = [ k_bin.mean() for k_bin in mean_px.k_bins ]
+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))
-    plt.errorbar(k_m, px_zt.P_m, yerr=yerr)
-    plt.xlim(0,0.8)
+    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]
@@ -167,3 +188,5 @@ def plot_px(mean_px, iz, it):
 plot_px(mean_4_4, iz=1, it=5)
 
 # %%
+
+# %%
diff --git a/notebooks/development/test_rebin.py b/notebooks/development/test_rebin.py
index 695be65..c54f242 100644
--- a/notebooks/development/test_rebin.py
+++ b/notebooks/development/test_rebin.py
@@ -17,7 +17,8 @@
 
 # %%
 # load PX measurement from different healpixels
-archive = data_healpix.HealpixPxArchive(fname)
+compute_window=False
+archive = data_healpix.HealpixPxArchive(fname, compute_window=compute_window)
 
 # %%
 # define new, coarser k bins
@@ -50,25 +51,19 @@ def sum_weights(pix, iz, it):
 
 
 # %%
-def compare(px_zt_1, px_zt_2):
-    k1 = [k_bin.k for k_bin in px_zt_1.k_bins]
+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 = [k_bin.mean() for k_bin in px_zt_2.k_bins]
-    p2 = px_zt_2.P_m    
-    plt.plot(k1, p1, 'o')
-    plt.plot(k2, p2, 'x')
-    plt.xlim(0,0.1)
+    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)<kmax], p1[np.abs(k1)<kmax], 'o')
+    plt.plot(k2[np.abs(k2)<kmax], p2[np.abs(k2)<kmax], 'x')
+    plt.xlim(0,kmax)
 
 
 # %%
 compare(test_px_zt, new_px_zt)
 
-# %%
-new_px_zt.k_bins[-1].mean()
-
-# %%
-test_px_zt.k_bins[512].k
-
 # %% [markdown]
 # ## Rebin in theta
 
@@ -96,12 +91,12 @@ def compare(test_px_z, new_px_z, new_it, plot_V_m=False):
         plt.plot(k_m, new_px_z.list_px_zt[new_it].V_m)
     else:
         plt.plot(k_m, new_px_z.list_px_zt[new_it].P_m)
-    plt.xlim([0,0.5])
-    #plt.ylim([-0.01, 0.02])
+        plt.ylim([-0.1, 0.2])
+        plt.xlim([0,0.5])
 
 
 # %%
-compare(test_px_z, new_px_z, 7)
+compare(test_px_z, new_px_z, 4)
 
 # %%
 compare(test_px_z, new_px_z, 3, plot_V_m=True)
diff --git a/notebooks/development/test_window.py b/notebooks/development/test_window.py
new file mode 100644
index 0000000..2e11d75
--- /dev/null
+++ b/notebooks/development/test_window.py
@@ -0,0 +1,93 @@
+# %% [markdown]
+# # Compute window matrix
+#
+# ### Consider the impact of rebinning, and of averaging over healpixels 
+
+# %%
+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) 
+
+# %% [markdown]
+# ### Computing the window matrices takes some time, and use a lot of memory
+
+# %%
+# 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]
+# ### Play with the window of a particular (z,t) bin, of a healpixel
+
+# %%
+# find a good PX measurement to work with its window matrix
+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)
+iz=1
+it=4
+for pix in archive.list_hp:
+    if sum_weights(pix, iz, it)>0:
+        print(pix)
+        ipix = pix
+        break
+
+# %%
+test_px_zt = archive.list_px[pix].list_px_z[iz].list_px_zt[it]
+F_m = test_px_zt.F_m
+W_m = test_px_zt.W_m
+z_bin = test_px_zt.z_bin
+t_bin = test_px_zt.t_bin
+k_bins = test_px_zt.k_bins
+L_A = archive.L_A
+sig_A = archive.sig_A
+start = time.time()
+
+# %%
+R2_m = px_window.compute_R2_m(k_bins, sig_A)
+
+# %%
+# %timeit px_window.compute_V_m(W_m, R2_m, L_A)
+
+# %%
+# %timeit px_window.compute_U_mn(W_m, R2_m, L_A)
+
+# %%
+new_px_zt = px_window.Px_zt_w.from_unnormalized(z_bin, t_bin, k_bins, F_m, W_m, L_A, sig_A, True)
+
+# %%
+new_px_zt.U_mn
+
+# %%
+plt.imshow(new_px_zt.U_mn, vmax=0.01)
+plt.colorbar()
+
+# %%
+zoom=50
+plt.imshow(new_px_zt.U_mn, vmax=0.01)
+plt.xlim(0,zoom)
+plt.ylim(0,zoom)
+plt.colorbar()
+
+# %% [markdown]
+# ### Compute the average window, over healpixels
+
+# %%
+V_m, P_m, U_mn = archive.get_mean_window_zt(iz=0, it=4)
+
+# %%