16184789126

Committed 10 Jul 2025 02:36AM UTC coverage: 65.194% (-0.06%) from 65.25%

Build # 16184789126

Build Type

Pull #1094

gihtub

Committed by

web-flow

Commit Message

Update pylintrc

Pull Request Pull Request #1094: Update utils.py with bootstrap covariance estimate

Run Details

1451 of 2805 branches covered (51.73%)

Branch coverage included in aggregate %.

8 of 19 new or added lines in 2 files covered. (42.11%)

1 existing line in 1 file now uncovered.

6783 of 9825 relevant lines covered (69.04%)

2.07 hits per line

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

26.84

/py/picca/utils.py

"""This module provides general functions.

These are:
    - userprint
    - compute_cov
    - smooth_cov
    - smooth_cov_wick
    - compute_ang_max
    - shuffle_distrib_forests
    - unred
See the respective docstrings for more details
"""
import sys
import numpy as np
import fitsio
import scipy.interpolate as interpolate
import iminuit


def userprint(*args, **kwds):
    """Defines an extension of the print function.

    Args:
        *args: arguments passed to print
        **kwargs: keyword arguments passed to print
    """
    print(*args, **kwds)
    sys.stdout.flush()


def compute_cov(xi, weights):
    """Computes the covariance matrix using the subsampling technique

    Args:
        xi: array of floats
            Correlation function measburement in each helpix
        weights: array of floats
            Weights on the correlation function measurement

    Returns:
        The covariance matrix
    """

    mean_xi = (xi * weights).sum(axis=0)
    sum_weights = weights.sum(axis=0)
    w = sum_weights > 0.
    mean_xi[w] /= sum_weights[w]

    meanless_xi_times_weight = weights * (xi - mean_xi)

    userprint("Computing cov...")

    covariance = meanless_xi_times_weight.T.dot(meanless_xi_times_weight)
    sum_weights_squared = sum_weights * sum_weights[:, None]
    w = sum_weights_squared > 0.
    covariance[w] /= sum_weights_squared[w]

    return covariance


def compute_cov_boot(xi, weights, nboots=10000, seed=121567):
    """Computes the covariance matrix using the bootstrap technique

    Args:
        xi: 2D np.array of shape (nhpx, ndata)
            Correlation function measurement in each helpix
        weights: 2D np.array of shape (nhpx, ndata)
            Weights on the correlation function measurement

    Returns:
        The covariance matrix
    """
    nhpx, ndata = xi.shape
    boot_xis = np.empty((nboots, ndata))
    rnst = np.random.default_rng(seed)

    for i in range(nboots):
        idx = rnst.choice(nhpx, size=nhpx)
        wei = weights[idx]
        boot_xis[i] = np.sum(wei * xi[idx], axis=0) / wei.sum(0)

    return np.cov(boot_xis, rowvar=False)


def smooth_cov(xi,
               weights,
               r_par,
               r_trans,
               delta_r_trans=4.0,
               delta_r_par=4.0,
               covariance=None,
               per_r_par=False):
    """Smoothes the covariance matrix

    Args:
        xi: array of floats
            Correlation function measburement in each helpix
        weights: array of floats
            Weights on the correlation function measurement
        r_par: array of floats
            Parallel distance of each pixel of xi (in Mpc/h)
        r_par: array of floats
            Transverse distance of each pixel of xi (in Mpc/h)
        delta_r_trans: float - default: 4.0
            Variation of the transverse distance between two pixels
        delta_r_par: float - default: 4.0
            Variation of the transverse distance between two pixels
        covariance: array of floats or None - defautl: None
            Covariance matrix. If None, it will be computed using the
            subsampling technique
        per_r_par: if true, average the correlation coef per
                  (r_par,delta_r_par,delta_r_trans)

    Returns:
        The smooth covariance matrix. If data is not correct, then print a
        warning and return the unsmoothed covariance matrix
    """
    if covariance is None:
        covariance = compute_cov(xi, weights)

    num_bins = covariance.shape[1]
    var = np.diagonal(covariance)
    if np.any(var == 0.):
        userprint('WARNING: data has some empty bins, impossible to smooth')
        userprint('WARNING: returning the unsmoothed covariance')
        return covariance

    correlation = covariance / np.sqrt(var * var[:, None])

    correlation_smooth = np.zeros([num_bins, num_bins])

    # add together the correlation from bins with similar separations in
    # parallel and perpendicular distances
    sum_correlation = {}
    counts_correlation = {}
    for index in range(num_bins):
        if per_r_par :
            index_r_par = int(r_par[index]/delta_r_par)
        userprint("\rsmoothing {}".format(index), end="")
        for index2 in range(index + 1, num_bins):
            index_delta_r_par = round(
                abs(r_par[index2] - r_par[index]) / delta_r_par)
            index_delta_r_trans = round(
                abs(r_trans[index] - r_trans[index2]) / delta_r_trans)
            if per_r_par :
                if (index_r_par, index_delta_r_par, index_delta_r_trans) not in sum_correlation:
                    sum_correlation[(index_r_par, index_delta_r_par, index_delta_r_trans)] = 0.
                    counts_correlation[(index_r_par, index_delta_r_par, index_delta_r_trans)] = 0

                sum_correlation[(index_r_par, index_delta_r_par,
                                 index_delta_r_trans)] += correlation[index, index2]
                counts_correlation[(index_r_par, index_delta_r_par, index_delta_r_trans)] += 1
            else :
                if (index_delta_r_par, index_delta_r_trans) not in sum_correlation:
                    sum_correlation[(index_delta_r_par, index_delta_r_trans)] = 0.
                    counts_correlation[(index_delta_r_par, index_delta_r_trans)] = 0

                sum_correlation[(index_delta_r_par,
                                 index_delta_r_trans)] += correlation[index, index2]
                counts_correlation[(index_delta_r_par, index_delta_r_trans)] += 1

    for index in range(num_bins):
        correlation_smooth[index, index] = 1.
        if per_r_par :
            index_r_par = int(r_par[index]/delta_r_par)
        for index2 in range(index + 1, num_bins):
            index_delta_r_par = round(
                abs(r_par[index2] - r_par[index]) / delta_r_par)
            index_delta_r_trans = round(
                abs(r_trans[index] - r_trans[index2]) / delta_r_trans)
            if per_r_par :
                correlation_smooth[index, index2] = (
                    sum_correlation[(index_r_par, index_delta_r_par, index_delta_r_trans)] /
                    counts_correlation[(index_r_par , index_delta_r_par, index_delta_r_trans)])
            else :
                correlation_smooth[index, index2] = (
                    sum_correlation[(index_delta_r_par, index_delta_r_trans)] /
                    counts_correlation[(index_delta_r_par, index_delta_r_trans)])
            correlation_smooth[index2, index] = correlation_smooth[index,
                                                                   index2]

    userprint("\n")
    covariance_smooth = correlation_smooth * np.sqrt(var * var[:, None])
    return covariance_smooth


