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

flip/data/test_covariance_reference_values.json

@@ -0,0 +1,145 @@
+{
+  "carreres23:velocity": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 3611379.1088774186,
+    "diag_mean": 72227.58217754838,
+    "entry_0_0": 72227.58217754838,
+    "entry_0_-1": 61263.88827893213,
+    "entry_mid_mid": 72227.58217754838
+  },
+  "adamsblake17plane:density": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 12.912536222207837,
+    "diag_mean": 0.25825072444415675,
+    "entry_0_0": 0.25894497176118253,
+    "entry_0_-1": -0.00012572887297247905,
+    "entry_mid_mid": 0.2579469829436334
+  },
+  "adamsblake17plane:velocity": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 3611379.108877419,
+    "diag_mean": 72227.58217754838,
+    "entry_0_0": 72227.5821775484,
+    "entry_0_-1": 62509.77032963803,
+    "entry_mid_mid": 72227.5821775484
+  },
+  "adamsblake17plane:full": {
+    "shape": [
+      100,
+      100
+    ],
+    "trace": 3611392.021413641,
+    "diag_mean": 36113.92021413641,
+    "entry_0_0": 0.25894497176118253,
+    "entry_0_-1": 0.11198990205834854,
+    "entry_mid_mid": 72227.5821775484
+  },
+  "adamsblake20:density": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 10.839977330113454,
+    "diag_mean": 0.21679954660226908,
+    "entry_0_0": 0.2174937939192948,
+    "entry_0_-1": 8.86951966632375e-05,
+    "entry_mid_mid": 0.21649580510174568
+  },
+  "adamsblake20:velocity": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 3611379.108877419,
+    "diag_mean": 72227.58217754838,
+    "entry_0_0": 72227.5821775484,
+    "entry_0_-1": 62509.77032963803,
+    "entry_mid_mid": 72227.5821775484
+  },
+  "adamsblake20:full": {
+    "shape": [
+      100,
+      100
+    ],
+    "trace": 3611389.9488547486,
+    "diag_mean": 36113.89948854749,
+    "entry_0_0": 0.2174937939192948,
+    "entry_0_-1": 0.2153532060784413,
+    "entry_mid_mid": 72227.5821775484
+  },
+  "ravouxcarreres:velocity": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 3611379.1088774186,
+    "diag_mean": 72227.58217754838,
+    "entry_0_0": 72227.58217754838,
+    "entry_0_-1": 61263.88827893213,
+    "entry_mid_mid": 72227.58217754838
+  },
+  "ravouxcarreres:density": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 10.818743417176453,
+    "diag_mean": 0.21637486834352906,
+    "entry_0_0": 0.21706911566055476,
+    "entry_0_-1": -0.00019596175905701665,
+    "entry_mid_mid": 0.21607112684300564
+  },
+  "ravouxcarreres:full": {
+    "shape": [
+      100,
+      100
+    ],
+    "trace": 3611389.9276208356,
+    "diag_mean": 36113.89927620836,
+    "entry_0_0": 0.21706911566055476,
+    "entry_0_-1": -0.8829953575001571,
+    "entry_mid_mid": 72227.58217754838
+  },
+  "lai22:density": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 24.619958402798297,
+    "diag_mean": 0.49239916805596595,
+    "entry_0_0": 0.49309341537299173,
+    "entry_0_-1": -0.0001843060148694274,
+    "entry_mid_mid": 0.4920954265554426
+  },
+  "lai22:velocity": {
+    "shape": [
+      50,
+      50
+    ],
+    "trace": 3611379.1088774186,
+    "diag_mean": 72227.58217754838,
+    "entry_0_0": 72227.58217754838,
+    "entry_0_-1": 61263.88827893213,
+    "entry_mid_mid": 72227.58217754838
+  },
+  "lai22:full": {
+    "shape": [
+      100,
+      100
+    ],
+    "trace": 3611403.728835821,
+    "diag_mean": 36114.03728835821,
+    "entry_0_0": 0.49309341537299173,
+    "entry_0_-1": -0.861244816639371,
+    "entry_mid_mid": 72227.58217754838
+  }
+}

flip/data/test_e2e_reference_values.json

@@ -0,0 +1,14 @@
+{
+  "e2e_density": {
+    "bs8": 0.7733310284360548
+  },
+  "e2e_velocity": {
+    "fs8": 0.1144194502321321,
+    "sigv": 352.4375360099262
+  },
+  "e2e_joint": {
+    "bs8": 0.7737580386847474,
+    "fs8": 0.2336828431986614,
+    "sigv": 337.7484435075482
+  }
+}

flip/data_vector/basic.py

@@ -7,13 +7,13 @@ import numpy as np
 from flip.covariance import CovMatrix
 from flip.utils import create_log
 
-from ..config import __use_jax__
+from .._config import __use_jax__
 from . import vector_utils
 
 if __use_jax__:
     try:
         import jax.numpy as jnp
