Gaussian likelihood with point mass marginalization (#573)

* Introducing Gaussian likelihood with point mass marginalization ConstGaussianPM
* Test improvements:
- Add test_compute_chisq_with_correction() for inv_cov_correction branch
- Add test_get_lens_statistic_not_found() for error handling
- Add test_get_src_statistic_not_found() for error handling
- Add test_compute_chisq_without_cholesky() to cover GaussFamily direct
  chi-squared path
- Use object.__setattr__() in test fixtures to bypass type checking for
  mock assignments
* consolidate ConstGaussian classes and remove unnecessary inheritance
- Update ConstGaussianPM to inherit directly from ConstGaussian
- Update all imports across codebase to use firecrown.likelihood.gaussian
- Fix mypy type ignore comment in test file

* consolidate point mass data into dataclass and improve error handling

- Add PointMassData dataclass to encapsulate all point mass attributes
- Replace individual _pm_* attributes with single _pm_data container
- Refactor _generate_maps, _cache_precomputed_data, and data access methods
- Improve _collect_data_vectors error handling with detailed missing attributes
- Consolidate _get_lens_statistic and _get_src_statistic into shared _get_statistic
- Add default constants for sigma_B and point_mass parameters
- Update tests to use new PointMassData structure

* make _generate_maps return None and inline caching logic

- Remove unused _cache_precomputed_data method
- Change _generate_maps return type from PointMassData to None
- Move data caching logic directly into _generate_maps
- Update early return to not return cached data
- Remove return value documentation from docstring

* Refactor point mass marginalization code for better organization and type safety

- Replace _pm_maps_ready boolean with pm_maps_ready property that checks _pm_data existence
- Make PointMassData fields optional with None defaults and add assert_prepared() method
- Modify _generate_maps() to set instance variables directly instead of returning PointMassData
- Add typing.cast() calls after assert_prepared() to satisfy mypy type checker
- Remove redundant state tracking and improve code maintainability
* Refactor likelihood module structure
Rename gaussian_pointmass.py to _gaussian_pointmass.py and
update imports to use new private module naming convention:
- Update internal imports to use _gaussian and _base modules
- Export ConstGaussianPM and TrivialStatistic from public
  firecrown.likelihood API
- Update all example and test imports to use public API
- Update Makefile pre-commit target to include docs build
- Refactor test fixtures to use public reset() method
  instead of protected _reset() for TwoPointTheory

This change aligns with the module reorganization where
implementation modules use underscore prefix while
maintaining a clean public API through __init__.py exports.

---------

Co-authored-by: Sandro Dias Pinto Vitenti <[email protected]>
Co-authored-by: Marc Paterno <[email protected]>

Author: 3 people

Makefile

@@ -251,7 +251,7 @@ clean: clean-coverage clean-docs clean-build  ## Remove all generated files (can
 
 ##@ Pre-commit
 
-pre-commit: format lint test-coverage  ## Run all pre-commit checks (format, lint, test with coverage)
+pre-commit: format lint test-coverage docs  ## Run all pre-commit checks (format, lint, test with coverage, docs)
 	@echo ""
 	@echo "✅ All pre-commit checks passed!"
 

examples/des_y1_3x2pt/factory_PM.py

@@ -0,0 +1,122 @@
+"""The DES Y1 3x2pt likelihood factory module with point mass marginalization."""
+
+import os
+
+import pyccl
+from firecrown.likelihood.factories import load_sacc_data
+from firecrown.likelihood import NamedParameters, TwoPoint, ConstGaussianPM
+import firecrown.likelihood.weak_lensing as wl
+import firecrown.likelihood.number_counts as nc
+from firecrown.modeling_tools import ModelingTools
+from firecrown.ccl_factory import CCLFactory
+
+
+# The likelihood used for DES Y1 3x2pt analysis is a Gaussian likelihood, which
+# necessitates providing a list of statistics that are each represented by a two-point
+# function. To construct the two-point function, a list of sources is required, with
+# each source being responsible for computing the theoretical prediction for a specific
+#  segment of the data. These sources are created using the build_likelihood function
+# and also contain a list of systematics. The systematics are classes that modify the
+# theoretical prediction and are also constructed in the build_likelihood function.
+def build_likelihood(params: NamedParameters) -> tuple[ConstGaussianPM, ModelingTools]:
+    """Build the DES Y1 3x2pt likelihood."""
+    # Creates a LAI systematic. This is a systematic that is applied to
+    # all weak-lensing probes. The `sacc_tracer` argument is used to identify the
+    # section of the SACC file that this systematic will be applied to. In this
+    # case we want to apply it to all weak-lensing probes, so we use the
+    # empty string.
+    lai_systematic = wl.LinearAlignmentSystematic(sacc_tracer="")
+
+    # Creating sources, each one maps to a specific section of a SACC file. In
+    # this case src0, src1, src2 and src3 describe weak-lensing probes. The
+    # sources are saved in a dictionary since they will be used by one or more
+    # two-point function.
+    sources: dict[str, wl.WeakLensing | nc.NumberCounts] = {}
+
+    for i in range(4):
+        # Each weak-lensing section has its own multiplicative bias. Parameters
+        # reflect this by using src{i}_ prefix.
+        mbias = wl.MultiplicativeShearBias(sacc_tracer=f"src{i}")
+
+        # We also include a photo-z shift bias (a constant shift in dndz). We
+        # also have a different parameter for each bin, so here again we use the
+        # src{i}_ prefix.
+        wl_pzshift = wl.PhotoZShift(sacc_tracer=f"src{i}")
+
+        # Now we can finally create the weak-lensing source that will compute the
+        # theoretical prediction for that section of the data, given the
+        # systematics.
+        sources[f"src{i}"] = wl.WeakLensing(
+            sacc_tracer=f"src{i}", systematics=[lai_systematic, mbias, wl_pzshift]
+        )
+
+    # Creating the number counting sources. There are five sources each one
+    # labeled by lens{i}.
+    for i in range(5):
+        # We also include a photo-z shift for the dndz.
+        nc_pzshift = nc.PhotoZShift(sacc_tracer=f"lens{i}")
+
+        # The source is created and saved (temporarily in the sources dict).
+        sources[f"lens{i}"] = nc.NumberCounts(
+            sacc_tracer=f"lens{i}", systematics=[nc_pzshift], derived_scale=True
+        )
+
+    # Now that we have all sources we can instantiate all the two-point
+    # functions. The weak-lensing sources have two "data types", for each one we
+    # create a new two-point function.
+    stats = {}
+    for stat, sacc_stat in [
+        ("xip", "galaxy_shear_xi_plus"),
+        ("xim", "galaxy_shear_xi_minus"),
+    ]:
+        # Creating all auto/cross-correlations two-point function objects for the
+        # weak-lensing probes.
+        for i in range(4):
+            for j in range(i, 4):
+                stats[f"{stat}_src{i}_src{j}"] = TwoPoint(
+                    source0=sources[f"src{i}"],
+                    source1=sources[f"src{j}"],
+                    sacc_data_type=sacc_stat,
+                )
+    # The following two-point function objects compute the cross correlations
+    # between the weak-lensing and number count sources.
+    for j in range(5):
+        for i in range(4):
+            stats[f"gammat_lens{j}_src{i}"] = TwoPoint(
+                source0=sources[f"lens{j}"],
+                source1=sources[f"src{i}"],
+                sacc_data_type="galaxy_shearDensity_xi_t",
+            )
+
+    # Finally the instances for the lensing auto-correlations are created.
+    for i in range(5):
+        stats[f"wtheta_lens{i}_lens{i}"] = TwoPoint(
+            source0=sources[f"lens{i}"],
+            source1=sources[f"lens{i}"],
+            sacc_data_type="galaxy_density_xi",
+        )
+
+    # Here we instantiate the actual likelihood. The statistics argument carry
+    # the order of the data/theory vector.
+    likelihood = ConstGaussianPM(statistics=list(stats.values()))
+
+    # We load the correct SACC file.
+    sacc_file = params.get_string("sacc_file")
+    # Translate envriornment variables, if needed.
+    sacc_file = os.path.expandvars(sacc_file)
+    sacc_data = load_sacc_data(sacc_file)
+
+    # The read likelihood method is called passing the loaded SACC file, the
+    # two-point functions will receive the appropriated sections of the SACC
+    # file and the sources their respective dndz.
+    likelihood.read(sacc_data)
+
+    modeling_tools = ModelingTools(ccl_factory=CCLFactory(require_nonlinear_pk=True))
+
+    # After reading in the sacc data, we apply the point mass marginalization.
+    cosmo = pyccl.CosmologyVanillaLCDM()
+    likelihood.compute_pointmass(cosmo)
+
+    # This script will be loaded by the appropriated connector. The framework
+    # will call the factory function that should return a Likelihood instance.
+    return likelihood, modeling_tools

firecrown/likelihood/__init__.py

@@ -39,6 +39,7 @@
     SourceSystematic,
     Statistic,
     Tracer,
+    TrivialStatistic,
 )
 from firecrown.likelihood._likelihood import (
     load_likelihood,
@@ -49,6 +50,7 @@
 from firecrown.likelihood._gaussian import ConstGaussian
 from firecrown.likelihood._gaussfamily import GaussFamily, State
 from firecrown.likelihood._student_t import StudentT
+from firecrown.likelihood._gaussian_pointmass import ConstGaussianPM
 
 # Two-point statistics
 from firecrown.likelihood._two_point import TwoPoint, TwoPointFactory
@@ -103,6 +105,7 @@
     "GaussFamily",
     "State",
     "StudentT",
+    "ConstGaussianPM",
     # Two-point statistics
     "TwoPoint",
     "TwoPointFactory",
@@ -120,6 +123,7 @@
     "SourceSystematic",
     # Base statistic class
     "Statistic",
+    "TrivialStatistic",
     # Subpackages
     "weak_lensing",
     "number_counts",

firecrown/likelihood/_gaussfamily.py

@@ -142,10 +142,14 @@ class GaussFamily(Likelihood):
     def __init__(
         self,
         statistics: Sequence[Statistic],
+        use_cholesky: bool = True,
     ) -> None:
         """Initialize the base class parts of a GaussFamily object.
 
         :param statistics: A list of statistics to be include in chisquared calculations
+        :param use_cholesky: Whether to use Cholesky decomposition for chi-squared
+            calculation. Set to False if covariance modifications make Cholesky
+            incompatible.
         """
         super().__init__()
         self.state: State = State.INITIALIZED
@@ -168,6 +172,7 @@ def __init__(
         self.cov_index_map: None | dict[int, int] = None
         self.theory_vector: None | npt.NDArray[np.double] = None
         self.data_vector: None | npt.NDArray[np.double] = None
+        self._use_cholesky_for_chisq: bool = use_cholesky
 
     @classmethod
     def create_ready(
@@ -392,6 +397,66 @@ def compute(
         )
         return self.get_data_vector(), self.compute_theory_vector(tools)
 
+    def _compute_chisq_cholesky(self, residuals: npt.NDArray[np.float64]) -> float:
+        """Compute chi-squared using Cholesky decomposition.
+
+        This is the numerically stable method for chi-squared calculation.
+
+        :param residuals: The residuals (data - theory)
+        :return: The chi-squared value
+        """
+        assert self.cholesky is not None
+        x = scipy.linalg.solve_triangular(self.cholesky, residuals, lower=True)
+        chisq = np.dot(x, x)
+        return float(chisq)
+
+    def _compute_chisq_direct(self, residuals: npt.NDArray[np.float64]) -> float:
+        """Compute chi-squared using direct inverse covariance multiplication.
+
+        This method is less numerically stable but necessary when covariance
+        modifications make Cholesky decomposition incompatible.
+
+        :param residuals: The residuals (data - theory)
+        :return: The chi-squared value
+        """
+        assert self.inv_cov is not None
+        chisq = residuals @ self.inv_cov @ residuals
+        return float(chisq)
+
+    def _get_theory_and_data(
+        self, tools: ModelingTools
+    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
+        """Get theory and data vectors.
+
+        :param tools: The ModelingTools to use for theory calculation
+        :return: Tuple of (theory_vector, data_vector)
+        """
+        theory_vector: npt.NDArray[np.float64]
+        data_vector: npt.NDArray[np.float64]
+        try:
+            theory_vector = self.compute_theory_vector(tools)
+            data_vector = self.get_data_vector()
+        except NotImplementedError:
+            data_vector, theory_vector = self.compute(tools)
+
+        assert len(data_vector) == len(theory_vector)
+        return theory_vector, data_vector
+
+    def compute_chisq_impl(self, residuals: npt.NDArray[np.float64]) -> float:
+        """Implementation of chi-squared calculation.
+
+        This method can be overridden by subclasses that need different
+        chi-squared calculation strategies. By default, it uses either
+        Cholesky decomposition (more stable) or direct inverse covariance
+        multiplication, depending on the _use_cholesky_for_chisq flag.
+
+        :param residuals: The residuals (data - theory)
+        :return: The chi-squared value
+        """
+        if self._use_cholesky_for_chisq:
+            return self._compute_chisq_cholesky(residuals)
+        return self._compute_chisq_direct(residuals)
+
     @final
     @enforce_states(
         initial=[State.UPDATED, State.COMPUTED],
@@ -405,22 +470,9 @@ def compute_chisq(self, tools: ModelingTools) -> float:
             theory vector
         :return: the chi-squared
         """
-        theory_vector: npt.NDArray[np.float64]
-        data_vector: npt.NDArray[np.float64]
-        residuals: npt.NDArray[np.float64]
-        try:
-            theory_vector = self.compute_theory_vector(tools)
-            data_vector = self.get_data_vector()
-        except NotImplementedError:
-            data_vector, theory_vector = self.compute(tools)
-
-        assert len(data_vector) == len(theory_vector)
+        theory_vector, data_vector = self._get_theory_and_data(tools)
         residuals = np.array(data_vector - theory_vector, dtype=np.float64)
-
-        assert self.cholesky is not None
-        x = scipy.linalg.solve_triangular(self.cholesky, residuals, lower=True)
-        chisq = np.dot(x, x)
-        return chisq
+        return self.compute_chisq_impl(residuals)
 
     @enforce_states(
         initial=[State.READY, State.UPDATED, State.COMPUTED],

firecrown/likelihood/_gaussian.py

@@ -9,7 +9,11 @@
 
 
 class ConstGaussian(GaussFamily):
-    """A Gaussian log-likelihood with a constant covariance matrix."""
+    """Base class for constant covariance Gaussian likelihoods.
+
+    Provides shared implementations of compute_loglike and make_realization_vector
+    for all constant covariance Gaussian likelihood variants.
+    """
 
     def compute_loglike(self, tools: ModelingTools) -> float:
         """Compute the log-likelihood.