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

flip/__init__.py

@@ -17,4 +17,4 @@ except:
 
 __flip_dir_path__ = os.path.dirname(__file__)
 
-__version__ = "1.2.1"
+__version__ = "1.2.2"

flip/data_vector/__init__.py

@@ -1,4 +1,4 @@
 """Init file of the flip.data_vector package."""
 
-from . import cosmo_utils, galaxypv_vectors, snia_vectors
+from . import cosmo_utils, galaxypv_vectors, gridding, snia_vectors, vector_utils
 from .basic import *

flip/data_vector/basic.py

@@ -43,8 +43,11 @@ class DataVector(abc.ABC):
         _kind (str): One of "velocity", "density" or "cross".
     """
 
-    _free_par = []
     _kind = ""  # 'velocity', 'density' or 'cross'
+    _needed_keys = []
+    _free_par = []
+    _number_dimension_observation_covariance = 0
+    _parameters_observation_covariance = []
 
     @property
     def conditional_free_par(self):
@@ -109,16 +112,6 @@ class DataVector(abc.ABC):
         """
         pass
 
-    def _check_keys(self, data):
-        """Validate that `data` contains all required keys.
-
-        Raises:
-            ValueError: When a required key is missing.
-        """
-        for k in self.needed_keys:
-            if k not in data:
-                raise ValueError(f"{k} field is needed in data")
-
     def __init__(self, data, covariance_observation=None, **kwargs):
         """Initialize data vector with data and optional observation covariance.
 
@@ -129,6 +122,8 @@ class DataVector(abc.ABC):
         """
         self._covariance_observation = covariance_observation
         self._check_keys(data)
+        self._number_datapoints = len(data[self.needed_keys[0]])
+        self.check_covariance_observation()
         self._data = copy.copy(data)
         self._kwargs = kwargs
 
@@ -138,6 +133,31 @@ class DataVector(abc.ABC):
         if jax_installed:
             self.give_data_and_variance_jit = jit(self.give_data_and_variance)
 
+    def check_covariance_observation(self):
+        if self._covariance_observation is not None:
+            if self._covariance_observation.shape != (
+                self._number_dimension_observation_covariance * self._number_datapoints,
+                self._number_dimension_observation_covariance * self._number_datapoints,
+            ):
+                raise ValueError(
+                    f"Observation covariance matrix should be {self._number_dimension_observation_covariance}N "
+                    f"x {self._number_dimension_observation_covariance}N"
+                )
+            log.add(
+                f"Loading observation covariance matrix, "
+                f"expecting {self._parameters_observation_covariance} parameters."
+            )
+
+    def _check_keys(self, data):
+        """Validate that `data` contains all required keys.
+
+        Raises:
+            ValueError: When a required key is missing.
+        """
+        for k in self.needed_keys:
+            if k not in data:
+                raise ValueError(f"{k} field is needed in data")
+
     def get_masked_data_and_cov(self, bool_mask):
         """Return masked data and corresponding masked observation covariance.
 
@@ -189,6 +209,9 @@ class DataVector(abc.ABC):
 class Dens(DataVector):
     _kind = "density"
     _needed_keys = ["density", "density_error"]
+    _free_par = []
+    _number_dimension_observation_covariance = 1
+    _parameters_observation_covariance = ["density"]
 
     def give_data_and_variance(self, *args):
         """Return density data and diagonal variance from `density_error`.
@@ -196,12 +219,21 @@ class Dens(DataVector):
         Returns:
             tuple: (density, density_error^2).
         """
+
+        if self._covariance_observation is not None:
+            return self._data["density"], self._covariance_observation
         return self._data["density"], self._data["density_error"] ** 2
 
+    def __init__(self, data, covariance_observation=None):
+        super().__init__(data, covariance_observation=covariance_observation)
+
 
 class DirectVel(DataVector):
     _kind = "velocity"
     _needed_keys = ["velocity"]
+    _free_par = []
+    _number_dimension_observation_covariance = 1
+    _parameters_observation_covariance = ["velocity"]
 
     @property
     def conditional_needed_keys(self):
@@ -253,66 +285,12 @@ class DirectVel(DataVector):
                 self._covariance_observation = velocity_variance
 
 
-class DensVel(DataVector):
-    _kind = "cross"
-
-    @property
-    def needed_keys(self):
-        return self.densities.needed_keys + self.velocities.needed_keys
-
-    @property
-    def free_par(self):
-        return self.densities.free_par + self.velocities.free_par
-
-    def give_data_and_variance(self, *args):
-        data_density, density_variance = self.densities.give_data_and_variance(*args)
-        data_velocity, velocity_variance = self.velocities.give_data_and_variance(*args)
-        data = jnp.hstack((data_density, data_velocity))
-        variance = jnp.hstack((density_variance, velocity_variance))
-        return data, variance
-
-    def __init__(self, density_vector, velocity_vector):
-        self.densities = density_vector
-        self.velocities = velocity_vector
-
-        if self.velocities._covariance_observation is not None:
-            raise NotImplementedError(
-                "Velocity with observed covariance + density not implemented yet"
-            )
-
-        if jax_installed:
-            self.give_data_and_variance_jit = jit(self.give_data_and_variance)
-
-    def compute_covariance(self, model, power_spectrum_dict, **kwargs):
-
-        coords_dens = np.vstack(
-            (
-                self.densities.data["ra"],
-                self.densities.data["dec"],
-                self.densities.data["rcom_zobs"],
-            )
-        )
-
-        coords_vel = np.vstack(
-            (
-                self.velocities.data["ra"],
-                self.velocities.data["dec"],
-                self.velocities.data["rcom_zobs"],
-            )
-        )
-        return CovMatrix.init_from_flip(
-            model,
-            "full",
-            power_spectrum_dict,
-            coordinates_density=coords_dens,
-            coordinates_velocity=coords_vel,
-            **kwargs,
-        )
-
-
 class VelFromHDres(DataVector):
+    _kind = "velocity"
     _needed_keys = ["dmu", "zobs"]
     _free_par = ["M_0"]
+    _number_dimension_observation_covariance = 1
+    _parameters_observation_covariance = ["dmu"]
 
     @property
     def conditional_needed_keys(self):
@@ -331,25 +309,23 @@ class VelFromHDres(DataVector):
             distance_modulus_difference_to_velocity * self._data["dmu"]
             - distance_modulus_difference_to_velocity * parameter_values_dict["M_0"]
         )
-        if self._covariance_observation is None and "dmu_error" in self._data:
-            velocity_error = (
-                distance_modulus_difference_to_velocity * self._data["dmu_error"]
-            )
-            return velocity, velocity_error**2
 
-        elif self._covariance_observation is not None:
-            J = jnp.diag(self._distance_modulus_difference_to_velocity)
-            velocity_variance = J @ self._covariance_observation @ J.T
-            return velocity, velocity_variance
+        if self._covariance_observation is None:
+            velocity_variance = (
+                distance_modulus_difference_to_velocity * self._data["dmu_error"]
+            ) ** 2
         else:
-            raise ValueError(
-                "Cannot compute velocity variance without dmu_error or covariance_observation"
+            conversion_matrix = jnp.diag(distance_modulus_difference_to_velocity)
+
+            velocity_variance = (
+                conversion_matrix @ self._covariance_observation @ conversion_matrix.T
             )
 
+        return velocity, velocity_variance
+
     def __init__(
         self, data, covariance_observation=None, velocity_estimator="full", **kwargs
     ):
-        # Compute conversion using provided input data, not uninitialized self._data
 
         self.velocity_estimator = velocity_estimator
 
@@ -360,6 +336,8 @@ class VelFromIntrinsicScatter(DataVector):
     _kind = "velocity"
     _needed_keys = ["zobs"]
     _free_par = ["sigma_M"]
+    _number_dimension_observation_covariance = 0
+    _parameters_observation_covariance = []
 
     def give_data_and_variance(self, parameter_values_dict):
         distance_modulus_difference_to_velocity = (
@@ -389,3 +367,60 @@ class VelFromIntrinsicScatter(DataVector):
     def __init__(self, data, velocity_estimator="full"):
         super().__init__(data)
         self.velocity_estimator = velocity_estimator
+
+
+class DensVel(DataVector):
+    _kind = "cross"
+
+    @property
+    def needed_keys(self):
+        return self.densities.needed_keys + self.velocities.needed_keys
+
+    @property
+    def free_par(self):
+        return self.densities.free_par + self.velocities.free_par
+
+    def give_data_and_variance(self, *args):
+        data_density, density_variance = self.densities.give_data_and_variance(*args)
+        data_velocity, velocity_variance = self.velocities.give_data_and_variance(*args)
+        data = jnp.hstack((data_density, data_velocity))
+        variance = jnp.hstack((density_variance, velocity_variance))
+        return data, variance
+
+    def __init__(self, density_vector, velocity_vector):
+        self.densities = density_vector
+        self.velocities = velocity_vector
+
+        if self.velocities._covariance_observation is not None:
+            raise NotImplementedError(
+                "Velocity with observed covariance + density not implemented yet"
+            )
+
+        if jax_installed:
+            self.give_data_and_variance_jit = jit(self.give_data_and_variance)
+
+    def compute_covariance(self, model, power_spectrum_dict, **kwargs):
+
+        coords_dens = np.vstack(
+            (
+                self.densities.data["ra"],
+                self.densities.data["dec"],
+                self.densities.data["rcom_zobs"],
+            )
+        )
+
+        coords_vel = np.vstack(
+            (
+                self.velocities.data["ra"],
+                self.velocities.data["dec"],
+                self.velocities.data["rcom_zobs"],
+            )
+        )
+        return CovMatrix.init_from_flip(
+            model,
+            "full",
+            power_spectrum_dict,
+            coordinates_density=coords_dens,
+            coordinates_velocity=coords_vel,
+            **kwargs,
+        )

flip/data_vector/galaxypv_vectors.py

@@ -1,3 +1,5 @@
+from flip.utils import create_log
+
 from .._config import __use_jax__
 from . import vector_utils
 from .basic import DataVector
@@ -19,10 +21,15 @@ else:
 
     jax_installed = False
 
+log = create_log()
+
 
 class VelFromLogDist(DataVector):
     _kind = "velocity"
     _needed_keys = ["eta"]
+    _free_par = []
+    _number_dimension_observation_covariance = 1
+    _parameters_observation_covariance = ["eta"]
 
     @property
     def conditional_needed_keys(self):
@@ -36,6 +43,22 @@ class VelFromLogDist(DataVector):
             cond_keys += ["eta_error"]
         return self._needed_keys + cond_keys
 
+    def __init__(
+        self,
+        data,
+        covariance_observation=None,
+        velocity_estimator="full",
+    ):
+        """Initialize velocity from log-distance `eta`.
+
+        Args:
+            data (dict): Must include `eta` and optionally `eta_error`.
+            covariance_observation (ndarray|None): Observed covariance.
+            velocity_estimator (str): Estimator name, default `"full"`.
+        """
+        self.velocity_estimator = velocity_estimator
+        super().__init__(data, covariance_observation=covariance_observation)
+
     def give_data_and_variance(self, parameter_values_dict, *args):
         """Return velocity and variance for log-distance based estimator.
 
