Merge pull request #17 from aidos-lab/spherical-grid-for-3D-directions

Spherical grid for 3 d directions

Author: ErnstRoell

dect/directions.py

@@ -109,3 +109,46 @@ def generate_multiview_directions(num_thetas: int, d: int):
         multiview_dirs.append(v)
 
     return torch.hstack(multiview_dirs)
+
+
+def generate_spherical_grid_directions(num_thetas: int, num_phis: int, d: int = 3):
+    """
+    Generates a smooth spherical grid of directions on the unit sphere in 3D using
+    latitude–longitude (θ, φ) style sampling.
+
+    The directions are parameterized by θ (polar angle, [0, π]) and φ (azimuthal angle, [0, 2π)),
+    and returned as a tensor of shape [3, num_thetas * num_phis], with each column a unit vector.
+
+    Parameters
+    ----------
+    num_thetas: int
+        Number of θ samples (from 0 to π, inclusive).
+    num_phis: int
+        Number of φ samples (from 0 to 2π, exclusive).
+    d: int
+        Must be 3, as spherical coordinates are for 3D.
+
+    Returns
+    -------
+    Tensor of shape [3, num_thetas * num_phis] containing unit vectors on the sphere.
+    """
+    assert d == 3, "Spherical coordinates are only defined for d=3."
+
+    # Removes both poles.
+    theta = torch.linspace(0, torch.pi, num_thetas + 2)[1:-1]
+
+    # Removes one endpoint.
+    phi = torch.linspace(0, 2 * torch.pi, num_phis + 1)[
+        :-1
+    ]  # Induces endpoint=False behavior.
+
+    phi_grid, theta_grid = torch.meshgrid(
+        phi, theta, indexing="ij"
+    )  # shape [num_phis, num_thetas]
+    sin_theta = torch.sin(theta_grid)
+    x = sin_theta * torch.cos(phi_grid)
+    y = sin_theta * torch.sin(phi_grid)
+    z = torch.cos(theta_grid)
+    dirs = torch.stack([x, y, z], dim=0)  # [3, num_phis, num_thetas]
+    dirs = dirs.reshape(3, -1)  # [3, num_thetas*num_phis]
+    return dirs

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "pdoc>=15.0.1,<16",
     "torch-geometric>=2.6.1,<3",
     "geotorch>=0.3.0,<0.4",
+    "pyvista>=0.46.3",
 ]
 
 [dependency-groups]