Added (Hyperhemi)sphericalGridDistribution

Author: FlorianPfaff

pyrecest/_backend/__init__.py

@@ -198,6 +198,7 @@ def get_backend_name():
         "vmap",
         "gammaln",
         "round",
+        "array_equal",
         # For Riemannian score-based SDE
         "log1p"
     ],

pyrecest/_backend/jax/__init__.py

@@ -73,6 +73,7 @@
     vstack,
     where,
     zeros_like,
+    # For pyrecest
     diag,
     diff,
     apply_along_axis,
@@ -139,6 +140,7 @@
     linspace,
     ones,
     round,
+    array_equal,
     # For Riemannian score-based SDE
     log1p,
 )

pyrecest/_backend/numpy/__init__.py

@@ -98,6 +98,7 @@
     roll,
     dstack,
     round,
+    array_equal,
     # For Riemannian score-based SDE
     log1p,
 )

pyrecest/_backend/pytorch/__init__.py

@@ -72,6 +72,8 @@
     # For Riemannian score-based SDE
     log1p,
 )
+from torch import equal as array_equal  # For PyRecEst
+
 from torch import broadcast_tensors as broadcast_arrays
 from torch import repeat_interleave as repeat
 from torch.special import gammaln as _gammaln

pyrecest/distributions/abstract_grid_distribution.py

@@ -2,11 +2,12 @@
 import warnings
 from abc import abstractmethod
 
-import numpy as np
 from beartype import beartype
 
 from .abstract_distribution_type import AbstractDistributionType
 
+# pylint: disable=redefined-builtin,no-name-in-module,no-member
+from pyrecest.backend import mean, abs, any
 
 class AbstractGridDistribution(AbstractDistributionType):
     # pylint: disable=too-many-positional-arguments
