Complex spectrum computations in python

[Dream4Leo]

Is there a way to create python equivalent type of dr::Complex<UnpolarizedSpectrum> ?
I tried dr.cuda.Complex2f(mi.Spectrum(0), mi.Spectrum(1)) and it says:
TypeError: Complex2f constructor expects: 1 or 2 values of type "float", a matching list/tuple, or a NumPy/PyTorch/TF/Jax array.

I am using cuda_spectral variant of mitsuba, having searched through docs but could not find a python example.
Edit:
I also wonder the correct way in drjit to implement something like dr.sum(dr.exp(a[:,None]*b), axis=1) where a is Spectrum and b is a drjit float array. I tried to manually unroll the spectral dimension like s[0]=dr.sum(dr.exp(a[0]*b)), but it throws RuntimeError: jitc_var_eval(r3749): placeholder variables are used to record computation symbolically and cannot be scheduled for evaluation.

[rtabbara]

Hi @Dream4Leo,

Is there a way to create python equivalent type of dr::Complex ?

No, we currently only have Python bindings for Complex2f or Complex2f64.

With regards to your second question, what you have should work (omitting the slicing operation), although I've only tested this out on Dr.Jit 1.0.1. I suspect you're using an older version of Dr.Jit as the terminology for placeholder variables is outdated.

If you're still encountering issues, it might be worth providing a bit more code with regards to how a and b are defined.

[Dream4Leo]

Thank you for your information. For the complex numbers I guess I have to compute real and imaginary parts seperately, although it definitely affects code readibility.

For the second question, I am indeed using an older Dr.Jit version (with Mitsuba 3.5.2).
My intention is to compute an outer product c(n_wvls,16) from a(n_wvls) and b(16), then reduce it on the second axis to s(n_wvls).
Relevant codes are:

class mybsdf(mi.BSDF):
    
    def __init__(self, props):
        mi.BSDF.__init__(self, props)
        self.b = dr.cuda.Float(np.ones(16))
    
    def eval(self, ctx, si, wo, active):
        a = dr.two_pi / si.wavelengths
        s = dr.zeros(mi.cuda_spectral.Spectrum)
        s[0] = dr.sum(dr.exp(a[0]*self.b))
        ## s = dr.sum(dr.exp(a[:,None]*self.b), axis=1) not applicable

Switching to evaluated modes leads to new error:
RuntimeError: jit_var_mul(r2858, r27): operands have incompatible sizes! (sizes: 4194304, 16)
the former size seems to be the hidden dimension for parallelism in a.

Would like to know if there is a workaround before updating to the new version (need to recompile variants).

Edit: below is a minimum example to replicate the above errors:

import numpy as np

import drjit as dr
dr.set_flag(dr.JitFlag.VCallRecord, False)
dr.set_flag(dr.JitFlag.LoopRecord, False)

import mitsuba as mi
mi.set_variant('cuda_spectral')

class mybsdf(mi.BSDF):
    
  def __init__(self, props):
    mi.BSDF.__init__(self, props)
    self.b = dr.cuda.Float(np.ones(16))
  
  def sample(self, ctx, si, sample1, sample2):
    bs = mi.BSDFSample3f()
    bs.wo = mi.warp.square_to_cosine_hemisphere(sample2)
    bs.pdf = mi.warp.square_to_cosine_hemisphere_pdf(bs.wo)
    
    val = self.eval(ctx, si, bs.wo) / bs.pdf
    return (bs, val)

  def eval(self, ctx, si, wo, active = True):
    a = dr.two_pi / si.wavelengths
    s = dr.zeros(mi.cuda_spectral.Spectrum)
    s[0] = dr.sum(dr.exp(a[0]*self.b))
    ## s = dr.sum(dr.exp(a[:,None]*self.b), axis=1) not applicable
    return s
  
  def pdf(self, ctx, si, wo, active):    
    pdf = mi.warp.square_to_cosine_hemisphere_pdf(wo)
    return dr.select(active, pdf, 0)


mi.register_bsdf('mybsdf', lambda props: mybsdf(props))

my_bsdf = mi.load_dict({
  'type': 'mybsdf'
})

scene = mi.load_dict({
  'type': 'scene',
  'integrator': {
      'type': 'path'
  },
  'sphere' : {
    'type': 'sphere',
    'bsdf': my_bsdf
  },
  'sensor': {
    'type': 'perspective'
  }
})

image = mi.render(scene)

[huanxi0319]

It looks like you're working with the Mitsuba 3 rendering framework, which uses the Dr.Jit (Differentiable Runtimes for Just-In-Time Compilation) library for its just-in-time compilation and automatic differentiation capabilities. The error messages you're encountering suggest that there are some specific requirements for constructing Complex2f objects and using certain operations on drjit arrays.

Let's address each of your issues:

1. Constructing Complex2f Objects

The Complex2f type in Dr.Jit is designed to represent complex numbers with two floating-point components. The error message indicates that the constructor expects either 1 or 2 float values, a matching list/tuple, or a NumPy/PyTorch/TF/Jax array.

If you want to create a Complex2f object from mi.Spectrum objects, you need to ensure that the Spectrum objects can be converted to float values. Here’s how you can do it:

import mitsuba as mi
import drjit as dr

# Initialize the spectrum objects
spectrum_real = mi.Spectrum(0)
spectrum_imag = mi.Spectrum(1)

# Convert the spectra to float values
real_part = float(spectrum_real)
imag_part = float(spectrum_imag)

# Create a Complex2f object
complex_value = dr.cuda.Complex2f(real_part, imag_part)

2. Using where and is in drjit Arrays

In Dr.Jit, you can use the dr.where function to perform conditional operations on arrays. The is keyword in Python is not used for array operations; instead, you would use comparison operators and dr.where.

Here’s an example of how to use dr.where:

import drjit as dr

# Example arrays
a = dr.array([1.0, 2.0, 3.0])
b = dr.array([4.0, 5.0, 6.0])

# Condition: a > 2
condition = a > 2

# Use dr.where to select elements based on the condition
result = dr.where(condition, a, b)  # If condition is True, take a, otherwise take b

print(result)  # Output: [4.0, 5.0, 3.0]

3. Handling Spectral Dimensions in drjit

When working with spectral dimensions, you need to be careful about how you handle the data. The error you encountered suggests that you might be trying to perform operations that are not supported directly on the spectral dimension.

To sum over the spectral dimension, you can use the dr.sum function, but you need to ensure that the dimensions match correctly. Here’s an example of how you might do this:

import mitsuba as mi
import drjit as dr

# Initialize the spectral arrays
a = mi.Spectrum([1.0, 2.0, 3.0])  # Example spectral array
b = mi.Spectrum([4.0, 5.0, 6.0])  # Example spectral array

# Expand the dimensions to allow broadcasting
a_expanded = dr.tile(a, (b.size, 1))  # Tile a to match the size of b
b_expanded = dr.tile(b, (a.size, 1)).T  # Tile b and transpose to match the shape of a

# Perform the element-wise multiplication and sum over the spectral dimension
result = dr.sum(dr.exp(a_expanded * b_expanded), axis=1)

print(result)

Summary

• Constructing Complex2f: Ensure that the input values are float and convert Spectrum objects to float if necessary.
• Using dr.where: Use dr.where for conditional operations on arrays.
• Handling Spectral Dimensions: Use dr.tile to expand dimensions and dr.sum to sum over the spectral dimension, ensuring that the dimensions match correctly.

If you have more specific requirements or if these examples don't fully address your needs, please provide more details, and I can help further.