[❔ other question] How to record a Path

[DoeringChristian]

Hi,

Is there a way to record a path, traced in a drjit loop? I want to implement a bdpt and want to record all the paths traced from the light as well as the view. My idea was to have an array and scatter the path vertices into it:

depth = mi.UInt32(0)
vpath_p = dr.zeros(mi.Point3f, shape=(max_depth * n_rays))
loop = mi.Loop("Path Tracing", state=lambda: (active, ray, depth, vpath_p))

while loop(active):
    si = scene.ray_intersect(ray)
    dr.scatter(vpath_p, si.p, depth * n_rays)
    depth += 1

This does not work however and the program fails with "loop state variables have inconsistent size". When leaving the vpath_p variable out of the loop declaration It doesn't work either and all values are 0. I also tried using Tensors but this also didn't work.

Thanks for your help.

[Speierers]

Hi @DoeringChristian,

Something like this should be possible. And indeed you will need to leave the vpath_p variable out of the loop state as variables affected by dr.scatter inside loop are considered as side effects and are handled differently in the JIT.

I don't understand why your second solution doesn't work. Could you maybe provide me with a reproducer that I can run on my side?

[DoeringChristian]

My current code (using dr.scatter):

import mitsuba as mi
import drjit as dr
import matplotlib.pyplot as plt

mi.set_variant("cuda_ad_rgb")

class Simple(mi.SamplingIntegrator):
    def __init__(self, props=mi.Properties()):
        super().__init__(props)
        self.max_depth = props.get("max_depth")
        self.rr_depth = props.get("rr_depth")

    def sample(self, scene: mi.Scene, sampler: mi.Sampler, ray_: mi.RayDifferential3f, medium: mi.Medium = None, active: mi.Bool = True):
        bsdf_ctx = mi.BSDFContext()

        ray = mi.Ray3f(ray_)
        depth = mi.UInt32(0)
        active = mi.Bool(active)

        n_rays = dr.shape(ray.o)[1]
        vpathf = dr.ones(mi.Spectrum, shape=(self.max_depth * n_rays))
        vpathp = dr.zeros(mi.Point3f, shape=(self.max_depth * n_rays))
        vpathL = dr.zeros(mi.Spectrum, shape=(self.max_depth * n_rays))

        prev_si = dr.zeros(mi.SurfaceInteraction3f)

        loop = mi.Loop(name="Path Tracing", state=lambda: (
            sampler, ray, depth, active, prev_si))

        loop.set_max_iterations(self.max_depth)

        while loop(active):
            si: mi.SurfaceInteraction3f = scene.ray_intersect(
                ray, ray_flags=mi.RayFlags.All, coherent=dr.eq(depth, 0))

            bsdf: mi.BSDF = si.bsdf(ray)

            # Direct emission

            ds = mi.DirectionSample3f(scene, si=si, ref=prev_si)

            L = ds.emitter.eval(si)

            active_next = (depth + 1 < self.max_depth) & si.is_valid()

            # BSDF Sampling
            bsdf_smaple, bsdf_val = bsdf.sample(
                bsdf_ctx, si, sampler.next_1d(), sampler.next_2d(), active_next)

            # Scatter into path variables
            f = bsdf_val

            dr.scatter(vpathf, f, depth * n_rays)
            dr.scatter(vpathp, si.p, depth * n_rays)
            dr.scatter(vpathL, L, depth * n_rays)

            # Update Loop variables
            ray = si.spawn_ray(si.to_world(bsdf_smaple.wo))
            prev_si = dr.detach(si, True)

            # Stopping criterion (russian roulettte)

            active_next &= dr.neq(dr.max(f), 0)

            active = active_next
            depth += 1

        L = dr.gather(mi.Spectrum, vpathL, dr.full(mi.UInt32, 1, n_rays))
        p = dr.gather(mi.Spectrum, vpathp, dr.full(mi.UInt32, 0, n_rays))
        f = dr.gather(mi.Spectrum, vpathf, dr.full(mi.UInt32, 0, n_rays))
        print(f"n_rays = {n_rays}")
        print(f"L = {L}")
        print(f"shape(L): {dr.shape(L)}")
        print(f"dr.max(L): {dr.max(dr.max(L))}")
        print(f"dr.max(p): {dr.max(dr.max(p))}")
        print(f"dr.max(f): {dr.max(dr.max(f))}")
        return (L, dr.neq(depth, 0), [])

