"""Chi-square calculations for LiLit likelihood computations.
This module provides a unified interface for different chi-square calculation methods
used in cosmological likelihood analysis, including exact, Gaussian, Hamimeche & Lewis,
and LoLLiPoP approximations.
"""
from enum import Enum
import numpy as np
[docs]
class ChiSquareMethod(Enum):
"""Available chi-square calculation methods."""
EXACT = "exact"
GAUSSIAN = "gaussian"
CORRELATED_GAUSSIAN = "correlated_gaussian"
HAMIMECHE_LEWIS = "hl"
LOLLIPOP = "lollipop"
[docs]
class ChiSquareCalculator:
"""Unified interface for chi-square calculations in likelihood analysis.
This class provides a factory pattern for different chi-square calculation methods,
making it easy to switch between approximations and maintain consistent interfaces.
"""
[docs]
@staticmethod
def calculate(
method: ChiSquareMethod | str, data: np.ndarray, coba: np.ndarray, **kwargs
) -> np.ndarray | list:
"""Calculate chi-square using the specified method.
Parameters:
method (ChiSquareMethod or str):
The chi-square calculation method to use.
data (ndarray):
The covariance matrix of the data.
coba (ndarray):
The covariance matrix on the MCMC step.
**kwargs:
Method-specific parameters.
Returns:
ndarray | list: Chi-square values, format depends on method.
Raises:
ValueError: If method is not supported or required parameters are missing.
"""
if isinstance(method, str):
try:
method = ChiSquareMethod(method)
except ValueError:
raise ValueError(f"Unsupported chi-square method: {method}")
method_map = {
ChiSquareMethod.EXACT: ChiSquareCalculator._calculate_exact,
ChiSquareMethod.GAUSSIAN: ChiSquareCalculator._calculate_gaussian,
ChiSquareMethod.CORRELATED_GAUSSIAN: (
ChiSquareCalculator._calculate_correlated_gaussian
),
ChiSquareMethod.HAMIMECHE_LEWIS: ChiSquareCalculator._calculate_HL,
ChiSquareMethod.LOLLIPOP: ChiSquareCalculator._calculate_LoLLiPoP,
}
return method_map[method](data, coba, **kwargs)
@staticmethod
def _calculate_exact(
data: np.ndarray,
coba: np.ndarray,
N: int,
lmin: int,
lmax: int,
fsky: float,
**kwargs,
) -> np.ndarray:
"""Computes proper chi-square term for the exact likelihood case.
Parameters:
data (ndarray):
The covariance matrix of the data.
coba (ndarray):
The covariance matrix on the MCMC step.
N (int):
Number of fields.
lmin (int):
The minimum multipole to consider.
lmax (int):
The maximum multipole to consider.
fsky (float):
The fraction of the sky if a unique number is provided. Otherwise, it is
the geometrical mean of the fraction of the sky for each field pair. In
other words an effective fraction of the sky.
modes (ndarray, optional):
Effective number of modes (degrees of freedom) per element of the last
axis. If given, it replaces the per-multipole (2*ell + 1) * fsky weight,
which makes the exact (Wishart) algebra agnostic to whether the last axis
indexes multipoles or bandpowers. For bandpowers the natural choice is
nu_b = (2*ell_center + 1) * delta_ell * fsky = ((lmax_b+1)^2 - lmin_b^2) *
fsky, i.e. the exact sum of (2*ell + 1) over the band.
"""
# Import here to avoid circular imports
from ..functions import get_reduced_covariances
modes = kwargs.get("modes", None)
if modes is None:
ell = np.arange(lmin, lmax + 1, 1)
modes = (2 * ell + 1) * fsky
modes = np.asarray(modes, dtype=float)
if N != 1:
reduced_data = get_reduced_covariances(data, lmin, lmax)
reduced_coba = get_reduced_covariances(coba, lmin, lmax)
M_ℓ = list(map(np.linalg.solve, reduced_coba, reduced_data))
return modes * np.array([
np.trace(M) - np.linalg.slogdet(M)[1] - M.shape[0] for M in M_ℓ
])
else:
ratio = data[0, 0, :] / coba[0, 0, :]
return modes * (ratio - np.log(ratio) - 1)
@staticmethod
def _calculate_gaussian(
data: np.ndarray,
coba: np.ndarray,
N: int,
mask: np.ndarray,
inverse_covariance: list[np.ndarray],
lmin: int,
lmax: int,
**kwargs,
) -> list | np.ndarray:
"""Computes proper chi-square term for the Gaussian likelihood case.
Parameters:
data (ndarray):
The covariance matrix of the data.
coba (ndarray):
The covariance matrix on the MCMC step.
N (int):
Number of fields.
mask (ndarray):
Mask corresponding to both excluded probes and excluded multipole ranges.
inverse_covariance (list(ndarray)):
Inverse of the covariance matrices for each multipole.
lmin (int):
The minimum multipole to consider.
lmax (int):
The maximum multipole to consider.
"""
# Import here to avoid circular imports
from ..functions import get_reduced_data_vectors
if N != 1:
reduced_coba = get_reduced_data_vectors(N, coba, mask, lmin, lmax)
reduced_data = get_reduced_data_vectors(N, data, mask, lmin, lmax)
return [
(reduced_coba[j] - reduced_data[j])
@ inverse_covariance[j]
@ (reduced_coba[j] - reduced_data[j])
for j in range(lmax + 1 - lmin)
]
else:
diff = coba[0, 0, :] - data[0, 0, :]
return diff**2 * np.array(inverse_covariance)[:, 0, 0]
@staticmethod
def _calculate_correlated_gaussian(
data: np.ndarray, coba: np.ndarray, inverse_covariance: list[np.ndarray], **kwargs
) -> float:
"""Computes proper chi-square term for the correlated Gaussian likelihood case.
Parameters:
data (ndarray):
The covariance matrix of the data.
coba (ndarray):
The covariance matrix on the MCMC step.
inverse_covariance (list(ndarray)):
Inverse of the covariance matrices for each multipole.
"""
diff = coba[0, 0, :] - data[0, 0, :]
return diff @ inverse_covariance @ diff
@staticmethod
def _calculate_HL(
data: np.ndarray,
coba: np.ndarray,
fidu: np.ndarray,
offset: np.ndarray,
inverse_covariance: list[np.ndarray],
**kwargs,
) -> float:
"""Computes proper chi-square term for the Hamimeche & Lewis likelihood case.
Parameters:
data (ndarray):
The covariance matrix of the data.
coba (ndarray):
The covariance matrix on the MCMC step.
fidu (ndarray):
The fiducial covariance matrix.
offset (ndarray):
The offset array for the calculation.
inverse_covariance (list(ndarray)):
Inverse of the covariance matrices for each multipole.
"""
data_offset = data[0, 0, :] + offset[0, 0, :]
coba_offset = coba[0, 0, :] + offset[0, 0, :]
M = np.array(data_offset / coba_offset)
g = np.sign(M - 1) * np.sqrt(2 * (M - np.log(M) - 1))
weighted_g = g * (fidu[0, 0, :] + offset[0, 0, :])
return weighted_g @ inverse_covariance @ weighted_g
@staticmethod
def _calculate_LoLLiPoP(
data: np.ndarray,
coba: np.ndarray,
fidu: np.ndarray,
offset: np.ndarray,
inverse_covariance: list[np.ndarray],
**kwargs,
) -> float:
"""Computes proper chi-square term for the LoLLiPoP likelihood case.
Parameters:
data (ndarray):
The covariance matrix of the data.
coba (ndarray):
The covariance matrix on the MCMC step.
fidu (ndarray):
The fiducial covariance matrix.
offset (ndarray):
The offset array for the calculation.
inverse_covariance (list(ndarray)):
Inverse of the covariance matrices for each multipole.
"""
data_offset = data[0, 0, :] + offset[0, 0, :]
coba_offset = coba[0, 0, :] + offset[0, 0, :]
M = np.array(data_offset / coba_offset)
g = (
np.sign(M)
* np.sign(np.abs(M) - 1)
* np.sqrt(2 * (np.abs(M) - np.log(np.abs(M)) - 1))
)
weighted_g = g * (fidu[0, 0, :] + offset[0, 0, :])
return weighted_g @ inverse_covariance @ weighted_g
# Legacy function wrappers for backward compatibility
[docs]
def get_chi_exact(
N: int,
data: np.ndarray,
coba: np.ndarray,
lmin: int,
lmax: int,
fsky: float,
) -> np.ndarray:
"""Legacy wrapper for exact chi-square calculation."""
return ChiSquareCalculator.calculate(
ChiSquareMethod.EXACT, data, coba, N=N, lmin=lmin, lmax=lmax, fsky=fsky
)
[docs]
def get_chi_gaussian(
N: int,
data: np.ndarray,
coba: np.ndarray,
mask: np.ndarray,
inverse_covariance: list[np.ndarray],
lmin: int,
lmax: int,
) -> list | np.ndarray:
"""Legacy wrapper for Gaussian chi-square calculation."""
return ChiSquareCalculator.calculate(
ChiSquareMethod.GAUSSIAN,
data,
coba,
N=N,
mask=mask,
inverse_covariance=inverse_covariance,
lmin=lmin,
lmax=lmax,
)
[docs]
def get_chi_HL(
data: np.ndarray,
coba: np.ndarray,
fidu: np.ndarray,
offset: np.ndarray,
inverse_covariance: list[np.ndarray],
) -> float:
"""Legacy wrapper for Hamimeche & Lewis chi-square calculation."""
return ChiSquareCalculator.calculate(
ChiSquareMethod.HAMIMECHE_LEWIS,
data,
coba,
fidu=fidu,
offset=offset,
inverse_covariance=inverse_covariance,
)
[docs]
def get_chi_LoLLiPoP(
data: np.ndarray,
coba: np.ndarray,
fidu: np.ndarray,
offset: np.ndarray,
inverse_covariance: list[np.ndarray],
) -> float:
"""Legacy wrapper for LoLLiPoP chi-square calculation."""
return ChiSquareCalculator.calculate(
ChiSquareMethod.LOLLIPOP,
data,
coba,
fidu=fidu,
offset=offset,
inverse_covariance=inverse_covariance,
)
__all__ = [
"ChiSquareCalculator",
"ChiSquareMethod",
# Legacy functions for backward compatibility
"get_chi_exact",
"get_chi_gaussian",
"get_chi_correlated_gaussian",
"get_chi_HL",
"get_chi_LoLLiPoP",
]