8649536934

Committed 11 Apr 2024 03:45PM UTC coverage: 66.24% (-1.4%) from 67.631%

Build # 8649536934

Build Type

Pull #1056

gihtub

Committed by

corentinravoux

Commit Message

changing tests

Pull Request Pull Request #1056: Cross exposure

Run Details

1764 of 3059 branches covered (57.67%)

Branch coverage included in aggregate %.

128 of 247 new or added lines in 9 files covered. (51.82%)

3 existing lines in 2 files now uncovered.

6102 of 8816 relevant lines covered (69.22%)

2.08 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

93.29

/py/picca/delta_extraction/data_catalogues/desi_healpix.py

"""This module defines the class DesiData to load DESI data
"""
import logging
import multiprocessing
import os

import fitsio
import numpy as np
from picca.delta_extraction.astronomical_objects.desi_pk1d_forest import DesiPk1dForest
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 DesiHealpix(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)

        if self.analysis_type == "PK 1D":
            DesiPk1dForest.update_class_variables()

    def get_filename(self, survey, healpix):
        """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
        """
        input_directory = f'{self.input_directory}/{survey}/dark'
        coadd_name = "spectra" if self.use_non_coadded_spectra else "coadd"
        filename = (
            f"{input_directory}/{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, False

    def read_data(self):
        """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
        """
        grouped_catalogue = self.catalogue.group_by(["HEALPIX", "SURVEY"])

        is_mock = False
        forests_by_targetid = {}

        arguments = []
        for group in grouped_catalogue.groups:
            healpix, survey = group["HEALPIX", "SURVEY"][0]

            filename, is_mock_aux = self.get_filename(survey, healpix)
            if is_mock_aux:
                is_mock = True

            arguments.append((filename, group))

        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.analysis_type,
                                           self.use_non_coadded_spectra,
                                           self.keep_single_exposures,
                                           self.logger), arguments)
                for index, output_imap in enumerate(imap_it):
                    forests_by_targetid_aux = output_imap[0]
                    if self.use_non_coadded_spectra & self.keep_single_exposures:
                        # Change the dictionary key to prevent coadding.
                        # exposures on different files.
                        forests_by_targetid_aux= {f"{index}_{key}": items
                                                  for key, items in forests_by_targetid_aux.items()}
                    # Merge each dict to master forests_by_targetid
                    merge_new_forest(forests_by_targetid,
                                     forests_by_targetid_aux)

        else:
            reader = DesiHealpixFileHandler(self.analysis_type,
                                            self.use_non_coadded_spectra,
                                            self.keep_single_exposures,
                                            self.logger)
            num_data = 0
            for index, this_arg in enumerate(arguments):
                forests_by_targetid_aux, num_data_aux = reader(this_arg)
                if self.use_non_coadded_spectra & self.keep_single_exposures:
                    # Change the dictionary key to prevent coadding
                    # exposures on different files.
                    forests_by_targetid_aux= {f"{index}_{key}": items
                                              for key, items in forests_by_targetid_aux.items()}
                merge_new_forest(forests_by_targetid, forests_by_targetid_aux)
                num_data += num_data_aux
                self.logger.progress(f"Read {index} of {len(arguments)}. "
                                     f"num_data: {num_data}")
        if self.use_non_coadded_spectra & self.keep_single_exposures:
            forests_by_targetid = verify_exposures_shape(forests_by_targetid)

        if len(forests_by_targetid) == 0:
            raise DataError("No quasars found, stopping here")
        self.forests = list(forests_by_targetid.values())

        return is_mock