def smooth_cov_wick(filename, wick_filename, results_filename):
    """Smoothes the Wick covariance matrix, saves it to out_file

    Model the missing correlation in the Wick computation
    with an exponential

    Args:
        filename: str
            Correlation function (produced by picca_cf, picca_xcf)
        wick_filename: str
            Wick covariance matrix (produced by picca_wick, picca_xwick)
        results_filename: str
            Filename where the smoothed covariance matrix will be saved
    """
    # load subsampling covariance
    hdul = fitsio.FITS(filename)
    xi = np.array(hdul[2]['DA'][:])
    weights = np.array(hdul[2]['WE'][:])
    header = hdul[1].read_header()
    num_bins_r_par = header['NP']
    num_bins_r_trans = header['NT']
    hdul.close()

    covariance = compute_cov(xi, weights)

    num_bins = xi.shape[1]
    var = np.diagonal(covariance)
    if np.any(var == 0.):
        userprint('WARNING: data has some empty bins, impossible to smooth')
        userprint('WARNING: returning')
        return

    correlation = covariance / np.sqrt(var * var[:, None])
    correlation1d = correlation.reshape(num_bins * num_bins)

    # load Wick covariance
    hdul = fitsio.FITS(wick_filename)
    covariance_wick = np.array(hdul[1]['CO'][:])
    hdul.close()

    var_wick = np.diagonal(covariance_wick)
    if np.any(var_wick == 0.):
        userprint('WARNING: Wick covariance has bins with var = 0')
        userprint('WARNING: returning')
        return

    correlation_wick = covariance_wick / np.sqrt(var_wick * var_wick[:, None])
    correlation_wick1d = correlation_wick.reshape(num_bins * num_bins)

    # difference between 1d correlations
    delta_correlation1d = correlation1d - correlation_wick1d

    # indices
    index = np.arange(num_bins)
    index_r_trans = index % num_bins_r_trans
    index_r_par = index // num_bins_r_trans
    index_delta_r_trans2d = abs(index_r_trans - index_r_trans[:, None])
    index_delta_r_par2d = abs(index_r_par - index_r_par[:, None])
    index_delta_r_trans1d = index_delta_r_trans2d.reshape(num_bins * num_bins)
    index_delta_r_par1d = index_delta_r_par2d.reshape(num_bins * num_bins)

    # compute the reduced correlation
    reduced_delta_correlation1d = np.zeros(num_bins)
    for index in range(0, num_bins):
        userprint("\rsmoothing {}".format(index), end="")
        reduced_delta_correlation1d[index] = np.mean(delta_correlation1d[
            (index_delta_r_par1d == index_r_par[index]) &
            (index_delta_r_trans1d == index_r_trans[index])])
    reduced_delta_correlation = reduced_delta_correlation1d.reshape(
        num_bins_r_par, num_bins_r_trans)
    userprint("")

    #### fit for length and amp at each delta_r_par
    def corrfun(index_delta_r_par, index_delta_r_trans, length, amp):
        """Models the missing correlation in the Wick computation
        with an exponential

        Args:
            index_delta_r_par: int
                Index associated with the separation in parallel distance
            index_delta_r_trans: int
                Index associated with the separation in transverse distance
            length: float
                Characteristic length of the exponential
            amp: float
                Amplitude of the exponential

        Returns:
            The model with the issed correlation
        """
        r = np.sqrt(
            float(index_delta_r_trans)**2 +
            float(index_delta_r_par)**2) - float(index_delta_r_par)
        return amp * np.exp(-r / length)

    def chi2(length, amp, index_delta_r_par):
        """Computes the chi2 for a given set of parameters in the model

        Args:
            length: float
                Characteristic length of the exponential
            amp: float
                Amplitude of the exponential
            index_delta_r_par: int
                Index associated with the separation in parallel distance

        Returns:
            The chi2 value
        """
        chi2 = 0.
        index_delta_r_par = int(index_delta_r_par)
        for index_delta_r_trans in range(1, num_bins_r_trans):
            chi = (reduced_delta_correlation[index_delta_r_par,
                                             index_delta_r_trans] -
                   corrfun(index_delta_r_par, index_delta_r_trans, length, amp))
            chi2 += chi**2
        chi2 = chi2 * num_bins_r_par * num_bins
        return chi2

    length_fit = np.zeros(num_bins_r_par)
    amp_fit = np.zeros(num_bins_r_par)
    for index_delta_r_par in range(num_bins_r_par):
        minimizer = iminuit.Minuit(chi2,
                                   length=5.,
                                   amp=1.,
                                   index_delta_r_par=index_delta_r_par)
        minimizer.fixed['index_delta_r_par']=True
        minimizer.errors['length'] = 0.2
        minimizer.errors['amp'] = 0.2
        minimizer.errordef = 1.
        minimizer.print_level = 1
        minimizer.limits['length'] = (1., 400.)
        minimizer.migrad()
        length_fit[index_delta_r_par] = minimizer.values['length']
        amp_fit[index_delta_r_par] = minimizer.values['amp']

    #### hybrid covariance from wick + fit
    covariance_smooth = np.sqrt(var * var[:, None])

    cor0 = reduced_delta_correlation1d[index_r_trans == 0]
    for index in range(num_bins):
        userprint("\rupdating {}".format(index), end="")
        for index2 in range(index + 1, num_bins):
            index_delta_r_par = index_delta_r_par2d[index, index2]
            index_delta_r_trans = index_delta_r_trans2d[index, index2]
            newcov = correlation_wick[index, index2]
            if index_delta_r_trans == 0:
                newcov += cor0[index_delta_r_par]
            else:
                newcov += corrfun(index_delta_r_par, index_delta_r_trans,
                                  length_fit[index_delta_r_par],
                                  amp_fit[index_delta_r_par])
            covariance_smooth[index, index2] *= newcov
            covariance_smooth[index2, index] *= newcov

    userprint("\n")

    results = fitsio.FITS(results_filename, 'rw', clobber=True)
    results.write([covariance_smooth], names=['CO'], extname='COR')
    results.close()
    userprint(results_filename, ' written')