mi.register_integrator("integrator", lambda props: Simple(props))

scene = mi.cornell_box()
scene['integrator']['type'] = 'integrator'
scene['integrator']['max_depth'] = 16
scene['integrator']['rr_depth'] = 2
scene['sensor']['sampler']['sample_count'] = 1
scene['sensor']['film']['width'] = 1024
scene['sensor']['film']['height'] = 1024
scene = mi.load_dict(scene)

img = mi.render(scene)

plt.imshow(img ** (1. / 2.2))
plt.axis("off")
plt.show()

At the end I'm trying to extract the values again so it should only show the lamp of the scene but as max(L.x) shows all recorded values are 0. Weirdly the f-values for the first vertex are 0 as well and all other further f-values are 1 to which they have been initialized. Maybe it is a problem with another part of the code.

[DoeringChristian]

Ok I found my mistake, it was a fault of my own. Thanks for your help.

[merlinND]

Hello @DoeringChristian,

Could you please describe the mistake so that users running into the same situation can avoid it as well?

[DoeringChristian]

Of course. Basically I got confused about how I wanted to encode the path vertices into the array. To fix this I decided the encoding of the path vertices to be in this form: [depth, ray_index] A ray index is needed:

idx = dr.arange(mi.UInt32, n_rays)

Then when scattering the vertices the index should be calculated as follows:

dr.scatter(vpathp, si.p, depth * n_rays + self.idx)

And gathering the vertices back like this:

p=dr.gather(mi.Point3f, vpathp, depth * n_rays + self.idx)

I have written a simple class to abstract it:

@dataclass
class PVert:
    f: mi.Spectrum
    L: mi.Spectrum
    p: mi.Point3d

class Path:
    idx: mi.UInt32
    f: mi.Spectrum
    L: mi.Spectrum
    p: mi.Point3d

    def __init__(self, n_rays: int, max_depth: int):
        self.n_rays = n_rays
        self.max_depth = max_depth
        self.idx = dr.arange(mi.UInt32, n_rays)

        self.f = dr.ones(mi.Spectrum, shape=(self.max_depth * self.n_rays))
        self.L = dr.zeros(mi.Spectrum, shape=(self.max_depth * self.n_rays))
        self.p = dr.zeros(mi.Spectrum, shape=(self.max_depth * self.n_rays))

    def __setitem__(self, depth: mi.UInt32, value: PVert):
        dr.scatter(self.f, value.f, depth * self.n_rays + self.idx)
        dr.scatter(self.L, value.L, depth * self.n_rays + self.idx)
        dr.scatter(self.p, value.p, depth * self.n_rays + self.idx)

    def __getitem__(self, depth: mi.UInt32) -> PVert:
        return PVert(
            f=dr.gather(mi.Spectrum, self.f, depth * self.n_rays + self.idx),
            L=dr.gather(mi.Spectrum, self.L, depth * self.n_rays + self.idx),
            p=dr.gather(mi.Point3f, self.p, depth * self.n_rays + self.idx)
        )

Which can be used like this in the integrator: (As you mentioned the path should not be in the loop state)

path = Path(n_rays, self.max_depth)
...
while loop(active):
    ...
    path[depth] = PVert(f, L, si.p)
    ...
    depth += 1
...
vert = path[mi.UInt32(0)]
return (vert.f, mi.Bool(True), [])

[Speierers]

Thanks for this in detail @DoeringChristian!

