Repository: igmhub/picca (language: python)

Issue:
Possible issue with rebinning
@andreicuceu noted that there is potentially an issue with the way we are doing the rebinning.

Original message by Andrei:
I think there's a small bug in picca delta extraction when defining the wavelength grid. These are true continuum deltas for the same forest, run with both picca and qsonic. You can see that for mocks the grid is slightly shifted compared to the data, so the bin centers are at 3599.9, 3600.7, etc.  instead of 3600.0, 3600.8, etc.. QSOnic has the correct grid, while picca shifts everything by 0.1A to the data grid. This is not a big deal, but it would be good to fix both the bug in picca and the shifted grid in QQ.

![issue](https://github.com/igmhub/picca/assets/9192975/966f4734-2d6f-4a56-972b-e65dcb77aaca)



Discussion / hints:
This is not a PICCA issue, that PICCA behaves as expected, but rather a quick quasars issue. While we wait for QuickQuasars to be fixed, it suffices to specify lambda min = 3599.9 and  lambda max = 5499.9 in the [data] section in the delta extraction config
![issue_check](https://github.com/igmhub/picca/assets/9192975/76824c4b-3186-4cdf-bfde-a6b0e7bfeff3)


Fix patch:
diff --git a/py/picca/delta_extraction/data_catalogues/desi_healpix_fast.py b/py/picca/delta_extraction/data_catalogues/desi_healpix_fast.py
new file mode 100644
index 000000000..0e75d2a38
--- /dev/null
+++ b/py/picca/delta_extraction/data_catalogues/desi_healpix_fast.py
@@ -0,0 +1,266 @@
+"""This module defines the class DesiData to load DESI data
+"""
+import logging
+import multiprocessing
+import time
+import itertools
+
+import fitsio
+import numpy as np
+
+from picca.delta_extraction.astronomical_objects.desi_forest import DesiForest
+from picca.delta_extraction.data_catalogues.desi_data import (  # pylint: disable=unused-import
+    DesiData,
+    DesiDataFileHandler,
+    accepted_options,
+    defaults,
+    merge_new_forest,
+    verify_exposures_shape,
+)
+from picca.delta_extraction.errors import DataError
+
+
+class DesiHealpixFast(DesiData):
+    """Reads the spectra from DESI using healpix mode and formats its data as a
+    list of Forest instances.
+
+    Methods
+    -------
+    (see DesiData in py/picca/delta_extraction/data_catalogues/desi_data.py)
+    __init__
+    __parse_config
+    get_filename
+    read_data
+    read_file
+
+    Attributes
+    ----------
+    (see DesiData in py/picca/delta_extraction/data_catalogues/desi_data.py)
+
+    logger: logging.Logger
+    Logger object
+
+    num_processors: int
+    Number of processors to be used for parallel reading
+    """
+
+    def __init__(self, config):
+        """Initialize class instance
+
+        Arguments
+        ---------
+        config: configparser.SectionProxy
+        Parsed options to initialize class
+        """
+        self.logger = logging.getLogger(__name__)
+
+        super().__init__(config)
+
+        #TODO: remove exception when this is implemented
+        if self.analysis_type == "PK 1D":
+            raise NotImplementedError("fast healpix reading is not implemented for PK 1D analyses")
+
+    def get_filename(self, survey, healpix, coadd_name):
+        """Get the name of the file to read
+
+        Arguments
+        ---------
+        survey: str
+        Name of the survey (sv, sv1, sv2, sv3, main, special)
+
+        healpix: int
+        Healpix of observations
+
+        Return
+        ------
+        filename: str
+        The name of the file to read
+
+        is_mock: bool
+        False, as we are reading DESI data
+        """
+        filename = (
+            f"{self.input_directory}/{survey}/dark/{healpix//100}/{healpix}/{coadd_name}-{survey}-"
+            f"dark-{healpix}.fits")
+        # TODO: not sure if we want the dark survey to be hard coded
+        # in here, probably won't run on anything else, but still
+        return filename
+
+    def read_data(self, is_mock=False):
+        """Read the data.
+
+        Method used to read healpix-based survey data.
+
+        Return
+        ------
+        is_mock: bool
+        False for DESI data and True for mocks
+
+        Raise
+        -----
+        DataError if no quasars were found
+        """
+        coadd_name = "spectra" if self.use_non_coadded_spectra else "coadd"
+
+        grouped_catalogue = self.catalogue.group_by(["HEALPIX", "SURVEY"])
+
+        arguments = [
+            (self.get_filename(group["SURVEY"][0], group["HEALPIX"][0], coadd_name), group)
+            for group in grouped_catalogue.groups
+        ]
+
+        self.logger.info(f"reading data from {len(arguments)} files")
+        if self.num_processors > 1:
+            context = multiprocessing.get_context('fork')
+            with context.Pool(processes=self.num_processors) as pool:
+                imap_it = pool.imap(
+                    DesiHealpixFileHandler(self.logger), arguments)
+                t0 = time.time()
+                self.logger.progress("Merging threads")
+
+
+                self.forests = combine_results(imap_it)
+                t1 = time.time()
+                self.logger.progress(f"Time spent meerging threads: {t1-t0}")
+
+        if len(self.forests) == 0:
+            raise DataError("No quasars found, stopping here")
+
+        return is_mock
+
+
+# Class to read in parallel
+# Seems lightweight to copy all these 3 arguments
+class DesiHealpixFileHandler():
+    """File handler for class DesiHealpix
+
+    Methods
+    -------
+    (see DesiDataFileHandler in py/picca/delta_extraction/data_catalogues/desi_data.py)
+    read_file
+
+    Attributes
+    ----------
+    (see DesiDataFileHandler in py/picca/delta_extraction/data_catalogues/desi_data.py)
+    """
+    def __init__(self, logger):
+        """Initialize file handler
+
+        Arguments
+        ---------
+        logger: logging.Logger
+        Logger object from parent class. Trying to initialize it here
+        without copying failed data_tests.py
+        """
+        # The next line gives failed tests
+        # self.logger = logging.getLogger(__name__)
+        self.logger = logger
+
+    def __call__(self, args):
+        """Call method read_file. Note imap can be called with
+        only one argument, hence tuple as argument.
+
+        Arguments
+        ---------
+        args: tuple
+        Arguments to be passed to read_file. Should contain a string with the
+        filename and a astropy.table.Table with the quasar catalogue
+        """
+        return self.read_file(*args)
+
+    def read_file(self, filename, catalogue):
+        """Read the spectra and formats its data as Forest instances.
+
+        Arguments
+        ---------
+        filename: str
+        Name of the file to read
+
+        catalogue: astropy.table.Table
+        The quasar catalogue fragment associated with this file
+
+        Returns:
+        ---------
+        forests_by_targetid: dict
+        Dictionary were forests are stored.
+
+        num_data: int
+        The number of instances loaded
+
+        Raise
+        -----
+        DataError if the analysis type is PK 1D and resolution data is not present
+        """
+        try:
+            hdul = fitsio.FITS(filename)
+        except IOError:
+            self.logger.warning(f"Error reading '{filename}'. Ignoring file")
+            return {}, 0
+        # Read targetid from fibermap to match to catalogue later
+        fibermap = hdul['FIBERMAP'].read()
+
+        colors = ["B", "R"]
+        if "Z_FLUX" in hdul:
+            colors.append("Z")
+
+        # read wavelength
+        wave = np.hstack([hdul[f"{color}_WAVELENGTH"].read() for color in colors])
+        log_lambda = np.log10(wave)
+        # read flux and ivar
+        flux_colors = {
+            color: hdul[f"{color}_FLUX"].read() for color in colors
+        }
+        ivar_colors = {
+            color: (hdul[f"{color}_IVAR"].read() * (hdul[f"{color}_MASK"].read() == 0))
+            for color in colors
+        }
+
+        # Loop over quasars in catalogue fragment
+        forests = []
+        for row in catalogue:
+            # Find which row in tile contains this quasar
+            # It should be there by construction
+            targetid = row["TARGETID"]
+            w_t = np.where(fibermap["TARGETID"] == targetid)[0]
+            if len(w_t) == 0:
+                self.logger.warning(
+                    f"Error reading {targetid}. Ignoring object")
+                continue
+
+            w_t = w_t[0]
+            # Construct DesiForest instance
+            # Fluxes from the different spectrographs will be coadded
+            flux = np.hstack([item[w_t] for item in flux_colors.values()])
+            ivar = np.hstack([item[w_t] for item in ivar_colors.values()])
+
+            args = {
+                "flux": flux,
+                "ivar": ivar,
+                "targetid": targetid,
+                "ra": row['RA'],
+                "dec": row['DEC'],
+                "z": row['Z'],
+                "log_lambda": log_lambda,
+            }
+            forest = DesiForest(**args)
+            forest.rebin()
+            forests.append(forest)
+
+        return forests
+
+
+def combine_results(lists):
+    """Combine the content of all the lists into a single one
+
+    Arguments
+    ---------
+    lists: list of lists
+    The content to be merged
+
+    Return
+    ------
+    combined_list: list
+    The combined list
+    """
+    # Combine a list of lists into a single list
+    return list(itertools.chain(*lists))
diff --git a/py/picca/delta_extraction/data_catalogues/desisim_mocks_fast.py b/py/picca/delta_extraction/data_catalogues/desisim_mocks_fast.py
new file mode 100644
index 000000000..b6ef81be7
--- /dev/null
+++ b/py/picca/delta_extraction/data_catalogues/desisim_mocks_fast.py
@@ -0,0 +1,87 @@
+"""This module defines the class DesiData to load DESI data
+"""
+import logging
+
+from picca.delta_extraction.data_catalogues.desi_healpix_fast import DesiHealpixFast
+from picca.delta_extraction.data_catalogues.desi_healpix_fast import (# pylint: disable=unused-import
+    defaults, accepted_options)
+
+class DesisimMocksFast(DesiHealpixFast):
+    """Reads the spectra from DESI using healpix mode and formats its data as a
+    list of Forest instances.
+
+    Should work for both data and mocks. This is specified using the 'mode'
+    keyword. It is required to set the in_nside member.
+
+    Methods
+    -------
+    (see DesiHealpix in py/picca/delta_extraction/data_catalogues/desi_healpix.py)
+    __init__
+    get_filename
+
+    Attributes
+    ----------
+    (see DesiHealpix in py/picca/delta_extraction/data_catalogues/desi_healpix.py)
+
+    in_nside: 16
+    Parameter in_nside to compute the healpix indexes
+
+    logger: logging.Logger
+    Logger object
+    """
+
+    def __init__(self, config):
+        """Initialize class instance
+
+        Arguments
+        ---------
+        config: configparser.SectionProxy
+        Parsed options to initialize class
+        """
+        self.logger = logging.getLogger(__name__)
+
+        # overwrite value for mocks
+        self.in_nside = 16
+        config["in_nside"] = str(self.in_nside)
+
+        super().__init__(config)
+
+    def get_filename(self, survey, healpix, coadd_name):
+        """Get the name of the file to read
+
+        Arguments
+        ---------
+        survey: str
+        Name of the survey (sv, sv1, sv2, sv3, main, special)
+
+        healpix: int
+        Healpix of observations
+
+        Return
+        ------
+        filename: str
+        The name of the file to read
+
+        is_mock: bool
+        True, as we are reading mocks
+        """
+        filename = (
+            f"{self.input_directory}/{healpix//100}/{healpix}/spectra-"
+            f"{self.in_nside}-{healpix}.fits")
+        return filename
+
+    def read_data(self, is_mock=True):
+        """Read the data.
+
+        Method used to read healpix-based survey data.
+
+        Return
+        ------
+        is_mock: bool
+        False for DESI data and True for mocks
+
+        Raise
+        -----
+        DataError if no quasars were found
+        """
+        return super().read_data(is_mock=is_mock)


Test patch:
diff --git a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Delta/delta-74.fits.gz b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Delta/delta-74.fits.gz
index 6d23ff885..a75603bd9 100644
Binary files a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Delta/delta-74.fits.gz and b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Delta/delta-74.fits.gz differ
diff --git a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes.fits.gz b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes.fits.gz
index 88207c5ca..3a02ad531 100644
Binary files a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes.fits.gz and b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes.fits.gz differ
diff --git a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration1.fits.gz b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration1.fits.gz
index 6601636fe..cec41cbcc 100644
Binary files a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration1.fits.gz and b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration1.fits.gz differ
diff --git a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration2.fits.gz b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration2.fits.gz
index 54d619033..0f86468c0 100644
Binary files a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration2.fits.gz and b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration2.fits.gz differ
diff --git a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration3.fits.gz b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration3.fits.gz
index 74d05b689..d12941734 100644
Binary files a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration3.fits.gz and b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration3.fits.gz differ
diff --git a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration4.fits.gz b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration4.fits.gz
index da74f33c9..1d2668e03 100644
Binary files a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration4.fits.gz and b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/delta_attributes_iteration4.fits.gz differ
diff --git a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/rejection_log.fits.gz b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/rejection_log.fits.gz
index 98b3c91df..81163d3b1 100644
Binary files a/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/rejection_log.fits.gz and b/py/picca/tests/delta_extraction/data/delta_extraction_lya_desi_mocks/Log/rejection_log.fits.gz differ
diff --git a/py/picca/tests/delta_extraction/data/delta_lin_fast.ini b/py/picca/tests/delta_extraction/data/delta_lin_fast.ini
new file mode 100644
index 000000000..04de3dae4
--- /dev/null
+++ b/py/picca/tests/delta_extraction/data/delta_lin_fast.ini
@@ -0,0 +1,13 @@
+[general]
+out dir = $THIS_DIR/results/delta_extraction_lin_fast/
+overwrite = True
+
+[data]
+type = DesiHealpixFast
+input directory = $THIS_DIR/data/
+save format = BinTableHDU
+catalogue = $THIS_DIR/data/QSO_cat_fuji_dark_healpix.fits.gz
+wave solution = lin
+
+[expected flux]
+type = Dr16ExpectedFlux
diff --git a/py/picca/tests/delta_extraction/data/delta_lin_image_fast.ini b/py/picca/tests/delta_extraction/data/delta_lin_image_fast.ini
new file mode 100644
index 000000000..ac4a98fe1
--- /dev/null
+++ b/py/picca/tests/delta_extraction/data/delta_lin_image_fast.ini
@@ -0,0 +1,13 @@
+[general]
+out dir = $THIS_DIR/results/delta_extraction_lin_image_fast/
+overwrite = True
+
+[data]
+type = DesiHealpixFast
+save format = ImageHDU
+input directory = $THIS_DIR/data/
+catalogue = $THIS_DIR/data/QSO_cat_fuji_dark_healpix.fits.gz
+wave solution = lin
+
+[expected flux]
+type = Dr16ExpectedFlux
diff --git a/py/picca/tests/delta_extraction/data/delta_lya_desi_mocks.ini b/py/picca/tests/delta_extraction/data/delta_lya_desi_mocks.ini
index 2f77370f9..efcc379bd 100644
--- a/py/picca/tests/delta_extraction/data/delta_lya_desi_mocks.ini
+++ b/py/picca/tests/delta_extraction/data/delta_lya_desi_mocks.ini
@@ -8,6 +8,8 @@ type = DesisimMocks
 save format = BinTableHDU
 catalogue = $THIS_DIR/data/desi_mock_test_catalogue.fits
 input directory = $THIS_DIR/data
+lambda min = 3599.9
+lambda max = 5499.9
 
 [corrections]
 num corrections = 0
diff --git a/py/picca/tests/delta_extraction/data/delta_lya_desi_mocks_fast.ini b/py/picca/tests/delta_extraction/data/delta_lya_desi_mocks_fast.ini
new file mode 100644
index 000000000..d185998c7
--- /dev/null
+++ b/py/picca/tests/delta_extraction/data/delta_lya_desi_mocks_fast.ini
@@ -0,0 +1,19 @@
+[general]
+overwrite = True
+out dir = $THIS_DIR/results/delta_extraction_lya_desi_mocks_fast/
+
+[data]
+wave solution = lin
+type = DesisimMocksFast
+save format = BinTableHDU
+catalogue = $THIS_DIR/data/desi_mock_test_catalogue.fits
+input directory = $THIS_DIR/data
+lambda min = 3599.9
+lambda max = 5499.9
+
+
+[corrections]
+num corrections = 0
+
+[expected flux]
+type = Dr16ExpectedFlux
diff --git a/py/picca/tests/delta_extraction/profiler.py b/py/picca/tests/delta_extraction/profiler.py
new file mode 100644
index 000000000..bd38952f2
--- /dev/null
+++ b/py/picca/tests/delta_extraction/profiler.py
@@ -0,0 +1,34 @@
+import os
+import sys
+import cProfile
+import argparse
+import pstats
+
+
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+os.environ["THIS_DIR"] = THIS_DIR
+
+PICCA_BIN = THIS_DIR.split("py/picca")[0]+"bin/"
+sys.path.append(PICCA_BIN)
+
+from picca_delta_extraction import main
+
+parser = argparse.ArgumentParser(
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    description=('Compute the delta field '
+                 'from a list of spectra'))
+
+parser.add_argument(
+    'config_file',
+    type=str,
+    default=None,
+    help=
+    ('Configuration file. To learn about all the available options '
+     'check the configuration tutorial in '
+     '$PICCA/tutorials/delta_extraction/picca_delta_extraction_tutorial.ipynb'
+    ))
+
+args = parser.parse_args()
+cProfile.run('main(args)', 'results/profile_output')
+p = pstats.Stats('results/profile_output')
+p.sort_stats('cumulative').print_stats()
diff --git a/py/picca/tests/delta_extraction/scripts_tests.py b/py/picca/tests/delta_extraction/scripts_tests.py
index 8bd10c86a..07f90c675 100644
--- a/py/picca/tests/delta_extraction/scripts_tests.py
+++ b/py/picca/tests/delta_extraction/scripts_tests.py
@@ -99,6 +99,14 @@ def test_delta_lin(self):
 
         self.run_delta_extraction(config_file, out_dir, test_dir)
 
+    def test_delta_lin_fast(self):
+        """End-to-end test using 'LYA' setup with  a linear wavelenth solution"""
+        config_file = "{}/data/delta_lin_fast.ini".format(THIS_DIR)
+        out_dir = "{}/results/delta_extraction_lin".format(THIS_DIR)
+        test_dir = "{}/data/delta_extraction_lin".format(THIS_DIR)
+
+        self.run_delta_extraction(config_file, out_dir, test_dir)
+
     def test_delta_lin_image(self):
         """End-to-end test using 'LYA' linear setup storing data as image"""
         config_file = "{}/data/delta_lin_image.ini".format(THIS_DIR)
@@ -107,6 +115,14 @@ def test_delta_lin_image(self):
 
         self.run_delta_extraction(config_file, out_dir, test_dir)
 
+    def test_delta_lin_image_fast(self):
+        """End-to-end test using 'LYA' linear setup storing data as image"""
+        config_file = "{}/data/delta_lin_image_fast.ini".format(THIS_DIR)
+        out_dir = "{}/results/delta_extraction_lin_image".format(THIS_DIR)
+        test_dir = "{}/data/delta_extraction_lin_image".format(THIS_DIR)
+
+        self.run_delta_extraction(config_file, out_dir, test_dir)
+
     def test_delta_lin_pk1d(self):
         """End-to-end test using 'LYA' setup with a linear wavelenth solution and Pk1D Forests.
         """
@@ -141,6 +157,15 @@ def test_delta_lya_desi_mocks(self):
 
         self.run_delta_extraction(config_file, out_dir, test_dir)
 
+    def test_delta_lya_desi_mocks_fast(self):
+        """End-to-end test using 'LYA' setup without corrections or masking,
+        for desi mocks"""
+        config_file = "{}/data/delta_lya_desi_mocks_fast.ini".format(THIS_DIR)
+        out_dir = "{}/results/delta_extraction_lya_desi_mocks".format(THIS_DIR)
+        test_dir = "{}/data/delta_extraction_lya_desi_mocks".format(THIS_DIR)
+
+        self.run_delta_extraction(config_file, out_dir, test_dir)
+
 
 if __name__ == '__main__':
     unittest.main()