Next

Author: FlorianPfaff

pyrecest/distributions/hypersphere_subset/hyperhemispherical_grid_distribution.py

@@ -1,4 +1,3 @@
-from matplotlib.pyplot import grid
 from .abstract_hypersphere_subset_grid_distribution import (
     AbstractHypersphereSubsetGridDistribution,
 )
@@ -20,16 +19,16 @@ class HyperhemisphericalGridDistribution(
     AbstractHypersphereSubsetGridDistribution, AbstractHyperhemisphericalDistribution
 ):
 
-    def __init__(self, grid_, grid_values, enforce_pdf_nonnegative=True):
+    def __init__(self, grid, grid_values, enforce_pdf_nonnegative=True):
         # Do not test norm precisely, only quick test if any coordinate exceeds 1.
-        assert all(abs(grid_) <= 1 + 1e-12), (
+        assert all(abs(grid) <= 1 + 1e-12), (
             "Grid points must not lie outside the unit ball."
         )
         assert all(
-            grid_[:, -1] >= -1e-12
+            grid[:, -1] >= -1e-12
         ), "Always using upper hemisphere (along last dimension)."
 
-        super().__init__(grid_, grid_values, enforce_pdf_nonnegative)
+        super().__init__(grid, grid_values, enforce_pdf_nonnegative)
 
     # ------------------------------------------------------------------
     # Basic functionality
@@ -59,9 +58,9 @@ def to_full_sphere(self):
         hyperspherical distribution normalized.
         """
         from .hyperspherical_grid_distribution import HypersphericalGridDistribution
-        grid_ = vstack((self.grid, -self.grid))
+        grid_full = vstack((self.grid, -self.grid))
         grid_values_ = 0.5 * concatenate((self.grid_values, self.grid_values))
-        hgd = HypersphericalGridDistribution(grid_, grid_values_)
+        hgd = HypersphericalGridDistribution(grid_full, grid_values_)
         return hgd
 
     def plot(self):
@@ -119,8 +118,8 @@ def get_closest_point(self, xs):
             raise ValueError("xs must be a 1D or 2D array.")
 
         # Distances to each grid point and its antipode.
-        diff1 = xs[:, None, :] - grid[None, :, :]  # (batch, n_grid, dim)
-        diff2 = xs[:, None, :] + grid[None, :, :]  # (batch, n_grid, dim)
+        diff1 = xs[:, None, :] - self.grid[None, :, :]  # (batch, n_grid, dim)
+        diff2 = xs[:, None, :] + self.grid[None, :, :]  # (batch, n_grid, dim)
 
         dists1 = linalg.norm(diff1, axis=2)  # (batch, n_grid)
         dists2 = linalg.norm(diff2, axis=2)  # (batch, n_grid)
@@ -189,8 +188,8 @@ def multiply(self: "HyperhemisphericalGridDistribution", other: "Hyperhemispheri
     @staticmethod
     def eq_point_set_upper_half(dim, n_points):
         ls = LeopardiSampler()
-        grid_, _ = ls.get_grid(n_points * 2, dim)
-        grid_upper_half = grid_[:grid_.shape[0]//2]
+        grid, _ = ls.get_grid(n_points * 2, dim)
+        grid_upper_half = grid[:grid.shape[0]//2]
         # To have upper half along last dim instead of first
         grid_upper_half[:, [0, -1]] = grid_upper_half[:, [-1, 0]]
         return grid_upper_half