Fix hyperhemisphere component state assignment

src/pyrecest/filters/hyperhemisphere_cart_prod_particle_filter.py

@@ -4,7 +4,7 @@
 from beartype import beartype
 
 # pylint: disable=no-name-in-module,no-member
-from pyrecest.backend import empty, ones
+from pyrecest.backend import empty, ones, ones_like
 from pyrecest.distributions import AbstractHypersphericalDistribution
 from pyrecest.distributions.cart_prod.hyperhemisphere_cart_prod_dirac_distribution import (
     HyperhemisphereCartProdDiracDistribution,
@@ -81,13 +81,6 @@ def set_state(self, new_state):
                 )
             self.filter_state = new_state
             return
-        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
@@ -163,5 +156,8 @@ def filter_state(self, new_state):
             if isinstance(new_state, AbstractHypersphericalDistribution):
                 samples[samples[:, -1] < 0] = -samples[samples[:, -1] < 0]
             self._filter_state.d = samples.reshape(self.filter_state.d.shape)
+            self._filter_state.w = (
+                ones_like(self._filter_state.w) / self._filter_state.w.shape[0]
+            )
         else:
             AbstractParticleFilter.filter_state.fset(self, new_state)

tests/filters/test_hyperhemisphere_cart_prod_component_state_assignment.py

@@ -0,0 +1,80 @@
+import unittest
+
+import numpy as np
+
+import pyrecest.backend  # pylint: disable=no-name-in-module,no-member
+from pyrecest.backend import array  # pylint: disable=no-name-in-module,no-member
+from pyrecest.distributions.cart_prod.hyperhemisphere_cart_prod_dirac_distribution import (
+    HyperhemisphereCartProdDiracDistribution,
+)
+from pyrecest.distributions.hypersphere_subset.hyperhemispherical_watson_distribution import (
+    HyperhemisphericalWatsonDistribution,
+)
+from pyrecest.filters.hyperhemisphere_cart_prod_particle_filter import (
+    HyperhemisphereCartProdParticleFilter,
+)
+
+
+@unittest.skipIf(
+    pyrecest.backend.__backend_name__  # pylint: disable=no-name-in-module,no-member
+    in ("jax", "pytorch"),
+    reason="Backend not supported",
+)
+class HyperhemisphereCartProdComponentStateAssignmentTest(unittest.TestCase):
+    @staticmethod
+    def _filter_with_nonuniform_weights():
+        particle_filter = HyperhemisphereCartProdParticleFilter(4, 2, 2)
+        particles = array(
+            [
+                [0.0, 0.0, 1.0, 0.0, 0.0, 1.0],
+                [1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0, 1.0, 0.0],
+                [0.0, 0.0, 1.0, 1.0, 0.0, 0.0],
+            ]
+        )
+        state = HyperhemisphereCartProdDiracDistribution(
+            particles,
+            w=array([0.7, 0.1, 0.1, 0.1]),
+            dim_hemisphere=2,
+            n_hemispheres=2,
+        )
+        particle_filter.set_state(state)
+        return particle_filter
+
+    @staticmethod
+    def _component_distribution():
+        return HyperhemisphericalWatsonDistribution(
+            array([0.0, 0.0, 1.0]), 2.0
+        )
+
+    def test_set_state_expands_component_distribution_across_product(self):
+        particle_filter = self._filter_with_nonuniform_weights()
+
+        particle_filter.set_state(self._component_distribution())
+
+        self.assertEqual(particle_filter.filter_state.d.shape, (4, 6))
+        self.assertEqual(
+            particle_filter.filter_state.as_component_array().shape, (4, 2, 3)
+        )
+        np.testing.assert_allclose(
+            particle_filter.filter_state.w, np.full(4, 0.25)
+        )
+        self.assertTrue(
+            np.all(
+                np.asarray(particle_filter.filter_state.as_component_array())[..., -1]
+                >= 0.0
+            )
+        )
+
+    def test_filter_state_distribution_assignment_resets_weights(self):
+        particle_filter = self._filter_with_nonuniform_weights()
+
+        particle_filter.filter_state = self._component_distribution()
+
+        np.testing.assert_allclose(
+            particle_filter.filter_state.w, np.full(4, 0.25)
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()