-        from jax import jit
+        from jax import jit, random
         from jax.experimental.sparse import BCOO
 
         jax_installed = True
@@ -25,6 +25,7 @@ if __use_jax__:
 else:
 
     import numpy as jnp
+    from numpy import random
 
     jax_installed = False
 
@@ -32,43 +33,100 @@ log = create_log()
 
 
 class DataVector(abc.ABC):
+    """Abstract base for data vectors used in fits.
+
+    Provides common storage, key validation, optional JAX acceleration,
+    covariance-aware masking, and covariance construction helpers.
+
+    Attributes:
+        _free_par (list[str]): Model parameters this vector depends on.
+        _kind (str): One of "velocity", "density" or "cross".
+    """
+
     _free_par = []
     _kind = ""  # 'velocity', 'density' or 'cross'
 
     @property
     def conditional_free_par(self):
+        """Conditional extra parameters required by this vector.
+
+        Returns:
+            list[str]: Parameter names required depending on data content.
+        """
         return []
 
     @property
     def free_par(self):
+        """All free parameters for this vector.
+
+        Returns:
+            list[str]: Base plus conditional parameters.
+        """
         return self._free_par + self.conditional_free_par
 
     @property
     def kind(self):
+        """Return the data vector type.
+
+        Returns:
+            str: "velocity", "density" or "cross".
+        """
         return self._kind
 
     @property
     def conditional_needed_keys(self):
+        """Keys conditionally required in input `data`.
+
+        Returns:
+            list[str]: Extra keys required depending on configuration.
+        """
         return []
 
     @property
     def needed_keys(self):
+        """All required keys for this data vector.
+
+        Returns:
+            list[str]: Static plus conditional keys.
+        """
         return self._needed_keys + self.conditional_needed_keys
 
     @property
     def data(self):
+        """Access the underlying data dictionary.
+
+        Returns:
+            dict: Data fields as arrays.
+        """
         return self._data
 
     @abc.abstractmethod
     def give_data_and_variance(self, **kwargs):
+        """Return data vector and its variance/covariance.
+
+        Returns:
+            tuple: (data_array, variance_or_cov).
+        """
         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.