# Class to read in parallel
# Seems lightweight to copy all these 3 arguments
class DesiHealpixFileHandler(DesiDataFileHandler):
    """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 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()
        # First read all wavelength, flux, ivar, mask, and resolution
        # from this file
        spectrographs_data = {}
        colors = ["B", "R"]
        if "Z_FLUX" in hdul:
            colors.append("Z")
        else:
            self.logger.warning(
                f"Missing Z band from {filename}. Ignoring color.")

        hdul_truth = None
        reso_from_truth = False
        if self.analysis_type == "PK 1D" and any(f"{c}_RESOLUTION" not in hdul for c in colors):
            self.logger.debug(
                    "no resolution in files, reading from truth files"
                )
            basename_truth = os.path.basename(filename).replace(
                            'spectra-', 'truth-')
            pathname_truth = os.path.dirname(filename)
            filename_truth = f"{pathname_truth}/{basename_truth}"
            if os.path.exists(filename_truth):
                hdul_truth = fitsio.FITS(filename_truth)
                reso_from_truth = True

        def _read_resolution(color):
            if f"{color}_RESOLUTION" in hdul:
                return hdul[f"{color}_RESOLUTION"].read()
            if hdul_truth is not None:
                return hdul_truth[f"{color}_RESOLUTION"].read()

            raise DataError(
                    f"Error while reading {color} band from "
                    f"{filename}. Analysis type is 'PK 1D', "
                    "but file does not contain HDU "
                    f"'{color}_RESOLUTION'")

        for color in colors:
            spec = {}
            try:
                spec["WAVELENGTH"] = hdul[f"{color}_WAVELENGTH"].read()
                spec["FLUX"] = hdul[f"{color}_FLUX"].read()
                spec["IVAR"] = (hdul[f"{color}_IVAR"].read() *
                                (hdul[f"{color}_MASK"].read() == 0))
                w = np.isnan(spec["FLUX"]) | np.isnan(spec["IVAR"])
                for key in ["FLUX", "IVAR"]:
                    spec[key][w] = 0.

                if self.analysis_type == "PK 1D":
                    spec["RESO"] = _read_resolution(color)

                spectrographs_data[color] = spec
            except OSError:
                self.logger.warning(
                    f"Error while reading {color} band from {filename}. "
                    "Ignoring color.")
        hdul.close()
        if hdul_truth is not None:
            hdul_truth.close()

        forests_by_targetid, num_data = self.format_data(
            catalogue,
            spectrographs_data,
            fibermap["TARGETID"],
            reso_from_truth=reso_from_truth)

        return forests_by_targetid, num_data

1	"""This module defines the class DesiData to load DESI data
2	"""
3	import logging	3✔
4	import multiprocessing	3✔
5	import os	3✔
6
7	import fitsio	3✔
8	import numpy as np	3✔
9	from picca.delta_extraction.astronomical_objects.desi_pk1d_forest import DesiPk1dForest	3✔
10	from picca.delta_extraction.data_catalogues.desi_data import ( # pylint: disable=unused-import	3✔
11	DesiData,
12	DesiDataFileHandler,
13	accepted_options,
14	defaults,
15	merge_new_forest,
16	verify_exposures_shape,
17	)
18	from picca.delta_extraction.errors import DataError	3✔
19
20
21	class DesiHealpix(DesiData):	3✔
22	"""Reads the spectra from DESI using healpix mode and formats its data as a
23	list of Forest instances.
24
25	Methods
26	-------
27	(see DesiData in py/picca/delta_extraction/data_catalogues/desi_data.py)
28	__init__
29	__parse_config
30	get_filename
31	read_data
32	read_file
33
34	Attributes
35	----------
36	(see DesiData in py/picca/delta_extraction/data_catalogues/desi_data.py)
37
38	logger: logging.Logger
39	Logger object
40
41	num_processors: int
42	Number of processors to be used for parallel reading
43	"""
44
45	def __init__(self, config):	3✔
46	"""Initialize class instance
47
48	Arguments
49	---------
50	config: configparser.SectionProxy
51	Parsed options to initialize class
52	"""
53	self.logger = logging.getLogger(__name__)	3✔
54
55	super().__init__(config)	3✔
56
57	if self.analysis_type == "PK 1D":	3✔
58	DesiPk1dForest.update_class_variables()	3✔
59
60	def get_filename(self, survey, healpix):	3✔
61	"""Get the name of the file to read
62
63	Arguments
64	---------
65	survey: str
66	Name of the survey (sv, sv1, sv2, sv3, main, special)
67
68	healpix: int
69	Healpix of observations
70
71	Return
72	------
73	filename: str
74	The name of the file to read
75
76	is_mock: bool
77	False, as we are reading DESI data
78	"""
79	input_directory = f'{self.input_directory}/{survey}/dark'	3✔
80	coadd_name = "spectra" if self.use_non_coadded_spectra else "coadd"	3✔
81	filename = (	3✔
82	f"{input_directory}/{healpix//100}/{healpix}/{coadd_name}-{survey}-"
83	f"dark-{healpix}.fits")
84	# TODO: not sure if we want the dark survey to be hard coded
85	# in here, probably won't run on anything else, but still
86	return filename, False	3✔
87
88	def read_data(self):	3✔
89	"""Read the data.
90
91	Method used to read healpix-based survey data.
92
93	Return
94	------
95	is_mock: bool
96	False for DESI data and True for mocks
97
98	Raise
99	-----
100	DataError if no quasars were found
101	"""
102	grouped_catalogue = self.catalogue.group_by(["HEALPIX", "SURVEY"])	3✔
103
104	is_mock = False	3✔
105	forests_by_targetid = {}	3✔
106
107	arguments = []	3✔
108	for group in grouped_catalogue.groups:	3✔
109	healpix, survey = group["HEALPIX", "SURVEY"][0]	3✔
110
111	filename, is_mock_aux = self.get_filename(survey, healpix)	3✔
112	if is_mock_aux:	3✔
113	is_mock = True	3✔
114
115	arguments.append((filename, group))	3✔
116
117	self.logger.info(f"reading data from {len(arguments)} files")	3✔
118	if self.num_processors > 1:	3✔
119	context = multiprocessing.get_context('fork')	3✔
120	with context.Pool(processes=self.num_processors) as pool:	3✔
121	imap_it = pool.imap(	3✔
122	DesiHealpixFileHandler(self.analysis_type,
123	self.use_non_coadded_spectra,
124	self.keep_single_exposures,
125	self.logger), arguments)
126	for index, output_imap in enumerate(imap_it):	3✔
127	forests_by_targetid_aux = output_imap[0]	3✔
128	if self.use_non_coadded_spectra & self.keep_single_exposures:	3!
129	# Change the dictionary key to prevent coadding.
130	# exposures on different files.
NEW 131	forests_by_targetid_aux= {f"{index}_{key}": items	×
132	for key, items in forests_by_targetid_aux.items()}
133	# Merge each dict to master forests_by_targetid
134	merge_new_forest(forests_by_targetid,	3✔
135	forests_by_targetid_aux)
136
137	else:
138	reader = DesiHealpixFileHandler(self.analysis_type,	3✔
139	self.use_non_coadded_spectra,
140	self.keep_single_exposures,
141	self.logger)
142	num_data = 0	3✔
143	for index, this_arg in enumerate(arguments):	3✔
144	forests_by_targetid_aux, num_data_aux = reader(this_arg)	3✔
145	if self.use_non_coadded_spectra & self.keep_single_exposures:	3!
146	# Change the dictionary key to prevent coadding
147	# exposures on different files.
NEW 148	forests_by_targetid_aux= {f"{index}_{key}": items	×
149	for key, items in forests_by_targetid_aux.items()}
150	merge_new_forest(forests_by_targetid, forests_by_targetid_aux)	3✔
151	num_data += num_data_aux	3✔
152	self.logger.progress(f"Read {index} of {len(arguments)}. "	3✔
153	f"num_data: {num_data}")
154	if self.use_non_coadded_spectra & self.keep_single_exposures:	3!
NEW 155	forests_by_targetid = verify_exposures_shape(forests_by_targetid)	×
156
157	if len(forests_by_targetid) == 0:	3✔
158	raise DataError("No quasars found, stopping here")	3✔
159	self.forests = list(forests_by_targetid.values())	3✔
160
161	return is_mock	3✔
162
163
164	# Class to read in parallel
165	# Seems lightweight to copy all these 3 arguments
166	class DesiHealpixFileHandler(DesiDataFileHandler):	3✔
167	"""File handler for class DesiHealpix
168
169	Methods
170	-------
171	(see DesiDataFileHandler in py/picca/delta_extraction/data_catalogues/desi_data.py)
172	read_file
173
174	Attributes
175	----------
176	(see DesiDataFileHandler in py/picca/delta_extraction/data_catalogues/desi_data.py)
177	"""
178
179	def read_file(self, filename, catalogue):	3✔
180	"""Read the spectra and formats its data as Forest instances.
181
182	Arguments
183	---------
184	filename: str
185	Name of the file to read
186
187	catalogue: astropy.table.Table
188	The quasar catalogue fragment associated with this file
189
190	Returns:
191	---------
192	forests_by_targetid: dict
193	Dictionary were forests are stored.
194
195	num_data: int
196	The number of instances loaded
197
198	Raise
199	-----
200	DataError if the analysis type is PK 1D and resolution data is not present
201	"""
202	try:	3✔
203	hdul = fitsio.FITS(filename)	3✔
204	except IOError:	3✔
205	self.logger.warning(f"Error reading '{filename}'. Ignoring file")	3✔
206	return {}, 0	3✔
207	# Read targetid from fibermap to match to catalogue later
208	fibermap = hdul['FIBERMAP'].read()	3✔
209	# First read all wavelength, flux, ivar, mask, and resolution
210	# from this file
211	spectrographs_data = {}	3✔
212	colors = ["B", "R"]	3✔
213	if "Z_FLUX" in hdul:	3✔
214	colors.append("Z")	3✔
215	else:
216	self.logger.warning(	3✔
217	f"Missing Z band from {filename}. Ignoring color.")
218
219	hdul_truth = None	3✔
220	reso_from_truth = False	3✔
221	if self.analysis_type == "PK 1D" and any(f"{c}_RESOLUTION" not in hdul for c in colors):	3✔
222	self.logger.debug(	3✔
223	"no resolution in files, reading from truth files"
224	)
225	basename_truth = os.path.basename(filename).replace(	3✔
226	'spectra-', 'truth-')
227	pathname_truth = os.path.dirname(filename)	3✔
228	filename_truth = f"{pathname_truth}/{basename_truth}"	3✔
229	if os.path.exists(filename_truth):	3✔
230	hdul_truth = fitsio.FITS(filename_truth)	3✔
231	reso_from_truth = True	3✔
232
233	def _read_resolution(color):	3✔
234	if f"{color}_RESOLUTION" in hdul:	3✔
235	return hdul[f"{color}_RESOLUTION"].read()	3✔
236	if hdul_truth is not None:	3✔
237	return hdul_truth[f"{color}_RESOLUTION"].read()	3✔
238
239	raise DataError(	3✔
240	f"Error while reading {color} band from "
241	f"{filename}. Analysis type is 'PK 1D', "
242	"but file does not contain HDU "
243	f"'{color}_RESOLUTION'")
244
245	for color in colors:	3✔
246	spec = {}	3✔
247	try:	3✔
248	spec["WAVELENGTH"] = hdul[f"{color}_WAVELENGTH"].read()	3✔
249	spec["FLUX"] = hdul[f"{color}_FLUX"].read()	3✔
250	spec["IVAR"] = (hdul[f"{color}_IVAR"].read() *	3✔
251	(hdul[f"{color}_MASK"].read() == 0))
252	w = np.isnan(spec["FLUX"]) \| np.isnan(spec["IVAR"])	3✔
253	for key in ["FLUX", "IVAR"]:	3✔
254	spec[key][w] = 0.	3✔
255
256	if self.analysis_type == "PK 1D":	3✔
257	spec["RESO"] = _read_resolution(color)	3✔
258
259	spectrographs_data[color] = spec	3✔
260	except OSError:	3✔
261	self.logger.warning(	3✔
262	f"Error while reading {color} band from {filename}. "
263	"Ignoring color.")
264	hdul.close()	3✔
265	if hdul_truth is not None:	3✔
266	hdul_truth.close()	3✔
267
268	forests_by_targetid, num_data = self.format_data(	3✔
269	catalogue,
270	spectrographs_data,
271	fibermap["TARGETID"],
272	reso_from_truth=reso_from_truth)
273
274	return forests_by_targetid, num_data	3✔

igmhub / picca / 8649536934

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous