Add gausian frustum calculations.

pyrad/cameras/rays.py

@@ -7,6 +7,7 @@ from typing import Optional
 
 import torch
 from torchtyping import TensorType
+from pyrad.utils.math import Gaussians, conical_frustum_to_gaussian
 
 from pyrad.utils.misc import is_not_none
 
@@ -30,9 +31,29 @@ class Frustums:
     pixel_area: TensorType[..., 1]
 
     def get_positions(self) -> TensorType[..., 3]:
-        """Returns "center" position of frustum. Not weighted by mass."""
+        """Calulates "center" position of frustum. Not weighted by mass.
+
+        Returns:
+            TensorType[..., 3]: xyz positions.
+        """
         return self.origins + self.directions * (self.frustum_starts + self.frustum_ends) / 2
 
+    def get_gaussian_blob(self) -> Gaussians:
+        """Calculates of guassian approximation of conical frustum.
+
+        Resturns:
+            Gaussians: Conical frustums approximated by gaussian distribution.
+        """
+        # Cone radius is set such that the square pixel_area matches the cone area.
+        cone_radius = torch.sqrt(self.pixel_area) / 1.7724538509055159  # r = sqrt(pixel_area / pi)
+        return conical_frustum_to_gaussian(
+            origins=self.origins,
+            directions=self.directions,
+            starts=self.frustum_starts,
+            ends=self.frustum_ends,
+            radius=cone_radius,
+        )
+
     @classmethod
     def get_mock_frustum(cls) -> "Frustums":
         """Helper function to generate a placeholder frustum.
@@ -40,13 +61,12 @@ class Frustums:
         Returns:
             Frustums: A size 1 frustum with meaningless values.
         """
-        device = cls.origins.device
         return Frustums(
-            origins=torch.ones((1, 3), device=device),
-            directions=torch.ones((1, 3), device=device),
-            frustum_starts=torch.ones((1, 1), device=device),
-            frustum_ends=torch.ones((1, 1), device=device) + 1,
-            pixel_area=torch.ones((1, 1), device=device),
+            origins=torch.ones((1, 3)),
+            directions=torch.ones((1, 3)),
+            frustum_starts=torch.ones((1, 1)),
+            frustum_ends=torch.ones((1, 1)),
+            pixel_area=torch.ones((1, 1)),
         )
 
     def apply_masks(self, mask: TensorType) -> "Frustums":

pyrad/utils/math.py

@@ -1,5 +1,6 @@
 """ Math Helper Functions """
 
+from dataclasses import dataclass
 import torch
 from torchtyping import TensorType
 
