flipcosmo 1.0.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

flip/covariance/fisher.py

@@ -0,0 +1,307 @@
+import importlib
+
+from flip.utils import create_log
+
+from .._config import __use_jax__
+
+if __use_jax__:
+    try:
+        import jax.numpy as jnp
+
+        jax_installed = True
+
+    except ImportError:
+        import numpy as jnp
+
+        jax_installed = False
+else:
+    import numpy as jnp
+
+    jax_installed = False
+
+log = create_log()
+
+
+def inverse_covariance_inverse(covariance):
+    """Return the explicit inverse of a covariance matrix.
+
+    Args:
+        covariance (array-like): Positive-definite covariance matrix of shape `(N, N)`.
+
+    Returns:
+        numpy.ndarray: The matrix inverse `covariance^{-1}`.
+
+    Notes:
+        For Fisher computations, explicit inversion is acceptable on modest sizes,
+        but Cholesky-based solves are usually more stable for ill-conditioned matrices.
+    """
+    return jnp.linalg.inv(covariance)
+
+
+def inverse_covariance_cholesky(covariance):
+    """Return the explicit inverse of a covariance matrix.
+
+    Args:
+        covariance (array-like): Positive-definite covariance matrix of shape `(N, N)`.
+
+    Returns:
+        numpy.ndarray: The matrix inverse `covariance^{-1}`.
+
+    Notes:
+        For Fisher computations, explicit inversion is acceptable on modest sizes,
+        but Cholesky-based solves are usually more stable for ill-conditioned matrices.
+    """
+    c = jnp.linalg.inv(jnp.linalg.cholesky(covariance))
+    inverse = jnp.dot(c.T, c)
+    return inverse
+
+
+def inverse_covariance_cholesky_inverse(covariance):
+    """Return the explicit inverse of a covariance matrix.
+
+    Args:
+        covariance (array-like): Positive-definite covariance matrix of shape `(N, N)`.
+
+    Returns:
+        numpy.ndarray: The matrix inverse `covariance^{-1}`.
+
+    Notes:
+        For Fisher computations, explicit inversion is acceptable on modest sizes,
+        but Cholesky-based solves are usually more stable for ill-conditioned matrices.
+    """
+    c = jnp.linalg.inv(jnp.linalg.cholesky(covariance))
+    inverse = jnp.dot(c.T, c)
+    return inverse
+
+
+class FisherMatrix:
+    """Compute Fisher information matrix from covariance derivatives.
+
+    Builds the Fisher matrix $F_{ij} = \tfrac{1}{2}\,\mathrm{Tr}[C^{-1} \partial_i C\, C^{-1} \partial_j C]$
+    using model-specific derivative coefficients and a provided inverse covariance.
+
+    Attributes:
+        covariance (CovMatrix): Covariance model describing blocks gg/gv/vv.
+        inverse_covariance_sum (numpy.ndarray): Precomputed `C^{-1}` used in traces.
+        parameter_values_dict (dict): Parameter values for computing coefficients.
+        free_par (list[str]): Combined free parameters from covariance and data.
+    """
+
+    _default_fisher_properties = {
+        "inversion_method": "inverse",
+        "negative_log_likelihood": True,
+    }
+
+    def __init__(
+        self,
+        covariance=None,
+        inverse_covariance_sum=None,
+        data_free_par=None,
+        parameter_values_dict=None,
+    ):
+
+        self.covariance = covariance
+        self.inverse_covariance_sum = inverse_covariance_sum
+        self.parameter_values_dict = parameter_values_dict
+        self.free_par = self.covariance.free_par[:]
+        if data_free_par is not None:
+            self.free_par += data_free_par
+
+    @classmethod
+    def init_from_covariance(
+        cls,
+        covariance,
+        data,
+        parameter_values_dict,
+        fisher_properties={},
+        covariance_prefactor_dict=None,
+    ):
+        """Initialize a FisherMatrix instance from a covariance and data model.
+
+        Ensures the covariance is in matrix form, prepares the `compute_covariance_sum`
+        functions, computes the data variance and the full covariance sum, and inverts
+        it according to the selected method.
+
+        Args:
+            covariance (CovMatrix): Covariance model to use.
+            data (object): Data model; called to obtain vector errors given parameters.
+            parameter_values_dict (dict): Parameter values for coefficient evaluation.
+            fisher_properties (dict, optional): Options including `inversion_method`.
+            covariance_prefactor_dict (dict, optional): Prefactors per block (gg/gv/vv).
+
+        Returns:
+            FisherMatrix: Ready-to-use Fisher matrix builder instance.
+        """
+        if covariance.matrix_form is False and covariance.emulator_flag is False:
+            covariance.compute_matrix_covariance()
+        if (
+            covariance.compute_covariance_sum is None
+            or covariance.compute_covariance_sum_jit is None
+        ):
+            covariance.init_compute_covariance_sum()
+
+        fisher_properties = {
+            **cls._default_fisher_properties,
+            **fisher_properties,
+        }
+
+        _, vector_variance = data.give_data_and_variance(
+            parameter_values_dict,
+        )
+        covariance_sum = covariance.compute_covariance_sum(
+            parameter_values_dict,
+            vector_variance,
+            covariance_prefactor_dict=covariance_prefactor_dict,
+        )
+
+        inverse_covariance_sum = eval(
+            f"inverse_covariance_{fisher_properties['inversion_method']}"
+        )(covariance_sum)
+
+        return cls(
+            covariance=covariance,
+            inverse_covariance_sum=inverse_covariance_sum,
+            data_free_par=data.free_par,
+            parameter_values_dict=parameter_values_dict,
+        )
+
+    def compute_covariance_derivative(
+        self,
+        partial_coefficients_dict_param,
+    ):
+        """Assemble covariance derivative matrix for a single parameter.
+
+        Uses model-kind-specific blocks and partial derivative coefficients to form
+        $\partial_i C$ for the parameter `i`.
+
+        Args:
+            partial_coefficients_dict_param (dict): Coefficients for blocks `gg/gv/vv`.
+
+        Returns:
+            numpy.ndarray: Covariance derivative matrix $\partial_i C$.
+
+        Raises:
+            ValueError: If the covariance model kind is unsupported.
+        """
+
+        if self.covariance.model_kind == "density":
+            covariance_derivative_sum = jnp.sum(
+                jnp.array(
+                    [
+                        partial_coefficients_dict_param["gg"][i] * cov
+                        for i, cov in enumerate(self.covariance.covariance_dict["gg"])
+                    ]
+                ),
+                axis=0,
+            )
+
+        elif self.covariance.model_kind == "velocity":
+            covariance_derivative_sum = jnp.sum(
+                jnp.array(
+                    [
+                        partial_coefficients_dict_param["vv"][i] * cov
+                        for i, cov in enumerate(self.covariance.covariance_dict["vv"])
+                    ]
+                ),
+                axis=0,
+            )
+
+        elif self.covariance.model_kind in ["density_velocity", "full"]:
+            number_densities = self.covariance.number_densities
+            number_velocities = self.covariance.number_velocities
+
+            if self.covariance.model_kind == "density_velocity":
+                covariance_derivative_sum_gv = jnp.zeros(
+                    (number_densities, number_velocities)
+                )
+            elif self.covariance.model_kind == "full":
+                covariance_derivative_sum_gv = jnp.sum(
+                    [
+                        partial_coefficients_dict_param["gv"][i] * cov
+                        for i, cov in enumerate(self.covariance.covariance_dict["gv"])
+                    ],
+                    axis=0,
+                )
+            covariance_derivative_sum_gg = jnp.sum(
+                [
+                    partial_coefficients_dict_param["gg"][i] * cov
+                    for i, cov in enumerate(self.covariance.covariance_dict["gg"])
+                ],
+                axis=0,
+            )
+
+            covariance_derivative_sum_vv = jnp.sum(
+                [
+                    partial_coefficients_dict_param["vv"][i] * cov
+                    for i, cov in enumerate(self.covariance.covariance_dict["vv"])
+                ],
+                axis=0,
+            )
+            covariance_derivative_sum_vg = covariance_derivative_sum_gv.T
+
+            covariance_derivative_sum = jnp.block(
+                [
+                    [covariance_derivative_sum_gg, covariance_derivative_sum_gv],
+                    [covariance_derivative_sum_vg, covariance_derivative_sum_vv],
+                ]
+            )
+        else:
+            log.add("Wrong model type in the loaded covariance.")
+
+        return covariance_derivative_sum
+
+    def compute_fisher_matrix(self, covariance_prefactor_dict=None):
+        """Compute the Fisher matrix using covariance derivatives.
+
+        Returns:
+            tuple[list[str], numpy.ndarray]: Parameter names list and Fisher matrix
+            of shape `(n_params, n_params)`.
+        """
+        coefficients = importlib.import_module(
+            f"flip.covariance.analytical.{self.covariance.model_name}.fisher_terms"
+        )
+        partial_coefficients_dict = coefficients.get_partial_derivative_coefficients(
+            self.covariance.model_kind,
+            self.parameter_values_dict,
+            variant=self.covariance.variant,
+            covariance_prefactor_dict=covariance_prefactor_dict,
+        )
+        parameter_name_list = []
+        covariance_derivative_sum_list = []
+
+        for (
+            parameter_name,
+            partial_coefficients_dict_param,
+        ) in partial_coefficients_dict.items():
+            parameter_name_list.append(parameter_name)
+            covariance_derivative_sum_list.append(
+                jnp.dot(
+                    self.inverse_covariance_sum,
+                    self.compute_covariance_derivative(
+                        partial_coefficients_dict_param,
+                    ),
+                )
+            )
+
+        fisher_matrix_size = len(partial_coefficients_dict.keys())
+        fisher_matrix = jnp.zeros((fisher_matrix_size, fisher_matrix_size))
+        tri_i, tri_j = jnp.triu_indices_from(fisher_matrix)
+
+        for i, j in zip(tri_i, tri_j):
+            if jax_installed:
+                fisher_matrix = fisher_matrix.at[i, j].set(
+                    0.5
+                    * jnp.trace(
+                        covariance_derivative_sum_list[i]
+                        @ covariance_derivative_sum_list[j]
+                    )
+                )
+                fisher_matrix = fisher_matrix.at[j, i].set(fisher_matrix[i, j])
+            else:
+                fisher_matrix[i][j] = 0.5 * jnp.trace(
+                    covariance_derivative_sum_list[i]
+                    @ covariance_derivative_sum_list[j]
+                )
+                fisher_matrix[j, i] = fisher_matrix[i, j]
+
+        return parameter_name_list, fisher_matrix

