How to sample sggx distribution in Python?

[Vrroom]

Summary

How do I sample from the SGGX distribution from Python. I'm trying to implement the SpongeCake BSDF, thus needing to sample micronormals from the SGGX distribution.

I was trying:

>>> mi.sggx_sample(sh, mi.Point2f([0.5, 0.5]), dr.llvm.ad.ArrayXf([[1, 1, 1, 0, 0, 0]]))

Mitsuba complains that:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: sggx_sample(): incompatible function arguments. The following argument types are supported:
    1. (sh_frame: mitsuba.llvm_ad_rgb.Frame3f, sample: mitsuba.llvm_ad_rgb.Point2f, s: drjit::Array<drjit::DiffArray<drjit::LLVMArray<float> >, 6ul>) -> mitsuba.llvm_ad_rgb.Normal3f
    2. (sh_frame: mitsuba.llvm_ad_rgb.Vector3f, sample: mitsuba.llvm_ad_rgb.Point2f, s: drjit::Array<drjit::DiffArray<drjit::LLVMArray<float> >, 6ul>) -> mitsuba.llvm_ad_rgb.Normal3f

Invoked with: Frame[
  s = [[1, -0, -0]],
  t = [[-0, 1, -0]],
  n = [[0, 0, 1]]
], [[0.5, 0.5]], [[1.0],
 [1.0],
 [1.0],
 [0.0],
 [0.0],
 [0.0]]

How do I construct this specific type: drjit::Array<drjit::DiffArray<drjit::LLVMArray<float> >, 6ul> in Python.

System configuration

System information:

  OS: Ubuntu 20.04.5 LTS
  CPU: Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz
  GPU: NVIDIA A10
  Python: 3.11.4 (main, Jul  5 2023, 13:45:01) [GCC 11.2.0]
  NVidia driver: 525.85.12
  CUDA: 11.8.89
  LLVM: 12.0.0

  Dr.Jit: 0.4.4
  Mitsuba: 3.5.0
     Is custom build? False
     Compiled with: GNU 10.2.1
     Variants:
        scalar_rgb
        scalar_spectral
        cuda_ad_rgb
        llvm_ad_rgb

Description

See the summary

Steps to reproduce

NA

[Vrroom]

I ported it here (for anyone interested):

import mitsuba as mi
import drjit as dr

def safe_rsqrt (x) :
    return dr.rsqrt(dr.maximum(x, 0.))

def sggx_sample(sh_frame, sample, s_mat) :
    k, j, i = 0, 1, 2
    m = mi.Matrix3f (sh_frame.s, sh_frame.t, sh_frame.n)
    m = dr.transpose(m);
    s2 = m @ s_mat @ dr.transpose(m)
    inv_sqrt_s_ii = safe_rsqrt(s2[i, i])
    tmp = dr.safe_sqrt(s2[j, j] * s2[i, i] - s2[j, i] * s2[j, i])
    m_k = mi.Vector3f(dr.safe_sqrt(dr.abs(dr.det(s2))) / tmp, 0., 0.);
    m_j = mi.Vector3f(-inv_sqrt_s_ii * (s2[k, i] * s2[j, i] - s2[k, j] * s2[i, i]) / tmp, inv_sqrt_s_ii * tmp, 0.)
    m_i = inv_sqrt_s_ii * mi.Vector3f(s2[k, i], s2[j, i], s2[i, i])
    uvw = mi.warp.square_to_cosine_hemisphere(sample)
    return sh_frame.to_world(dr.normalize(uvw.x * m_k + uvw.y * m_j + uvw.z * m_i))

[njroussel]

Hi @Vrroom

This is the issue which is addressed in https://github.com/mitsuba-renderer/mitsuba3/pull/1052.

As mentionned in that PR, we're undergoing a large rewrite of the Python bindings which should be released soon. In the meantime, you can use the fix suggested in that PR in your own fork. Alternatively, I think your Python port of the function is just fine too :)

Note to self: I'm keeping this open until the nanobind port is merged.

[Vrroom]

Thanks for the reply @njroussel !!

[rtabbara]

Added in Mitsuba 3.6.0 — see mi.sggx_sample