micmac.likelihood.sampling module

class SamplingFunctions(nside, lmax, nstokes, frequency_array, freq_inverse_noise, pos_special_freqs=[0, -1], spv_nodes_b=None, freq_noise_c_ell=None, mask=None, n_components=3, lmin=2, n_iter=8, limit_iter_cg=200, limit_iter_cg_eta=200, tolerance_CG=1e-10, atol_CG=1e-08, bin_ell_distribution=None)[source]

Bases: MixingMatrix

Sampling functions object Contains all the functions needed for the sampling step of the Gibbs sampler

Parameters:
nside : int

nside of the input frequency maps

lmax : int

maximum multipole for the spherical harmonics transforms and harmonic domain objects,

nstokes : int

number of Stokes parameters

frequency_array : array[float]

array of frequencies, in GHz

freq_inverse_noise : array[float]

array of inverse noise for each frequency, in uK^-2

pos_special_freqs : list[int] (optional)

indexes of the special frequencies in the frequency array respectively for synchrotron and dust, default is [0,-1] for first and last frequencies

spv_nodes_b : list[dictionaries] (optional)

tree for the spatial variability, to generate from a yaml file, default None in principle set up by get_nodes_b

freq_noise_c_ell : array[float] of dimensions [frequencies, frequencies, lmax+1-lmin] or [frequencies, frequencies, lmax] (in which case it will be cut to lmax+1-lmin) (optional)

optional, noise power spectra for each frequency, in uK^2, dimensions, default None

mask : None or array[float] of dimensions [n_pix] (optional)

mask to use in the sampling ; if not given, no mask is used, default None Note: the mask WILL NOT be applied to the input maps, it will be only used for the propagated noise covariance

n_components : int (optional)

number of components for the mixing matrix, default 3

lmin : int (optional)

minimum multipole for the spherical harmonics transforms and harmonic domain objects, default 2

n_iter : int (optional)

number of iterations the spherical harmonics transforms (for map2alm transformations), default 8

limit_iter_cg : int (optional)

maximum number of iterations for the conjugate gradient for the CMB map sampling, default 200

limit_iter_cg_eta : int (optional)

maximum number of iterations for the conjugate gradient for eta maps sampling, default 200

tolerance_CG : float (optional)

tolerance for the conjugate gradient, default 1e-8

atol_CG : float (optional)

absolute tolerance for the conjugate gradient, default 1e-8

bin_ell_distribution : array[int] (optional)

array of the bounds of the bins if binned Inverse Wishart considered, of size nbins+1, default None for no binning

__init__(nside, lmax, nstokes, frequency_array, freq_inverse_noise, pos_special_freqs=[0, -1], spv_nodes_b=None, freq_noise_c_ell=None, mask=None, n_components=3, lmin=2, n_iter=8, limit_iter_cg=200, limit_iter_cg_eta=200, tolerance_CG=1e-10, atol_CG=1e-08, bin_ell_distribution=None)[source]

Sampling functions object Contains all the functions needed for the sampling step of the Gibbs sampler

Parameters:
nside : int

nside of the input frequency maps

lmax : int

maximum multipole for the spherical harmonics transforms and harmonic domain objects,

nstokes : int

number of Stokes parameters

frequency_array : array[float]

array of frequencies, in GHz

freq_inverse_noise : array[float]

array of inverse noise for each frequency, in uK^-2

pos_special_freqs : list[int] (optional)

indexes of the special frequencies in the frequency array respectively for synchrotron and dust, default is [0,-1] for first and last frequencies

spv_nodes_b : list[dictionaries] (optional)

tree for the spatial variability, to generate from a yaml file, default None in principle set up by get_nodes_b

freq_noise_c_ell : array[float] of dimensions [frequencies, frequencies, lmax+1-lmin] or [frequencies, frequencies, lmax] (in which case it will be cut to lmax+1-lmin) (optional)

optional, noise power spectra for each frequency, in uK^2, dimensions, default None

mask : None or array[float] of dimensions [n_pix] (optional)

mask to use in the sampling ; if not given, no mask is used, default None Note: the mask WILL NOT be applied to the input maps, it will be only used for the propagated noise covariance

n_components : int (optional)

number of components for the mixing matrix, default 3

lmin : int (optional)

minimum multipole for the spherical harmonics transforms and harmonic domain objects, default 2

n_iter : int (optional)

number of iterations the spherical harmonics transforms (for map2alm transformations), default 8

limit_iter_cg : int (optional)

maximum number of iterations for the conjugate gradient for the CMB map sampling, default 200

limit_iter_cg_eta : int (optional)

maximum number of iterations for the conjugate gradient for eta maps sampling, default 200

tolerance_CG : float (optional)

tolerance for the conjugate gradient, default 1e-8

atol_CG : float (optional)

absolute tolerance for the conjugate gradient, default 1e-8

bin_ell_distribution : array[int] (optional)

array of the bounds of the bins if binned Inverse Wishart considered, of size nbins+1, default None for no binning

bin_and_reproject_red_c_ell(red_cov_matrix)[source]

Bin and reproject the power spectrum

Parameters:
c_ells_to_reproject : array[float] of dimensions [lmax+1, nstokes, nstokes]

power spectrum to bin and reproject

Returns:

binned_reprojected_c_ells – Binned and reprojected power spectrum

Return type:

array[float] of dimensions [lmax+1, nstokes, nstokes]

property f_sky

Return fraction of the observed sky

get_band_limited_maps(input_map)[source]

Get band limited maps from input maps between lmin and lmax

Parameters:
input_map : array[float] of dimensions [nstokes, n_pix]

input maps to be band limited

Returns:

band limited maps

Return type:

