Smoothing Discontinuities for Fully Differentiable Ray Tracing

In a previous work [2], we introduce a smoothing technique that aims at smoothing the discontinuities caused by Ray Tracing in the hope to make optimization through gradient descent easier. This technique was originally implemented in [5], and we now aim to reproduce a similar implementation inside DiffeRT, but with less flexibility to reduce code complexity.

We will not go into details of how the method actually work, nor the actual implementation details, but you can always read the source code or the original paper [2] for more information.

Important

This tutorial is still TODO, but you can find premise of it below.

Smoothing Function

The smoothing function can be any smooth function \(s: (x;\alpha) \in \mathbb{R} \times \mathbb{R}^+ \mapsto s(x;\alpha) \in [0;1]\), with \(s \in C^1\), with the following properties on the smoothing factor \(\alpha\):

\[ \lim_{\alpha\rightarrow\infty} s(x;\alpha) = \theta(x),\]

where

\[\begin{split}\theta(x) = \begin{cases} 1, &\text{if }x>0,\\ 0, &\text{otherwise.}\end{cases}\end{split}\]

In practice, additional constraints are imposed to \(s\), but they are not detailed here. Evaluating \(s(0;\alpha)\) is considered to be undefined behavior, even if it returns actual values (e.g., \(\frac{1}{2}\)), as it might change in the future.

Upstream functions using smoothing

Currently, only a limited set of functions allow returning a smoothed-out value instead of a hard boolean decision. The most important function is TriangleScene.compute_paths. The rest are downstream functions used by the latter, and can be identified with their optional smoothing_factor argument.

Boolean equivalents

The following boolean comparisons are evaluated with the following equivalent real-value functions:

Table 1 Real-valued equivalents to common boolean operators

boolean	real-valued
x | y	jnp.minimum(x, y)
x & y	jnp.maximum(x, y)
x > y	s(x - y)
x < y	s(y - x)
x >= y	s(x - y)
x <= y	s(y - x)

Note the lack of strict equality operator, as it is (currently) not required by upstream functions.

Examples

In this section, we show usage examples of the smoothing techniques.

But first, we need to import a few packages (the import cell is hidden by default to increase readability).

Show code cell source

Hide code cell source

import jax.numpy as jnp
from differt.geometry import (
    TriangleMesh,
    path_lengths,
)
from differt.plotting import draw_image, reuse, set_backend
from differt.scene import (
    TriangleScene,
)

Coverage Map

In the cells below, we construct a basic indoor scenario with one blocking object in the center. Then, we compute a rough approximate of the path gain for many receiver positions and different values of smoothing factor (see slider).

set_backend(
    "plotly"
)  # Our scene is simple, and Plotly is the best backend for online interactive plots :-)

mesh = TriangleMesh.box(length=6, width=4, height=2).set_face_colors(
    jnp.asarray([
        0.45,
        0.27,
        0.0,
    ])
).translate(
    jnp.asarray([
        0.0,
        0.0,
        1.0,
    ])
) + TriangleMesh.plane(
    jnp.asarray([0.0, 0.0, 1.0]),
    normal=jnp.asarray([1.0, 0.0, 0.0]),
    side_length=2.0,
).set_face_colors(jnp.asarray([0.57, 0.57, 0.57]))
scene = TriangleScene(transmitters=jnp.asarray([1.0, 0.0, 1.5]), mesh=mesh)
scene.plot()

Show code cell source

Hide code cell source

# Our scene can be simplified to quadrilaterals,
# so informing the code of that matter will make it run faster
scene = scene.set_assume_quads()
batch = (
    50,
    50,
)  # Warning: a too large batch could easily cause OOM issues,
#    or you may want to reduce the 'chunk_size' value below.
z0 = 0.5  # The z coordinate of the receivers
scene_grid = scene.with_receivers_grid(*batch, height=z0)

# Only needed for plotting purposes
x, y, _ = jnp.unstack(scene_grid.receivers, axis=-1)

with reuse() as fig:
    scene.plot()
    offset = len(fig.data)

    smoothing_factors = jnp.logspace(-1, 5, 20)

    G_min = G_max = None

    for smoothing_factor in [None, *smoothing_factors]:
        P = jnp.zeros(batch)

        for order in range(3):
            for paths in scene_grid.compute_paths(
                order=order, chunk_size=1_000, smoothing_factor=smoothing_factor
            ):
                # Path loss as r^2, 0.5 reflection coefficient
                P += paths.reduce(
                    lambda path_vertices: (
                        0.5 ** (path_vertices.shape[-2] - 2)
                        / (path_lengths(path_vertices) ** 2)
                    ),
                    axis=-1,
                )

        G_dB = 10 * jnp.log10(P)

        if G_min is None and G_max is None:
            isfinite = jnp.isfinite(G_dB)
            G_min = float(G_dB.min(where=isfinite, initial=+jnp.inf))
            G_max = float(G_dB.max(where=isfinite, initial=-jnp.inf))

        draw_image(
            G_dB,
            x=x[0, :],
            y=y[:, 0],
            z0=z0,
            colorbar={"title": "Gain (dB)"},
            colorscale="viridis",
            cmin=G_min,
            cmax=G_max,
            visible=False,
        )

    # Put 'smoothing_factor=None' trace at the latest position
    fig.data = [*fig.data[:offset], *fig.data[offset + 1 :], fig.data[offset]]

    steps = []

    for i, value in enumerate([*smoothing_factors, None]):
        step = {
            "method": "update",
            "args": [
                {
                    "visible": [True] * offset
                    + [False] * (len(smoothing_factors) + 1),
                },
            ],
            "label": f"{value:.2e}" if value is not None else "no smoothing",
        }
        step["args"][0]["visible"][offset + i] = True  # Show CM
        steps.append(step)

    sliders = [
        {
            "active": 0,
            "currentvalue": {"prefix": "Smoothing factor: "},
            "pad": {"t": 50},
            "steps": steps,
        }
    ]

    fig.data[offset].visible = True

    fig.update_layout(
        height=600,
        sliders=sliders,
    )
fig