Representation of complex vectors

[pingpongballz]

Hi,

I wanted to ask if Mitsuba/drjit has an efficient way of representing complex vectors? From the API documentation of both of them, it doesn't seem like it.

The problem is as follows. Given a collection of ray propagation, the scattered Electric field is given by $\vec{E_s} = 2ik\iint_{rays} \hat{r} \times (\hat{r} \times ( \hat{n} \times \hat{H_i} ) e^{ik(t - \hat{r} \cdot \vec{r'})}ds'$, where $\hat{r}$ is the unit camera direction, $\vec{r'}$ is the hit location of the rays, and $\hat{H_i}$ is the initial ray polarisation. The result is a 3D complex vector.

I have broken down the integral into $\vec{E_s} = 2ik\iint_{rays} \hat{r} \times (\hat{r} \times ( \hat{n} \times \hat{H_i} ) \Re(e^{ik(t - \hat{r} \cdot \vec{r'})})ds' - 2k\iint_{rays} \hat{r} \times (\hat{r} \times ( \hat{n} \times \hat{H_i} ) \Im(e^{ik(t - \hat{r} \cdot \vec{r'})})ds'$, and treated each Imaginary and Real term as its own vector.

So $\Re (\vec{E_s} ) = - 2k\iint_{rays} \hat{r} \times (\hat{r} \times ( \hat{n} \times \hat{H_i} ) \Im(e^{ik(t - \hat{r} \cdot \vec{r'})})ds'$
and $\Im (\vec{E_s} ) = 2k\iint_{rays} \hat{r} \times (\hat{r} \times ( \hat{n} \times \hat{H_i} ) \Re(e^{ik(t - \hat{r} \cdot \vec{r'})})ds'$.

Below is a snippet of the code. Everything sorta works 😅

import mitsuba as mi
import drjit as dr
import numpy as np
import matplotlib.pyplot as plt


frequency = 1e9  
wavelength = 3e8 / frequency
k = 2 * np.pi / wavelength  

mi.set_variant('llvm_ad_mono')

scene = mi.load_file('scenes/test.xml')

def PO_integral(si, ray, pol, r_dir, ds_prime):

    # Incident wave vector (opposite to ray direction)
    k_inc = -ray.d
    distance = si.t
        
    H_inc = dr.cross(k_inc, pol)
    
    phase = dr.exp(dr.scalar.Complex2f(1j) * k * (distance - dr.dot(r_dir, si.p)))
        
    E_imag = 2 * k * dr.cross(cam_dir, dr.cross(cam_dir, dr.cross(si.n, H_inc))) * dr.real(phase) * ds_prime
    E_real = -2 * k * dr.cross(cam_dir, dr.cross(cam_dir, dr.cross(si.n, H_inc))) * dr.imag(phase) * ds_prime

    return E_real, E_imag

# Camera origin in world space
cam_origin = mi.Point3f(0, 1, 1)

# Camera view direction in world space
cam_dir = dr.normalize(mi.Vector3f(0, -1, -1))

# Camera width and height in world space
cam_width  = 0.1
cam_height = 0.1

resolution = 256
ds_prime = (cam_width*cam_height)/(resolution*resolution)
# Image pixel resolution
image_res = (resolution, resolution)

pol = dr.zeros(mi.Vector3f, resolution*resolution)
pol.x = 1.0

# Construct a grid of 2D coordinates
x, y = dr.meshgrid(
    dr.linspace(mi.Float, -cam_width  / 2,   cam_width / 2, image_res[0]),
    dr.linspace(mi.Float, -cam_height / 2,  cam_height / 2, image_res[1])
)

# Ray origin in local coordinates
ray_origin_local = mi.Vector3f(x, y, 0)

# Ray origin in world coordinates
ray_origin = mi.Frame3f(cam_dir).to_world(ray_origin_local) + cam_origin

ray = mi.Ray3f(o=ray_origin, d=cam_dir)

si = scene.ray_intersect(ray)

result_real = mi.Vector3f(0)
result_imag = mi.Vector3f(0)

@dr.syntax
def compute(result_real, result_imag, si, ray):
    if si.is_valid():

        result_real, result_imag = PO_integral(si, ray, pol, cam_dir, ds_prime)

    return result_real, result_imag

result_real, result_imag = compute(result_real, result_imag, si, ray)

result_real_sum = dr.sum(result_real, axis = 1)
result_imag_sum = dr.sum(result_imag, axis = 1)

I was wondering if the complex version of Vector3f can somehow exist or be implemented? Because I am also not sure what problems can arise from this Real and Imaginary decomposition of the integral in the context of inverse rendering. Once again, thank you!

[njroussel]

Hi @pingpongballz

Dr.Jit has a Complex2f type (which Mitsuba aliases as mi.Complex2f), but we don't have nested structured like a Vector[Complex2f, 3]. I think this could be added fairly easily - it's just a different template specialization in C++, you would just need to add bindings for it in Python. Another solution is to rely on Dr.Jit's handling of Pytrees, where you would define a custom dataclass or just handle your vectors as Python tuples/lists.

Because I am also not sure what problems can arise from this Real and Imaginary decomposition of the integral in the context of inverse rendering.

I have no idea. Ultimately, under the hood these are still decomposed in two arrays, the Complex2f type is just a convenience wrapper to handle common arithmetic operations by correctly applying (or not) the operation on either real or imaginary components.

[pingpongballz]

Thank you! 😊