array[float] of dimensions [nstokes, n_pix]

get_binned_c_ells(c_ells_to_bin)[source]

Bin the power spectrum to get the binned power spectrum

Parameters:
c_ells_to_bin : array[float] of dimensions [lmax+1, nstokes, nstokes]

power spectrum to bin

Returns:

binned_c_ells – Binned power spectrum

Return type:

array[float] of dimensions [number_bins, nstokes, nstokes]

get_binned_conditional_proba_C_from_r_wBB(r_param, red_sigma_ell, theoretical_red_cov_r1_tensor, theoretical_red_cov_r0_total)[source]

Compute log-proba of C parametrized by r_param only from the BB power spectrum with binning. The parametrisation is given by: C(r) = r * theoretical_red_cov_r1_tensor + theoretical_red_cov_r0_total

The associated log proba is:

-1/2 (tr sigma_ell C(r)^-1) - 1/2 log det C(r)

Parameters:
r_param : float

parameter of the covariance C to be sampled

red_sigma_ell : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

covariance matrices in harmonic domain

theoretical_red_cov_r1_tensor : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

tensor mode covariance matrices in harmonic domain

theoretical_red_cov_r0_total : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

scalar mode covariance matrices in harmonic domain

Returns:

log-proba of C parametrized by r_param

Return type:

float

get_binned_inverse_wishart_sampling_from_c_ells_v2(sigma_ell, PRNGKey, old_sample=None, acceptance_posdef=False)[source]

Solve sampling step 3: inverse Wishart distribution with C

sigma_ell is expected to be exactly the parameter of the inverse Wishart (so it should NOT be multiplied by 2*ell+1 if it is thought as a power spectrum)

Compute a matrix sample following an inverse Wishart distribution. The 3 steps follow Gupta & Nagar (2000):

  1. Sample n = 2*ell - p + 2*q_prior independent Gaussian vectors with covariance (sigma_ell)^{-1}

  2. Compute their outer product to form a matrix of dimension n_stokes*n_stokes ; which gives us a sample following the Wishart distribution

  3. Invert this matrix to obtain the final result: a matrix sample following an inverse Wishart distribution

Also assumes the monopole and dipole to be 0

Parameters:
sigma_ell : initial power spectrum which will define the parameter matrix of the inverse Wishart distribution ; must be of dimension [n_correlations, lmax+1]

PRNGKey : random key for JAX PNRG

Returns:

Matrices following an inverse Wishart distribution, of dimensions [lmin

Return type:

lmax, nstokes, nstokes]

get_binned_inverse_wishart_sampling_from_c_ells_v3(sigma_ell, PRNGKey, old_sample=None, acceptance_posdef=False)[source]

Solve sampling step 3: inverse Wishart distribution with C

sigma_ell is expected to be exactly the parameter of the inverse Wishart (so it should NOT be multiplied by 2*ell+1 if it is thought as a power spectrum)

Compute a matrix sample following an inverse Wishart distribution. The 3 steps follow Gupta & Nagar (2000):

  1. Sample n = 2*ell - p + 2*q_prior independent Gaussian vectors with covariance (sigma_ell)^{-1}

  2. Compute their outer product to form a matrix of dimension n_stokes*n_stokes ; which gives us a sample following the Wishart distribution

  3. Invert this matrix to obtain the final result: a matrix sample following an inverse Wishart distribution

Also assumes the monopole and dipole to be 0

Parameters:
sigma_ell : initial power spectrum which will define the parameter matrix of the inverse Wishart distribution ; must be of dimension [n_correlations, lmax+1]

PRNGKey : random key for JAX PNRG

Returns:

Matrices following an inverse Wishart distribution, of dimensions [lmin

Return type:

lmax, nstokes, nstokes]

get_binned_red_c_ells_v2(red_c_ells_to_bin)[source]

Bin the power spectrum to get the binned power spectrum

Parameters:
red_c_ells_to_bin : power spectrum to bin ; must be of dimension [lmax+1, nstokes, nstokes]

Return type:

Binned power spectrum, of dimension [number_bins, nstokes, nstokes]

get_binned_red_c_ells_v3(red_c_ells_to_bin)[source]

Bin the power spectrum to get the binned power spectrum

Parameters:
red_c_ells_to_bin : power spectrum to bin ; must be of dimension [lmax+1, nstokes, nstokes]

Return type:

Binned power spectrum, of dimension [number_bins, nstokes, nstokes]

get_binned_red_c_ells_v4(red_c_ells_to_bin)[source]

Bin the power spectrum to get the binned power spectrum

Parameters:
red_c_ells_to_bin : power spectrum to bin ; must be of dimension [lmax+1, nstokes, nstokes]

Return type:

Binned power spectrum, of dimension [number_bins, nstokes, nstokes]

get_cond_unobserved_patches()[source]

Get boolean condition on the free Bf indices corresponding to patches within the mask

get_cond_unobserved_patches_from_indices(indices)[source]

Get boolean condition on the free Bf indices corresponding to patches within the mask

get_cond_unobserved_patches_from_indices_optimized(indices)[source]

Get boolean condition on the free Bf indices corresponding to patches within the mask

get_conditional_proba_C_from_previous_sample(red_sigma_ell, red_cov_matrix_sampled)[source]
Compute log-proba of C. The associated log proba is:

-1/2 (tr sigma_ell C^-1) - 1/2 log det C

Parameters:
red_sigma_ell : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

covariance matrices in harmonic domain

red_cov_matrix_sampled : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

previous sampled covariance matrices in harmonic domain

Returns:

log-proba of C parametrized by r_param

Return type:

float

get_conditional_proba_C_from_r(r_param, red_sigma_ell, theoretical_red_cov_r1_tensor, theoretical_red_cov_r0_total)[source]