@@ -53,31 +76,19 @@ class VelFromLogDist(DataVector):
 
         velocity = log_distance_to_velocity * self._data["eta"]
 
-        if self._covariance_observation is not None:
-            J = jnp.diag(log_distance_to_velocity)
-            velocity_variance = J @ self._covariance_observation @ J.T
-            return velocity, velocity_variance
+        if self._covariance_observation is None:
+            velocity_variance = (
+                log_distance_to_velocity * self._data["eta_error"]
+            ) ** 2
 
-        return velocity, (log_distance_to_velocity * self._data["eta_error"]) ** 2
+        else:
+            conversion_matrix = jnp.diag(log_distance_to_velocity)
 
-    def __init__(
-        self,
-        data,
-        covariance_observation=None,
-        velocity_estimator="full",
-    ):
-        """Initialize velocity from log-distance `eta`.
+            velocity_variance = (
+                conversion_matrix @ self._covariance_observation @ conversion_matrix.T
+            )
 
-        Args:
-            data (dict): Must include `eta` and optionally `eta_error`.
-            covariance_observation (ndarray|None): Observed covariance.
-            velocity_estimator (str): Estimator name, default `"full"`.
-        """
-        self.velocity_estimator = velocity_estimator
-        super().__init__(
-            data,
-            covariance_observation=covariance_observation,
-        )
+        return velocity, velocity_variance
 
 
 class VelFromTullyFisher(DataVector):
@@ -85,6 +96,8 @@ class VelFromTullyFisher(DataVector):
     _kind = "velocity"
     _needed_keys = ["zobs", "logW", "m_mean", "rcom_zobs"]
     _free_par = ["a", "b"]
+    _number_dimension_observation_covariance = 2
+    _parameters_observation_covariance = ["logW", "m_mean"]
 
     @property
     def conditional_needed_keys(self):
@@ -98,6 +111,39 @@ class VelFromTullyFisher(DataVector):
             cond_keys += ["e_logW", "e_m_mean"]
         return cond_keys
 