Just as a side note, you should be able to use DRJIT_STRUCT in Python to add support for your custom data-structures to Dr.Jit. Once this is done, functions like dr.gather, dr.scatter, ... will understand how to work with it.

Here is a simple example (untested):

class MyStruct:
    def __init__(self):
        self.a = mi.Vector3f()
        self.b = mi.Float()

    DRJIT_STRUCT = { 'a' : mi.Vector3f, 'b' : mi.Float }

structs = dr.zeros(MyStruct, 100)
idx = dr.arange(mi.UInt32, 4)
subset = dr.gather(MyStruct, structs, idx) 
...

[DoeringChristian]

Sorry I what I have posted previously didn't work. Here an implementation that does:

class PVert:
    f: mi.Spectrum
    L: mi.Spectrum
    p: mi.Point3d
    DRJIT_STRUCT = {'f': mi.Spectrum, 'L': mi.Spectrum, 'p': mi.Point3f}

    def __init__(self, f=mi.Spectrum(), L=mi.Spectrum(), p=mi.Point3f()):
        self.f = f
        self.L = L
        self.p = p

class Path:
    idx: mi.UInt32
    vertices: PVert

    def __init__(self, n_rays: int, max_depth: int):
        self.n_rays = n_rays
        self.max_depth = max_depth
        self.idx = dr.arange(mi.UInt32, n_rays)

        self.vertices = dr.zeros(PVert, shape=(self.max_depth * self.n_rays))

    def __setitem__(self, depth: mi.UInt32, value: PVert):
        dr.scatter(self.vertices, value, depth * self.n_rays + self.idx)

    def __getitem__(self, depth: mi.UInt32) -> PVert:
        return dr.gather(PVert, self.vertices, depth * self.n_rays + self.idx)

I think a decorator that infers the DRJIT_STRUCT automatically would be useful.

[DoeringChristian]

Hi @Speierers,

I reopened this issue since this one is closely related to the original one, hope this is OK. Scattering path vertices works great now and I can access them later however accessing them inside another loop does not work. I get the following error:

Critical Dr.Jit compiler failure: jit_var_new_gather(): variable remains dirty after evaluation!

Is there a way to gather from an array in a loop this way? I haven't simplified the code to reproduce it but I hope it is not too long to put into an issue:

import mitsuba as mi
import drjit as dr
import matplotlib.pyplot as plt

mi.set_variant("cuda_ad_rgb")
# dr.set_log_level(dr.LogLevel.Debug)

def drjitstruct(cls):
    annotations = cls.__dict__.get('__annotations__', {})
    drjit_struct = {}
    for name, type in annotations.items():
        drjit_struct[name] = type
    cls.DRJIT_STRUCT = drjit_struct
    return cls

@drjitstruct
class PVert:
    f: mi.Spectrum
    L: mi.Spectrum
    p: mi.Point3f

    def __init__(self, f=mi.Spectrum(), L=mi.Spectrum(), p=mi.Point3f()):
        self.f = f
        self.L = L
        self.p = p

class Path:
    idx: mi.UInt32
    vertices: PVert

    def __init__(self, n_rays: int, max_depth: int):
        self.n_rays = n_rays
        self.max_depth = max_depth
        self.idx = dr.arange(mi.UInt32, n_rays)

        self.vertices = dr.zeros(PVert, shape=(self.max_depth * self.n_rays))

    def __setitem__(self, depth: mi.UInt32, value: PVert):
        dr.scatter(self.vertices, value, depth * self.n_rays + self.idx)

    def __getitem__(self, depth: mi.UInt32) -> PVert:
        return dr.gather(PVert, self.vertices, depth * self.n_rays + self.idx)