+
+        Args:
+            data (dict): Mapping of required fields to arrays.
+            covariance_observation (ndarray|None): Observation covariance matrix or None.
+            **kwargs: Extra configuration for subclasses.
+        """
         self._covariance_observation = covariance_observation
         self._check_keys(data)
         self._data = copy.copy(data)
@@ -80,11 +138,18 @@ class DataVector(abc.ABC):
         if jax_installed:
             self.give_data_and_variance_jit = jit(self.give_data_and_variance)
 
-    # TODO: deprecate this call
-    def __call__(self, *args):
-        return self.give_data_and_variance(*args)
-
     def get_masked_data_and_cov(self, bool_mask):
+        """Return masked data and corresponding masked observation covariance.
+
+        Args:
+            bool_mask (array-like): Boolean mask aligned with first data key length.
+
+        Returns:
+            tuple: (new_data_dict, new_cov) with covariance masked or None.
+
+        Raises:
+            ValueError: If mask length mismatches data length.
+        """
         if len(bool_mask) != len(self.data[self.needed_keys[0]]):
             raise ValueError("Boolean mask does not align with data")
         new_data = {k: v[bool_mask] for k, v in self._data.items()}
@@ -95,9 +160,19 @@ class DataVector(abc.ABC):
         return new_data, new_cov
 
     def compute_covariance(self, model, power_spectrum_dict, **kwargs):
+        """Build a `CovMatrix` for this vector and model.
+
+        Args:
+            model (str): Covariance model module under `flip.covariance`.
+            power_spectrum_dict (dict): Power spectra inputs for model.
+            **kwargs: Model-specific options.
+
+        Returns:
+            CovMatrix: Initialized covariance matrix object.
+        """
 
         coordinate_keys = importlib.import_module(
-            f"flip.covariance.{model}"
+            f"flip.covariance.analytical.{model}"
         )._coordinate_keys
 
         coords = np.vstack([self.data[k] for k in coordinate_keys])
@@ -116,6 +191,11 @@ class Dens(DataVector):
     _needed_keys = ["density", "density_error"]
 
     def give_data_and_variance(self, *args):
+        """Return density data and diagonal variance from `density_error`.
+
+        Returns:
+            tuple: (density, density_error^2).
+        """
         return self._data["density"], self._data["density_error"] ** 2
 
 
@@ -230,7 +310,7 @@ class DensVel(DataVector):
         )
 
 
-class VelFromHDres(DirectVel):
+class VelFromHDres(DataVector):
     _needed_keys = ["dmu", "zobs"]
     _free_par = ["M_0"]
 
@@ -248,54 +328,37 @@ class VelFromHDres(DirectVel):
             )
         )
         velocity = (
-            self._data["velocity"]
+            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
 
-        if self._covariance_observation is not None:
+        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
-        return velocity, self._data["velocity_error"] ** 2
+        else:
+            raise ValueError(
+                "Cannot compute velocity variance without dmu_error or covariance_observation"
+            )
 
     def __init__(
         self, data, covariance_observation=None, velocity_estimator="full", **kwargs
     ):
         # Compute conversion using provided input data, not uninitialized self._data
-        distance_modulus_difference_to_velocity = (
-            vector_utils.redshift_dependence_velocity(
-                data, velocity_estimator, **kwargs
-            )
-        )
+
         self.velocity_estimator = velocity_estimator
-        data = dict(data)  # shallow copy to avoid side-effects upstream
-        data["velocity"] = distance_modulus_difference_to_velocity * data["dmu"]
 
-        if covariance_observation is None and "dmu_error" in data:
-            data["velocity_error"] = (
-                distance_modulus_difference_to_velocity * data["dmu_error"]
-            )
         super().__init__(data, covariance_observation=covariance_observation)
 
 
-class FisherVelMesh(DataVector):
+class VelFromIntrinsicScatter(DataVector):
     _kind = "velocity"
-    _needed_keys = ["zobs", "ra", "dec", "rcom_zobs"]
-
-    def give_data_and_variance(self, parameter_values_dict):
-        variance = self.data["velocity_variance"]
-        return self._distance_modulus_difference_to_velocity**2 * variance
-
-    def __init__(self, data, velocity_estimator="full", **kwargs):
-        super().__init__(data)
-        self._distance_modulus_difference_to_velocity = (
-            vector_utils.redshift_dependence_velocity(
-                self._data, velocity_estimator, **kwargs
-            )
-        )
-class FisherVelFromHDres(DataVector):
-    _kind = "velocity"
-    _needed_keys = ["zobs", "ra", "dec", "rcom_zobs"]
+    _needed_keys = ["zobs"]
     _free_par = ["sigma_M"]
 
     def give_data_and_variance(self, parameter_values_dict):
@@ -304,71 +367,25 @@ class FisherVelFromHDres(DataVector):
                 self._data, self.velocity_estimator, **parameter_values_dict
             )
         )
+        if jax_installed:
+            key = random.PRNGKey(0)
+            distance_modulus = parameter_values_dict["sigma_M"] * random.normal(
+                key, (len(self._data["zobs"]),)
+            )
+        else:
+            distance_modulus = random.normal(
+                loc=0.0,
+                scale=parameter_values_dict["sigma_M"],
+                size=len(self._data["zobs"]),
+            )
 
         variance = parameter_values_dict["sigma_M"] ** 2
-        if "dmu_error" in self.data:
-            variance += self.data["dmu_error"] ** 2
-        return distance_modulus_difference_to_velocity**2 * variance
+
+        return (
+            distance_modulus_difference_to_velocity * distance_modulus,
+            distance_modulus_difference_to_velocity**2 * variance,
+        )
 
     def __init__(self, data, velocity_estimator="full"):
         super().__init__(data)
         self.velocity_estimator = velocity_estimator
-
-
-class FisherDens(DataVector):
-    _kind = "density"
-    _needed_keys = ["ra", "dec", "rcom_zobs"]
-    _free_par = []
-
-    def give_data_and_variance(self, parameter_values_dict):
-        variance = 0
-        if "density_error" in self.data:
-            variance += self.data["density_error"] ** 2
-        return variance
-
-    def __init__(self, data, velocity_estimator="full", **kwargs):
-        super().__init__(data)
-
-
-class FisherDensVel(DataVector):
-    _kind = "cross"
-
-    def give_data_and_variance(self, *args):
-        density_variance = self.densities.give_data_and_variance(*args)
-        velocity_variance = self.velocities.give_data_and_variance(*args)
-
-        variance = np.hstack((density_variance, velocity_variance))
-        return variance
-
-    def __init__(self, fisher_density, fisher_velocity):
-        self.densities = fisher_density
-        self.velocities = fisher_velocity
-
-        if self.velocities._covariance_observation is not None:
-            raise NotImplementedError("Vel with cov + density not implemented yet")
-
-    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/cosmo_utils.py

@@ -1,6 +1,24 @@
 def compute_hubble_norm_zobs(z, cosmo):
+    """Compute normalized Hubble parameter at redshift.
+
+    Args:
+        z (array-like): Redshift values.
+        cosmo: Astropy cosmology object with `efunc`.
+
+    Returns:
+        ndarray: `100 * E(z)`.
+    """
     return 100 * cosmo.efunc(z)
 
 
 def compute_rcom_zobs(z, cosmo):
+    """Compute comoving distance times little-h at redshift.
+
+    Args:
+        z (array-like): Redshift values.
+        cosmo: Astropy cosmology object with `comoving_distance` and `h`.
+
+    Returns:
+        ndarray: `Dm(z) * h` in Mpc/h.
+    """
     return cosmo.comoving_distance(z).value * cosmo.h