+    def __init__(
+        self,
+        data,
+        h,
+        covariance_observation=None,
+        velocity_estimator="full",
+    ):
+        """Initialize Tully–Fisher velocity vector.
+
+        Args:
+            data (dict): Includes `logW`, `m_mean`, redshifts and distances.
+            h (float): Little-h scaling for distances.
+            covariance_observation (ndarray|None): Optional observation covariance.
+            velocity_estimator (str): Estimator name.
+
+        Raises:
+            ValueError: If covariance shape is not `2N x 2N` when provided.
+        """
+        super().__init__(data, covariance_observation=covariance_observation)
+        self.velocity_estimator = velocity_estimator
+        self.h = h
+        self._host_matrix = None
+
+        if "host_group_id" in data:
+            self._host_matrix, self._data_to_group_mapping = (
+                vector_utils.compute_host_matrix(self._data["host_group_id"])
+            )
+            self._data = vector_utils.format_data_multiple_host(
+                self._data, self._host_matrix
+            )
+            if jax_installed:
+                self._host_matrix = BCOO.from_scipy_sparse(self._host_matrix)
+
     def compute_observed_distance_modulus(self, parameter_values_dict):
         """Compute observed distance modulus from Tully–Fisher relation.
 
@@ -159,11 +205,23 @@ class VelFromTullyFisher(DataVector):
             )
             variance_distance_modulus += parameter_values_dict["sigma_M"] ** 2
         else:
+            weights_observation_covariance = jnp.array(
+                [
+                    1.0,
+                    parameter_values_dict["a"],
+                ]
+            )
+            jacobian = jnp.kron(
+                weights_observation_covariance,
+                jnp.eye(self._number_datapoints),
+            )
             variance_distance_modulus = (
-                self._covariance_observation
-                + jnp.eye(self._covariance_observation.shape[0])
-                * parameter_values_dict["sigma_M"] ** 2
+                jacobian @ self._covariance_observation @ jacobian.T
+            )
+            variance_distance_modulus += (
+                jnp.eye(self._number_datapoints) * parameter_values_dict["sigma_M"] ** 2
             )
+
         return variance_distance_modulus
 
     def give_data_and_variance(self, parameter_values_dict):
@@ -190,10 +248,13 @@ class VelFromTullyFisher(DataVector):
                 * distance_modulus_difference_to_velocity**2
             )
         else:
-            A = self._init_A()
-            J = A[0] + parameter_values_dict["a"] * A[1]
-            J = jnp.diag(distance_modulus_difference_to_velocity) @ J
-            velocity_variance = J @ observed_distance_modulus_variance @ J.T
+            conversion_matrix = jnp.diag(distance_modulus_difference_to_velocity)
+
+            velocity_variance = (
+                conversion_matrix
+                @ observed_distance_modulus_variance
+                @ conversion_matrix.T
+            )
 
         velocities = (
             distance_modulus_difference_to_velocity
@@ -207,18 +268,25 @@ class VelFromTullyFisher(DataVector):
 
         return velocities, velocity_variance
 
-    def _init_A(self):
-        """Initialize design matrices for linear propagation with covariance.
+
+class VelFromFundamentalPlane(DataVector):
+    _kind = "velocity"
+    _needed_keys = ["zobs", "logRe", "logsig", "logI", "rcom_zobs"]
+    _free_par = ["a", "b", "c"]
+    _number_dimension_observation_covariance = 3
+    _parameters_observation_covariance = ["logRe", "logsig", "logI"]
+
+    @property
+    def conditional_needed_keys(self):
+        """Conditionally required keys when covariance is absent.
 
         Returns:
-            ndarray: Matrix A blocks.
+            list[str]: Includes `e_logRe`, `e_logsig`, `e_logI` when needed.
         """
-        N = len(self._data)
-        A = jnp.ones((2, N, 2 * N))
-        ij = jnp.ogrid[:N, : 2 * N]
-        for k in range(2):
-            A[k][ij[1] == 2 * ij[0] + k] = 1
-        return A
+        cond_keys = []
+        if self._covariance_observation is None:
+            cond_keys += ["e_logRe", "e_logsig", "e_logI"]
+        return cond_keys
 
     def __init__(
         self,
@@ -227,21 +295,20 @@ class VelFromTullyFisher(DataVector):
         covariance_observation=None,
         velocity_estimator="full",
     ):
-        """Initialize Tully–Fisher velocity vector.
+        """Initialize Fundamental Plane velocity vector.
 
         Args:
-            data (dict): Includes `logW`, `m_mean`, redshifts and distances.
+            data (dict): Includes `logRe`, `logsig`, `logI`, redshifts and distances.
             h (float): Little-h scaling for distances.
             covariance_observation (ndarray|None): Optional observation covariance.
             velocity_estimator (str): Estimator name.
 
         Raises:
-            ValueError: If covariance shape is not `2N x 2N` when provided.
+            ValueError: If covariance shape is not `3N x 3N` when provided.
         """
         super().__init__(data, covariance_observation=covariance_observation)
         self.velocity_estimator = velocity_estimator
         self.h = h
-        self._A = None
         self._host_matrix = None
 
         if "host_group_id" in data:
@@ -254,29 +321,6 @@ class VelFromTullyFisher(DataVector):
             if jax_installed:
                 self._host_matrix = BCOO.from_scipy_sparse(self._host_matrix)
 
-        if self._covariance_observation is not None:
-            if self._covariance_observation.shape != (2 * len(data), 2 * len(data)):
-                raise ValueError("Cov should be 2N x 2N")
-
-
-class VelFromFundamentalPlane(DataVector):
-
-    _kind = "velocity"
-    _needed_keys = ["zobs", "logRe", "logsig", "logI", "rcom_zobs"]
-    _free_par = ["a", "b", "c"]
-
-    @property
-    def conditional_needed_keys(self):
-        """Conditionally required keys when covariance is absent.
-
-        Returns:
-            list[str]: Includes `e_logRe`, `e_logsig`, `e_logI` when needed.
-        """
-        cond_keys = []
-        if self._covariance_observation is None:
-            cond_keys += ["e_logRe", "e_logsig", "e_logI"]
-        return cond_keys
-
     def compute_observed_distance_modulus(self, parameter_values_dict):
         """Compute observed distance modulus from Fundamental Plane relation.
 