class Simple(mi.SamplingIntegrator):
    def __init__(self, props=mi.Properties()):
        super().__init__(props)
        self.max_depth = props.get("max_depth")
        self.rr_depth = props.get("rr_depth")

    # record a path
    def record(self, scene: mi.Scene, sampler: mi.Sampler, ray: mi.Ray3f, active) -> Path:
        ray = mi.Ray3f(ray)
        bsdf_ctx = mi.BSDFContext()

        prev_si = dr.zeros(mi.SurfaceInteraction3f)

        n_rays = dr.shape(ray.o)[1]

        path = Path(n_rays, self.max_depth)
        depth = mi.UInt32(0)
        active = mi.Bool(active)

        loop = mi.Loop(name="Path Record", state=lambda: (
            sampler, ray, depth, active, prev_si))

        loop.set_max_iterations(self.max_depth)

        while loop(active):
            si: mi.SurfaceInteraction3f = scene.ray_intersect(
                ray, ray_flags=mi.RayFlags.All, coherent=dr.eq(depth, 0))

            bsdf: mi.BSDF = si.bsdf(ray)

            ds = mi.DirectionSample3f(scene, si=si, ref=prev_si)

            L = ds.emitter.eval(si)

            active_next = (depth + 1 < self.max_depth) & si.is_valid()

            bsdf_sample, f = bsdf.sample(
                bsdf_ctx, si, sampler.next_1d(), sampler.next_2d(), active_next)

            ray = si.spawn_ray(si.to_world(bsdf_sample.wo))

            path[depth] = PVert(f, L, si.p)

            prev_si = dr.detach(si, True)

            active_next &= dr.neq(dr.max(f), 0)

            active = active_next
            depth += 1

        return path

    def sample(self, scene: mi.Scene, sampler: mi.Sampler, ray: mi.RayDifferential3f, medium: mi.Medium = None, active: mi.Bool = True):

        lray, lweight, emitter = scene.sample_emitter_ray(
            1., sampler.next_1d(), sampler.next_2d(), sampler.next_2d(), active)

        #lpath = self.record(scene, sampler, lray, True)

        vpath = self.record(scene, sampler, ray, True)

        depth = mi.UInt32(0)
        active = mi.Bool(active)
        L = mi.Spectrum(0.)
        f = mi.Spectrum(1.)
        loop = mi.Loop("Defferd lighting",
                       state=lambda: (sampler, depth, L, f, active))
        loop.set_max_iterations(self.max_depth)

        while loop(active):
            vert = vpath[depth]
            L += vert.L * f
            f *= vert.f
            depth += 1

        #vert = vpath[mi.UInt32(0)]
        return (L, mi.Bool(True), [])

mi.register_integrator("integrator", lambda props: Simple(props))

scene = mi.cornell_box()
scene['integrator']['type'] = 'integrator'
scene['integrator']['max_depth'] = 16
scene['integrator']['rr_depth'] = 2
scene['sensor']['sampler']['sample_count'] = 1
scene['sensor']['film']['width'] = 1024
scene['sensor']['film']['height'] = 1024
scene = mi.load_dict(scene)

img = mi.render(scene)

plt.imshow(img ** (1. / 2.2))
plt.axis("off")
plt.show()

Thanks for your Help.

[DoeringChristian]

Sorry for reopening the issue after all the fix was quite easy, but maybe if someone has a similar problem this information could be useful. I found out how I could fix it: The path vertices need to be evaluated with dr.eval(vpath.vertices) I assume evaluating the variable will clear it of all side effects which was the reason the error has been triggered. Thinking about it makes sense since the threads need to be synchronized before it is safe to read from the array again.

[merlinND]

Thanks for the update @DoeringChristian!

@Speierers, could it be a bug? As far as I know we should never need explicit calls to dr.eval().

[Speierers]

If I understand correctly you have a first loop that scatters into some buffer and a second loop that gathers from that same buffer. Indeed this will require some syncronisation as otherwise would would have no guaranty that the sideeffects took place before entering the second loop.

Another solution would be to recompute those values in the second loop, or arrange your code differently so that everything is done in a single loop.

[DoeringChristian]