@@ -23,10 +24,10 @@ def __init__(
             not grid_type == "custom" or grid is not None
         )  # if grid_type is custom, grid needs to be given
         assert (
-            grid is None or np.size(grid) == 0 or grid.shape[0] == grid_values.shape[0]
+            grid is None or grid.shape == () or grid.shape[0] == grid_values.shape[0]
         )
         assert (
-            grid is None or np.size(grid) == 0 or grid.ndim == 1 or grid.shape[1] == dim
+            grid is None or grid.shape == () or grid.ndim == 1 or grid.shape[1] == dim
         )
         if grid is None or grid.ndim > 1 and grid.shape[0] < grid.shape[1]:
             warnings.warn(
@@ -37,7 +38,7 @@ def __init__(
         self.grid = grid
         self.enforce_pdf_nonnegative = enforce_pdf_nonnegative
         # Overwrite with more descriptive parameterization
-        self.grid_density_description = {"n_grid_values": np.size(grid_values)}
+        self.grid_density_description = {"n_grid_values": grid_values.shape[0], "grid_type": grid_type}
 
     def pdf(self, xs):
         # Use nearest neighbor interpolation by default
@@ -47,7 +48,7 @@ def pdf(self, xs):
     @property
     def n_grid_points(self):
         # Overwrite if grid_values contains values that are not used as grid values
-        return np.size(self.grid_values)
+        return self.grid_values.shape[0]
 
     @abstractmethod
     def get_closest_point(self, xs):
@@ -61,19 +62,19 @@ def integrate(self, integration_boundaries=None):
         assert (
             integration_boundaries is None
         ), "Custom integration boundaries are currently not supported"
-        return self.get_manifold_size() * np.mean(self.grid_values)
+        return self.get_manifold_size() * mean(self.grid_values)
 
     def normalize_in_place(self, tol=1e-4, warn_unnorm=True):
         int_val = self.integrate()
-        if np.any(self.grid_values < 0):
+        if any(self.grid_values < 0):
             warnings.warn(
                 "Warning: There are negative values. This usually points to a user error."
             )
-        elif np.abs(int_val) < 1e-200:
+        elif abs(int_val) < 1e-200:
             raise ValueError(
                 "Sum of grid values is too close to zero, this usually points to a user error."
             )
-        elif np.abs(int_val - 1) > tol:
+        elif abs(int_val - 1) > tol:
             if warn_unnorm:
                 warnings.warn(
                     "Warning: Grid values apparently do not belong to a normalized density. Normalizing..."

pyrecest/distributions/hypersphere_subset/abstract_hypersphere_subset_grid_distribution.py

@@ -0,0 +1,79 @@
+import warnings
+
+from ..abstract_grid_distribution import AbstractGridDistribution 
+from .abstract_hypersphere_subset_distribution import AbstractHypersphereSubsetDistribution
+from .abstract_hyperhemispherical_distribution import AbstractHyperhemisphericalDistribution
+from .abstract_hyperspherical_distribution import AbstractHypersphericalDistribution
+from .von_mises_fisher_distribution import VonMisesFisherDistribution
+from .bingham_distribution import BinghamDistribution
+from .hyperspherical_mixture import HypersphericalMixture
+from .watson_distribution import WatsonDistribution
+from beartype import beartype
+
+# pylint: disable=redefined-builtin,no-name-in-module,no-member
+from pyrecest.backend import array_equal, argmax, sum
+
+class AbstractHypersphereSubsetGridDistribution(AbstractGridDistribution, AbstractHypersphereSubsetDistribution):
+    
+    def __init__(self, grid, grid_values, enforce_pdf_nonnegative=True):
+        # Check size consistency
+        if grid.shape[0] != grid_values.shape[0]:
+            raise ValueError("Grid size must match number of grid values.")
+            
+        AbstractGridDistribution.__init__(self, grid_values, grid_type = "unknown", grid=grid, dim=grid.shape[1],  enforce_pdf_nonnegative=enforce_pdf_nonnegative)     
+        AbstractHypersphereSubsetDistribution.__init__(self, dim=grid.shape[1])
+        self.normalize() 
+
+    def mean_direction(self):
+        warnings.warn("For hyperhemispheres, this function yields the mode and not the mean.", UserWarning)
+        # If we took the mean, it would be biased toward [0;...;0;1]
+        # because the lower half is considered inexistant.
+        index_max = argmax(self.grid_values)
+        mu = self.get_grid_point(index_max)
+        return mu
+
+    def moment(self):
+        weights = self.grid_values / sum(self.grid_values) # (N,)
+        
+        weighted_grid = self.get_grid() * weights 
+
+        C = weighted_grid * (self.get_grid().T @ self.get_grid())
+        return C
+
+    @beartype
+    def multiply(self: "AbstractHypersphereSubsetGridDistribution", other: "AbstractHypersphereSubsetGridDistribution") -> "AbstractHypersphereSubsetGridDistribution":
+        # Check for grid compatibility
+        if not array_equal(self.get_grid(), other.get_grid()):
+            raise ValueError("Can only multiply for equal grids. Grids are incompatible.")
+        
+        # Delegates multiplication logic to AbstractGridDistribution
+        return super().multiply(other)
+
+    @staticmethod
+    def from_distribution(distribution, no_of_grid_points, grid_type='healpix'):
+        # Import here to avoid circular imports
+        from .hyperhemispherical_grid_distribution import HyperhemisphericalGridDistribution
+        from .hyperspherical_grid_distribution import HypersphericalGridDistribution
+        # pylint: disable=too-many-boolean-expressions
+        if isinstance(distribution, AbstractHyperhemisphericalDistribution):
+            fun = distribution.pdf
+        elif (isinstance(distribution, (WatsonDistribution, BinghamDistribution)) or 
+              (isinstance(distribution, VonMisesFisherDistribution) and distribution.mu[-1] == 0) or 
+              (isinstance(distribution, HypersphericalMixture) and
+                len(distribution.dists) == 2 and all(w == 0.5 for w in distribution.w) and
+                array_equal(distribution.dists[1].mu, -distribution.dists[0].mu))):
+            def fun(x):
+                return 2 * distribution.pdf(x)
+        elif isinstance(distribution, HypersphericalGridDistribution):
+            raise ValueError('Converting a HypersphericalGridDistribution to a HyperhemisphericalGridDistribution is not supported')
+        elif isinstance(distribution, AbstractHypersphericalDistribution):
+            warnings.warn('Approximating a hyperspherical distribution on a hemisphere. The density may not be symmetric. Double check if this is intentional.',
+                          UserWarning)
+            def fun(x):
+                return 2 * distribution.pdf(x)
+        else:
+            raise ValueError('Distribution currently not supported.')
+            
+        sgd = HyperhemisphericalGridDistribution.from_function(fun, no_of_grid_points, distribution.dim, grid_type)
+        return sgd
+

pyrecest/distributions/hypersphere_subset/bingham_distribution.py

@@ -2,7 +2,7 @@
 
 # pylint: disable=redefined-builtin,no-name-in-module,no-member
 # pylint: disable=no-name-in-module,no-member
-from pyrecest.backend import abs, all, argsort, diag, exp, eye, linalg, max, sum, zeros
+from pyrecest.backend import abs, all, argsort, diag, exp, eye, linalg, max, sum, zeros, array
 from scipy.integrate import quad
 from scipy.special import iv
 
@@ -20,7 +20,7 @@ def __init__(self, Z, M):
         assert all(Z[:-1] <= Z[1:]), "Values in Z have to be ascending"
 
         # Verify that M is orthogonal
-        epsilon = 0.001
+        epsilon = array(0.001)
         assert max(abs(M @ M.T - eye(self.dim + 1))) < epsilon, "M is not orthogonal"
 
         self.Z = Z