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)

[merlinND]

Hello @Dream4Leo,

Would like to know if there is a workaround before updating to the new version

I don't think that we can provide support for older versions of DrJit and Mitsuba.