Custom Normalmap BSDF (python) DrJit Evaluation Symbolic Variable

[imanolooo]

Hi all,

I'm trying to implement a custom version of the normalmap bsdf plugin in python (I have it in c++ working fine). I've started by implementing in python the mitsuba3 version of this plugin but I'm stucked translating two lines from c++ to python. Maybe is something trivial, but I'm very newie with DrJIT and I'm not fully understanding what is the problem (sorry in advance if it is the case).

In the sample function, there are these three lines in c++ to evaluate the result of the nested plugin sample:

auto [bs, weight] = m_nested_bsdf->sample(ctx, perturbed_si,
                                                  sample1, sample2, active);
        active &= dr::any(unpolarized_spectrum(weight) != 0.f);
        if (dr::none_or<false>(active))
            return { bs, 0.f };

After checking that none_or function is not available in python DrJIT AP, I have translated it as:

        bs, weight = self.m_nested_bsdf.sample(ctx, perturbed_si, sample1, sample2, active)
        active &= dr.any(mi.unpolarized_spectrum(weight) != 0.0)
        if not dr.any(active):
            return bs, 0.0

It worked fine using the scalar_rgb variant of Mitsuba. But when I activate LLVM_ad_rgb I'm getting the following error:

File "/mnt/c/Users/imanol.munoz/workspace/Manvir/pedret_def/code/final/MyNormalMap.py", line 38, in sample
    if not dr.any(active):
           ~~~~~~^^^^^^^^
RuntimeError: drjit.any(<drjit.llvm.ad.Bool>): jit_var_schedule(r1475): not permitted.

You performed an operation that tried to evalute a *symbolic*
variable which is not permitted.

I thought that maybe the problem was the if was a regular python control statement, and I've substituted it to:

    weight = dr.if_stmt(
            dr.any(active), 
            weight,
            0.0
        )

But I'm getting the same error... Any clue on how I should solve this?

Best,

Imanol

[imanolooo]

I've solved it using dr.select instead of dr.any and if statements. Seems to be working.

def sample(self, ctx, si, sample1, sample2, active):
        # Sample the nested BSDF with a perturbed shading frame
        perturbed_si = mi.SurfaceInteraction3f(si)
        perturbed_si.sh_frame = self.frame(si, active)
        perturbed_si.wi = perturbed_si.to_local(si.wi)

        bs, weight = self.m_nested_bsdf.sample(ctx, perturbed_si, sample1, sample2, active)

        # Update the active mask based on the weight
        # Symbolic update of the active mask based on weight
        weight_nonzero = mi.unpolarized_spectrum(weight) != 0.0
        active = active & weight_nonzero  # Combine active mask with weight non-zero mask

        # Mask out weight where active is false
        weight = dr.select(active, weight, 0.0)  # Symbolic masking using dr.select

        # Transform the sampled 'wo' back to the original frame and verify orientation
        perturbed_wo = perturbed_si.to_world(bs.wo)
        active &= (mi.Frame3f.cos_theta(bs.wo) * 
                mi.Frame3f.cos_theta(perturbed_wo)) > 0.0
        #bs.wo = perturbed_wo
        bs.wo = dr.select(active, perturbed_wo, bs.wo)

        return bs, weight & active

[merlinND]

dr::none_or<false> is a helper to avoid horizontal reductions when running in the JIT variants like cuda_ad_rgb. It means "dr::none(...) in scalar variants, or false in JIT variants".
This means the code inside of the if block needs to be masked to get correct results in JIT variants.

Since custom Python plugins are expected to always run in JIT variants, there's no equivalent to dr::none_or<>. You can either remove the if and use masks / selects (like you did), or you @dr.syntax / dr.if_stmt.