Jes I have two loops one scattering and one gathering the path vertices. The idea was I could record some light paths and then connect them to view paths in a second loop (I guess similar to photon mapping). I thought I could save some computation time but I don't think it is possible since I would need to save the surface interactions which can't be scattered/gathered since they hold poiters to the shape. Also if I did it the way I thought the rays wouldn't be stochastically independent anymore and without interpolation I would get some artifacts.

[saeedhd96]

Thanks for this very helpful issue here.

I am having a problem when taking this to the next level. I can now successfuly record a path thanks to the notes here, however, during path recording, I would like to have a second recorded loop and do some operation.

For example, in any of the loops above, Path Record or Defferd lighting, I would like to have another recorded loop.

Below is an example of this:

I have a loop in inner_loop() in which I update a struct SampleHolder as a side effect. there is a print statement at the end of this loop. When I run this loop, I can successfully get the printed result. However, when I put this into another recorded loop, it crashes with RuntimeError: jit_var_eval(): variable r224 remains dirty after evaluation!.

import mitsuba as mi
mi.set_variant('cuda_ad_rgb')
import drjit as dr
import torch

# Coppied from above code
def drjitstruct(cls):
    annotations = cls.__dict__.get('__annotations__', {})
    drjit_struct = {}
    for name, type in annotations.items():
        drjit_struct[name] = type
    cls.DRJIT_STRUCT = drjit_struct
    return cls

# New DRJIT struct, similar to PathVert class above
@drjitstruct
class Sample:
    wi: mi.Vector3f

    def __init__(self, wi = mi.Vector3f()):
        self.wi = wi

# A holder of a sample object which can assign to it with alterations to index
#Similar to Path class above
class SampleHolder:
    idx: mi.UInt32
    samples: Sample

    def __init__(self, n_rays: int, M: int):
        self.n_rays = n_rays
        self.M = M
        self.idx = dr.arange(mi.UInt32, n_rays) * self.M
        self.samples = dr.zeros(Sample, shape=(self.M * self.n_rays))

    def __setitem__(self, sampled_id: mi.UInt32, value):
        dr.scatter(self.samples, value, sampled_id  + self.idx)

    def __getitem__(self, sampled_id: mi.UInt32):
        return dr.gather(Sample, self.samples, sampled_id  + self.idx)

n_rays = 100
M = 32

def loop1(n_rays, M):
    m = mi.UInt32(0)
    samplesM_2D = mi.Vector2f(torch.randn(n_rays*M, 2).cuda())

    sample_holder = SampleHolder(n_rays, M)
    indices = dr.arange(mi.UInt, 0, n_rays) * (M)

    loop = mi.Loop(name="MySAMPLEHolder",
                state=lambda: (m, indices))

    while loop(m<M):
        samplesM_2D_batch = dr.gather(type(samplesM_2D), samplesM_2D, indices)

        w_batch =  mi.warp.square_to_uniform_hemisphere(samplesM_2D_batch)
        batch = Sample(mi.Vector3f(w_batch))
        sample_holder[m] = batch

        indices += 1
        m+=1

    #Error happens here
    print(sample_holder.samples.wi)

print('just one loop')
loop1(n_rays, M)
print('success')

i = mi.UInt32(0)
loop = mi.Loop(name="outerloop",
            state=lambda: (i))

print('inside another loop')
while loop(i<10):
    loop1(n_rays, M)
    i+=1
print('success')

Can I ask what I am doing wrong?

[njroussel]

Hi @saeedhd96

I quickly tried reproducing your issue, which I could. I then change the script to no longer turn Sample into a DRJIT_STRUCT and it seemed to work as expected. I'm not quite sure what changed since the original posts in this thread. Could you confirm this?

[brabbitdousha]

@njroussel @Speierers Hi, is there any way to record a limited path? For example, max_depth is 10, but I only want a path with depth <=4 to be recorded, since I can't use if in mi.loop, is there any way to achieve this?

[njroussel]

@brabbitdousha You can record every path, and when the path ends you additionally scatter to a Bool array that indicates if it's end is shorter or equal to 4. With that you should have everything you need.