def compute_ang_max(cosmo, r_trans_max, z_min, z_min2=None):
    """Computes the maximum anglular separation the correlation should be
    calculated to.

    This angle is given by the maximum transverse separation and the fiducial
    cosmology

    Args:
        comso: constants.Cosmo
            Fiducial cosmology
        r_trans_max: float
            Maximum transverse separation
        z_min: float
            Minimum redshift of the first set of data
        z_min2: float or None - default: None
            Minimum redshift of the second set of data. If None, use z_min

    Returns:
        The maximum anglular separation
    """
    if z_min2 is None:
        z_min2 = z_min

    r_min = cosmo.get_dist_m(z_min)
    r_min2 = cosmo.get_dist_m(z_min2)

    if r_min + r_min2 < r_trans_max:
        ang_max = np.pi
    else:
        ang_max = 2. * np.arcsin(r_trans_max / (r_min + r_min2))

    return ang_max


def shuffle_distrib_forests(data, seed):
    """Shuffles the distribution of forests by assiging the angular
        positions from another forest

    Args:
        data: dict
            A dictionary with the data. Keys are the healpix numbers of each
            spectrum. Values are lists of delta instances.
        seed: int
            Seed for the given realization of the shuffle

    Returns:
        The shuffled catalogue
    """

    userprint(("INFO: Shuffling the forests angular position with seed "
               "{}").format(seed))

    data_info = {}
    param_list = [
        'ra', 'dec', 'x_cart', 'y_cart', 'z_cart', 'cos_dec', 'thingid'
    ]
    data_info['healpix'] = []
    for param in param_list:
        data_info[param] = []

    for healpix, deltas in data.items():
        for delta in deltas:
            data_info['healpix'].append(healpix)
            for param in param_list:
                data_info[param].append(getattr(delta, param))

    np.random.seed(seed)
    new_index = np.arange(len(data_info['ra']))
    np.random.shuffle(new_index)

    index = 0
    data_shuffled = {}
    for deltas in data.values():
        for delta in deltas:
            for param in param_list:
                setattr(delta, param, data_info[param][new_index[index]])
            if not data_info['healpix'][new_index[index]] in data_shuffled:
                data_shuffled[data_info['healpix'][new_index[index]]] = []
            data_shuffled[data_info['healpix'][new_index[index]]].append(delta)
            index += 1

    return data_shuffled


