Improving mesh in flip

flip/data_vector/__init__.py

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

flip/data_vector/basic.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from flip.covariance import CovMatrix
+from flip.data_vector import mesh
 from flip.utils import create_log
 
 from .._config import __use_jax__
@@ -285,6 +286,95 @@ def __init__(self, data, covariance_observation=None):
                 self._covariance_observation = velocity_variance
 
 
+class DensMesh(Dens):
+    _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`.
+
+        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)
+
+    @classmethod
+    def init_from_catalog(
+        cls,
+        data_position_sky,
+        rcom_max,
+        grid_size,
+        grid_type,
+        kind,
+        **kwargs,
+    ):
+        grid = mesh.grid_data_density(
+            data_position_sky,
+            rcom_max,
+            grid_size,
+            grid_type,
+            kind,
+            **kwargs,
+        )
+
+        return cls(grid)
+
+
+class DirectVelMesh(DirectVel):
+    _kind = "velocity"
+    _needed_keys = ["velocity", "velocity_error"]
+    _free_par = []
+    _number_dimension_observation_covariance = 1
+    _parameters_observation_covariance = ["velocity"]
+
+    def give_data_and_variance(self, *args):
+        """Return velocity data and diagonal variance from `velocity_error`.
+
+        Returns:
+            tuple: (velocity, velocity_error^2).
+        """
+
+        if self._covariance_observation is not None:
+            return self._data["velocity"], self._covariance_observation
+        return self._data["velocity"], self._data["velocity_error"] ** 2
+
+    def __init__(self, data, covariance_observation=None):
+        super().__init__(data, covariance_observation=covariance_observation)
+
+    @classmethod
+    def init_from_catalog(
+        cls,
+        data_position_sky,
+        data,
+        rcom_max,
+        grid_size,
+        grid_type,
+        kind,
+        **kwargs,
+    ):
+        grid_velocity = mesh.grid_data_velocity(
+            data_position_sky,
+            rcom_max,
+            grid_size,
+            grid_type,
+            kind,
+            data["velocity_error"] ** 2,
+            velocity=data["velocity"],
+            **kwargs,
+        )
+
+        return cls(grid_velocity)
+
+
 class VelFromHDres(DataVector):
     _kind = "velocity"
     _needed_keys = ["dmu", "zobs"]