diff --git a/mflike/mflike.py b/mflike/mflike.py index e3cf08d..898d4e3 100644 --- a/mflike/mflike.py +++ b/mflike/mflike.py @@ -24,7 +24,7 @@ import numpy as np from cobaya.conventions import data_path, packages_path_input -from cobaya.likelihoods.base_classes import InstallableLikelihood +from cobaya.likelihoods.base_classes import InstallableLikelihood, _fast_chi_square from cobaya.log import LoggedError from cobaya.tools import are_different_params_lists @@ -45,6 +45,8 @@ class MFLike(InstallableLikelihood): top_hat_band: dict systematics_template: dict + _fast_chi_squared = _fast_chi_square() + def initialize(self): # Set default values to data member not initialized via yaml file self.l_bpws = None @@ -183,7 +185,8 @@ def loglike(self, cl, **params_values_nocosmo): ps_vec = self._get_power_spectra(cl, **params_values_nocosmo) delta = self.data_vec - ps_vec - logp = -0.5 * (delta @ self.inv_cov @ delta) + # logp = -0.5 * (delta @ self.inv_cov @ delta) + logp = -0.5 * self._fast_chi_squared(self.inv_cov, delta) logp += self.logp_const self.log.debug( f"Log-likelihood value computed = {logp} (Χ² = {-2 * (logp - self.logp_const)})" @@ -222,7 +225,7 @@ def prepare_data(self): except AttributeError: raise KeyError("You must provide a list of default cuts") - # Translation betwen TEB and sacc C_ell types + # Translation between TEB and sacc C_ell types pol_dict = {"T": "0", "E": "e", "B": "b"} ppol_dict = { "TT": "tt", @@ -234,7 +237,6 @@ def prepare_data(self): "BE": "eb", "TB": "tb", "BT": "tb", - "BB": "bb", } def get_cl_meta(spec): @@ -369,6 +371,11 @@ def get_sacc_names(pol, exp_1, exp_2): ws = s_b.get_bandpower_windows(ind_b) else: ws = s.get_bandpower_windows(ind) + # pre-compute the actual slices of the weights that are needed + nonzeros = np.array([np.nonzero(ws.weight[:, i])[0][[0, -1]] for i in range(ws.weight.shape[1])]) + ws.nonzeros = [slice(i[0], i[1] + 1) for i in nonzeros] + ws.sliced_weights = [np.ascontiguousarray(ws.weight[ws.nonzeros[i], i]) + for i in range(len(nonzeros))] if self.l_bpws is None: # The assumption here is that bandpower windows @@ -454,18 +461,22 @@ def _get_power_spectra(self, cl, **params_values_nocosmo): Dls = {s: cl[s][self.l_bpws] for s, _ in self.lcuts.items()} DlsObs = self.ThFo.get_modified_theory(Dls, **params_values_nocosmo) + return self._get_ps_vec(DlsObs) + + def _get_ps_vec(self, DlsObs): ps_vec = np.zeros_like(self.data_vec) for m in self.spec_meta: p = m["pol"] - i = m["ids"] - w = m["bpw"].weight.T + w = m["bpw"] # If symmetrize = False, the (ET, exp1, exp2) spectrum # will have the flag m["hasYX_xsp"] = True. # In this case, the power spectrum # is computed as DlsObs["te", m["t2"], m["t1"]], to associate # T --> exp2, E --> exp1 dls_obs = DlsObs[p, m["t2"], m["t1"]] if m["hasYX_xsp"] else DlsObs[p, m["t1"], m["t2"]] - clt = w @ dls_obs - ps_vec[i] = clt + for i, nonzero, weights in zip(m["ids"], w.nonzeros, w.sliced_weights): + ps_vec[i] = weights @ dls_obs[nonzero] + # can check against unoptimized version + # assert np.allclose(ps_vec[m["ids"]], np.dot(w.weight.T, dls_obs)) return ps_vec diff --git a/pyproject.toml b/pyproject.toml index c24fd61..5277cb4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ classifiers = [ "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12" ] requires-python = ">=3.9.0" dependencies = [