Laplacian loss during vertex position optimization

[sapo17]

Hello!

First of all, thank you for the amazing work and for providing a great platform for asking questions.

The other day I came across a tutorial Optimizing a mesh using a Differentiable Renderer, where they show mesh optimization using Kaolin.

During optimization, they also calculate the laplacian loss. Combining kaolin** and mitsuba 3, I wanted to test a similar approach, however unfortunately I could not manage to do it. Has anyone tried out a similar approach? What would be the best/easiest way to calculate the laplacian loss (in mitsuba) during vertex position optimization?

I would appreciate any hints/tips--thank you!

** Similar to the linked tutorial, I tried to calculate laplacian loss (using uniform_laplacian), however in my case, the vertex positions (and actually all the optimization steps) are handled by mitsuba. The optimization loop has a lot of similarities to the provided mitsuba tutorial.

[njroussel]

Hi @sapo17

This is indeed a bit complicated in Mitsuba. Here are a few things that come to mind:

• Mitsuba/Dr.Jit doesn't have matrix-matrix multiplications for arbitrary sizes. You would need to implement something like matrix_mutiply(a: mi.TensorXf, b: mi:TensorXf) using Dr.Jit operations only (to preserve automatic differentiation).
• The matrix initialization in unfirom_laplacian is awful to do in Mitsuba. Since it's only computed once per mesh/topology you can safely build it with some other framework and convert it to a Mitsuba/DrJit type.
• When using mi.traverse() on a Mesh, you'll get vertex_positions, faces as flattened arrays. Internally the vertex_positions parameter is a bit un-conventionial. We duplicate a vertex for every face it is a part of, this saves us some indexing/lookups at different areas in the code. You can iterate through faces in steps of 3 (we only support triangle meshes) to gather the 3 indices which can then be used to do a lookup in the vertex_positions array. I would recommend trying this out with some basic shape like a cube, to understand the memory layout.

[sapo17]

Hello @njroussel, thank you very much for your detailed answer! I will give this a shot and update this post if I get successful.

[DoeringChristian]

Hi,
sorry for posting on a closed issue but I might have some implementations of matrix multiplication and calculating the mesh Laplacian using Dr.Jit.

def matmul(a: mi.TensorXf, b: mi.TensorXf) -> mi.TensorXf:
    assert len(a.shape) == 2
    assert a.shape[1] == b.shape[0]

    N = a.shape[0]
    M = b.shape[1]
    K = b.shape[0]

    i, j = dr.arange(mi.UInt, N), dr.arange(mi.UInt, M)
    j, i = dr.meshgrid(j, i)

    print(f"{i=}")
    print(f"{j=}")

    k = mi.UInt(0)
    dst = mi.Float(0)

    loop = mi.Loop(name="matmul", state=lambda: (k, dst))
    loop.set_max_iterations(K)

    while loop(k < K):
        dst += dr.gather(mi.Float, a.array, i * a.shape[1] + k) * dr.gather(
            mi.Float, b.array, k * b.shape[1] + j
        )

        k += 1

    return mi.TensorXf(dst, shape=(N, M))


def laplacian_uniform(verts: mi.Float, faces: mi.UInt) -> mi.TensorXf:
    verts: mi.Point3f = dr.unravel(mi.Point3f, verts)
    faces: mi.Vector3u = dr.unravel(mi.Vector3u, faces)

    n_verts = dr.shape(verts)[-1]

    adj = dr.zeros(mi.Float, n_verts * n_verts)

    w = -1

    def scatter_weights(i, j):
        dr.scatter(adj, w, i * n_verts + j)
        dr.scatter(adj, w, j * n_verts + i)

    scatter_weights(faces.x, faces.y)
    scatter_weights(faces.y, faces.z)
    scatter_weights(faces.x, faces.z)

    i = dr.arange(mi.UInt, n_verts)
    j = dr.arange(mi.UInt, n_verts)
    i, j = dr.meshgrid(i, j)

    dr.scatter_reduce(
        dr.ReduceOp.Add,
        adj,
        -dr.gather(mi.Float, adj, i * n_verts + j),
        i * n_verts + i,
    )

    return mi.TensorXf(adj, shape=(n_verts, n_verts))

This implementation only calculates the uniform Laplacian however I think this would be extendable to calculate the cosine Laplacian matrix.
I don't think this is a memory efficient implementation but I tested it against the largesteps Laplace matrix calculation implementation and it seems to work. I would recommend you to test it yourself non the less.

[sapo17]

I will give this a shot, thank you very much!

[DoeringChristian]

Btw. it seems that there is also some work being done to implement the largesteps paper in Dr.Jit here. They might come up with a better implementation that uses sparse matrix representations and has an inverse matrix solver, which would be necessary to implement that paper.

[sapo17]

Wow, this is a very nice tutorial, thank you very much for your effort! I look forward to reading and trying it.

[GuangyanCai]

Thanks for your code! However, when I tried to use it for a MLP, I realized that recorded loops don't support differentiation. So here is a modified version that support differentiation while maintaining decent speed:

import drjit as dr
import mitsuba as mi
from timeit import default_timer

mi.set_variant('cuda_ad_rgb')

def matmul(a: mi.TensorXf, b: mi.TensorXf) -> mi.TensorXf:
    # Check conditions
    assert len(a.shape) == 2
    assert a.shape[1] == b.shape[0]

    # Matrix sizes
    N = a.shape[0]
    M = b.shape[1]
    K = b.shape[0]

    # Indices of the final matrix c repeat K times
    i, j = dr.arange(mi.UInt, N), dr.arange(mi.UInt, M)
    i, j = dr.meshgrid(i, j, indexing='ij')
    i, j = dr.repeat(i, K), dr.repeat(j, K)

    # [0, 1, ..., K - 1] repeated N * M times
    offset = dr.tile(dr.arange(mi.UInt, K), N * M)

    # Compute [a[0][0] * b[0][0], a[0][1] * b[1][0], ... ]
    tmp = dr.gather(mi.Float, a.array, i * K + offset) * dr.gather(mi.Float, b.array, offset * M + j)

    # Compute [c[0][0], c[0][1], ... ]
    c = dr.zeros(mi.TensorXf, shape=(N, M))
    dr.scatter_reduce(dr.ReduceOp.Add, c.array, tmp, dr.repeat(dr.arange(mi.UInt, N * M), K))
    return c

def rand(shape):
    shape = mi.UInt(shape)
    return mi.TensorXf(mi.PCG32(dr.prod(shape)[0]).next_float32(), shape=shape)

N = 100
K = 1000
M = 10000

a = rand((N, K))
b = rand((K, M))
dr.eval()

start = default_timer()
c = matmul(a, b)
dr.eval()

dr.sync_thread()

duration = default_timer() - start

c_drjit = c.torch()
print('drjit result:')
print(c_drjit)

c_torch = a.torch() @ b.torch()
print('pytorch result:')
print(c_torch)

print(f'\ndrjit matmul time: {duration}')

print('avg diff: ')
print((c_drjit - c_torch).mean())