flip/{fit_utils.py → covariance/fit_utils.py}

@@ -7,8 +7,7 @@ import snsim
 import snutils
 from astropy.cosmology import FlatLambdaCDM
 
-from flip import fitter
-from flip.covariance import covariance
+from flip.covariance import covariance, fitter
 from flip.utils import create_log
 
 log = create_log()
@@ -26,6 +25,26 @@ def fit_density_minuit(
     additional_parameters_values=(),
     maximum_number_coordinates=None,
 ):
+    """Fit density field parameters with Minuit.
+
+    Loads a grid from parquet, builds a density covariance via `flip.covariance`,
+    constructs the likelihood, runs Minuit, and pickle-dumps fit results.
+
+    Args:
+        parameter_dict (dict): Parameter specs with values/errors/limits.
+        model_name (str): Covariance model name.
+        likelihood_type (str): Likelihood variant key.
+        likelihood_properties (dict): Likelihood options.
+        parameter_fit (tuple): `(grid_name, power_spectrum_dict, name_out_fit, str_fit)`.
+        overwrite (bool): Overwrite existing pickle output.
+        size_batch (int): Batch size for covariance generation.
+        number_worker (int): Parallel workers for covariance generation.
+        additional_parameters_values (tuple): Extra parameters for covariance.
+        maximum_number_coordinates (int, optional): Early stop if grid too large.
+
+    Returns:
+        None: Writes results to `name_out_fit`.
+    """
     grid_name = parameter_fit[0]
     power_spectrum_dict = parameter_fit[1]
     name_out_fit = parameter_fit[2]
