Pixel Kernels (3/2): RGBD Kernels (#152)

While working on the image kernels, I realized that it might be useful
to have a unified class that sample the RGBD values jointly, so I'm
introducing this utility class that takes in an color kernel (from
#147) and a depth kernel (from
#149) to create an RGBD kernel,
which supports sample and logpdf computation.

This is a general class that's agnostic of the type of RGB and depth
kernels -- it'll pass in the arguments it receives to _both_ of the
kernels during `sample` and `logpdf` call. Since both color and depth
kernels have `*args, **kwargs` as part of their function signature, they
should simply ignore additional arguments that's not relevant to them.

(I also just realized that I was referring to these classes as "kernels"
but the actual class names are `*Distributions`... maybe I should fix
these in a future PR.)

I'm submitting this PR to `main` for now. Once I'm wrapping up my local
changes, I'm going to work on resolving merge conflicts with `gen3d`.

## Test Plan

Similar to the previous PRs, I've added some unit tests to make sure
that the kernels roughly have the behaviors that we expected:

```bash
pytest tests/gen3d/test_pixel_rgbd_kernels.py
```

Author: horizon-blue

src/b3d/chisight/gen3d/pixel_kernels/pixel_color_kernels.py

@@ -129,15 +129,15 @@ def logpdf_per_channel(
 class MixturePixelColorDistribution(PixelColorDistribution):
     """A distribution that generates the color of a pixel from a mixture of a
     truncated Laplace distribution centered around the latent color (inlier
-    branch) and a uniform distribution (outlier branch). The mixture is
-    controlled by the outlier_prob parameter. The support of the
+    branch) and a uniform distribution (occluded branch). The mixture is
+    controlled by the occluded_prob parameter. The support of the
     distribution is ([0, 1]^3).
     """
 
     color_scale: float
 
     @property
-    def _outlier_dist(self) -> PixelColorDistribution:
+    def _occluded_dist(self) -> PixelColorDistribution:
         return UniformPixelColorDistribution()
 
     @property
@@ -146,35 +146,35 @@ def _inlier_dist(self) -> PixelColorDistribution:
 
     @property
     def _mixture_dists(self) -> tuple[PixelColorDistribution, PixelColorDistribution]:
-        return (self._outlier_dist, self._inlier_dist)
+        return (self._occluded_dist, self._inlier_dist)
 
-    def _get_mix_ratio(self, outlier_prob: float) -> FloatArray:
-        return jnp.array((outlier_prob, 1 - outlier_prob))
+    def _get_mix_ratio(self, occluded_prob: float) -> FloatArray:
+        return jnp.array((occluded_prob, 1 - occluded_prob))
 
     def sample(
         self,
         key: PRNGKey,
         latent_color: FloatArray,
-        outlier_prob: float,
+        occluded_prob: float,
         *args,
         **kwargs,
     ) -> FloatArray:
         return PythonMixtureDistribution(self._mixture_dists).sample(
-            key, self._get_mix_ratio(outlier_prob), [(), (latent_color,)]
+            key, self._get_mix_ratio(occluded_prob), [(), (latent_color,)]
         )
 
     def logpdf_per_channel(
         self,
         observed_color: FloatArray,
         latent_color: FloatArray,
-        outlier_prob: float,
+        occluded_prob: float,
         *args,
         **kwargs,
     ) -> FloatArray:
         # Since the mixture model class does not keep the per-channel information,
         # we have to redefine this method to allow for testing
         logprobs = []
-        for dist, prob in zip(self._mixture_dists, self._get_mix_ratio(outlier_prob)):
+        for dist, prob in zip(self._mixture_dists, self._get_mix_ratio(occluded_prob)):
             logprobs.append(
                 dist.logpdf_per_channel(observed_color, latent_color) + jnp.log(prob)
             )
@@ -192,7 +192,7 @@ class FullPixelColorDistribution(PixelColorDistribution):
     else:
         color ~ mixture(
             [uniform(0, 1), truncated_laplace(latent_color; color_scale)],
-            [outlier_prob, 1 - outlier_prob]
+            [occluded_prob, 1 - occluded_prob]
         )
 
     Constructor args:
@@ -203,9 +203,9 @@ class FullPixelColorDistribution(PixelColorDistribution):
     - `latent_color`: 3-array.  If no latent point hits the pixel, should contain
         3 negative values.  If a latent point hits the pixel, should contain the point's
         color as an RGB value in [0, 1]^3.
-    - `color_outlier_prob`: float.  If a latent point hits the pixel, should contain
+    - `color_occluded_prob`: float.  If a latent point hits the pixel, should contain
         the probability associated with that point that the generated color is
-        an outlier.  If no latent point hits the pixel, this value is ignored.
+        an occluded.  If no latent point hits the pixel, this value is ignored.
 
     Distribution support:
         - An RGB value in [0, 1]^3.
@@ -225,7 +225,7 @@ def sample(
         self,
         key: PRNGKey,
         latent_color: FloatArray,
-        outlier_prob: FloatArray,
+        occluded_prob: FloatArray,
         *args,
         **kwargs,
     ) -> FloatArray:
@@ -236,14 +236,14 @@ def sample(
             # sample args
             key,
             latent_color,
-            outlier_prob,
+            occluded_prob,
         )
 
     def logpdf_per_channel(
         self,
         observed_color: FloatArray,
         latent_color: FloatArray,
-        outlier_prob: float,
+        occluded_prob: float,
         *args,
         **kwargs,
     ) -> FloatArray:
@@ -254,5 +254,5 @@ def logpdf_per_channel(
             # logpdf args
             observed_color,
             latent_color,
-            outlier_prob,
+            occluded_prob,
         )

src/b3d/chisight/gen3d/pixel_kernels/pixel_depth_kernels.py

@@ -114,7 +114,7 @@ class MixturePixelDepthDistribution(PixelDepthDistribution):
     """A distribution that generates the depth of a pixel from
     mixture(
         [delta(-1), uniform(near, far), laplace(latent_depth; depth_scale)],
-        [depth_nonreturn_prob, (1 - depth_nonreturn_prob) * outlier_prob, remaining_prob]
+        [depth_nonreturn_prob, (1 - depth_nonreturn_prob) * occluded_prob, remaining_prob]
     )
 
     The support of the distribution is [near, far] ∪ { "nonreturn" }.
@@ -129,7 +129,7 @@ def _nonreturn_dist(self) -> PixelDepthDistribution:
         return DeltaDistribution(DEPTH_NONRETURN_VAL)
 
     @property
-    def _outlier_dist(self) -> PixelDepthDistribution:
+    def _occluded_dist(self) -> PixelDepthDistribution:
         return UniformPixelDepthDistribution(self.near, self.far)
 
     @property
@@ -141,47 +141,47 @@ def _inlier_dist(self) -> PixelDepthDistribution:
     @property
     def _mixture_dist(self) -> PythonMixtureDistribution:
         return PythonMixtureDistribution(
-            (self._nonreturn_dist, self._outlier_dist, self._inlier_dist)
+            (self._nonreturn_dist, self._occluded_dist, self._inlier_dist)
         )
 
     def _get_mix_ratio(
-        self, depth_nonreturn_prob: float, outlier_prob: float
+        self, occluded_prob: float, depth_nonreturn_prob: float
     ) -> FloatArray:
         return jnp.array(
             (
                 depth_nonreturn_prob,
-                (1 - depth_nonreturn_prob) * outlier_prob,
-                (1 - depth_nonreturn_prob) * (1 - outlier_prob),
+                (1 - depth_nonreturn_prob) * occluded_prob,
+                (1 - depth_nonreturn_prob) * (1 - occluded_prob),
             )
         )
 
     def sample(
         self,
         key: PRNGKey,
         latent_depth: float,
+        occluded_prob: float,
         depth_nonreturn_prob: float,
-        outlier_prob: float,
         *args,
         **kwargs,
     ) -> float:
         return self._mixture_dist.sample(
             key,
-            self._get_mix_ratio(depth_nonreturn_prob, outlier_prob),
+            self._get_mix_ratio(occluded_prob, depth_nonreturn_prob),
             [(), (), (latent_depth,)],
         )
 
     def logpdf(
         self,
         observed_depth: float,
         latent_depth: float,
+        occluded_prob: float,
         depth_nonreturn_prob: float,
-        outlier_prob: float,
         *args,
         **kwargs,
     ) -> float:
         return self._mixture_dist.logpdf(
             observed_depth,
-            self._get_mix_ratio(depth_nonreturn_prob, outlier_prob),
+            self._get_mix_ratio(occluded_prob, depth_nonreturn_prob),
             [(), (), (latent_depth,)],
         )
 
@@ -255,7 +255,7 @@ class FullPixelDepthDistribution(PixelDepthDistribution):
     else:
         mixture(
             [delta(-1), uniform(near, far), laplace(latent_depth; depth_scale)],
-            [depth_nonreturn_prob, (1 - depth_nonreturn_prob) * outlier_prob, remaining_prob]
+            [depth_nonreturn_prob, (1 - depth_nonreturn_prob) * occluded_prob, remaining_prob]
         )
     """
 
@@ -275,8 +275,8 @@ def sample(
         self,
         key: PRNGKey,
         latent_depth: FloatArray,
+        occluded_prob: FloatArray,
         depth_nonreturn_prob: float,
-        outlier_prob: FloatArray,
         *args,
         **kwargs,
     ) -> FloatArray:
@@ -287,16 +287,16 @@ def sample(
             # sample args
             key,
             latent_depth,
+            occluded_prob,
             depth_nonreturn_prob,
-            outlier_prob,
         )
 
     def logpdf(
         self,
         observed_depth: FloatArray,
         latent_depth: FloatArray,
+        occluded_prob: float,
         depth_nonreturn_prob: float,
-        outlier_prob: float,
         *args,
         **kwargs,
     ) -> FloatArray:
@@ -307,6 +307,6 @@ def logpdf(
             # logpdf args
             observed_depth,
             latent_depth,
+            occluded_prob,
             depth_nonreturn_prob,
-            outlier_prob,
         )

src/b3d/chisight/gen3d/pixel_kernels/pixel_rgbd_kernels.py

@@ -0,0 +1,37 @@
+import genjax
+import jax
+import jax.numpy as jnp
+from genjax import Pytree
+from genjax.typing import FloatArray, PRNGKey
+
+from b3d.chisight.gen3d.pixel_kernels.pixel_color_kernels import PixelColorDistribution
+from b3d.chisight.gen3d.pixel_kernels.pixel_depth_kernels import PixelDepthDistribution
+
+
+@Pytree.dataclass
+class PixelRGBDDistribution(genjax.ExactDensity):
+    color_kernel: PixelColorDistribution
+    depth_kernel: PixelDepthDistribution
+
+    def sample(
+        self, key: PRNGKey, latent_rgbd: FloatArray, *args, **kwargs
+    ) -> FloatArray:
+        keys = jax.random.split(key, 2)
+        observed_color = self.color_kernel.sample(
+            keys[0], latent_rgbd[:3], *args, **kwargs
+        )
+        observed_depth = self.depth_kernel.sample(
+            keys[1], latent_rgbd[3], *args, **kwargs
+        )
+        return jnp.append(observed_color, observed_depth)
+
+    def logpdf(
+        self, observed_rgbd: FloatArray, latent_rgbd: FloatArray, *args, **kwargs
+    ) -> float:
+        color_logpdf = self.color_kernel.logpdf(
+            observed_rgbd[:3], latent_rgbd[:3], *args, **kwargs
+        )
+        depth_logpdf = self.depth_kernel.logpdf(
+            observed_rgbd[3], latent_rgbd[3], *args, **kwargs
+        )
+        return color_logpdf + depth_logpdf

tests/gen3d/test_pixel_color_kernels.py

@@ -36,8 +36,8 @@ def generate_color_grid(n_grid_steps: int):
 sample_kernels_to_test = [
     (UniformPixelColorDistribution(), ()),
     (TruncatedLaplacePixelColorDistribution(0.1), ()),
-    (MixturePixelColorDistribution(0.3), (0.5,)),  # outlier_prob
-    (FullPixelColorDistribution(0.5), (0.3,)),  # outlier_prob
+    (MixturePixelColorDistribution(0.3), (0.5,)),  # occluded_prob
+    (FullPixelColorDistribution(0.5), (0.3,)),  # occluded_prob
 ]
 
 
@@ -80,15 +80,15 @@ def test_relative_logpdf():
     latent_color = -jnp.ones(3)  # use -1 to denote invalid pixel
     logpdf_1 = kernel.logpdf(obs_color, latent_color, 0.2)
     logpdf_2 = kernel.logpdf(obs_color, latent_color, 0.8)
-    # the logpdf should be the same because the outlier probability is not used
+    # the logpdf should be the same because the occluded probability is not used
     # in the case when no color hit the pixel
     assert jnp.allclose(logpdf_1, logpdf_2)
 
     # case 2: a color hit the pixel, but the color is not close to the observed color
     latent_color = jnp.array([1.0, 0.5, 0.0])
     logpdf_3 = kernel.logpdf(obs_color, latent_color, 0.2)
     logpdf_4 = kernel.logpdf(obs_color, latent_color, 0.8)
-    # the pixel should be more likely to be an outlier
+    # the pixel should be more likely to be an occluded
     assert logpdf_3 < logpdf_4
 
     # case 3: a color hit the pixel, and the color is close to the observed color

tests/gen3d/test_pixel_depth_kernels.py

@@ -22,15 +22,15 @@
     (
         MixturePixelDepthDistribution(near, far, 0.15),
         (
+            0.5,  # occluded_prob
             0.23,  # depth_nonreturn_prob
-            0.5,  # outlier_prob
         ),
     ),
     (
         FullPixelDepthDistribution(near, far, 0.5),
         (
+            0.3,  # occluded_prob
             0.1,  # depth_nonreturn_prob
-            0.3,  # outlier_prob
         ),
     ),
 ]
@@ -81,26 +81,26 @@ def test_relative_logpdf():
     obs_depth = DEPTH_NONRETURN_VAL
     latent_depth = DEPTH_NONRETURN_VAL
     depth_nonreturn_prob = 0.2
-    logpdf_1 = kernel.logpdf(obs_depth, latent_depth, depth_nonreturn_prob, 0.8)
+    logpdf_1 = kernel.logpdf(obs_depth, latent_depth, 0.8, depth_nonreturn_prob)
     assert logpdf_1 == jnp.log(depth_nonreturn_prob)
 
     latent_depth = -1.0  # no depth information from latent
-    logpdf_2 = kernel.logpdf(obs_depth, latent_depth, depth_nonreturn_prob, 0.8)
+    logpdf_2 = kernel.logpdf(obs_depth, latent_depth, 0.8, depth_nonreturn_prob)
     # nonreturn obs cannot be generates from latent that is not nonreturn
     assert logpdf_2 == jnp.log(UNEXPLAINED_DEPTH_NONRETURN_PROB)
 
     # case 2: valid depth is observed, but latent depth is far from the observed depth
     obs_depth = 10.0
     latent_depth = 0.01
-    logpdf_3 = kernel.logpdf(obs_depth, latent_depth, depth_nonreturn_prob, 0.9)
-    logpdf_4 = kernel.logpdf(obs_depth, latent_depth, depth_nonreturn_prob, 0.1)
-    # the pixel should be more likely to be an outlier
+    logpdf_3 = kernel.logpdf(obs_depth, latent_depth, 0.9, depth_nonreturn_prob)
+    logpdf_4 = kernel.logpdf(obs_depth, latent_depth, 0.1, depth_nonreturn_prob)
+    # the pixel should be more likely to be an occluded
     assert logpdf_3 > logpdf_4
 
     # case 3: valid depth is observed, but latent depth is close from the observed depth
     obs_depth = 6.0
     latent_depth = 6.01
-    logpdf_5 = kernel.logpdf(obs_depth, latent_depth, depth_nonreturn_prob, 0.9)
-    logpdf_6 = kernel.logpdf(obs_depth, latent_depth, depth_nonreturn_prob, 0.1)
+    logpdf_5 = kernel.logpdf(obs_depth, latent_depth, 0.9, depth_nonreturn_prob)
+    logpdf_6 = kernel.logpdf(obs_depth, latent_depth, 0.1, depth_nonreturn_prob)
     # the pixel should be more likely to be an inliner
     assert logpdf_5 < logpdf_6