# pylint: disable=invalid-name,locally-disabled
# we keep variable names since this function is adopted from elsewhere
def unred(wave, ebv, R_V=3.1, LMC2=False, AVGLMC=False):
    """
    https://github.com/sczesla/PyAstronomy
    in /src/pyasl/asl/unred
    """

    # Convert to inverse microns
    x = 10000. / wave
    curve = x * 0.

    # Set some standard values:
    x0 = 4.596
    gamma = 0.99
    c3 = 3.23
    c4 = 0.41
    c2 = -0.824 + 4.717 / R_V
    c1 = 2.030 - 3.007 * c2

    if LMC2:
        x0 = 4.626
        gamma = 1.05
        c4 = 0.42
        c3 = 1.92
        c2 = 1.31
        c1 = -2.16
    elif AVGLMC:
        x0 = 4.596
        gamma = 0.91
        c4 = 0.64
        c3 = 2.73
        c2 = 1.11
        c1 = -1.28

    # Compute UV portion of A(lambda)/E(B-V) curve using FM fitting function and
    # R-dependent coefficients
    xcutuv = np.array([10000.0 / 2700.0])
    xspluv = 10000.0 / np.array([2700.0, 2600.0])

    iuv = np.where(x >= xcutuv)[0]
    N_UV = iuv.size
    iopir = np.where(x < xcutuv)[0]
    Nopir = iopir.size
    if N_UV > 0:
        xuv = np.concatenate((xspluv, x[iuv]))
    else:
        xuv = xspluv

    yuv = c1 + c2 * xuv
    yuv = yuv + c3 * xuv**2 / ((xuv**2 - x0**2)**2 + (xuv * gamma)**2)
    yuv = yuv + c4 * (0.5392 * (np.maximum(xuv, 5.9) - 5.9)**2 + 0.05644 *
                      (np.maximum(xuv, 5.9) - 5.9)**3)
    yuv = yuv + R_V
    yspluv = yuv[0:2]  # save spline points

    if N_UV > 0:
        curve[iuv] = yuv[2::]  # remove spline points

    # Compute optical portion of A(lambda)/E(B-V) curve
    # using cubic spline anchored in UV, optical, and IR
    xsplopir = np.concatenate(
        ([0], 10000.0 /
         np.array([26500.0, 12200.0, 6000.0, 5470.0, 4670.0, 4110.0])))
    ysplir = np.array([0.0, 0.26469, 0.82925]) * R_V / 3.1
    ysplop = np.array(
        (np.polyval([-4.22809e-01, 1.00270, 2.13572e-04][::-1], R_V),
         np.polyval([-5.13540e-02, 1.00216, -7.35778e-05][::-1], R_V),
         np.polyval([7.00127e-01, 1.00184, -3.32598e-05][::-1], R_V),
         np.polyval([1.19456, 1.01707, -5.46959e-03, 7.97809e-04,
                     -4.45636e-05][::-1], R_V)))
    ysplopir = np.concatenate((ysplir, ysplop))

    if Nopir > 0:
        tck = interpolate.splrep(np.concatenate((xsplopir, xspluv)),
                                 np.concatenate((ysplopir, yspluv)),
                                 s=0)
        curve[iopir] = interpolate.splev(x[iopir], tck)

    #Now apply extinction correction to input flux vector
    curve *= ebv
    corr = 1. / (10.**(0.4 * curve))

    return corr