@@ -340,11 +384,24 @@ class VelFromFundamentalPlane(DataVector):
             )
             variance_distance_modulus += parameter_values_dict["sigma_M"] ** 2
         else:
+            weights_observation_covariance = jnp.array(
+                [
+                    1.0,
+                    parameter_values_dict["a"],
+                    parameter_values_dict["b"],
+                ]
+            )
+            jacobian = jnp.kron(
+                weights_observation_covariance,
+                jnp.eye(self._number_datapoints),
+            )
             variance_distance_modulus = (
-                self._covariance_observation
-                + jnp.eye(self._covariance_observation.shape[0])
-                * parameter_values_dict["sigma_M"] ** 2
+                jacobian @ self._covariance_observation @ jacobian.T
             )
+            variance_distance_modulus += (
+                jnp.eye(self._number_datapoints) * parameter_values_dict["sigma_M"] ** 2
+            )
+
         return variance_distance_modulus
 
     def give_data_and_variance(self, parameter_values_dict):
@@ -371,14 +428,13 @@ class VelFromFundamentalPlane(DataVector):
                 * distance_modulus_difference_to_velocity**2
             )
         else:
-            A = self._init_A()
-            J = (
-                A[0]
-                + parameter_values_dict["a"] * A[1]
-                + parameter_values_dict["b"] * A[2]
+            conversion_matrix = jnp.diag(distance_modulus_difference_to_velocity)
+
+            velocity_variance = (
+                conversion_matrix
+                @ observed_distance_modulus_variance
+                @ conversion_matrix.T
             )
-            J = jnp.diag(distance_modulus_difference_to_velocity) @ J
-            velocity_variance = J @ observed_distance_modulus_variance @ J.T
 
         velocities = (
             distance_modulus_difference_to_velocity
@@ -391,54 +447,3 @@ class VelFromFundamentalPlane(DataVector):
             )
 
         return velocities, velocity_variance
-
-    def _init_A(self):
-        """Initialize design matrices for linear propagation with covariance.
-
-        Returns:
-            ndarray: Matrix A blocks.
-        """
-        N = len(self._data)
-        A = jnp.ones((3, N, 3 * N))
-        ij = jnp.ogrid[:N, : 3 * N]
-        for k in range(3):
-            A[k][ij[1] == 3 * ij[0] + k] = 1
-        return A
-
-    def __init__(
-        self,
-        data,
-        h,
-        covariance_observation=None,
-        velocity_estimator="full",
-    ):
-        """Initialize Fundamental Plane velocity vector.
-
-        Args:
-            data (dict): Includes `logRe`, `logsig`, `logI`, redshifts and distances.
-            h (float): Little-h scaling for distances.
-            covariance_observation (ndarray|None): Optional observation covariance.
-            velocity_estimator (str): Estimator name.
-
-        Raises:
-            ValueError: If covariance shape is not `3N x 3N` when provided.
-        """
-        super().__init__(data, covariance_observation=covariance_observation)
-        self.velocity_estimator = velocity_estimator
-        self.h = h
-        self._A = None
-        self._host_matrix = None
-
-        if "host_group_id" in data:
-            self._host_matrix, self._data_to_group_mapping = (
-                vector_utils.compute_host_matrix(self._data["host_group_id"])
-            )
-            self._data = vector_utils.format_data_multiple_host(
-                self._data, self._host_matrix
-            )
-            if jax_installed:
-                self._host_matrix = BCOO.from_scipy_sparse(self._host_matrix)
-
-        if self._covariance_observation is not None:
-            if self._covariance_observation.shape != (3 * len(data), 3 * len(data)):
-                raise ValueError("Cov should be 3N x 3N")

flip/data_vector/snia_vectors.py

@@ -1,3 +1,5 @@
+from flip.utils import create_log
+
 from .._config import __use_jax__
 from . import vector_utils
 from .basic import DataVector
@@ -19,11 +21,15 @@ else:
 
     jax_installed = False
 
+log = create_log()
+
 
 class VelFromSALTfit(DataVector):
     _kind = "velocity"
     _needed_keys = ["zobs", "mb", "x1", "c", "rcom_zobs"]
     _free_par = ["alpha", "beta", "M_0", "sigma_M"]
+    _number_dimension_observation_covariance = 3
+    _parameters_observation_covariance = ["mb", "x1", "c"]
 
     @property
     def conditional_needed_keys(self):
@@ -49,6 +55,42 @@ class VelFromSALTfit(DataVector):
             _cond_fpar += ["gamma"]
         return _cond_fpar
 
+    def __init__(
+        self,
+        data,
+        h,
+        covariance_observation=None,
+        velocity_estimator="full",
+        mass_step=10,
+    ):
+        """Initialize SN Ia velocity vector from SALT2 fits.
+
+        Args:
+            data (dict): Includes SALT2 parameters and cosmology fields.
+            h (float): Little-h scaling for distances.
+            covariance_observation (ndarray|None): Optional observation covariance.
+            velocity_estimator (str): Estimator name.
+            mass_step (float): Threshold for host mass step correction.
+
+        Raises:
+            ValueError: If covariance shape is not adapted
+        """
+        super().__init__(data, covariance_observation=covariance_observation)
+        self.velocity_estimator = velocity_estimator
+        self.h = h
+        self._host_matrix = None
+        self._mass_step = mass_step
+
+        if "host_group_id" in data:
+            self._host_matrix, self._data_to_group_mapping = (
+                vector_utils.compute_host_matrix(self._data["host_group_id"])
+            )
+            self._data = vector_utils.format_data_multiple_host(
+                self._data, self._host_matrix
+            )
+            if jax_installed:
+                self._host_matrix = BCOO.from_scipy_sparse(self._host_matrix)
+
     def compute_observed_distance_modulus(self, parameter_values_dict):
         """Compute observed distance modulus from SALT2 fit parameters.
 
@@ -122,11 +164,24 @@ class VelFromSALTfit(DataVector):
             )
             variance_distance_modulus += parameter_values_dict["sigma_M"] ** 2
         else:
+            weights_observation_covariance = jnp.array(
+                [
+                    1.0,
+                    parameter_values_dict["alpha"],
+                    -parameter_values_dict["beta"],
+                ]
+            )
+            jacobian = jnp.kron(
+                weights_observation_covariance,
+                jnp.eye(self._number_datapoints),
+            )
             variance_distance_modulus = (
-                self._covariance_observation
-                + jnp.eye(self._covariance_observation.shape[0])
-                * parameter_values_dict["sigma_M"] ** 2
+                jacobian @ self._covariance_observation @ jacobian.T
             )
+            variance_distance_modulus += (
+                jnp.eye(self._number_datapoints) * parameter_values_dict["sigma_M"] ** 2
+            )
+
         return variance_distance_modulus
 
     def give_data_and_variance(self, parameter_values_dict):
@@ -152,14 +207,13 @@ class VelFromSALTfit(DataVector):
                 * distance_modulus_difference_to_velocity**2
             )
         else:
-            A = self._init_A()
-            J = (
-                A[0]
-                + parameter_values_dict["alpha"] * A[1]
-                - parameter_values_dict["beta"] * A[2]
+            conversion_matrix = jnp.diag(distance_modulus_difference_to_velocity)
+
+            velocity_variance = (
+                conversion_matrix
+                @ observed_distance_modulus_variance
+                @ conversion_matrix.T
             )
-            J = jnp.diag(distance_modulus_difference_to_velocity) @ J
-            velocity_variance = J @ observed_distance_modulus_variance @ J.T
 
         velocities = (
             distance_modulus_difference_to_velocity
@@ -172,58 +226,3 @@ class VelFromSALTfit(DataVector):
             )
 
         return velocities, velocity_variance
-
-    def _init_A(self):
-        """Initialize design matrices for linear covariance propagation.
-
-        Returns:
-            ndarray: Matrix A blocks.
-        """
-        N = len(self._data)
-        A = jnp.ones((3, N, 3 * N))
-        ij = jnp.ogrid[:N, : 3 * N]
-        for k in range(3):
-            A[k][ij[1] == 3 * ij[0] + k] = 1
-        return A
-
-    def __init__(
-        self,
-        data,
-        h,
-        covariance_observation=None,
-        velocity_estimator="full",
-        mass_step=10,
-    ):
-        """Initialize SN Ia velocity vector from SALT2 fits.
-
-        Args:
-            data (dict): Includes SALT2 parameters and cosmology fields.
-            h (float): Little-h scaling for distances.
-            covariance_observation (ndarray|None): Optional observation covariance.
-            velocity_estimator (str): Estimator name.
-            mass_step (float): Threshold for host mass step correction.
-
-        Raises:
-            ValueError: If covariance shape is not `3N x 3N` when provided.
-        """
-        super().__init__(data, covariance_observation=covariance_observation)
-        self.velocity_estimator = velocity_estimator
-        self.h = h
-        self._A = None
-        self._host_matrix = None
-        self._mass_step = mass_step
-
-        if "host_group_id" in data:
-            self._host_matrix, self._data_to_group_mapping = (
-                vector_utils.compute_host_matrix(self._data["host_group_id"])
-            )
-            self._data = vector_utils.format_data_multiple_host(
-                self._data, self._host_matrix
-            )
-            if jax_installed:
-                self._host_matrix = BCOO.from_scipy_sparse(self._host_matrix)
-
-        if self._covariance_observation is not None:
-            if self._covariance_observation.shape != (3 * len(data), 3 * len(data)):
-                raise ValueError("Cov should be 3N x 3N")
-            self._A = self._init_A()

{flipcosmo-1.2.1.dist-info → flipcosmo-1.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flipcosmo
-Version: 1.2.1
+Version: 1.2.2
 Summary:  Field Level Inference Package
 Author-email: Corentin Ravoux <corentin.ravoux.research@gmail.com>, Bastien Carreres <bastien.carreres@duke.edu>
 License-Expression: MIT
@@ -30,7 +30,7 @@ Requires-Dist: sphinx-tabs; extra == "docs"
 Requires-Dist: sphinx-autoapi; extra == "docs"
 Dynamic: license-file
 
-<img src="docs/_static/flip_logo.webp" width=350>
+<img src="https://github.com/corentinravoux/flip/blob/main/docs/_static/flip_logo.webp?raw=true" width=350>
 
 # flip: Field Level Inference Package
 

{flipcosmo-1.2.1.dist-info → flipcosmo-1.2.2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 docs/conf.py,sha256=oQH2FVsFdgs7Z2z0qyORgjUz1Wn35TS-oyd3GaQ8Uos,4694
-flip/__init__.py,sha256=nvvWo-nfwNPaATqooXCW0p6EPun32rOQ385ilNb2H_A,419
+flip/__init__.py,sha256=FAXygNnmz1DnB-Y4vBQgjDolJ3I8sLUxgWgXORWGbrE,419
 flip/_config.py,sha256=wf09C_i7q-G8v3ADk63WtrSEZyXQhy4zNzQ9AY3Qq9U,19
 flip/utils.py,sha256=mkLEW-qDWYRSpycnF6MxWUtTOAqogKZb708I8ys14p0,5462
 flip/covariance/__init__.py,sha256=bzxPuQOFUX7YfFb5WmoWE-ykQWK7sLh_zQHyRuvF2ng,233
@@ -64,7 +64,6 @@ flip/data/__init__.py,sha256=8OBjUqbJBny0jMSCxCyKVi4-MNjR6vg3ZTGCGX8Vgf4,72
 flip/data/data_density.parquet,sha256=_156O6Q05FQDqQj9QaStQBpiS1eWngTHN_QdBsfos2M,14061
 flip/data/data_velocity.parquet,sha256=m5WkL3LRzwHKl-AgxMVnhn_8IaB5extm1Fipo3fKR1c,181039
 flip/data/data_window_density.parquet,sha256=vSF7GGJzIkzGNTnhksFa-zJlEs9i5THugEzUl52si2o,40372
-flip/data/gridding.py,sha256=_3PHazY-KbxQyHKWxQguB4LGJ1yU-uHNcQqzQSk5O8c,30864
 flip/data/load_data_test.py,sha256=DmHn3gJTDtGqJCGObnpvhKm0YDGlhTM8kJuHIZUxZ9Q,2871
 flip/data/power_spectrum_mm.txt,sha256=TltaM8gxTnyUR3eLM1GHf6HCctRA1F6cHmuhjRYkj8M,50085
 flip/data/power_spectrum_mt.txt,sha256=S9fuSnNoW9qwJL7xuRUoVSz_EzmgEJENy3td8u9FkXU,50085
@@ -72,11 +71,12 @@ flip/data/power_spectrum_tt.txt,sha256=Ws3b8XtzZJ0aNKT7g3PRb0NQWDtvpFBTU6nCi0G5Z
 flip/data/style.mplstyle,sha256=2og3sSZU3SqBBB1Jm4alodk3jJF-PMyf7XZqCvItwdA,2341
 flip/data/test_covariance_reference_values.json,sha256=O6Z-QeDBMjWQUKpLqbhwilGfO-fQaGGRjAdfHuq_WvQ,3360
 flip/data/test_e2e_reference_values.json,sha256=qfqJDGL2CCiR4xiiYW6SX4_FdzawXQLodFXQREmT_Y0,256
-flip/data_vector/__init__.py,sha256=ZUGmpWl3D_PvZSmQNte4L-sW7trfL-SJj7jtzeysnwM,129
-flip/data_vector/basic.py,sha256=e5hsvSB19JA5QDeaNMdAQndLbDqcrNxx6zdCKeTTM5o,12362
+flip/data_vector/__init__.py,sha256=RINRUs2aQF4-erLqSxb9AvdnixbuVPWLhAw5VAlrfgE,153
+flip/data_vector/basic.py,sha256=L4ZN3PUMLLoAjBW1SNq-Pa_y-4OMRnzvcIKxhjI6TP0,13792
 flip/data_vector/cosmo_utils.py,sha256=WvrQmWM0lPkdy1Ak8NMe3IzauRcBV9bubAQ3YGCdqQo,624
-flip/data_vector/galaxypv_vectors.py,sha256=VI897662S5s7g6iwAsQssvKK6ydBto1LBmvjBez6mfA,14890
-flip/data_vector/snia_vectors.py,sha256=IzM2nH7y1y9puWGugLfLUJURpRksBNg2w5XZxypoWPY,7978
+flip/data_vector/galaxypv_vectors.py,sha256=apwCslFNU6Xa2IJ00M5QlzORmXUFptod4Qdd1Ic5NgE,14958
+flip/data_vector/gridding.py,sha256=_3PHazY-KbxQyHKWxQguB4LGJ1yU-uHNcQqzQSk5O8c,30864
+flip/data_vector/snia_vectors.py,sha256=-RrIqUJMulsYy9U29wSR1S-DUc1TR8uazKy8Xsk056s,7857
 flip/data_vector/vector_utils.py,sha256=P0AfdyYUU5F6ClgrhyASmqs3ai6PsKmB_o_KnHlWVJY,6715
 flip/power_spectra/__init__.py,sha256=9XDOGCo4YFXg0QdQ6gEC66oVaB1Q_NtCmqMn8PZs30w,118
 flip/power_spectra/class_engine.py,sha256=MWKvSFKOK242JRpc4wNxbT_1ldH9c00umVW75A4px-A,5964
@@ -85,7 +85,7 @@ flip/power_spectra/generator.py,sha256=X6AAsuiYWXNjfJDdYGd0Pna1wgXH0_eNEZefGNcxZ
 flip/power_spectra/models.py,sha256=bRR5lIQHyoRtE0lqQKjiPxCdQNqwK1ZG5zrjMPQmzoo,3544
 flip/power_spectra/pyccl_engine.py,sha256=D8PTQaB43L6A0VhiQs55kuk8t6uOcDIxwLNW6i8Ja6w,3527
 flip/simulation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flipcosmo-1.2.1.dist-info/licenses/LICENSE,sha256=jQfEFGXoNPb52M0jMfnWZtz8reTMYXItOx9-wevuHWg,1071
+flipcosmo-1.2.2.dist-info/licenses/LICENSE,sha256=jQfEFGXoNPb52M0jMfnWZtz8reTMYXItOx9-wevuHWg,1071
 scripts/flip_compute_correlation_model.py,sha256=9IhGD2rctHIuE_FeOuTze22EEqGVWUqFzoTle1HNQ6k,1710
 scripts/flip_compute_power_spectra.py,sha256=Ez36BS8bmxPs0iSeHVGIhzXcohlRb8BG6NxsioN7S0Y,1298
 scripts/flip_fisher_forecast_velocity.py,sha256=coQ_YdsSouq-4AGXExwojr1fIC_tO0wz4-3HU9ZYTco,1838
@@ -103,7 +103,7 @@ test/test_e2e_density.py,sha256=vVaRphhfGdHxKknzVOYbH3YMGma7BUk7o2PCqebHSr0,1833
 test/test_e2e_joint.py,sha256=GfQkm2_NGPXiZWRxmbk1xX0Q8JXtfue59klDD6tHOBE,2386
 test/test_e2e_velocity.py,sha256=FebMsWwzqvfhkest5kkXTdkW6Me-q1gTypcU6Qc4Kh4,1796
 test/test_likelihood_inversions.py,sha256=RE2IiCtBPKJzlvEaFQysYh1nht9wTQo2Bfsfa0hChaY,830
-flipcosmo-1.2.1.dist-info/METADATA,sha256=WxnK7RJ3Nxyc0gaH7IipR1SepPDVo4AUaWsVuMZj8iE,3616
-flipcosmo-1.2.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-flipcosmo-1.2.1.dist-info/top_level.txt,sha256=9MRA5iw3NarUSIfyQLBZB_b_44iHdvuTW9CjqvMi8_4,44
-flipcosmo-1.2.1.dist-info/RECORD,,
+flipcosmo-1.2.2.dist-info/METADATA,sha256=Z6Q7Fxv9C7H8oXHxO4bTu8ejeHGQEb4aOMlqtH35HS4,3674
+flipcosmo-1.2.2.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+flipcosmo-1.2.2.dist-info/top_level.txt,sha256=9MRA5iw3NarUSIfyQLBZB_b_44iHdvuTW9CjqvMi8_4,44
+flipcosmo-1.2.2.dist-info/RECORD,,