3750956619

Build Type

push

github-actions

Committed by Ignasi Pérez-Ràfols

Commit Message

linted code

Run Details

1940 of 3641 branches covered (53.28%)

Branch coverage included in aggregate %.

20 of 20 new or added lines in 2 files covered. (100.0%)

7405 of 11412 relevant lines covered (64.89%)

1.95 hits per line

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

48.77

/py/picca/pk1d/compute_pk1d.py

"""This module defines functions and variables required for the analysis of
the 1D power spectra

This module provides with one clas (Pk1D) and several functions:
    - split_forest
    - rebin_diff_noise
    - fill_masked_pixels
    - compute_pk_raw
    - compute_pk_noise
    - compute_correction_reso
See the respective docstrings for more details
"""
import numpy as np
from numpy.fft import rfft, rfftfreq

from picca import constants
from picca.utils import userprint


def split_forest(num_parts,
                 pixel_step,
                 lambda_or_log_lambda,
                 delta,
                 exposures_diff,
                 ivar,
                 first_pixel_index,
                 abs_igm="LYA",
                 reso_matrix=None,
                 linear_binning=False):
    """Split the forest in n parts

    Arguments
    ---------
    num_parts: int
    Number of parts

    pixel_step: float
    Variation of the wavelength (or log(wavelength)) between two pixels

    lambda_or_log_lambda: array of float
    Wavelength (in Angs) (or its log, but needs to be consistent with
    delta_lambda_or_log_lambda)

    delta: array of float
    Mean transmission fluctuation (delta field)

    exposures_diff: array of float
    Semidifference between two customized coadded spectra obtained from
    weighted averages of the even-number exposures, for the first
    spectrum, and of the odd-number exposures, for the second one

    ivar: array of float
    Inverse variances

    first_pixel_index: int
    Index of the first pixel in the forest

    abs_igm: string - Default: "LYA"
    Name of the absorption in picca.constants defining the
    redshift of the forest pixels

    reso_matrix: 2d-array of float
    The resolution matrix used for corrections

    linear_binning: bool
    assume linear wavelength binning, lambda_or_log_lambda vectors will be
    actual lambda (else they will be log_lambda)

    Return
    ------
    mean_z_array: array of float
    Array with the mean redshift the parts of the forest
    lambda_or_log_lambda_array: Array with logarith of the wavelength for the
        parts of the forest

    delta_array: array of float
    Array with the deltas for the parts of the forest
    exposures_diff_array: Array with the exposures_diff for the parts of
    the forest

    ivar_array: array of float
    Array with the ivar for the parts of the forest
    """
    lambda_or_log_lambda_limit = [lambda_or_log_lambda[first_pixel_index]]
    num_bins = (len(lambda_or_log_lambda) - first_pixel_index) // num_parts

    mean_z_array = []
    lambda_or_log_lambda_array = []
    delta_array = []
    exposures_diff_array = []
    ivar_array = []
    if reso_matrix is not None:
        reso_matrix_array = []

    for index in range(1, num_parts):
        lambda_or_log_lambda_limit.append(
            lambda_or_log_lambda[num_bins * index + first_pixel_index])

    lambda_or_log_lambda_limit.append(
        lambda_or_log_lambda[len(lambda_or_log_lambda) - 1] + 0.1 * pixel_step)

    for index in range(num_parts):
        selection = (
            (lambda_or_log_lambda >= lambda_or_log_lambda_limit[index]) &
            (lambda_or_log_lambda < lambda_or_log_lambda_limit[index + 1]))

        lambda_or_log_lambda_part = lambda_or_log_lambda[selection].copy()
        lambda_abs_igm = constants.ABSORBER_IGM[abs_igm]

        if linear_binning:
            mean_z = np.mean(lambda_or_log_lambda_part) / lambda_abs_igm - 1.0
        else:
            mean_z = np.mean(10**
                             lambda_or_log_lambda_part) / lambda_abs_igm - 1.0

        if reso_matrix is not None:
            reso_matrix_part = reso_matrix[:, selection].copy()

        mean_z_array.append(mean_z)
        lambda_or_log_lambda_array.append(lambda_or_log_lambda_part)
        delta_array.append(delta[selection].copy())
        if exposures_diff is not None:
            exposures_diff_array.append(exposures_diff[selection].copy())
        else:  # fill exposures_diff_array with zeros
            userprint('WARNING: exposure_diff is None, filling with zeros.')
            exposures_diff_array.append(np.zeros(delta[selection].shape))
        ivar_array.append(ivar[selection].copy())
        if reso_matrix is not None:
            reso_matrix_array.append(reso_matrix_part)

    out = [
        mean_z_array, lambda_or_log_lambda_array, delta_array,
        exposures_diff_array, ivar_array
    ]
    if reso_matrix is not None:
        out.append(reso_matrix_array)
    return out