1	"""This module provides general functions.
2
3	These are:
4	- userprint
5	- compute_cov
6	- smooth_cov
7	- smooth_cov_wick
8	- compute_ang_max
9	- shuffle_distrib_forests
10	- unred
11	See the respective docstrings for more details
12	"""
13	import sys	3✔
14	import numpy as np	3✔
15	import fitsio	3✔
16	import scipy.interpolate as interpolate	3✔
17	import iminuit	3✔
18
19
20	def userprint(args, *kwds):	3✔
21	"""Defines an extension of the print function.
22
23	Args:
24	*args: arguments passed to print
25	**kwargs: keyword arguments passed to print
26	"""
27	print(args, *kwds)	3✔
28	sys.stdout.flush()	3✔
29
30
31	def compute_cov(xi, weights):	3✔
32	"""Computes the covariance matrix using the subsampling technique
33
34	Args:
35	xi: array of floats
36	Correlation function measburement in each helpix
37	weights: array of floats
38	Weights on the correlation function measurement
39
40	Returns:
41	The covariance matrix
42	"""
43
44	mean_xi = (xi * weights).sum(axis=0)	3✔
45	sum_weights = weights.sum(axis=0)	3✔
46	w = sum_weights > 0.	3✔
47	mean_xi[w] /= sum_weights[w]	3✔
48
49	meanless_xi_times_weight = weights * (xi - mean_xi)	3✔
50
51	userprint("Computing cov...")	3✔
52
53	covariance = meanless_xi_times_weight.T.dot(meanless_xi_times_weight)	3✔
54	sum_weights_squared = sum_weights * sum_weights[:, None]	3✔
55	w = sum_weights_squared > 0.	3✔
56	covariance[w] /= sum_weights_squared[w]	3✔
57
58	return covariance	3✔
59
60
61	def compute_cov_boot(xi, weights, nboots=10000, seed=121567):	3✔
62	"""Computes the covariance matrix using the bootstrap technique
63
64	Args:
65	xi: 2D np.array of shape (nhpx, ndata)
66	Correlation function measurement in each helpix
67	weights: 2D np.array of shape (nhpx, ndata)
68	Weights on the correlation function measurement
69
70	Returns:
71	The covariance matrix
72	"""
NEW 73	nhpx, ndata = xi.shape	×
NEW 74	boot_xis = np.empty((nboots, ndata))	×
NEW 75	rnst = np.random.default_rng(seed)	×
76
NEW 77	for i in range(nboots):	×
NEW 78	idx = rnst.choice(nhpx, size=nhpx)	×
NEW 79	wei = weights[idx]	×
NEW 80	boot_xis[i] = np.sum(wei * xi[idx], axis=0) / wei.sum(0)	×
81
NEW 82	return np.cov(boot_xis, rowvar=False)	×
83
84
85	def smooth_cov(xi,	3✔
86	weights,
87	r_par,
88	r_trans,
89	delta_r_trans=4.0,
90	delta_r_par=4.0,
91	covariance=None,
92	per_r_par=False):
93	"""Smoothes the covariance matrix
94
95	Args:
96	xi: array of floats
97	Correlation function measburement in each helpix
98	weights: array of floats
99	Weights on the correlation function measurement
100	r_par: array of floats
101	Parallel distance of each pixel of xi (in Mpc/h)
102	r_par: array of floats
103	Transverse distance of each pixel of xi (in Mpc/h)
104	delta_r_trans: float - default: 4.0
105	Variation of the transverse distance between two pixels
106	delta_r_par: float - default: 4.0
107	Variation of the transverse distance between two pixels
108	covariance: array of floats or None - defautl: None
109	Covariance matrix. If None, it will be computed using the
110	subsampling technique
111	per_r_par: if true, average the correlation coef per
112	(r_par,delta_r_par,delta_r_trans)
113
114	Returns:
115	The smooth covariance matrix. If data is not correct, then print a
116	warning and return the unsmoothed covariance matrix
117	"""
118	if covariance is None:	3!
UNCOV 119	covariance = compute_cov(xi, weights)	×
120
121	num_bins = covariance.shape[1]	3✔
122	var = np.diagonal(covariance)	3✔
123	if np.any(var == 0.):	3!
124	userprint('WARNING: data has some empty bins, impossible to smooth')	×
125	userprint('WARNING: returning the unsmoothed covariance')	×
126	return covariance	×
127
128	correlation = covariance / np.sqrt(var * var[:, None])	3✔
129
130	correlation_smooth = np.zeros([num_bins, num_bins])	3✔
131
132	# add together the correlation from bins with similar separations in
133	# parallel and perpendicular distances
134	sum_correlation = {}	3✔
135	counts_correlation = {}	3✔
136	for index in range(num_bins):	3✔
137	if per_r_par :	3!
138	index_r_par = int(r_par[index]/delta_r_par)	×
139	userprint("\rsmoothing {}".format(index), end="")	3✔
140	for index2 in range(index + 1, num_bins):	3✔
141	index_delta_r_par = round(	3✔
142	abs(r_par[index2] - r_par[index]) / delta_r_par)
143	index_delta_r_trans = round(	3✔
144	abs(r_trans[index] - r_trans[index2]) / delta_r_trans)
145	if per_r_par :	3!
146	if (index_r_par, index_delta_r_par, index_delta_r_trans) not in sum_correlation:	×
147	sum_correlation[(index_r_par, index_delta_r_par, index_delta_r_trans)] = 0.	×
148	counts_correlation[(index_r_par, index_delta_r_par, index_delta_r_trans)] = 0	×
149
150	sum_correlation[(index_r_par, index_delta_r_par,	×
151	index_delta_r_trans)] += correlation[index, index2]
152	counts_correlation[(index_r_par, index_delta_r_par, index_delta_r_trans)] += 1	×
153	else :
154	if (index_delta_r_par, index_delta_r_trans) not in sum_correlation:	3✔
155	sum_correlation[(index_delta_r_par, index_delta_r_trans)] = 0.	3✔
156	counts_correlation[(index_delta_r_par, index_delta_r_trans)] = 0	3✔
157
158	sum_correlation[(index_delta_r_par,	3✔
159	index_delta_r_trans)] += correlation[index, index2]
160	counts_correlation[(index_delta_r_par, index_delta_r_trans)] += 1	3✔
161
162	for index in range(num_bins):	3✔
163	correlation_smooth[index, index] = 1.	3✔
164	if per_r_par :	3!
165	index_r_par = int(r_par[index]/delta_r_par)	×
166	for index2 in range(index + 1, num_bins):	3✔
167	index_delta_r_par = round(	3✔
168	abs(r_par[index2] - r_par[index]) / delta_r_par)
169	index_delta_r_trans = round(	3✔
170	abs(r_trans[index] - r_trans[index2]) / delta_r_trans)
171	if per_r_par :	3!
172	correlation_smooth[index, index2] = (	×
173	sum_correlation[(index_r_par, index_delta_r_par, index_delta_r_trans)] /
174	counts_correlation[(index_r_par , index_delta_r_par, index_delta_r_trans)])
175	else :
176	correlation_smooth[index, index2] = (	3✔
177	sum_correlation[(index_delta_r_par, index_delta_r_trans)] /
178	counts_correlation[(index_delta_r_par, index_delta_r_trans)])
179	correlation_smooth[index2, index] = correlation_smooth[index,	3✔
180	index2]
181
182	userprint("\n")	3✔
183	covariance_smooth = correlation_smooth * np.sqrt(var * var[:, None])	3✔
184	return covariance_smooth	3✔
185
186
187	def smooth_cov_wick(filename, wick_filename, results_filename):	3✔
188	"""Smoothes the Wick covariance matrix, saves it to out_file
189
190	Model the missing correlation in the Wick computation
191	with an exponential
192
193	Args:
194	filename: str
195	Correlation function (produced by picca_cf, picca_xcf)
196	wick_filename: str
197	Wick covariance matrix (produced by picca_wick, picca_xwick)
198	results_filename: str
199	Filename where the smoothed covariance matrix will be saved
200	"""
201	# load subsampling covariance
202	hdul = fitsio.FITS(filename)	×
203	xi = np.array(hdul[2]['DA'][:])	×
204	weights = np.array(hdul[2]['WE'][:])	×
205	header = hdul[1].read_header()	×
206	num_bins_r_par = header['NP']	×
207	num_bins_r_trans = header['NT']	×
208	hdul.close()	×
209
210	covariance = compute_cov(xi, weights)	×
211
212	num_bins = xi.shape[1]	×
213	var = np.diagonal(covariance)	×
214	if np.any(var == 0.):	×
215	userprint('WARNING: data has some empty bins, impossible to smooth')	×
216	userprint('WARNING: returning')	×
217	return	×
218
219	correlation = covariance / np.sqrt(var * var[:, None])	×
220	correlation1d = correlation.reshape(num_bins * num_bins)	×
221
222	# load Wick covariance
223	hdul = fitsio.FITS(wick_filename)	×
224	covariance_wick = np.array(hdul[1]['CO'][:])	×
225	hdul.close()	×
226
227	var_wick = np.diagonal(covariance_wick)	×
228	if np.any(var_wick == 0.):	×
229	userprint('WARNING: Wick covariance has bins with var = 0')	×
230	userprint('WARNING: returning')	×
231	return	×
232
233	correlation_wick = covariance_wick / np.sqrt(var_wick * var_wick[:, None])	×
234	correlation_wick1d = correlation_wick.reshape(num_bins * num_bins)	×
235
236	# difference between 1d correlations
237	delta_correlation1d = correlation1d - correlation_wick1d	×
238
239	# indices
240	index = np.arange(num_bins)	×
241	index_r_trans = index % num_bins_r_trans	×
242	index_r_par = index // num_bins_r_trans	×
243	index_delta_r_trans2d = abs(index_r_trans - index_r_trans[:, None])	×
244	index_delta_r_par2d = abs(index_r_par - index_r_par[:, None])	×
245	index_delta_r_trans1d = index_delta_r_trans2d.reshape(num_bins * num_bins)	×
246	index_delta_r_par1d = index_delta_r_par2d.reshape(num_bins * num_bins)	×
247
248	# compute the reduced correlation
249	reduced_delta_correlation1d = np.zeros(num_bins)	×
250	for index in range(0, num_bins):	×
251	userprint("\rsmoothing {}".format(index), end="")	×
252	reduced_delta_correlation1d[index] = np.mean(delta_correlation1d[	×
253	(index_delta_r_par1d == index_r_par[index]) &
254	(index_delta_r_trans1d == index_r_trans[index])])
255	reduced_delta_correlation = reduced_delta_correlation1d.reshape(	×
256	num_bins_r_par, num_bins_r_trans)
257	userprint("")	×
258
259	#### fit for length and amp at each delta_r_par
260	def corrfun(index_delta_r_par, index_delta_r_trans, length, amp):	×
261	"""Models the missing correlation in the Wick computation
262	with an exponential
263
264	Args:
265	index_delta_r_par: int
266	Index associated with the separation in parallel distance
267	index_delta_r_trans: int
268	Index associated with the separation in transverse distance
269	length: float
270	Characteristic length of the exponential
271	amp: float
272	Amplitude of the exponential
273
274	Returns:
275	The model with the issed correlation
276	"""
277	r = np.sqrt(	×
278	float(index_delta_r_trans)**2 +
279	float(index_delta_r_par)**2) - float(index_delta_r_par)
280	return amp * np.exp(-r / length)	×
281
282	def chi2(length, amp, index_delta_r_par):	×
283	"""Computes the chi2 for a given set of parameters in the model
284
285	Args:
286	length: float
287	Characteristic length of the exponential
288	amp: float
289	Amplitude of the exponential
290	index_delta_r_par: int
291	Index associated with the separation in parallel distance
292
293	Returns:
294	The chi2 value
295	"""
296	chi2 = 0.	×
297	index_delta_r_par = int(index_delta_r_par)	×
298	for index_delta_r_trans in range(1, num_bins_r_trans):	×
299	chi = (reduced_delta_correlation[index_delta_r_par,	×
300	index_delta_r_trans] -
301	corrfun(index_delta_r_par, index_delta_r_trans, length, amp))
302	chi2 += chi**2	×
303	chi2 = chi2 * num_bins_r_par * num_bins	×
304	return chi2	×
305
306	length_fit = np.zeros(num_bins_r_par)	×
307	amp_fit = np.zeros(num_bins_r_par)	×
308	for index_delta_r_par in range(num_bins_r_par):	×
309	minimizer = iminuit.Minuit(chi2,	×
310	length=5.,
311	amp=1.,
312	index_delta_r_par=index_delta_r_par)
313	minimizer.fixed['index_delta_r_par']=True	×
314	minimizer.errors['length'] = 0.2	×
315	minimizer.errors['amp'] = 0.2	×
316	minimizer.errordef = 1.	×
317	minimizer.print_level = 1	×
318	minimizer.limits['length'] = (1., 400.)	×
319	minimizer.migrad()	×
320	length_fit[index_delta_r_par] = minimizer.values['length']	×
321	amp_fit[index_delta_r_par] = minimizer.values['amp']	×
322
323	#### hybrid covariance from wick + fit
324	covariance_smooth = np.sqrt(var * var[:, None])	×
325
326	cor0 = reduced_delta_correlation1d[index_r_trans == 0]	×
327	for index in range(num_bins):	×
328	userprint("\rupdating {}".format(index), end="")	×
329	for index2 in range(index + 1, num_bins):	×
330	index_delta_r_par = index_delta_r_par2d[index, index2]	×
331	index_delta_r_trans = index_delta_r_trans2d[index, index2]	×
332	newcov = correlation_wick[index, index2]	×
333	if index_delta_r_trans == 0:	×
334	newcov += cor0[index_delta_r_par]	×
335	else:
336	newcov += corrfun(index_delta_r_par, index_delta_r_trans,	×
337	length_fit[index_delta_r_par],
338	amp_fit[index_delta_r_par])
339	covariance_smooth[index, index2] *= newcov	×
340	covariance_smooth[index2, index] *= newcov	×
341
342	userprint("\n")	×
343
344	results = fitsio.FITS(results_filename, 'rw', clobber=True)	×
345	results.write([covariance_smooth], names=['CO'], extname='COR')	×
346	results.close()	×
347	userprint(results_filename, ' written')	×
348
349
350	def compute_ang_max(cosmo, r_trans_max, z_min, z_min2=None):	3✔
351	"""Computes the maximum anglular separation the correlation should be
352	calculated to.
353
354	This angle is given by the maximum transverse separation and the fiducial
355	cosmology
356
357	Args:
358	comso: constants.Cosmo
359	Fiducial cosmology
360	r_trans_max: float
361	Maximum transverse separation
362	z_min: float
363	Minimum redshift of the first set of data
364	z_min2: float or None - default: None
365	Minimum redshift of the second set of data. If None, use z_min
366
367	Returns:
368	The maximum anglular separation
369	"""
370	if z_min2 is None:	3✔
371	z_min2 = z_min	3✔
372
373	r_min = cosmo.get_dist_m(z_min)	3✔
374	r_min2 = cosmo.get_dist_m(z_min2)	3✔
375
376	if r_min + r_min2 < r_trans_max:	3!
377	ang_max = np.pi	×
378	else:
379	ang_max = 2. * np.arcsin(r_trans_max / (r_min + r_min2))	3✔
380
381	return ang_max	3✔
382
383
384	def shuffle_distrib_forests(data, seed):	3✔
385	"""Shuffles the distribution of forests by assiging the angular
386	positions from another forest
387
388	Args:
389	data: dict
390	A dictionary with the data. Keys are the healpix numbers of each
391	spectrum. Values are lists of delta instances.
392	seed: int
393	Seed for the given realization of the shuffle
394
395	Returns:
396	The shuffled catalogue
397	"""
398
399	userprint(("INFO: Shuffling the forests angular position with seed "	×
400	"{}").format(seed))
401
402	data_info = {}	×
403	param_list = [	×
404	'ra', 'dec', 'x_cart', 'y_cart', 'z_cart', 'cos_dec', 'thingid'
405	]
406	data_info['healpix'] = []	×
407	for param in param_list:	×
408	data_info[param] = []	×
409
410	for healpix, deltas in data.items():	×
411	for delta in deltas:	×
412	data_info['healpix'].append(healpix)	×
413	for param in param_list:	×
414	data_info[param].append(getattr(delta, param))	×
415
416	np.random.seed(seed)	×
417	new_index = np.arange(len(data_info['ra']))	×
418	np.random.shuffle(new_index)	×
419
420	index = 0	×
421	data_shuffled = {}	×
422	for deltas in data.values():	×
423	for delta in deltas:	×
424	for param in param_list:	×
425	setattr(delta, param, data_info[param][new_index[index]])	×
426	if not data_info['healpix'][new_index[index]] in data_shuffled:	×
427	data_shuffled[data_info['healpix'][new_index[index]]] = []	×
428	data_shuffled[data_info['healpix'][new_index[index]]].append(delta)	×
429	index += 1	×
430
431	return data_shuffled	×
432
433
434	# pylint: disable=invalid-name,locally-disabled
435	# we keep variable names since this function is adopted from elsewhere
436	def unred(wave, ebv, R_V=3.1, LMC2=False, AVGLMC=False):	3✔
437	"""
438	https://github.com/sczesla/PyAstronomy
439	in /src/pyasl/asl/unred
440	"""
441
442	# Convert to inverse microns
443	x = 10000. / wave	×
444	curve = x * 0.	×
445
446	# Set some standard values:
447	x0 = 4.596	×
448	gamma = 0.99	×
449	c3 = 3.23	×
450	c4 = 0.41	×
451	c2 = -0.824 + 4.717 / R_V	×
452	c1 = 2.030 - 3.007 * c2	×
453
454	if LMC2:	×
455	x0 = 4.626	×
456	gamma = 1.05	×
457	c4 = 0.42	×
458	c3 = 1.92	×
459	c2 = 1.31	×
460	c1 = -2.16	×
461	elif AVGLMC:	×
462	x0 = 4.596	×
463	gamma = 0.91	×
464	c4 = 0.64	×
465	c3 = 2.73	×
466	c2 = 1.11	×
467	c1 = -1.28	×
468
469	# Compute UV portion of A(lambda)/E(B-V) curve using FM fitting function and
470	# R-dependent coefficients
471	xcutuv = np.array([10000.0 / 2700.0])	×
472	xspluv = 10000.0 / np.array([2700.0, 2600.0])	×
473
474	iuv = np.where(x >= xcutuv)[0]	×
475	N_UV = iuv.size	×
476	iopir = np.where(x < xcutuv)[0]	×
477	Nopir = iopir.size	×
478	if N_UV > 0:	×
479	xuv = np.concatenate((xspluv, x[iuv]))	×
480	else:
481	xuv = xspluv	×
482
483	yuv = c1 + c2 * xuv	×
484	yuv = yuv + c3 * xuv2 / ((xuv2 - x02)2 + (xuv * gamma)**2)	×
485	yuv = yuv + c4 * (0.5392 * (np.maximum(xuv, 5.9) - 5.9)*2 + 0.05644	×
486	(np.maximum(xuv, 5.9) - 5.9)**3)
487	yuv = yuv + R_V	×
488	yspluv = yuv[0:2] # save spline points	×
489
490	if N_UV > 0:	×
491	curve[iuv] = yuv[2::] # remove spline points	×
492
493	# Compute optical portion of A(lambda)/E(B-V) curve
494	# using cubic spline anchored in UV, optical, and IR
495	xsplopir = np.concatenate(	×
496	([0], 10000.0 /
497	np.array([26500.0, 12200.0, 6000.0, 5470.0, 4670.0, 4110.0])))
498	ysplir = np.array([0.0, 0.26469, 0.82925]) * R_V / 3.1	×
499	ysplop = np.array(	×
500	(np.polyval([-4.22809e-01, 1.00270, 2.13572e-04][::-1], R_V),
501	np.polyval([-5.13540e-02, 1.00216, -7.35778e-05][::-1], R_V),
502	np.polyval([7.00127e-01, 1.00184, -3.32598e-05][::-1], R_V),
503	np.polyval([1.19456, 1.01707, -5.46959e-03, 7.97809e-04,
504	-4.45636e-05][::-1], R_V)))
505	ysplopir = np.concatenate((ysplir, ysplop))	×
506
507	if Nopir > 0:	×
508	tck = interpolate.splrep(np.concatenate((xsplopir, xspluv)),	×
509	np.concatenate((ysplopir, yspluv)),
510	s=0)
511	curve[iopir] = interpolate.splev(x[iopir], tck)	×
512
513	#Now apply extinction correction to input flux vector
514	curve *= ebv	×
515	corr = 1. / (10.*(0.4 curve))	×
516
517	return corr	×

igmhub / picca / 16184789126

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