Compute log-proba of C parametrized by r_param. The parametrisation is given by: C(r) = r * theoretical_red_cov_r1_tensor + theoretical_red_cov_r0_total

The associated log proba is:

-1/2 (tr sigma_ell C(r)^-1) - 1/2 log det C(r)

Parameters:
r_param : float

parameter of the covariance C to be sampled

red_sigma_ell : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

covariance matrices in harmonic domain

theoretical_red_cov_r1_tensor : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

tensor mode covariance matrices in harmonic domain

theoretical_red_cov_r0_total : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

scalar mode covariance matrices in harmonic domain

Returns:

log-proba of C parametrized by r_param

Return type:

float

get_conditional_proba_C_from_r_wBB(r_param, red_sigma_ell, theoretical_red_cov_r1_tensor, theoretical_red_cov_r0_total)[source]

Compute log-proba of C parametrized by r_param only from the BB power spectrum. The parametrisation is given by: C(r) = r * theoretical_red_cov_r1_tensor + theoretical_red_cov_r0_total

The associated log proba is:

-1/2 (tr sigma_ell C(r)^-1) - 1/2 log det C(r)

Parameters:
r_param : float

parameter of the covariance C to be sampled

red_sigma_ell : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

covariance matrices in harmonic domain

theoretical_red_cov_r1_tensor : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

tensor mode covariance matrices in harmonic domain

theoretical_red_cov_r0_total : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

scalar mode covariance matrices in harmonic domain

Returns:

log-proba of C parametrized by r_param

Return type:

float

get_conditional_proba_correction_likelihood_JAX_pixel(old_params_mixing_matrix, new_params_mixing_matrix, inverse_term, component_eta_maps, red_cov_approx_matrix_sqrt, inverse_term_x_Capprox_root=None)[source]

Get conditional probability of correction term in the likelihood from the full mixing matrix, assuming the difference between the old and new mixing matrix is small