@@ -67,3 +68,98 @@ def components_from_spherical_harmonics(levels: int, directions: TensorType[...,
         components[..., 24] = 0.4425326924449826 * (xx * (xx - 3 * yy) - yy * (3 * xx - yy))
 
     return components
+
+
+@dataclass
+class Gaussians:
+    """Stores Gaussians
+
+    Args:
+        mean (TensorType[..., "dim"]): Mean of multivariate Gaussian
+        cov (TensorType[..., "dim", "dim"]): Covariance of multivariate Gaussian.
+    """
+
+    mean: TensorType[..., "dim"]
+    cov: TensorType[..., "dim", "dim"]
+
+
+def compute_3d_gaussian(
+    directions: TensorType[..., 3],
+    means: TensorType[..., 3],
+    dir_variance: TensorType[..., 1],
+    radius_variance: TensorType[..., 1],
+) -> Gaussians:
+    """Compute guassian along ray.
+
+    Args:
+        directions (TensorType[..., 3]): Axis of Gaussian.
+        means (TensorType[..., 3]): Mean of Gaussian.
+        dir_variance (TensorType[..., 3]): Variance along direction axis.
+        radius_variance (TensorType[..., 3]): Variance tangent to direction axis.
+
+    Returns:
+        Gaussians: Oriented 3D gaussian.
+    """
+
+    dir_outer_product = directions[..., :, None] * directions[..., None, :]
+    eye = torch.eye(directions.shape[-1])
+    dir_mag_sq = torch.clamp(torch.sum(directions**2, axis=-1, keepdim=True), min=1e-10)
+    null_outer_product = eye - directions[..., :, None] * (directions / dir_mag_sq)[..., None, :]
+    dir_cov_diag = dir_variance[..., None] * dir_outer_product[..., :, :]
+    radius_cov_diag = radius_variance[..., None] * null_outer_product[..., :, :]
+    cov = dir_cov_diag + radius_cov_diag
+    return Gaussians(mean=means, cov=cov)
+
+
+def cylinder_to_gaussian(
+    origins: TensorType[..., 3],
+    directions: TensorType[..., 3],
+    starts: TensorType[..., 1],
+    ends: TensorType[..., 1],
+    radius: TensorType[..., 1],
+) -> Gaussians:
+    """Approximates cylinders with a Gaussian distributions.
+
+    Args:
+        origins (TensorType[..., 3]): Origins of cylinders.
+        directions (TensorType[..., 3]): Direction (axis) of cylinders.
+        starts (TensorType[..., 1]): Start of cylinders.
+        ends (TensorType[..., 1]): End of cylinders.
+        radius (TensorType[..., 1]): Radii of cylinders.
+
+    Returns:
+        Gaussians: Approximation of cylinders
+    """
+    means = origins + directions * ((starts + ends) / 2.0)
+    dir_variance = (ends - starts) ** 2 / 12
+    radius_variance = radius**2 / 4.0
+    return compute_3d_gaussian(directions, means, dir_variance, radius_variance)
+
+
+def conical_frustum_to_gaussian(
+    origins: TensorType[..., 3],
+    directions: TensorType[..., 3],
+    starts: TensorType[..., 1],
+    ends: TensorType[..., 1],
+    radius: TensorType[..., 1],
+) -> Gaussians:
+    """Approximates conical frustums with a Gaussian distributions.
+
+    Uses stable parameterization described in mip-NeRF publication.
+
+    Args:
+        origins (TensorType[..., 3]): Origins of cones.
+        directions (TensorType[..., 3]): Direction (axis) of frustums.
+        starts (TensorType[..., 1]): Start of conical frustums.
+        ends (TensorType[..., 1]): End of conical frustums.
+        radius (TensorType[..., 1]): Radii of cone a distance of 1 from the origin.
+
+    Returns:
+        Gaussians: Approximation of conical frustums
+    """
+    mu = (starts + ends) / 2.0
+    hw = (ends - starts) / 2.0
+    means = origins + directions * (mu + (2.0 * mu * hw**2.0) / (3.0 * mu**2.0 + hw**2.0))
+    dir_variance = (hw**2) / 3 - (4 / 15) * ((hw**4 * (12 * mu**2 - hw**2)) / (3 * mu**2 + hw**2) ** 2)
+    radius_variance = radius**2 * ((mu**2) / 4 + (5 / 12) * hw**2 - 4 / 15 * (hw**4) / (3 * mu**2 + hw**2))
+    return compute_3d_gaussian(directions, means, dir_variance, radius_variance)

tests/cameras/test_rays.py

@@ -0,0 +1,74 @@
+"""
+Test the ray classes.
+"""
+
+import pytest
+import torch
+from pyrad.cameras.rays import Frustums
+
+
+def test_frustum_get_position():
+    """Test position calculation"""
+
+    origin = torch.Tensor([0, 1, 2])[None, ...]
+    direction = torch.Tensor([0, 1, 0])[None, ...]
+    frustum_start = torch.Tensor([2])[None, ...]
+    frustum_end = torch.Tensor([3])[None, ...]
+
+    target_position = torch.Tensor([0, 3.5, 2])[None, ...]
+
+    frustum = Frustums(
+        origins=origin,
+        directions=direction,
+        frustum_starts=frustum_start,
+        frustum_ends=frustum_end,
+        pixel_area=torch.ones((1, 1)),
+    )
+
+    positions = frustum.get_positions()
+    assert positions == pytest.approx(target_position, abs=1e-6)
+
+
+def test_frustum_get_gaussian_blob():
+    """Test gaussian blob calculation"""
+
+    frustum = Frustums(
+        origins=torch.ones((5, 3)),
+        directions=torch.ones((5, 3)),
+        frustum_starts=torch.ones((5, 1)),
+        frustum_ends=torch.ones((5, 1)),
+        pixel_area=torch.ones((5, 1)),
+    )
+
+    gaussian_blob = frustum.get_gaussian_blob()
+    assert gaussian_blob.mean.shape == (5, 3)
+    assert gaussian_blob.cov.shape == (5, 3, 3)
+
+
+def test_frustum_apply_masks():
+    """Test masking frustum"""
+    frustum = Frustums(
+        origins=torch.ones((5, 3)),
+        directions=torch.ones((5, 3)),
+        frustum_starts=torch.ones((5, 1)),
+        frustum_ends=torch.ones((5, 1)),
+        pixel_area=torch.ones((5, 1)),
+    )
+
+    mask = torch.tensor([False, True, False, True, True], dtype=torch.bool)
+    frustum = frustum.apply_masks(mask)
+
+    assert frustum.origins.shape == (3, 3)
+    assert frustum.directions.shape == (3, 3)
+    assert frustum.frustum_starts.shape == (3, 1)
+    assert frustum.frustum_ends.shape == (3, 1)
+    assert frustum.pixel_area.shape == (3, 1)
+
+
+def test_get_mock_frustum():
+    """Test creation of mock frustum"""
+    Frustums.get_mock_frustum()
+
+
+if __name__ == "__main__":
+    test_frustum_get_gaussian_blob()