def rebin_diff_noise(pixel_step, lambda_or_log_lambda, exposures_diff):
    """Rebin the semidifference between two customized coadded spectra to
    construct the noise array

    Note that inputs can be either linear or log-lambda spaced units (but
    pixel_step and lambda_or_log_lambda need the same unit)

    The rebinning is done by combining 3 of the original pixels into analysis
    pixels.

    Arguments
    ---------
    pixel_step: float
    Variation of the logarithm of the wavelength between two pixels
    for linear binnings this would need to be the wavelength difference

    lambda_or_log_lambda: array of float
    Array containing the logarithm of the wavelengths (in Angs)
    for linear binnings this would need to be just wavelength

    exposures_diff: array of float
    Semidifference between two customized coadded spectra obtained from
    weighted averages of the even-number exposures, for the first
    spectrum, and of the odd-number exposures, for the second one

    Return
    ------
    noise: array of float
    The noise array
    """
    rebin = 3
    if exposures_diff.size < rebin:
        userprint("Warning: exposures_diff.size too small for rebin")
        return exposures_diff
    rebin_delta_lambda_or_log_lambda = rebin * pixel_step

    # rebin not mixing pixels separated by masks
    bins = np.floor((lambda_or_log_lambda - lambda_or_log_lambda.min()) /
                    rebin_delta_lambda_or_log_lambda + 0.5).astype(int)

    rebin_exposure_diff = np.bincount(bins.astype(int), weights=exposures_diff)
    rebin_counts = np.bincount(bins.astype(int))
    w = (rebin_counts > 0)
    if len(rebin_counts) == 0:
        userprint("Error: exposures_diff size = 0 ", exposures_diff)
    rebin_exposure_diff = rebin_exposure_diff[w] / np.sqrt(rebin_counts[w])

    # now merge the rebinned array into a noise array
    noise = np.zeros(exposures_diff.size)
    for index in range(len(exposures_diff) // len(rebin_exposure_diff) + 1):
        length_max = min(len(exposures_diff),
                         (index + 1) * len(rebin_exposure_diff))
        noise[index *
              len(rebin_exposure_diff):length_max] = rebin_exposure_diff[:(
                  length_max - index * len(rebin_exposure_diff))]
        # shuffle the array before the next iteration
        np.random.shuffle(rebin_exposure_diff)

    return noise


def fill_masked_pixels(delta_lambda_or_log_lambda, lambda_or_log_lambda, delta,
                       exposures_diff, ivar, no_apply_filling):
    """Fills the masked pixels with zeros

    Note that inputs can be either linear or log-lambda spaced units (but
    delta_lambda_or_log_lambda and lambda_or_log_lambda need the same unit);
    spectrum needs to be uniformly binned in that unit

    Arguments
    ---------
    delta_lambda_or_log_lambda: float
    Variation of the logarithm of the wavelength between two pixels

    lambda_or_log_lambda: array of float
    Array containing the logarithm of the wavelengths (in Angs)

    delta: array of float
    Mean transmission fluctuation (delta field)

    exposures_diff: array of float
    Semidifference between two customized coadded spectra obtained from
    weighted averages of the even-number exposures, for the first
    spectrum, and of the odd-number exposures, for the second one

    ivar: array of float
    Array containing the inverse variance

    no_apply_filling: boolean
    If True, then return the original arrays

    Return
    ------
    lambda_or_log_lambda_new: array of float
    Array containing the wavelength (in Angs) or its logarithm depending on the
    wavelength solution used

    delta_new: array of float
    Mean transmission fluctuation (delta field)

    exposures_diff_new: array of float
    Semidifference between two customized coadded spectra obtained from
    weighted averages of the even-number exposures, for the first spectrum,
    and of the odd-number exposures, for the second one

    ivar_new: array of float
    Array containing the inverse variance

    num_masked_pixels: array of int
    Number of masked pixels
    """
    if no_apply_filling:
        return lambda_or_log_lambda, delta, exposures_diff, ivar, 0

    lambda_or_log_lambda_index = lambda_or_log_lambda.copy()
    lambda_or_log_lambda_index -= lambda_or_log_lambda[0]
    lambda_or_log_lambda_index /= delta_lambda_or_log_lambda
    lambda_or_log_lambda_index += 0.5
    lambda_or_log_lambda_index = np.array(lambda_or_log_lambda_index, dtype=int)
    index_all = range(lambda_or_log_lambda_index[-1] + 1)
    index_ok = np.in1d(index_all, lambda_or_log_lambda_index)

    delta_new = np.zeros(len(index_all))
    delta_new[index_ok] = delta

    lambda_or_log_lambda_new = np.array(index_all, dtype=float)
    lambda_or_log_lambda_new *= delta_lambda_or_log_lambda
    lambda_or_log_lambda_new += lambda_or_log_lambda[0]

    exposures_diff_new = np.zeros(len(index_all))
    if exposures_diff is not None:
        exposures_diff_new[index_ok] = exposures_diff

    ivar_new = np.zeros(len(index_all), dtype=float)
    ivar_new[index_ok] = ivar

    num_masked_pixels = len(index_all) - len(lambda_or_log_lambda_index)

    return (lambda_or_log_lambda_new, delta_new, exposures_diff_new, ivar_new,
            num_masked_pixels)


def compute_pk_raw(delta_lambda_or_log_lambda, delta, linear_binning=False):
    """Compute the raw power spectrum

    Arguments
    ---------
    delta_lambda_or_log_lambda: float
    Variation of (the logarithm of) the wavelength between two pixels

    delta: array of float
    Mean transmission fluctuation (delta field)

    linear_binning: bool
    If set then inputs need to be in AA, outputs will be 1/AA else inputs will
    be in log(AA) and outputs in s/km

    Return
    ------
    k: array of float
    The Fourier modes the Power Spectrum is measured on

    pk: array of float
    The Power Spectrum
    """
    if linear_binning:  # spectral length in AA
        length_lambda = (delta_lambda_or_log_lambda * len(delta))
    else:  # spectral length in km/s
        length_lambda = (delta_lambda_or_log_lambda * constants.SPEED_LIGHT *
                         np.log(10.) * len(delta))

    # make 1D FFT
    num_pixels = len(delta)
    fft_delta = rfft(delta)

    # compute power spectrum
    pk = (fft_delta.real**2 + fft_delta.imag**2) * length_lambda / num_pixels**2
    k = 2 * np.pi * rfftfreq(num_pixels, length_lambda / num_pixels)

    return k, pk


def compute_pk_noise(delta_lambda_or_log_lambda,
                     ivar,
                     exposures_diff,
                     run_noise,
                     num_noise_exposures=10,
                     linear_binning=False):
    """Compute the noise power spectrum

    Two noise power spectrum are computed: one using the pipeline noise and
    another one using the noise derived from exposures_diff

    Arguments
    ---------
    delta_lambda_or_log_lambda: float
    Variation of the logarithm of the wavelength between two pixels

    ivar: array of float
    Array containing the inverse variance

    exposures_diff: array of float
    Semidifference between two customized coadded spectra obtained from
    weighted averages of the even-number exposures, for the first
    spectrum, and of the odd-number exposures, for the second one

    run_noise: boolean
    If False the noise power spectrum using the pipeline noise is not
    computed and an array filled with zeros is returned instead

    num_noise_exposures: int
    Number of exposures to average for noise power estimate

    Return
    ------
    pk_noise: array of float
    The noise Power Spectrum using the pipeline noise

    pk_diff: array of float
    The noise Power Spectrum using the noise derived from exposures_diff
    """
    num_pixels = len(ivar)
    num_bins_fft = num_pixels // 2 + 1

    pk_noise = np.zeros(num_bins_fft)
    error = np.zeros(num_pixels)
    w = ivar > 0
    error[w] = 1.0 / np.sqrt(ivar[w])

    if run_noise:
        for _ in range(num_noise_exposures):
            delta_exp = np.zeros(num_pixels)
            delta_exp[w] = np.random.normal(0., error[w])
            _, pk_exp = compute_pk_raw(delta_lambda_or_log_lambda,
                                       delta_exp,
                                       linear_binning=linear_binning)
            pk_noise += pk_exp

        pk_noise /= float(num_noise_exposures)

    _, pk_diff = compute_pk_raw(delta_lambda_or_log_lambda,
                                exposures_diff,
                                linear_binning=linear_binning)

    return pk_noise, pk_diff


def compute_correction_reso(delta_pixel, mean_reso, k):
    """Computes the resolution correction

    Arguments
    ---------
    delta_pixel: float
    Variation of the logarithm of the wavelength between two pixels
    (in km/s or Ang depending on the units of k submitted)

    mean_reso: float
    Mean resolution of the forest

    k: array of float
    Fourier modes

    Return
    ------
    correction: array of float
    The resolution correction
    """
    num_bins_fft = len(k)
    correction = np.ones(num_bins_fft)

    pixelization_factor = np.sinc(k * delta_pixel / (2 * np.pi))**2

    correction *= np.exp(-(k * mean_reso)**2)
    correction *= pixelization_factor
    return correction


def compute_correction_reso_matrix(reso_matrix, k, delta_pixel, num_pixel):
    """Computes the resolution correction based on the resolution matrix using linear binning

    Arguments
    ---------
    delta_pixel: float
    Variation of the logarithm of the wavelength between two pixels
    (in km/s or Ang depending on the units of k submitted)

    num_pixel: int
    Length  of the spectrum in pixels

    mean_reso: float
    Mean resolution of the forest

    k: array of float
    Fourier modes

    Return
    ------
    correction: array of float
    The resolution correction
    """
    #this allows either computing the power for each pixel seperately or for the mean
    reso_matrix = np.atleast_2d(reso_matrix)

    w2_arr = []
    #first compute the power in the resmat for each pixel, then average
    for resmat in reso_matrix:
        r = np.append(resmat, np.zeros(num_pixel - resmat.size))
        k_resmat, w2 = compute_pk_raw(delta_pixel, r, linear_binning=True)
        try:
            assert np.all(k_resmat == k)
        except AssertionError as error:
            raise Exception(
                "for some reason the resolution matrix correction has "
                "different k scaling than the pk") from error
        w2 /= w2[0]
        w2_arr.append(w2)

    w_res2 = np.mean(w2_arr, axis=0)
    pixelization_factor = np.sinc(k * delta_pixel / (2 * np.pi))**2

    # the following assumes that the resolution matrix is storing the actual
    # resolution convolved with the pixelization kernel along each matrix axis
    correction = w_res2
    correction /= pixelization_factor

    return correction


class Pk1D:
    """Class to represent the 1D Power Spectrum for a given forest

    Class Methods
    -------------
    from_fitsio

    Methods
    -------
    __init__

    Attributes
    ----------
    ra: float
    Right-ascension of the quasar (in radians).

    dec: float
    Declination of the quasar (in radians).

    z_qso: float
    Redshift of the quasar.

    plate: int
    Plate number of the observation.

    fiberid: int
    Fiberid of the observation.

    mjd: int
    Modified Julian Date of the observation.

    mean_snr: float
    Mean signal-to-noise ratio in the forest

    mean_reso: float
    Mean resolution of the forest

    mean_z: float
    Mean redshift of the forest

    num_masked_pixels: int
    Number of masked pixels

    k: array of float
    Fourier modes

    pk_raw: array of float
    Raw power spectrum

    pk_noise: array of float
    Noise power spectrum for the different Fourier modes

    correction_reso: array of float
    Resolution correction

    pk: array of float
    Power Spectrum

    pk_diff: array of float or None
    Power spectrum of exposures_diff for the different Fourier modes
    """
    def __init__(self,
                 ra,
                 dec,
                 z_qso,
                 mean_z,
                 plate,
                 mjd,
                 fiberid,
                 mean_snr,
                 mean_reso,
                 k,
                 pk_raw,
                 pk_noise,
                 correction_reso,
                 pk,
                 num_masked_pixels,
                 pk_diff=None):
        """Initialize instance

        Arguments
        ---------
        ra: float
        Right-ascension of the quasar (in radians).

        dec: float
        Declination of the quasar (in radians).

        z_qso: float
        Redshift of the quasar.

        mean_z: float
        Mean redshift of the forest

        plate: int
        Plate number of the observation.

        mjd: int
        Modified Julian Date of the observation.

        fiberid: int
        Fiberid of the observation.

        mean_snr: float
        Mean signal-to-noise ratio in the forest

        mean_reso: float
        Mean resolution of the forest

        k: array of float
        Fourier modes

        pk_raw: array of float
        Raw power spectrum

        pk_noise: array of float
        Noise power spectrum for the different Fourier modes

        correction_reso: array of float
        Resolution correction

        pk: array of float
        Power Spectrum

        num_masked_pixels: int
        Number of masked pixels

        pk_diff: array of float or None - default: None
        Power spectrum of exposures_diff for the different Fourier modes
        """
        self.ra = ra
        self.dec = dec
        self.z_qso = z_qso
        self.mean_z = mean_z
        self.mean_snr = mean_snr
        self.mean_reso = mean_reso
        self.num_masked_pixels = num_masked_pixels

        self.plate = plate
        self.mjd = mjd
        self.fiberid = fiberid
        self.k = k
        self.pk_raw = pk_raw
        self.pk_noise = pk_noise
        self.correction_reso = correction_reso
        self.pk = pk
        self.pk_diff = pk_diff

    @classmethod
    def from_fitsio(cls, hdu):
        """Read the 1D Power Spectrum from fits file

        Arguments
        ---------
        hdu: Header Data Unit
        Header Data Unit where the 1D Power Spectrum is read

        Return
        ------
        An intialized instance of Pk1D
        """

        header = hdu.read_header()

        ra = header['RA']
        dec = header['DEC']
        z_qso = header['Z']
        mean_z = header['MEANZ']
        mean_reso = header['MEANRESO']
        mean_snr = header['MEANSNR']
        plate = header['PLATE']
        mjd = header['MJD']
        fiberid = header['FIBER']
        num_masked_pixels = header['NBMASKPIX']

        data = hdu.read()
        k = data['k'][:]
        pk = data['Pk'][:]
        pk_raw = data['Pk_raw'][:]
        pk_noise = data['Pk_noise'][:]
        correction_reso = data['cor_reso'][:]
        pk_diff = data['Pk_diff'][:]

        return cls(ra, dec, z_qso, mean_z, plate, mjd, fiberid, mean_snr,
                   mean_reso, k, pk_raw, pk_noise, correction_reso, pk,
                   num_masked_pixels, pk_diff)

1	"""This module defines functions and variables required for the analysis of
2	the 1D power spectra
3
4	This module provides with one clas (Pk1D) and several functions:
5	- split_forest
6	- rebin_diff_noise
7	- fill_masked_pixels
8	- compute_pk_raw
9	- compute_pk_noise
10	- compute_correction_reso
11	See the respective docstrings for more details
12	"""
13	import numpy as np	3✔
14	from numpy.fft import rfft, rfftfreq	3✔
15
16	from picca import constants	3✔
17	from picca.utils import userprint	3✔
18
19
20	def split_forest(num_parts,	3✔
21	pixel_step,
22	lambda_or_log_lambda,
23	delta,
24	exposures_diff,
25	ivar,
26	first_pixel_index,
27	abs_igm="LYA",
28	reso_matrix=None,
29	linear_binning=False):
30	"""Split the forest in n parts
31
32	Arguments
33	---------
34	num_parts: int
35	Number of parts
36
37	pixel_step: float
38	Variation of the wavelength (or log(wavelength)) between two pixels
39
40	lambda_or_log_lambda: array of float
41	Wavelength (in Angs) (or its log, but needs to be consistent with
42	delta_lambda_or_log_lambda)
43
44	delta: array of float
45	Mean transmission fluctuation (delta field)
46
47	exposures_diff: array of float
48	Semidifference between two customized coadded spectra obtained from
49	weighted averages of the even-number exposures, for the first
50	spectrum, and of the odd-number exposures, for the second one
51
52	ivar: array of float
53	Inverse variances
54
55	first_pixel_index: int
56	Index of the first pixel in the forest
57
58	abs_igm: string - Default: "LYA"
59	Name of the absorption in picca.constants defining the
60	redshift of the forest pixels
61
62	reso_matrix: 2d-array of float
63	The resolution matrix used for corrections
64
65	linear_binning: bool
66	assume linear wavelength binning, lambda_or_log_lambda vectors will be
67	actual lambda (else they will be log_lambda)
68
69	Return
70	------
71	mean_z_array: array of float
72	Array with the mean redshift the parts of the forest
73	lambda_or_log_lambda_array: Array with logarith of the wavelength for the
74	parts of the forest
75
76	delta_array: array of float
77	Array with the deltas for the parts of the forest
78	exposures_diff_array: Array with the exposures_diff for the parts of
79	the forest
80
81	ivar_array: array of float
82	Array with the ivar for the parts of the forest
83	"""
84	lambda_or_log_lambda_limit = [lambda_or_log_lambda[first_pixel_index]]	3✔
85	num_bins = (len(lambda_or_log_lambda) - first_pixel_index) // num_parts	3✔
86
87	mean_z_array = []	3✔
88	lambda_or_log_lambda_array = []	3✔
89	delta_array = []	3✔
90	exposures_diff_array = []	3✔
91	ivar_array = []	3✔
92	if reso_matrix is not None:	3!
93	reso_matrix_array = []	×
94
95	for index in range(1, num_parts):	3✔
96	lambda_or_log_lambda_limit.append(	3✔
97	lambda_or_log_lambda[num_bins * index + first_pixel_index])
98
99	lambda_or_log_lambda_limit.append(	3✔
100	lambda_or_log_lambda[len(lambda_or_log_lambda) - 1] + 0.1 * pixel_step)
101
102	for index in range(num_parts):	3✔
103	selection = (	3✔
104	(lambda_or_log_lambda >= lambda_or_log_lambda_limit[index]) &
105	(lambda_or_log_lambda < lambda_or_log_lambda_limit[index + 1]))
106
107	lambda_or_log_lambda_part = lambda_or_log_lambda[selection].copy()	3✔
108	lambda_abs_igm = constants.ABSORBER_IGM[abs_igm]	3✔
109
110	if linear_binning:	3!
111	mean_z = np.mean(lambda_or_log_lambda_part) / lambda_abs_igm - 1.0	×
112	else:
113	mean_z = np.mean(10**	3✔
114	lambda_or_log_lambda_part) / lambda_abs_igm - 1.0
115
116	if reso_matrix is not None:	3!
117	reso_matrix_part = reso_matrix[:, selection].copy()	×
118
119	mean_z_array.append(mean_z)	3✔
120	lambda_or_log_lambda_array.append(lambda_or_log_lambda_part)	3✔
121	delta_array.append(delta[selection].copy())	3✔
122	if exposures_diff is not None:	3!
123	exposures_diff_array.append(exposures_diff[selection].copy())	3✔
124	else: # fill exposures_diff_array with zeros
125	userprint('WARNING: exposure_diff is None, filling with zeros.')	×
126	exposures_diff_array.append(np.zeros(delta[selection].shape))	×
127	ivar_array.append(ivar[selection].copy())	3✔
128	if reso_matrix is not None:	3!
129	reso_matrix_array.append(reso_matrix_part)	×
130
131	out = [	3✔
132	mean_z_array, lambda_or_log_lambda_array, delta_array,
133	exposures_diff_array, ivar_array
134	]
135	if reso_matrix is not None:	3!
136	out.append(reso_matrix_array)	×
137	return out	3✔
138
139
140	def rebin_diff_noise(pixel_step, lambda_or_log_lambda, exposures_diff):	3✔
141	"""Rebin the semidifference between two customized coadded spectra to
142	construct the noise array
143
144	Note that inputs can be either linear or log-lambda spaced units (but
145	pixel_step and lambda_or_log_lambda need the same unit)
146
147	The rebinning is done by combining 3 of the original pixels into analysis
148	pixels.
149
150	Arguments
151	---------
152	pixel_step: float
153	Variation of the logarithm of the wavelength between two pixels
154	for linear binnings this would need to be the wavelength difference
155
156	lambda_or_log_lambda: array of float
157	Array containing the logarithm of the wavelengths (in Angs)
158	for linear binnings this would need to be just wavelength
159
160	exposures_diff: array of float
161	Semidifference between two customized coadded spectra obtained from
162	weighted averages of the even-number exposures, for the first
163	spectrum, and of the odd-number exposures, for the second one
164
165	Return
166	------
167	noise: array of float
168	The noise array
169	"""
170	rebin = 3	×
171	if exposures_diff.size < rebin:	×
172	userprint("Warning: exposures_diff.size too small for rebin")	×
173	return exposures_diff	×
174	rebin_delta_lambda_or_log_lambda = rebin * pixel_step	×
175
176	# rebin not mixing pixels separated by masks
177	bins = np.floor((lambda_or_log_lambda - lambda_or_log_lambda.min()) /	×
178	rebin_delta_lambda_or_log_lambda + 0.5).astype(int)
179
180	rebin_exposure_diff = np.bincount(bins.astype(int), weights=exposures_diff)	×
181	rebin_counts = np.bincount(bins.astype(int))	×
182	w = (rebin_counts > 0)	×
183	if len(rebin_counts) == 0:	×
184	userprint("Error: exposures_diff size = 0 ", exposures_diff)	×
185	rebin_exposure_diff = rebin_exposure_diff[w] / np.sqrt(rebin_counts[w])	×
186
187	# now merge the rebinned array into a noise array
188	noise = np.zeros(exposures_diff.size)	×
189	for index in range(len(exposures_diff) // len(rebin_exposure_diff) + 1):	×
190	length_max = min(len(exposures_diff),	×
191	(index + 1) * len(rebin_exposure_diff))
192	noise[index *	×
193	len(rebin_exposure_diff):length_max] = rebin_exposure_diff[:(
194	length_max - index * len(rebin_exposure_diff))]
195	# shuffle the array before the next iteration
196	np.random.shuffle(rebin_exposure_diff)	×
197
198	return noise	×
199
200
201	def fill_masked_pixels(delta_lambda_or_log_lambda, lambda_or_log_lambda, delta,	3✔
202	exposures_diff, ivar, no_apply_filling):
203	"""Fills the masked pixels with zeros
204
205	Note that inputs can be either linear or log-lambda spaced units (but
206	delta_lambda_or_log_lambda and lambda_or_log_lambda need the same unit);
207	spectrum needs to be uniformly binned in that unit
208
209	Arguments
210	---------
211	delta_lambda_or_log_lambda: float
212	Variation of the logarithm of the wavelength between two pixels
213
214	lambda_or_log_lambda: array of float
215	Array containing the logarithm of the wavelengths (in Angs)
216
217	delta: array of float
218	Mean transmission fluctuation (delta field)
219
220	exposures_diff: array of float
221	Semidifference between two customized coadded spectra obtained from
222	weighted averages of the even-number exposures, for the first
223	spectrum, and of the odd-number exposures, for the second one
224
225	ivar: array of float
226	Array containing the inverse variance
227
228	no_apply_filling: boolean
229	If True, then return the original arrays
230
231	Return
232	------
233	lambda_or_log_lambda_new: array of float
234	Array containing the wavelength (in Angs) or its logarithm depending on the
235	wavelength solution used
236
237	delta_new: array of float
238	Mean transmission fluctuation (delta field)
239
240	exposures_diff_new: array of float
241	Semidifference between two customized coadded spectra obtained from
242	weighted averages of the even-number exposures, for the first spectrum,
243	and of the odd-number exposures, for the second one
244
245	ivar_new: array of float
246	Array containing the inverse variance
247
248	num_masked_pixels: array of int
249	Number of masked pixels
250	"""
251	if no_apply_filling:	3!
252	return lambda_or_log_lambda, delta, exposures_diff, ivar, 0	×
253
254	lambda_or_log_lambda_index = lambda_or_log_lambda.copy()	3✔
255	lambda_or_log_lambda_index -= lambda_or_log_lambda[0]	3✔
256	lambda_or_log_lambda_index /= delta_lambda_or_log_lambda	3✔
257	lambda_or_log_lambda_index += 0.5	3✔
258	lambda_or_log_lambda_index = np.array(lambda_or_log_lambda_index, dtype=int)	3✔
259	index_all = range(lambda_or_log_lambda_index[-1] + 1)	3✔
260	index_ok = np.in1d(index_all, lambda_or_log_lambda_index)	3✔
261
262	delta_new = np.zeros(len(index_all))	3✔
263	delta_new[index_ok] = delta	3✔
264
265	lambda_or_log_lambda_new = np.array(index_all, dtype=float)	3✔
266	lambda_or_log_lambda_new *= delta_lambda_or_log_lambda	3✔
267	lambda_or_log_lambda_new += lambda_or_log_lambda[0]	3✔
268
269	exposures_diff_new = np.zeros(len(index_all))	3✔
270	if exposures_diff is not None:	3!
271	exposures_diff_new[index_ok] = exposures_diff	3✔
272
273	ivar_new = np.zeros(len(index_all), dtype=float)	3✔
274	ivar_new[index_ok] = ivar	3✔
275
276	num_masked_pixels = len(index_all) - len(lambda_or_log_lambda_index)	3✔
277
278	return (lambda_or_log_lambda_new, delta_new, exposures_diff_new, ivar_new,	3✔
279	num_masked_pixels)
280
281
282	def compute_pk_raw(delta_lambda_or_log_lambda, delta, linear_binning=False):	3✔
283	"""Compute the raw power spectrum
284
285	Arguments
286	---------
287	delta_lambda_or_log_lambda: float
288	Variation of (the logarithm of) the wavelength between two pixels
289
290	delta: array of float
291	Mean transmission fluctuation (delta field)
292
293	linear_binning: bool
294	If set then inputs need to be in AA, outputs will be 1/AA else inputs will
295	be in log(AA) and outputs in s/km
296
297	Return
298	------
299	k: array of float
300	The Fourier modes the Power Spectrum is measured on
301
302	pk: array of float
303	The Power Spectrum
304	"""
305	if linear_binning: # spectral length in AA	3!
306	length_lambda = (delta_lambda_or_log_lambda * len(delta))	×
307	else: # spectral length in km/s
308	length_lambda = (delta_lambda_or_log_lambda * constants.SPEED_LIGHT *	3✔
309	np.log(10.) * len(delta))
310
311	# make 1D FFT
312	num_pixels = len(delta)	3✔
313	fft_delta = rfft(delta)	3✔
314
315	# compute power spectrum
316	pk = (fft_delta.real2 + fft_delta.imag2) * length_lambda / num_pixels**2	3✔
317	k = 2 * np.pi * rfftfreq(num_pixels, length_lambda / num_pixels)	3✔
318
319	return k, pk	3✔
320
321
322	def compute_pk_noise(delta_lambda_or_log_lambda,	3✔
323	ivar,
324	exposures_diff,
325	run_noise,
326	num_noise_exposures=10,
327	linear_binning=False):
328	"""Compute the noise power spectrum
329
330	Two noise power spectrum are computed: one using the pipeline noise and
331	another one using the noise derived from exposures_diff
332
333	Arguments
334	---------
335	delta_lambda_or_log_lambda: float
336	Variation of the logarithm of the wavelength between two pixels
337
338	ivar: array of float
339	Array containing the inverse variance
340
341	exposures_diff: array of float
342	Semidifference between two customized coadded spectra obtained from
343	weighted averages of the even-number exposures, for the first
344	spectrum, and of the odd-number exposures, for the second one
345
346	run_noise: boolean
347	If False the noise power spectrum using the pipeline noise is not
348	computed and an array filled with zeros is returned instead
349
350	num_noise_exposures: int
351	Number of exposures to average for noise power estimate
352
353	Return
354	------
355	pk_noise: array of float
356	The noise Power Spectrum using the pipeline noise
357
358	pk_diff: array of float
359	The noise Power Spectrum using the noise derived from exposures_diff
360	"""
361	num_pixels = len(ivar)	3✔
362	num_bins_fft = num_pixels // 2 + 1	3✔
363
364	pk_noise = np.zeros(num_bins_fft)	3✔
365	error = np.zeros(num_pixels)	3✔
366	w = ivar > 0	3✔
367	error[w] = 1.0 / np.sqrt(ivar[w])	3✔
368
369	if run_noise:	3!
370	for _ in range(num_noise_exposures):	×
371	delta_exp = np.zeros(num_pixels)	×
372	delta_exp[w] = np.random.normal(0., error[w])	×
373	_, pk_exp = compute_pk_raw(delta_lambda_or_log_lambda,	×
374	delta_exp,
375	linear_binning=linear_binning)
376	pk_noise += pk_exp	×
377
378	pk_noise /= float(num_noise_exposures)	×
379
380	_, pk_diff = compute_pk_raw(delta_lambda_or_log_lambda,	3✔
381	exposures_diff,
382	linear_binning=linear_binning)
383
384	return pk_noise, pk_diff	3✔
385
386
387	def compute_correction_reso(delta_pixel, mean_reso, k):	3✔
388	"""Computes the resolution correction
389
390	Arguments
391	---------
392	delta_pixel: float
393	Variation of the logarithm of the wavelength between two pixels
394	(in km/s or Ang depending on the units of k submitted)
395
396	mean_reso: float
397	Mean resolution of the forest
398
399	k: array of float
400	Fourier modes
401
402	Return
403	------
404	correction: array of float
405	The resolution correction
406	"""
407	num_bins_fft = len(k)	3✔
408	correction = np.ones(num_bins_fft)	3✔
409
410	pixelization_factor = np.sinc(k * delta_pixel / (2 * np.pi))**2	3✔
411
412	correction = np.exp(-(k mean_reso)**2)	3✔
413	correction *= pixelization_factor	3✔
414	return correction	3✔
415
416
417	def compute_correction_reso_matrix(reso_matrix, k, delta_pixel, num_pixel):	3✔
418	"""Computes the resolution correction based on the resolution matrix using linear binning
419
420	Arguments
421	---------
422	delta_pixel: float
423	Variation of the logarithm of the wavelength between two pixels
424	(in km/s or Ang depending on the units of k submitted)
425
426	num_pixel: int
427	Length of the spectrum in pixels
428
429	mean_reso: float
430	Mean resolution of the forest
431
432	k: array of float
433	Fourier modes
434
435	Return
436	------
437	correction: array of float
438	The resolution correction
439	"""
440	#this allows either computing the power for each pixel seperately or for the mean
441	reso_matrix = np.atleast_2d(reso_matrix)	×
442
443	w2_arr = []	×
444	#first compute the power in the resmat for each pixel, then average
445	for resmat in reso_matrix:	×
446	r = np.append(resmat, np.zeros(num_pixel - resmat.size))	×
447	k_resmat, w2 = compute_pk_raw(delta_pixel, r, linear_binning=True)	×
448	try:	×
449	assert np.all(k_resmat == k)	×
450	except AssertionError as error:	×
451	raise Exception(	×
452	"for some reason the resolution matrix correction has "
453	"different k scaling than the pk") from error
454	w2 /= w2[0]	×
455	w2_arr.append(w2)	×
456
457	w_res2 = np.mean(w2_arr, axis=0)	×
458	pixelization_factor = np.sinc(k * delta_pixel / (2 * np.pi))**2	×
459
460	# the following assumes that the resolution matrix is storing the actual
461	# resolution convolved with the pixelization kernel along each matrix axis
462	correction = w_res2	×
463	correction /= pixelization_factor	×
464
465	return correction	×
466
467
468	class Pk1D:	3✔
469	"""Class to represent the 1D Power Spectrum for a given forest
470
471	Class Methods
472	-------------
473	from_fitsio
474
475	Methods
476	-------
477	__init__
478
479	Attributes
480	----------
481	ra: float
482	Right-ascension of the quasar (in radians).
483
484	dec: float
485	Declination of the quasar (in radians).
486
487	z_qso: float
488	Redshift of the quasar.
489
490	plate: int
491	Plate number of the observation.
492
493	fiberid: int
494	Fiberid of the observation.
495
496	mjd: int
497	Modified Julian Date of the observation.
498
499	mean_snr: float
500	Mean signal-to-noise ratio in the forest
501
502	mean_reso: float
503	Mean resolution of the forest
504
505	mean_z: float
506	Mean redshift of the forest
507
508	num_masked_pixels: int
509	Number of masked pixels
510
511	k: array of float
512	Fourier modes
513
514	pk_raw: array of float
515	Raw power spectrum
516
517	pk_noise: array of float
518	Noise power spectrum for the different Fourier modes
519
520	correction_reso: array of float
521	Resolution correction
522
523	pk: array of float
524	Power Spectrum
525
526	pk_diff: array of float or None
527	Power spectrum of exposures_diff for the different Fourier modes
528	"""
529	def __init__(self,	3✔
530	ra,
531	dec,
532	z_qso,
533	mean_z,
534	plate,
535	mjd,
536	fiberid,
537	mean_snr,
538	mean_reso,
539	k,
540	pk_raw,
541	pk_noise,
542	correction_reso,
543	pk,
544	num_masked_pixels,
545	pk_diff=None):
546	"""Initialize instance
547
548	Arguments
549	---------
550	ra: float
551	Right-ascension of the quasar (in radians).
552
553	dec: float
554	Declination of the quasar (in radians).
555
556	z_qso: float
557	Redshift of the quasar.
558
559	mean_z: float
560	Mean redshift of the forest
561
562	plate: int
563	Plate number of the observation.
564
565	mjd: int
566	Modified Julian Date of the observation.
567
568	fiberid: int
569	Fiberid of the observation.
570
571	mean_snr: float
572	Mean signal-to-noise ratio in the forest
573
574	mean_reso: float
575	Mean resolution of the forest
576
577	k: array of float
578	Fourier modes
579
580	pk_raw: array of float
581	Raw power spectrum
582
583	pk_noise: array of float
584	Noise power spectrum for the different Fourier modes
585
586	correction_reso: array of float
587	Resolution correction
588
589	pk: array of float
590	Power Spectrum
591
592	num_masked_pixels: int
593	Number of masked pixels
594
595	pk_diff: array of float or None - default: None
596	Power spectrum of exposures_diff for the different Fourier modes
597	"""
598	self.ra = ra	×
599	self.dec = dec	×
600	self.z_qso = z_qso	×
601	self.mean_z = mean_z	×
602	self.mean_snr = mean_snr	×
603	self.mean_reso = mean_reso	×
604	self.num_masked_pixels = num_masked_pixels	×
605
606	self.plate = plate	×
607	self.mjd = mjd	×
608	self.fiberid = fiberid	×
609	self.k = k	×
610	self.pk_raw = pk_raw	×
611	self.pk_noise = pk_noise	×
612	self.correction_reso = correction_reso	×
613	self.pk = pk	×
614	self.pk_diff = pk_diff	×
615
616	@classmethod	3✔
617	def from_fitsio(cls, hdu):	3✔
618	"""Read the 1D Power Spectrum from fits file
619
620	Arguments
621	---------
622	hdu: Header Data Unit
623	Header Data Unit where the 1D Power Spectrum is read
624
625	Return
626	------
627	An intialized instance of Pk1D
628	"""
629
630	header = hdu.read_header()	×
631
632	ra = header['RA']	×
633	dec = header['DEC']	×
634	z_qso = header['Z']	×
635	mean_z = header['MEANZ']	×
636	mean_reso = header['MEANRESO']	×
637	mean_snr = header['MEANSNR']	×
638	plate = header['PLATE']	×
639	mjd = header['MJD']	×
640	fiberid = header['FIBER']	×
641	num_masked_pixels = header['NBMASKPIX']	×
642
643	data = hdu.read()	×
644	k = data['k'][:]	×
645	pk = data['Pk'][:]	×
646	pk_raw = data['Pk_raw'][:]	×
647	pk_noise = data['Pk_noise'][:]	×
648	correction_reso = data['cor_reso'][:]	×
649	pk_diff = data['Pk_diff'][:]	×
650
651	return cls(ra, dec, z_qso, mean_z, plate, mjd, fiberid, mean_snr,	×
652	mean_reso, k, pk_raw, pk_noise, correction_reso, pk,
653	num_masked_pixels, pk_diff)

igmhub / picca / 3750956619

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