With notation C_approx instead of ilde{C}, the associated conditional probability is given by:

  • (eta ^t ( Id + C_approx^{1/2} N_c^{-1} C_approx^{1/2} )^{-1} eta

Or:

  • (eta^t ( Id + C_approx^{1/2} (E^t (B^t N^{-1} B)^{-1} E)^{-1} C_approx^{1/2}) eta

The computation proceeds as follows:

Knowing A^{-1} = ( Id + C_approx^{1/2} N_{c,old}^{-1} C_approx^{1/2} )^{-1} ; We compute ( Id + C_approx^{1/2} N_{c,new}^{-1} C_approx^{1/2} )^{-1} eta From: eta (A + C_approx^{1/2} (N_{c,new}^{-1} - N_{c,old}^{-1}) C_approx^{1/2})^{-1} eta

~= eta^t (A^{-1} - A^{-1} C_approx^{1/2} (N_{c,new}^{-1} - N_{c,old}^{-1}) C_approx^{1/2} A^{-1}) eta

Which holds if || N_{c,new}^{-1} - N_{c,old}^{-1} || is small enough

As we have:

  • inverse_term = A^{-1} eta

  • inverse_term_x_Capprox_root = C_approx^{1/2} A^{-1} eta (optional, otherwise it can be computed in the routine)

  • The B_{old} and B_{new} mixing matrices to reconstruct easily N_{c,old}^{-1} and N_{c,new}^{-1}

We can thus easily compute:

eta (A^{-1} - A^{-1} C_approx^{1/2} (N_{c,new}^{-1} - N_{c,old}^{-1}) C_approx^{1/2} A^{-1}) eta

Parameters:
old_params_mixing_matrix : array[float] of dimensions [component, frequencies]

old mixing matrix B_{old} to generate N_{c,old}

new_params_mixing_matrix : array[float] of dimensions [component, frequencies]

new mixing matrix B_{new} to generate N_{c,new}

inverse_term : array[float] of dimensions [component, n_pix]

previous inverse term computed with N_{c,old}

component_eta_maps : array[float] of dimensions [nstokes, n_pix]

set of eta maps

red_cov_approx_matrix : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

matrix square root of the covariance of the covariance matrice approx (C_approx) in harmonic domain

inverse_term_x_Capprox_root : array[float] of dimensions [component, n_pix] (optional)

precomputed product C_approx^{1/2} A^{-1} eta, if not given it will be recomputed here

Returns:

computation of the log-proba of the correction term to the likelihood

Return type:

float

get_conditional_proba_correction_likelihood_JAX_v2d(noise_CMB_component_, component_eta_maps, red_cov_approx_matrix_sqrt, first_guess=None, return_inverse=False, precond_func=None)[source]

Get conditional probability of correction term in the likelihood from the full mixing matrix

Noting C_approx = ilde{C}, the associated conditional probability is given by:

  • (eta ^t ( Id + C_approx^{1/2} N_c^{-1} C_approx^{1/2} )^{-1} eta

Or, with the developped version:

  • (eta^t ( Id + C_approx^{1/2} (E^t (B^t N^{-1} B)^{-1} E)^{-1} C_approx^{1/2}) eta

Parameters:
noise_CMB_component : array[float] of dimensions [n_pix]

pixel covariance matrix (B^t N^{-1} B)^{-1}

component_eta_maps : array[float] of dimensions [nstokes, n_pix]

set of eta maps

red_cov_approx_matrix_sqrt : array[float] of dimensions [component, component, n_pix]

square root of the approximate covariance matrix

first_guess : array[float] of dimensions [component, n_pix] (optional)

initial guess for the CG algorithm

return_inverse : bool (optional)

boolean to return the inverse term or not (default is False)

precond_func : function or None (optional)

preconditioning function to be used in the CG algorithm

Return type:

computation of log-proba correction term to the likelihood

get_conditional_proba_correction_likelihood_JAX_v2db(old_params_mixing_matrix, new_params_mixing_matrix, inverse_term, component_eta_maps, red_cov_approx_matrix_sqrt, inverse_term_x_Capprox_root=None)[source]

Get conditional probability of correction term in the likelihood from the full mixing matrix, assuming the difference between the old and new mixing matrix is small

With notation C_approx instead of ilde{C}, the associated conditional probability is given by:

  • (eta ^t ( Id + C_approx^{1/2} N_c^{-1} C_approx^{1/2} )^{-1} eta

Or:

  • (eta^t ( Id + C_approx^{1/2} (E^t (B^t N^{-1} B)^{-1} E)^{-1} C_approx^{1/2}) eta

The computation proceeds as follows:

Knowing A^{-1} = ( Id + C_approx^{1/2} N_{c,old}^{-1} C_approx^{1/2} )^{-1} ; We compute ( Id + C_approx^{1/2} N_{c,new}^{-1} C_approx^{1/2} )^{-1} eta From: eta (A + C_approx^{1/2} (N_{c,new}^{-1} - N_{c,old}^{-1}) C_approx^{1/2})^{-1} eta

~= eta^t (A^{-1} - A^{-1} C_approx^{1/2} (N_{c,new}^{-1} - N_{c,old}^{-1}) C_approx^{1/2} A^{-1}) eta

Which holds if || N_{c,new}^{-1} - N_{c,old}^{-1} || is small enough

As we have:

  • inverse_term = A^{-1} eta

  • inverse_term_x_Capprox_root = C_approx^{1/2} A^{-1} eta (optional, otherwise it can be computed in the routine)

  • The B_{old} and B_{new} mixing matrices to reconstruct easily N_{c,old}^{-1} and N_{c,new}^{-1}

We can thus easily compute:

eta (A^{-1} - A^{-1} C_approx^{1/2} (N_{c,new}^{-1} - N_{c,old}^{-1}) C_approx^{1/2} A^{-1}) eta

Parameters:
old_params_mixing_matrix : array[float] of dimensions [component, frequencies]

old mixing matrix B_{old} to generate N_{c,old}

new_params_mixing_matrix : array[float] of dimensions [component, frequencies]

new mixing matrix B_{new} to generate N_{c,new}

inverse_term : array[float] of dimensions [component, n_pix]

previous inverse term computed with N_{c,old}

component_eta_maps : array[float] of dimensions [nstokes, n_pix]

set of eta maps

red_cov_approx_matrix : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

matrix square root of the covariance of the covariance matrice approx (C_approx) in harmonic domain

inverse_term_x_Capprox_root : array[float] of dimensions [component, n_pix] (optional)

precomputed product C_approx^{1/2} A^{-1} eta, if not given it will be recomputed here

Returns:

computation of the log-proba of the correction term to the likelihood

Return type:

float

get_conditional_proba_mixing_matrix_v2b_JAX(new_params_mixing_matrix, full_data_without_CMB, red_cov_approx_matrix_sqrt, component_eta_maps=None, first_guess=None, biased_bool=False, precond_func=None)[source]

Get conditional probability of the conditional probability associated with the Bf parameters

With notation C_approx instead of ilde{C}, the associated conditional probability is given by:

  • (d - B_c s_c)^t N^{-1} B_f (B_f^t N^{-1} B_f)^{-1} B_f^t N^{-1} (d - B_c s_c) + eta^t ( Id + C_approx^{1/2} N_c^{-1} C_approx^{1/2} )^{-1} eta

Parameters:
new_params_mixing_matrix : array[float] of dimensions [nfreq-len(pos_special_frequencies), ncomp-1]

new Bf parameters of the mixing matrix to compute the log-proba

full_data_without_CMB : array[float] of dimensions [frequencies, n_pix]

data without from which the CMB (sample) was substracted, of dimension

red_cov_approx_matrix : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

matrix square root of the covariance of the covariance matrice approx (C_approx) in harmonic domain

component_eta_maps : array[float] of dimensions [nstokes, n_pix]

set of eta maps

first_guess : array[float] of dimensions [component, n_pix] (optional)

initial guess for the CG, default is None to use maps of 0

biased_bool : bool (optional)

indicate if the log-proba is biased, so computed without the correction, or not

precond_func : function or None (optional)

preconditioning function to be used in the CG algorithm

Returns:

  • log_proba (float) – computation of the conditional probability of the mixing matrix

  • inverse_term (array[float] of dimensions [component, n_pix]) – inverse term computed with N_{c,new} to be used in the CG algorithm

get_conditional_proba_mixing_matrix_v3_JAX(new_params_mixing_matrix, old_params_mixing_matrix, full_data_without_CMB, red_cov_approx_matrix_sqrt=None, component_eta_maps=None, first_guess=None, previous_inverse_x_Capprox_root=None, biased_bool=False)[source]

Get conditional probability of the conditional probability associated with the Bf parameters

Note that the difference between the old and new mixing matrix is assumed to be small

With notation C_approx instead of ilde{C}, the associated conditional probability is given by:

  • (d - B_c s_c)^t N^{-1} B_f (B_f^t N^{-1} B_f)^{-1} B_f^t N^{-1} (d - B_c s_c) + eta^t C_approx^{-1/2} ( C_approx^{-1} + N_c^{-1} )^{-1} C_approx^{-1/2} eta

Parameters:
old_params_mixing_matrix : array[float] of dimensions [component, frequencies]

old mixing matrix B_{old} to generate N_{c,old}

new_params_mixing_matrix : array[float] of dimensions [component, frequencies]

new mixing matrix B_{new} to generate N_{c,new}

full_data_without_CMB : array[float] of dimensions [frequencies, n_pix]

data without from which the CMB (sample) was substracted, of dimension

red_cov_approx_matrix : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

matrix square root of the covariance of the covariance matrice approx (C_approx) in harmonic domain

component_eta_maps : array[float] of dimensions [nstokes, n_pix]

set of eta maps

first_guess : array[float] of dimensions [component, n_pix] (optional)

previous inverse term computed with N_{c,old}, default is None to use maps of 0

previous_inverse_x_Capprox_root : array[float] of dimensions [component, n_pix] (optional)

C_approx^{1/2} A^{-1} eta, default None to have it recomputed

biased_bool : bool (optional)

indicate if the log-proba is biased, so computed without the correction, or not

Returns:

computation of the conditional probability of the mixing matrix

Return type:

float

get_conditional_proba_mixing_matrix_v3_pixel_JAX(new_params_mixing_matrix, old_params_mixing_matrix, full_data_without_CMB, red_cov_approx_matrix_sqrt, nside_patch, component_eta_maps=None, first_guess=None, previous_inverse_x_Capprox_root=None, biased_bool=False)[source]

Get conditional probability of the conditional probability associated with the Bf parameters

Note that the difference between the old and new mixing matrix is assumed to be small

With notation C_approx instead of ilde{C}, the associated conditional probability is given by:

  • (d - B_c s_c)^t N^{-1} B_f (B_f^t N^{-1} B_f)^{-1} B_f^t N^{-1} (d - B_c s_c) + eta^t C_approx^{-1/2} ( C_approx^{-1} + N_c^{-1} )^{-1} C_approx^{-1/2} eta

Parameters:
old_params_mixing_matrix : array[float] of dimensions [component, frequencies]

old mixing matrix B_{old} to generate N_{c,old}

new_params_mixing_matrix : array[float] of dimensions [component, frequencies]

new mixing matrix B_{new} to generate N_{c,new}

full_data_without_CMB : array[float] of dimensions [frequencies, n_pix]

data without from which the CMB (sample) was substracted, of dimension

red_cov_approx_matrix : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

matrix square root of the covariance of the covariance matrice approx (C_approx) in harmonic domain

nside_patch : int

nside of the parameter patch to retrieve in params

component_eta_maps : array[float] of dimensions [nstokes, n_pix]

set of eta maps

first_guess : array[float] of dimensions [component, n_pix] (optional)

previous inverse term computed with N_{c,old}, default is None to use maps of 0

previous_inverse_x_Capprox_root : array[float] of dimensions [component, n_pix] (optional)

C_approx^{1/2} A^{-1} eta, default None to have it recomputed

biased_bool : bool (optional)

indicate if the log-proba is biased, so computed without the correction, or not

Returns:

computation of the conditional probability of the mixing matrix for each patch

Return type:

array[float] of dimensions [max_length_patch]

get_conditional_proba_spectral_likelihood_JAX(complete_mixing_matrix, full_data_without_CMB)[source]

Get conditional probability of spectral likelihood from the full mixing matrix

The associated conditional probability is given by: - (d - B_c s_c)^t N^{-1} B_f (B_f^t N^{-1} B_f)^{-1} B_f^t N^{-1} (d - B_c s_c)

with d = full_data_without_CMB, B_c = complete_mixing_matrix, B_f = complete_mixing_matrix[:,1:,:] d is assumed to be band-limited

Parameters:
complete_mixing_matrix : array[float] of dimensions [component, frequencies]

complete mixing matrix

full_data_without_CMB : array[float] of dimensions [frequencies, n_pix]

data without from which the CMB (sample) was substracted

Returns:

computation of spectral likelihood

Return type:

array[float] of dimensions [n_pix]

get_conditional_proba_spectral_likelihood_JAX_pixel(complete_mixing_matrix, full_data_without_CMB)[source]

Get conditional probability of spectral likelihood from the full mixing matrix

The associated conditional probability is given by: - (d - B_c s_c)^t N^{-1} B_f (B_f^t N^{-1} B_f)^{-1} B_f^t N^{-1} (d - B_c s_c)

with d = full_data_without_CMB, B_c = complete_mixing_matrix, B_f = complete_mixing_matrix[:,1:,:] d is assumed to be band-limited

Parameters:
complete_mixing_matrix : array[float] of dimensions [component, frequencies]

complete mixing matrix

full_data_without_CMB : array[float] of dimensions [frequencies, n_pix]

data without from which the CMB (sample) was substracted

Returns:

computation of spectral likelihood per pixel

Return type:

array[float] of dimensions [n_pix]

get_fluctuating_term_maps_v2d(red_cov_matrix_sqrt, invBtinvNB, BtinvN_sqrt, jax_key_PNRG, map_random_realization_xi=None, map_random_realization_chi=None, initial_guess=Array([], shape=(0,), dtype=float64), precond_func=None)[source]

Sampling step 2: fluctuating term

Solve fluctuation term:

(Id + C^{1/2} N_c^{-1} C^{1/2}) C^{-1/2} zeta = xi + C^{1/2} N_c^{-1/2} chi

Or, fully developped:

(Id + C^{1/2} (E^t (B^t N^{-1} B)^{-1} E)^{-1} C^{1/2}) zeta = C^{1/2} C^{-1/2} xi + C^{1/2} (E^t (B^t N^{-1} B)^{-1} E)^{-1} E^t (B^t N^{-1} B)^{-1} B^t N^{-1/2} chi

This ensures: <zeta zeta^t> = (C^{-1} + N_c^{-1})^{-1} < zeta > = 0

Note C is assumed to be block diagonal

Parameters:
red_cov_matrix_sqrt : array[float] of dimension [lmin:lmax, nstokes, nstokes]

term C^{1/2}, matrix square root of CMB covariance matrices in harmonic domain

invBtinvNB : array[float] of dimension [component, component, n_pix]

matrix (B^t N^{-1} B)^{-1}

BtinvN_sqrt : array[float] of dimension [component, frequencies, n_pix]

matrix B^T N^{-1/2}

jax_key_PNRG : array[jnp.uint32]

random key for JAX PNRG

map_white_noise_xi : array[float] of dimensions [nstokes, n_pix] (optional)

set of maps 0 with mean and variance 1, which will be used to compute the fluctuation term

map_white_noise_chi : array[float] of dimensions [nfreq, nstokes, n_pix] (optional)

set of maps 0 with mean and variance 1, which will be used to compute the fluctuation term

initial_guess : array[float] of dimensions [nstokes, n_pix] (optional)

initial guess for the CG, default jnp.empty(0) (then set to 0)

precond_func : function (optional)

function preconditioner for the CG, default None

Returns:

fluctuating_map_z – Fluctuation maps for s_c sampling

Return type:

array[float] of dimensions [nstokes, n_pix]

get_inverse_gamma_sampling_from_c_ells(sigma_ell, PRNGKey)[source]

Solve sampling step 3: inverse Gamma distribution with C

sigma_ell is expected to be exactly the parameter of the inverse Gamma (so it should NOT be multiplied by 2*ell+1 if it is thought as a power spectrum)

Also assumes the monopole and dipole to be 0

Parameters:
sigma_ell : initial power spectrum which will define the parameter matrix of the inverse Gamma distribution ; must be of dimension [n_correlations, lmax+1]

PRNGKey : random key for JAX PNRG

Returns:

Matrices following an inverse Gamma distribution, of dimensions [lmin

Return type:

lmax, nstokes, nstokes]

get_inverse_wishart_sampling_from_c_ells(sigma_ell, PRNGKey, old_sample=None, acceptance_posdef=False)[source]

Solve sampling step 3: inverse Wishart distribution with C

sigma_ell is expected to be exactly the parameter of the inverse Wishart (so it should NOT be multiplied by 2*ell+1 if it is thought as a power spectrum)

Compute a matrix sample following an inverse Wishart distribution. The 3 steps follow Gupta & Nagar (2000):

  1. Sample n = 2*ell - p + 2*q_prior independent Gaussian vectors with covariance (sigma_ell)^{-1}

  2. Compute their outer product to form a matrix of dimension n_stokes*n_stokes ; which gives us a sample following the Wishart distribution

  3. Invert this matrix to obtain the final result: a matrix sample following an inverse Wishart distribution

Also assumes the monopole and dipole to be 0

Parameters:
sigma_ell : array[float] of dimensions [n_correlations, lmin:lmax+1]

initial power spectrum which will define the parameter matrix of the inverse Wishart distribution

PRNGKey : array[jnp.uint32]

random key for JAX PNRG

Returns:

array[float] of dimensions [lmin – Matrices following an inverse Wishart distribution

Return type:

lmax+1, nstokes, nstokes]

get_log_proba_non_centered_move_C(new_red_cov_matrix_sampled, old_red_cov_matrix_sampled, invBtinvNB, s_cML, s_c_sample, theoretical_red_cov_r1_tensor, theoretical_red_cov_r0_total)[source]

WARNING: Exploratory, to test further

get_log_proba_non_centered_move_C_from_r(new_r_sample, old_r_sample, invBtinvNB, s_cML, s_c_sample, theoretical_red_cov_r1_tensor, theoretical_red_cov_r0_total)[source]

WARNING: Exploratory, to test further

get_sampling_eta_v2(red_cov_approx_matrix_sqrt, invBtinvNB, BtinvN_sqrt, jax_key_PNRG, map_random_x=None, map_random_y=None, suppress_low_modes=True)[source]

Sampling step 1: eta maps Solve CG for eta term with formulation:

eta = ilde{C}^(1/2) N_c^{-1/2} x + y

Or, fully developped:

eta = ilde{C}^(1/2) ( (E (B^t N^{-1} B)^{-1} E) E (B^t N^{-1} B)^{-1} B^t N^{-1/2} x + ilde{C}^(-1/2) y )

This computation is performed assuming ilde{C} (or C_approx) is block diagonal

Parameters:
red_cov_approx_matrix_sqrt : array[float] of dimensions [lmax+1-lmin, nstokes, nstokes]

matrix square root of the covariance matrice ( ilde{C} or C_approx) in harmonic domain

invBtinvNB : array[float] of dimensions [component, component, n_pix]

matrix (B^t N^{-1} B)^{-1}

BtinvN_sqrt : array[float] of dimensions [component, frequencies, n_pix]

matrix B^T N^{-1/2}

jax_key_PNRG : array[jnp.uint32]

random key for JAX PNRG

map_random_x : array[float] of dimensions [nfreq, nstokes, n_pix] (optional)

set of maps 0 with mean and variance 1, which will be used to compute eta, default None (to have them computed in the routine)

map_random_y : array[float] of dimensions [nstokes, n_pix] (optional)

set of maps 0 with mean and variance 1, which will be used to compute eta, default None (to have them computed in the routine)

suppress_low_modes : bool (optional)

if True, suppress low modes in the CG between lmin and lmax, default True

Returns:

eta_maps – eta maps for Bf sampling

Return type:

array[float] of dimensions [nstokes, n_pix]

harmonic_marginal_probability(sample_Bf_r, noise_weighted_alm_data, theoretical_red_cov_r1_tensor, theoretical_red_cov_r0_total, red_cov_approx_matrix)[source]
Compute marginal probability of the full likelihood over s_c, given by:

-1/2 (d^t P d - s_{c,ML}^t (C^{-1} + N_c^{-1})^{-1} s_{c,ML} + ln | (C + N_c) (C_approx + N_c)^-1 |)

With:

P = N^{-1} - N^{-1} B (B^t N^{-1} B)^{-1} B^t N^{-1} In the routine, we don’t compute the d^t N^{-1} d part as it stays constant per sample, but only the second part

Here we denote C_approx = ilde{C}.

The data are assumed to be provided in the harmonic domain already noise weighted, and the covariance matrices are assumed to be in the harmonic domain as well. All harmonic covariance matrices are assumed to be block diagonal.

The routine will only work in harmonic domain, so any mixing matrix within sample_Bf_r will be averaged over the pixels.

The routine doesn’t explicitely retrieve C, but assume it to be parametrize by r, with C(r) = r * theoretical_red_cov_r1_tensor + theoretical_red_cov_r0_total

Parameters:
sample_Bf_r : array[float] of dimensions [nfreq-len(pos_special_frequencies)*(ncomp-1) + 1]

sample of the mixing matrix and r parameter

noise_weighted_alm_data : array[float] of dimensions [nfreq, nstokes]

noise weighted alms of the data, do N^{-1} d, given in harmonic domain

theoretical_red_cov_r1_tensor : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

tensor mode covariance matrices in harmonic domain

theoretical_red_cov_r0_total : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

scalar mode covariance matrices in harmonic domain

red_cov_approx_matrix : array[float] of dimensions [lmin:lmax, nstokes, nstokes]

approximate covariance matrices in harmonic domain

Returns:

log-proba computation of the marginal probability of the full likelihood over s_c

Return type:

float

property n_pix

Return number of pixels of 1 map, for a given Stokes parameter

property number_bins

Return number of bins

number_dof(bin_index)[source]

Return number of degrees of freedom for a given bin, for inverse Wishart distribution

project_bin_to_ell(binned_spectra)[source]

Project binned spectra by repeating the bin value for each ell

Parameters:
binned_spectra : array[float] of dimensions [number_bins, nstokes, nstokes]

binned power spectra

Returns:

ell_projection – dProjected power spectra

Return type:

array[float] of dimensions [lmax+1, nstokes, nstokes]

solve_generalized_wiener_filter_term_v2d(s_cML, red_cov_matrix_sqrt, invBtinvNB, initial_guess=Array([], shape=(0,), dtype=float64), precond_func=None)[source]
Solve Wiener filter term with CG:

(Id + C^{1/2} N_c^{-1} C^{1/2}) C^{-1/2} s_{c,WF} = C^{1/2} N_c^{-1} s_{c,ML}

This ensures:

s_{c,WF} = (C^{-1} + N_c^{-1})^{-1} N_c^{-1} s_{c,ML}

Note C is assumed to be block diagonal

Parameters:
s_cML : array[float] of dimensions [nstokes, n_pix]

Maximum Likelihood solution of component separation from input frequency maps

red_cov_matrix_sqrt : array[float] of dimension [lmin:lmax, nstokes, nstokes]

term C^{1/2}, matrix square root of CMB covariance matrices in harmonic domain

invBtinvNB : array[float] of dimension [component, component, n_pix]

matrix (B^t N^{-1} B)^{-1}

initial_guess : array[float] of dimensions [nstokes, n_pix] (optional)

initial guess for the CG, default jnp.empty(0) (then set to 0)

precond_func : function (optional)

function preconditioner for the CG, default None

Returns:

Wiener filter maps for s_c sampling

Return type:

array[float] of dimensions [nstokes, n_pix]

bounded_single_Metropolis_Hasting_step(random_PRNGKey, old_sample, step_size, log_proba, min_value, **model_kwargs)[source]

Single bounded Metropolis-Hasting step for a single parameter, with a Gaussian proposal distribution If the proposal is lower than the minimum value, the proposal is rejected

Parameters:
random_PRNGKey : array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample : array

old sample of the parameter

step_size : array[float]

standard deviation for the proposal distribution

log_proba : array[float]

log-probability function of the model

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

new sample of the parameter

Return type:

array

bounded_single_Metropolis_Hasting_step_numpyro(state, step_size, log_proba, min_value=-inf, **model_kwargs)[source]

Single Metropolis-Hasting step for a single parameter, with a Gaussian proposal distribution

Parameters:
random_PRNGKey : array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample : array

old sample of the parameter

step_size : array[float]

standard deviation for the proposal distribution

log_proba : array[float]

log-probability function of the model

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

new sample of the parameter

Return type:

array

multivariate_Metropolis_Hasting_step(random_PRNGKey, old_sample, covariance_matrix, log_proba, **model_kwargs)[source]

Metropolis-Hasting step for a multivariate parameter, with a multivariate Gaussian proposal distribution

Parameters:
random_PRNGKey : array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample : array

old sample of the parameter

covariance_matrix : array[float]

covariance matrix for the proposal distribution

log_proba : array[float]

log-probability function of the model

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

new sample of the parameter

Return type:

array

multivariate_Metropolis_Hasting_step_numpyro(state, covariance_matrix, log_proba, **model_kwargs)[source]

Metropolis-Hasting step for a multivariate parameter, with a multivariate Gaussian proposal distribution, from a Numpyro state

Parameters:
state : MHState

Numpyro state

covariance_matrix : array[float]

covariance matrix for the proposal distribution

log_proba : array[float]

log-probability function of the model

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

new sample of the parameter

Return type:

array

multivariate_Metropolis_Hasting_step_numpyro_bounded(state, covariance_matrix, log_proba, boundary, **model_kwargs)[source]

Metropolis-Hasting step for a multivariate parameter, with a multivariate Gaussian proposal distribution, from a Numpyro state

Parameters:
state : MHState

Numpyro state

covariance_matrix : array[float]

covariance matrix for the proposal distribution

log_proba : array[float]

log-probability function of the model

model_kwargs : dictionary

additional arguments for the log-probability function

boundary : array[float] of the same shape as the parameters

boundary for the parameters

Returns:

new sample of the parameter

Return type:

array

separate_single_MH_step_index_accelerated(random_PRNGKey, old_sample, step_size, log_proba, indexes_Bf, first_guess, **model_kwargs)[source]

Perform Metroplis-Hasting step for a given set of indexes given by indexes_patches_Bf

Parameters:
random_PRNGKey : array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample : array

old sample of the parameter

step_size : array[float]

standard deviation for the proposal distribution

log_proba : array[float]

log-probability function of the model

indexes_Bf : array[int]

indexes of the parameters to be sampled

first_guess : array[float]

first guess for the inverse term

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

  • latest_PRNGKey (array[jnp.uint32] of dimensions [2]) – latest random key

  • new_sample (array) – new sample of the parameter

  • new_inverse_term (array[float]) – array to speed-up computations

separate_single_MH_step_index_v2b(random_PRNGKey, old_sample, step_size, log_proba, indexes_Bf, **model_kwargs)[source]

Perform a set of single Metroplis-Hasting steps on a given set of indexes given by indexes_Bf, with a Gaussian proposal distribution. Each index will be sampled conditionned on each other.

Parameters:
random_PRNGKey : array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample : array

old sample of the parameter

step_size : array[float]

standard deviation for the proposal distribution

log_proba : array[float]

log-probability function of the model

indexes_Bf : array[int]

indexes of the parameters to be sampled

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

new sample of the parameter

Return type:

array

separate_single_MH_step_index_v3(random_PRNGKey, old_sample, step_size, log_proba, indexes_Bf, indexes_patches_Bf, size_patches, max_len_patches_Bf, len_indexes_Bf, **model_kwargs)[source]

Perform Metroplis-Hasting step for a given set of indexes given by indexes_patches_Bf

Parameters:
random_PRNGKey : array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample : array

old sample of the parameter

step_size : array[float]

standard deviation for the proposal distribution

log_proba : array[float]

log-probability function of the model

indexes_Bf : array[int]

indexes of the parameters to be sampled

indexes_patches_Bf : array[int]

indexes of the first patch of each parameter to be updated

size_patches : array[int]

size of all patches

max_len_patches_Bf : int

maximum length of the patches

len_indexes_Bf : int

maximum index of all possible Bf (not only the free ones)

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

  • latest_PRNGKey (array[jnp.uint32] of dimensions [2]) – latest random key

  • new_sample (array) – new sample of the parameter

separate_single_MH_step_index_v4_pixel(random_PRNGKey, old_sample, step_size, log_proba, indexes_Bf, indexes_patches_Bf, size_patches, max_len_patches_Bf, len_indexes_Bf, **model_kwargs)[source]

Perform Metroplis-Hasting step for a given set of indexes given by indexes_patches_Bf

Assumes all patches have the same disposition on the sky

random_PRNGKey: array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample: array

old sample of the parameter

step_size: array[float]

standard deviation for the proposal distribution

log_proba: array[float]

log-probability function of the model

indexes_Bf: array[int]

indexes of the parameters to be sampled

indexes_patches_Bf: array[int]

indexes of the first patch of each parameter to be updated

size_patches: array[int]

size of all patches

max_len_patches_Bf: int

maximum length of the patches

len_indexes_Bf: int

maximum index of all possible Bf (not only the free ones)

model_kwargs: dictionary

additional arguments for the log-probability function

Returns:

  • latest_PRNGKey (array[jnp.uint32] of dimensions [2]) – latest random key

  • new_sample (array) – new sample of the parameter

separate_single_MH_step_index_v4b_pixel(random_PRNGKey, old_sample, step_size, log_proba, indexes_Bf, indexes_patches_Bf, size_patches, max_len_patches_Bf, len_indexes_Bf, **model_kwargs)[source]

Perform Metroplis-Hasting step for a given set of indexes given by indexes_patches_Bf

Assumes all patches have the same disposition on the sky

Parameters:
random_PRNGKey : array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample : array

old sample of the parameter

step_size : array[float]

standard deviation for the proposal distribution

log_proba : array[float]

log-probability function of the model

indexes_Bf : array[int]

indexes of the parameters to be sampled

indexes_patches_Bf : array[int]

indexes of the first patch of each parameter to be updated

size_patches : array[int]

size of all patches

max_len_patches_Bf : int

maximum length of the patches

len_indexes_Bf : int

maximum index of all possible Bf (not only the free ones)

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

  • latest_PRNGKey (array[jnp.uint32] of dimensions [2]) – latest random key

  • new_sample (array) – new sample of the parameter

single_Metropolis_Hasting_step(random_PRNGKey, old_sample, step_size, log_proba, **model_kwargs)[source]

Single Metropolis-Hasting step for a single parameter, with a Gaussian proposal distribution

Parameters:
random_PRNGKey : array[jnp.uint32] of dimensions [2]

JAX random key to be splitted to generate the proposal and uniform distribution sample

old_sample : array

old sample of the parameter

step_size : array[float]

standard deviation for the proposal distribution

log_proba : array[float]

log-probability function of the model

model_kwargs : dictionary

additional arguments for the log-probability function

Returns:

new sample of the parameter

Return type:

array