Added cart prod pf

FlorianPfaff · FlorianPfaff · commit dcdb4efab4a1 · 2025-03-21T19:28:15.000+01:00
diff --git a/pyrecest/distributions/cart_prod/cart_prod_dirac_distribution.py b/pyrecest/distributions/cart_prod/cart_prod_dirac_distribution.py
@@ -0,0 +1,21 @@
+from ..hypersphere_subset.abstract_hyperhemispherical_distribution import (
+    AbstractHyperhemisphericalDistribution,
+)
+from ..hypersphere_subset.abstract_hypersphere_subset_dirac_distribution import (
+    AbstractHypersphereSubsetDiracDistribution,
+)
+from ..hypersphere_subset.abstract_hyperspherical_distribution import (
+    AbstractHypersphericalDistribution,
+)
+
+
+class HyperhemisphericalDiracDistribution(
+    AbstractHypersphereSubsetDiracDistribution, AbstractHyperhemisphericalDistribution
+):
+    pass
+
+
+class HypersphericalDiracDistribution(
+    AbstractHypersphereSubsetDiracDistribution, AbstractHypersphericalDistribution
+):
+    pass
diff --git a/pyrecest/distributions/cart_prod/hyperhemisphere_cart_prod_dirac_distribution.py b/pyrecest/distributions/cart_prod/hyperhemisphere_cart_prod_dirac_distribution.py
@@ -0,0 +1,38 @@
+import copy
+from collections.abc import Callable
+
+from ..abstract_dirac_distribution import AbstractDiracDistribution
+
+
+class HyperhemisphereCartProdDiracDistribution(AbstractDiracDistribution):
+    def __init__(self, d, w, dim_hemisphere, n_hemispheres):
+        """
+        Initialize a Dirac distribution with given Dirac locations and weights.
+
+        :param d: Dirac locations as a numpy array.
+        :param w: Weights of Dirac locations as a numpy array. If not provided, defaults to uniform weights.
+        """
+        super().__init__(d, w)
+        self.dim_hemisphere = dim_hemisphere
+        self.n_hemispheres = n_hemispheres
+        assert self.d.shape[-1] == (
+            (1 + dim_hemisphere) * n_hemispheres
+        ), "Dimension is not correct."
+        self.dim = dim_hemisphere * n_hemispheres
+
+    def apply_function_component_wise(
+        self, f: Callable, f_supports_multiple: bool = True
+    ):
+        """
+        Apply a function to the Dirac locations and return a new distribution.
+
+        :param f: Function to apply.
+        :returns: A new distribution with the function applied to the locations.
+        """
+        assert f_supports_multiple, "Function must support multiple inputs."
+        dist = copy.deepcopy(self)
+        for i in range(self.n_hemispheres):
+            dist.d[
+                i * self.dim_hemisphere : (i + 1) * self.dim_hemisphere  # noqa: E203 
+                ] = f(self.d[i])
+        return dist
diff --git a/pyrecest/filters/hyperhemisphere_cart_prod_particle_filter.py b/pyrecest/filters/hyperhemisphere_cart_prod_particle_filter.py
@@ -0,0 +1,113 @@
+from collections.abc import Callable
+
+import numpy as np
+
+# pylint: disable=no-name-in-module,no-member
+from pyrecest.backend import ones  # noqa: F821
+from pyrecest.distributions import AbstractHypersphericalDistribution
+from pyrecest.distributions.cart_prod.hyperhemisphere_cart_prod_dirac_distribution import (
+    HyperhemisphereCartProdDiracDistribution,
+)
+from pyrecest.distributions.hypersphere_subset.abstract_hyperhemispherical_distribution import (
+    AbstractHyperhemisphericalDistribution,
+)
+from pyrecest.distributions.hypersphere_subset.hyperhemispherical_watson_distribution import (
+    HyperhemisphericalWatsonDistribution,
+)
+
+from .abstract_particle_filter import AbstractParticleFilter
+from beartype import beartype
+
+class HyperhemisphereCartProdParticleFilter(AbstractParticleFilter):
+    def __init__(
+        self, n_particles: int | np.int32 | np.int64, dim_hemisphere, n_hemispheres
+    ) -> None:
+        """
+        Constructor
+
+        Parameters:
+        n_particles (int > 0): Number of particles to use
+        dim (int > 0): Dimension
+        """
+        initial_filter_state = HyperhemisphereCartProdDiracDistribution(
+            np.empty((n_particles, (dim_hemisphere + 1) * n_hemispheres)),
+            ones(n_particles) / n_particles,
+            dim_hemisphere,
+            n_hemispheres,
+        )
+        AbstractParticleFilter.__init__(self, initial_filter_state=initial_filter_state)
+
+    def set_state(self, new_state):
+        """
+        Sets the current system state
+
+        Parameters:
+        dist_ (HyperhemisphericalDiracDistribution): New state
+        """
+        assert isinstance(new_state, AbstractHyperhemisphericalDistribution)
+        if not isinstance(new_state, HyperhemisphereCartProdDiracDistribution):
+            new_state = HyperhemisphereCartProdDiracDistribution(
+                new_state.sample(self.filter_state.d.shape[0]),
+                w=ones(self.filter_state.d.shape[0]) / self.filter_state.d.shape[0],
+                dim_hemisphere=self.filter_state.dim_hemisphere,
+                n_hemispheres=self.filter_state.n_hemispheres,
+            )
+        self.filter_state = new_state
+
+    @beartype
+    def predict_nonlinear_each_part(
+        self,
+        f,
+        # Limit noise distribution to ones that support set_mode
+        noise_distribution: HyperhemisphericalWatsonDistribution | None = None,
+        function_is_vectorized: bool = True,
+        shift_instead_of_add: bool = True,
+    ):
+        """
+        Predicts the next state for each hyperhemisphere
+        """
+        assert function_is_vectorized, "Only vectorized functions are supported"
+        assert (
+            noise_distribution is None or noise_distribution.dim == self.filter_state.dim_hemisphere
+        ), "Noise dimension must match state dimension in Cartesian product"
+        assert shift_instead_of_add, "Only shifting is supported"
+        for i in range(self.filter_state.n_hemispheres):
+            # Apply the function to each hyperhemisphere
+            index_arr = range(
+                i * (self.filter_state.dim_hemisphere + 1),
+                (i + 1) * (self.filter_state.dim_hemisphere + 1),
+            )
+            # Consider only part of the state of the current hemisphere
+            curr_d = self.filter_state.d[:, index_arr]
+            d_fun_applied = f(curr_d)
+            # Add noise
+            if noise_distribution is not None:
+                for j in range(self.filter_state.d.shape[0]):
+                    # Set mean to transformed state
+                    # This will fail if set_mean is unavailable
+                    noise_distribution.set_mode(d_fun_applied[j])
+                    # Sample one noise vector centered at the transformed state to add noise
+                    self.filter_state.d[j, index_arr] = noise_distribution.sample(1)
+
+    @property
+    def filter_state(self):
+        return self._filter_state
+
+    @filter_state.setter
+    def filter_state(self, new_state):
+        if isinstance(
+            new_state,
+            (
+                AbstractHyperhemisphericalDistribution,
+                AbstractHypersphericalDistribution,
+            ),
+        ):
+            assert new_state.dim == self.filter_state.dim_hemisphere
+            samples = new_state.sample(
+                self._filter_state.d.shape[0] * self._filter_state.n_hemispheres
+            )
+            if isinstance(new_state, AbstractHypersphericalDistribution):
+                samples[samples[:, -1] < 0] = -samples[samples[:, -1] < 0]
+            self._filter_state.d = samples.reshape(self.filter_state.d.shape)
+        else:
+            AbstractParticleFilter.filter_state.fset(self, new_state)
diff --git a/pyrecest/tests/filters/test_hyperhemisphere_cart_prod_particle_filter.py b/pyrecest/tests/filters/test_hyperhemisphere_cart_prod_particle_filter.py
@@ -0,0 +1,83 @@
+import unittest
+
+import pyrecest.backend  # pylint: disable=no-name-in-module,no-member
+from pyrecest.backend import abs, array, sum  # pylint: disable=no-name-in-module,no-member,redefined-builtin
+from pyrecest.distributions import VonMisesFisherDistribution
+from pyrecest.filters.hyperhemisphere_cart_prod_particle_filter import (
+    HyperhemisphereCartProdParticleFilter,
+)
+
+
+class HyperHemisphereCartProdParticleFilterTest(unittest.TestCase):
+    def test_init(self):
+        n_particles = 1000
+        dim_hemisphere = 3
+        n_hemispheres = 2
+        pf = HyperhemisphereCartProdParticleFilter(
+            n_particles, dim_hemisphere, n_hemispheres
+        )
+        self.assertEqual(
+            pf.filter_state.d.shape,
+            (n_particles, (dim_hemisphere + 1) * n_hemispheres),
+        )
+
+    @unittest.skipIf(
+        pyrecest.backend.__name__ in ("pyrecest.jax", "pyrecest.pytorch"), # pylint: disable=no-name-in-module,no-member
+        reason="Backend not supported'",
+    )
+    def test_set_state(self):
+        n_particles = 1000
+        dim_hemisphere = 3
+        n_hemispheres = 2
+        pf = HyperhemisphereCartProdParticleFilter(
+            n_particles, dim_hemisphere, n_hemispheres
+        )
+        self.assertEqual(
+            pf.filter_state.d.shape,
+            (n_particles, (dim_hemisphere + 1) * n_hemispheres),
+        )
+        pf.filter_state = VonMisesFisherDistribution(array([0.0, 0.0, 0.0, 1.0]), 1.0)
+
+    @unittest.skipIf(
+        pyrecest.backend.__name__ in ("pyrecest.jax", "pyrecest.pytorch"),  # pylint: disable=no-name-in-module,no-member
+        reason="Backend not supported'",
+    )
+    def test_predict(self):
+        n_particles = 1000
+        dim_hemisphere = 3
+        n_hemispheres = 2
+        pf = HyperhemisphereCartProdParticleFilter(
+            n_particles, dim_hemisphere, n_hemispheres
+        )
+        pf.filter_state = VonMisesFisherDistribution(array([0.0, 0.0, 0.0, 1.0]), 1.0)
+
+        noise_distribution = VonMisesFisherDistribution(
+            array([0.0, 0.0, 0.0, 1.0]), 1.0
+        )
+
+        def identity_function(x):
+            return x
+
+        pf.predict_nonlinear_each_part(identity_function, noise_distribution)
+
+    @unittest.skipIf(
+        pyrecest.backend.__name__ in ("pyrecest.jax", "pyrecest.pytorch"),  # pylint: disable=no-name-in-module,no-member
+        reason="Backend not supported'",
+    )
+    def test_update(self):
+        n_particles = 1000
+        dim_hemisphere = 3
+        n_hemispheres = 2
+        pf = HyperhemisphereCartProdParticleFilter(
+            n_particles, dim_hemisphere, n_hemispheres
+        )
+        pf.filter_state = VonMisesFisherDistribution(array([0.0, 0.0, 0.0, 1.0]), 1.0)
+
+        def likelihood_function(x):
+            return abs(sum(x, axis=1))  # noqa: A001
+
+        pf.update_nonlinear_using_likelihood(likelihood_function)
+
+
+if __name__ == "__main__":
+    unittest.main()