@@ -45,10 +64,10 @@ def fit_density_minuit(
     for i in range(len(str_fit)):
         log.add(str_fit[i])
 
-    coordinates_density = np.array([grid["ra"], grid["dec"], grid["rcom"]])
+    coordinates_density = np.array([grid["ra"], grid["dec"], grid["rcom_zobs"]])
     data_density = {
         "density": np.array(grid["density"]),
-        "density_error": np.array(grid["density_err"]),
+        "density_error": np.array(grid["density_error"]),
     }
 
     covariance_fit = covariance.CovMatrix.init_from_flip(
@@ -96,6 +115,27 @@ def fit_density_interp_sigg_minuit(
     number_worker=32,
     maximum_number_coordinates=None,
 ):
+    """Fit density with 1D interpolation over `sigg` using Minuit.
+
+    Precomputes covariances over `interpolation_value_range`, interpolates during
+    likelihood evaluation, and runs Minuit.
+
+    Args:
+        parameter_dict (dict): Parameter specs.
+        model_name (str): Covariance model name.
+        likelihood_type (str): Likelihood variant key.
+        likelihood_properties (dict): Likelihood options.
+        interpolation_value_name (str): Interpolation parameter name.
+        interpolation_value_range (array-like): Grid of `sigg` values.
+        parameter_fit (tuple): `(grid_name, power_spectrum_dict, name_out_fit, str_fit)`.
+        overwrite (bool): Overwrite existing output.
+        size_batch (int): Covariance generation batch size.
+        number_worker (int): Parallel workers for covariance.
+        maximum_number_coordinates (int, optional): Early stop if grid too large.
+
+    Returns:
+        None: Writes results to `name_out_fit`.
+    """
     grid_name = parameter_fit[0]
     power_spectrum_dict = parameter_fit[1]
     name_out_fit = parameter_fit[2]
@@ -115,10 +155,10 @@ def fit_density_interp_sigg_minuit(
     for i in range(len(str_fit)):
         log.add(str_fit[i])
 
-    coordinates_density = np.array([grid["ra"], grid["dec"], grid["rcom"]])
+    coordinates_density = np.array([grid["ra"], grid["dec"], grid["rcom_zobs"]])
     data_density = {
         "density": np.array(grid["density"]),
-        "density_error": np.array(grid["density_err"]),
+        "density_error": np.array(grid["density_error"]),
     }
 
     covariance_list = []
@@ -172,6 +212,27 @@ def fit_velocity_true_minuit(
     size_batch=10_000,
     number_worker=32,
 ):
+    """Fit true peculiar velocities (no measurement error) with Minuit.
+
+    Extracts light curves from simulation, applies selection masks, builds the
+    velocity covariance with appropriate cosmological scaling, and runs Minuit.
+
+    Args:
+        parameter_dict (dict): Parameter specs.
+        likelihood_type (str): Likelihood variant key.
+        likelihood_properties (dict): Likelihood options.
+        zmin (float): Minimum redshift for selection.
+        z_simu (float): Simulation redshift for scaling.
+        photo_type (str): Phototyping mode for selection.
+        completeness_file (str): Path to completeness table.
+        parameter_fit (tuple): `(sim_name, fit_name, zmax, power_spectrum_dict, name_out_fit, str_fit)`.
+        overwrite (bool): Overwrite existing output.
+        size_batch (int): Covariance generation batch size.
+        number_worker (int): Parallel workers for covariance.
+
+    Returns:
+        None: Writes results to `name_out_fit`.
+    """
 
     sim_name = parameter_fit[0]
     fit_name = parameter_fit[1]
@@ -283,6 +344,29 @@ def fit_velocity_true_interp_sigu_minuit(
     size_batch=10_000,
     number_worker=32,
 ):
+    """Fit true velocities with 1D interpolation over `sigu` using Minuit.
+
+    Precomputes velocity covariances for multiple `sigu` values and interpolates
+    during likelihood evaluation.
+
+    Args:
+        parameter_dict (dict): Parameter specs.
+        likelihood_type (str): Likelihood variant key.
+        likelihood_properties (dict): Likelihood options.
+        interpolation_value_name (str): Name of interpolation parameter.
+        interpolation_value_range (array-like): Grid of `sigu` values.
+        zmin (float): Minimum redshift.
+        z_simu (float): Simulation redshift for scaling.
+        photo_type (str): Phototyping mode.
+        completeness_file (str): Path to completeness table.
+        parameter_fit (tuple): `(sim_name, fit_name, zmax, power_spectrum_dict_list, name_out_fit, str_fit)`.
+        overwrite (bool): Overwrite existing output.
+        size_batch (int): Covariance generation batch size.
+        number_worker (int): Parallel workers for covariance.
+
+    Returns:
+        None
+    """
 
     sim_name = parameter_fit[0]
     fit_name = parameter_fit[1]
@@ -400,6 +484,26 @@ def fit_velocity_estimated_minuit(
     size_batch=10_000,
     number_worker=32,
 ):
+    """Fit estimated peculiar velocities (with measurement error) using Minuit.
+
+    Similar to true velocity fit but uses observed velocity estimates with errors.
+
+    Args:
+        parameter_dict (dict): Parameter specs.
+        likelihood_type (str): Likelihood variant key.
+        likelihood_properties (dict): Likelihood options.
+        zmin (float): Minimum redshift.
+        z_simu (float): Simulation redshift for scaling.
+        photo_type (str): Phototyping mode.
+        completeness_file (str): Path to completeness table.
+        parameter_fit (tuple): `(sim_name, fit_name, zmax, power_spectrum_dict, name_out_fit, str_fit)`.
+        overwrite (bool): Overwrite existing output.
+        size_batch (int): Covariance generation batch size.
+        number_worker (int): Parallel workers for covariance.
+
+        Returns:
+            None
+    """
 
     sim_name = parameter_fit[0]
     fit_name = parameter_fit[1]
@@ -507,6 +611,28 @@ def fit_velocity_estimated_interp_sigu_minuit(
     size_batch=10_000,
     number_worker=32,
 ):
+    """Fit estimated velocities with 1D interpolation over `sigu` using Minuit.
+
+    Args follow the non-interpolated version, with additional interpolation controls.
+
+    Args:
+        parameter_dict (dict): Parameter specs.
+        likelihood_type (str): Likelihood variant key.
+        likelihood_properties (dict): Likelihood options.
+        interpolation_value_name (str): Interpolation parameter name.
+        interpolation_value_range (array-like): Grid of `sigu` values.
+        zmin (float): Minimum redshift.
+        z_simu (float): Simulation redshift for scaling.
+        photo_type (str): Phototyping mode.
+        completeness_file (str): Path to completeness.
+        parameter_fit (tuple): `(sim_name, fit_name, zmax, power_spectrum_dict_list, name_out_fit, str_fit)`.
+        overwrite (bool): Overwrite existing output.
+        size_batch (int): Covariance generation batch size.
+        number_worker (int): Parallel workers.
+
+    Returns:
+        None
+    """
 
     sim_name = parameter_fit[0]
     fit_name = parameter_fit[1]
@@ -623,6 +749,30 @@ def fit_full_velocity_estimated_minuit(
     additional_parameters_values=(),
     maximum_number_coordinates=None,
 ):
+    """Joint density-velocity fit with cross-terms using Minuit.
+
+    Builds full covariance blocks (gg, gv, vv), applies cosmological scaling
+    to velocity terms, and runs Minuit.
+
+    Args:
+        parameter_dict (dict): Parameter specs.
+        model_name (str): Covariance model name.
+        likelihood_type (str): Likelihood variant key.
+        likelihood_properties (dict): Likelihood options.
+        zmin (float): Minimum redshift for SN selection.
+        z_simu (float): Simulation redshift for scaling.
+        photo_type (str): Phototyping mode.
+        completeness_file (str): Path to completeness table.
+        parameter_fit (tuple): `(sim_name, fit_name, grid_name, zmax, power_spectrum_dict, name_out_fit, str_fit)`.
+        overwrite (bool): Overwrite existing output.
+        size_batch (int): Covariance generation batch size.
+        number_worker (int): Parallel workers for covariance.
+        additional_parameters_values (tuple): Extra parameters for covariance.
+        maximum_number_coordinates (int, optional): Early stop for large grids.
+
+    Returns:
+        None
+    """
 
     sim_name = parameter_fit[0]
     fit_name = parameter_fit[1]
@@ -646,10 +796,10 @@ def fit_full_velocity_estimated_minuit(
     for i in range(len(str_fit)):
         log.add(str_fit[i])
 
-    coordinates_density = np.array([grid["ra"], grid["dec"], grid["rcom"]])
+    coordinates_density = np.array([grid["ra"], grid["dec"], grid["rcom_zobs"]])
     data_density = {
         "density": np.array(grid["density"]),
-        "density_error": np.array(grid["density_err"]),
+        "density_error": np.array(grid["density_error"]),
     }
 
     fit = snsim.io_utils.open_fit(fit_name)
@@ -757,6 +907,31 @@ def fit_full_velocity_estimated_interp_sigu_minuit(
     additional_parameters_values=(),
     maximum_number_coordinates=None,
 ):
+    """Joint density-velocity fit with 1D interpolation over `sigu` using Minuit.
+
+    Args mirror the non-interpolated full fit plus interpolation controls.
+
+    Args:
+        parameter_dict (dict): Parameter specs.
+        model_name (str): Covariance model name.
+        likelihood_type (str): Likelihood variant key.
+        likelihood_properties (dict): Likelihood options.
+        interpolation_value_name (str): Interpolation parameter name.
+        interpolation_value_range (array-like): Grid of values.
+        zmin (float): Minimum redshift.
+        z_simu (float): Simulation redshift for scaling.
+        photo_type (str): Phototyping mode.
+        completeness_file (str): Path to completeness.
+        parameter_fit (tuple): `(sim_name, fit_name, grid_name, zmax, power_spectrum_dict_list, name_out_fit, str_fit)`.
+        overwrite (bool): Overwrite existing output.
+        size_batch (int): Covariance generation batch size.
+        number_worker (int): Parallel workers.
+        additional_parameters_values (tuple): Extra parameters for covariance.
+        maximum_number_coordinates (int, optional): Early stop.
+
+    Returns:
+        None
+    """
 
     sim_name = parameter_fit[0]
     fit_name = parameter_fit[1]
@@ -780,10 +955,10 @@ def fit_full_velocity_estimated_interp_sigu_minuit(
     for i in range(len(str_fit)):
         log.add(str_fit[i])
 
-    coordinates_density = np.array([grid["ra"], grid["dec"], grid["rcom"]])
+    coordinates_density = np.array([grid["ra"], grid["dec"], grid["rcom_zobs"]])
     data_density = {
         "density": np.array(grid["density"]),
-        "density_error": np.array(grid["density_err"]),
+        "density_error": np.array(grid["density_error"]),
     }
 
     fit = snsim.io_utils.open_fit(fit_name)