From d187a17f1abb2afd10c2330391c40a41d60f4aa3 Mon Sep 17 00:00:00 2001
From: Javier Sanchez <jsanchez@Javiers-MacBook-Air.local>
Date: Thu, 15 Feb 2024 12:32:20 -0500
Subject: [PATCH] some fixes

---
 augur/generate.py | 34 ++++++++++++++++++++++++----------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/augur/generate.py b/augur/generate.py
index 4c8bc8b..8c0e924 100644
--- a/augur/generate.py
+++ b/augur/generate.py
@@ -209,7 +209,7 @@ def generate_sacc_and_stats(config):
             sources[sacc_tracer].bias = bias
             sys_params[f'{sacc_tracer}_bias'] = bias
             sys_params[f'{sacc_tracer}_delta_z'] = delta_z[i]
-
+    
     # Read data vector combinations
     if 'statistics' not in config.keys():
         raise ValueError('statistics key is required in config file')
@@ -231,8 +231,18 @@ def generate_sacc_and_stats(config):
                 ells_here = ells[ells < ell_max]
             else:
                 ells_here = ells
-            for ell in ells_here:
-                S.add_data_point(key, (tr1, tr2), 0.0, ell=ell, error=1e30)
+            # Trying to add bandpower windows
+            ells_aux = np.arange(0, np.max(ells_here)+1)
+            wgt = np.zeros((len(ells_aux), len(ells_here)))
+            for i in range(len(ells_here)):
+                in_win = (ells_aux > ell_edges[i]) & (ells_aux < ell_edges[i+1])
+                wgt[in_win, i] = 1.0 
+            win = sacc.BandpowerWindow(ells_aux, wgt)
+            S.add_ell_cl(key, tr1, tr2,
+                         ells_here, np.zeros(len(ells_here)), window=win)
+            # # For some reason add_ell_cl does not update the sacc file properly on the fly...
+            # for i, ell in enumerate(ells_here):
+            #     S.add_data_point(key, (tr1, tr2), 0.0, ell=ell, error=1e30, window=win[i])
             # Now create TwoPoint objects for firecrown
             _aux_stat = TwoPoint(source0=sources[tr1], source1=sources[tr2],
                                  sacc_data_type=key)
@@ -264,19 +274,23 @@ def generate(config, return_all_outputs=False, write_sacc=True, force_read=True)
 
     Returns:
     -------
-    noise : float
-       Noise power for Cls for that particular tracer. That is 1/nbar for
-       number counts and e**2/nbar for weak lensing tracer.
+    lk : firecrown.likelihood.ConstGaussian
+        Likelihood object, only returned if `return_all_outputs` is True.
+    S : sacc.Sacc
+        Sacc object with fake data vector and covariance. It is only returned if `return_all_outputs` 
+        is True.
+    tools : firecrown.modeling.ModelingTools
+        Modeling tools, only returned if `return_all_outputs` is True.
+    sys_params : dict
+        Dictionary containing the modeling systematic parameters. It is only returned if `return_all_outputs`
+        is True.
 
-    Note:
-    -----
-    The input number_densities are #/arcmin.
-    The output units are in steradian.
     """
 
     config = parse_config(config)
     # Generate placeholders
     S, cosmo, stats, sys_params = generate_sacc_and_stats(config)
+
     # Generate likelihood object
     lk = ConstGaussian(statistics=stats)
     # Pass the correct binning/tracers