Skip to content

Commit

Permalink
Add support for AOVs to Python AD integrators.
Browse files Browse the repository at this point in the history 
This extends the existing Python AD integrator interface to support (differentiable) AOVs.

(This commit is a backport of f3b427e)

Co-authored-by: Sebastien Speierer <[email protected]>
Co-authored-by: Nicolas Roussel <[email protected]>
  • Loading branch information
@dvicini @Speierers @njroussel
3 people committed Jun 10, 2024
1 parent 4258c28 commit 3012b0a
Show file tree
Hide file tree
Showing 8 changed files with 81 additions and 42 deletions.
84 changes: 60 additions & 24 deletions src/python/python/ad/integrators/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,13 +73,14 @@ def render(self: mi.SamplingIntegrator,
ray, weight, pos = self.sample_rays(scene, sensor, sampler)

# Launch the Monte Carlo sampling process in primal mode
L, valid, _ = self.sample(
L, valid, aovs, _ = self.sample(
mode=dr.ADMode.Primal,
scene=scene,
sampler=sampler,
ray=ray,
depth=mi.UInt32(0),
δL=None,
δaovs=None,
state_in=None,
active=mi.Bool(True)
)
Expand All @@ -96,6 +97,7 @@ def render(self: mi.SamplingIntegrator,
value=L * weight,
weight=1.0,
alpha=dr.select(valid, mi.Float(1), mi.Float(0)),
aovs=aovs,
wavelengths=ray.wavelengths
)

Expand Down Expand Up @@ -130,7 +132,7 @@ def render_forward(self: mi.SamplingIntegrator,
ray, weight, pos = self.sample_rays(scene, sensor, sampler)

with dr.resume_grad():
L, valid, _ = self.sample(
L, valid, aovs, _ = self.sample(
mode=dr.ADMode.Forward,
scene=scene,
sampler=sampler,
Expand All @@ -147,6 +149,7 @@ def render_forward(self: mi.SamplingIntegrator,
value=L * weight,
weight=1,
alpha=dr.select(valid, mi.Float(1), mi.Float(0)),
aovs=aovs,
wavelengths=ray.wavelengths
)

Expand Down Expand Up @@ -182,7 +185,7 @@ def render_backward(self: mi.SamplingIntegrator,
ray, weight, pos = self.sample_rays(scene, sensor, sampler)

with dr.resume_grad():
L, valid, _ = self.sample(
L, valid, aovs, _ = self.sample(
mode=dr.ADMode.Backward,
scene=scene,
sampler=sampler,
Expand All @@ -202,6 +205,7 @@ def render_backward(self: mi.SamplingIntegrator,
value=L * weight,
weight=1,
alpha=dr.select(valid, mi.Float(1), mi.Float(0)),
aovs=aovs,
wavelengths=ray.wavelengths
)

Expand Down Expand Up @@ -372,6 +376,7 @@ def _splat_to_block(block: mi.ImageBlock,
value: mi.Spectrum,
weight: mi.Float,
alpha: mi.Float,
aovs: Sequence[mi.Float],
wavelengths: mi.Spectrum):
'''Helper function to splat values to a imageblock'''
if (dr.all(mi.has_flag(film.flags(), mi.FilmFlags.Special))):
Expand All @@ -382,13 +387,18 @@ def _splat_to_block(block: mi.ImageBlock,
block.put(pos, aovs)
del aovs
else:
block.put(
pos=pos,
wavelengths=wavelengths,
value=value,
weight=weight,
alpha=alpha
)
if mi.is_spectral:
rgb = mi.spectrum_to_srgb(value, wavelengths)
elif mi.is_monochromatic:
rgb = value.x
else:
rgb = value
if mi.has_flag(film.flags(), mi.FilmFlags.Alpha):
aovs = [rgb.x, rgb.y, rgb.z, alpha, weight] + aovs
else:
aovs = [rgb.x, rgb.y, rgb.z, weight] + aovs
block.put(pos, aovs)


def sample(self,
mode: dr.ADMode,
Expand All @@ -397,8 +407,9 @@ def sample(self,
ray: mi.Ray3f,
depth: mi.UInt32,
δL: Optional[mi.Spectrum],
δaovs: Optional[mi.Spectrum],
state_in: Any,
active: mi.Bool) -> Tuple[mi.Spectrum, mi.Bool]:
active: mi.Bool) -> Tuple[mi.Spectrum, mi.Bool, List[mi.Float]]:
"""
This function does the main work of differentiable rendering and
remains unimplemented here. It is provided by subclasses of the
Expand Down Expand Up @@ -465,6 +476,12 @@ def sample(self,
Output ``valid`` (``mi.Bool``):
Indicates whether the rays intersected a surface, which can be used
to compute an alpha channel.
Output ``aovs`` (``List[mi.Float]``):
Integrators may return one or more arbitrary output variables (AOVs).
The implementation has to guarantee that the number of returned AOVs
matches the length of self.aov_names().
"""

raise Exception('RBIntegrator does not provide the sample() method. '
Expand Down Expand Up @@ -555,7 +572,7 @@ def render_forward(self: mi.SamplingIntegrator,
ray, weight, pos = self.sample_rays(scene, sensor, sampler)

# Launch the Monte Carlo sampling process in primal mode (1)
L, valid, state_out = self.sample(
L, valid, aovs, state_out = self.sample(
mode=dr.ADMode.Primal,
scene=scene,
sampler=sampler.clone(),
Expand All @@ -567,13 +584,14 @@ def render_forward(self: mi.SamplingIntegrator,
)

# Launch the Monte Carlo sampling process in forward mode (2)
δL, valid_2, state_out_2 = self.sample(
δL, valid_2, δaovs, state_out_2 = self.sample(
mode=dr.ADMode.Forward,
scene=scene,
sampler=sampler,
ray=ray,
depth=mi.UInt32(0),
δL=None,
δaovs=None,
state_in=state_out,
active=mi.Bool(True)
)
Expand All @@ -590,15 +608,16 @@ def render_forward(self: mi.SamplingIntegrator,
value=δL * weight,
weight=1.0,
alpha=dr.select(valid_2, mi.Float(1), mi.Float(0)),
aovs=[δaov * weight for δaov in δaovs],
wavelengths=ray.wavelengths
)

# Perform the weight division and return an image tensor
film.put_block(block)

# Explicitly delete any remaining unused variables
del sampler, ray, weight, pos, L, valid, δL, valid_2, params, \
state_out, state_out_2, block
del sampler, ray, weight, pos, L, valid, aovs, δL, δaovs, \
valid_2, params, state_out, state_out_2, block

# Probably a little overkill, but why not.. If there are any
# DrJit arrays to be collected by Python's cyclic GC, then
Expand Down Expand Up @@ -682,7 +701,8 @@ def render_backward(self: mi.SamplingIntegrator,

def splatting_and_backward_gradient_image(value: mi.Spectrum,
weight: mi.Float,
alpha: mi.Float):
alpha: mi.Float,
aovs: Sequence[mi.Float]):
'''
Backward propagation of the gradient image through the sample
splatting and weight division steps.
Expand All @@ -699,6 +719,7 @@ def splatting_and_backward_gradient_image(value: mi.Spectrum,
value=value,
weight=weight,
alpha=alpha,
aovs=aovs,
wavelengths=ray.wavelengths
)

Expand All @@ -721,45 +742,53 @@ def splatting_and_backward_gradient_image(value: mi.Spectrum,
with dr.suspend_grad(pos, ray, weight):
L = dr.full(mi.Spectrum, 1.0, dr.width(ray))
dr.enable_grad(L)

aovs = []
for _ in self.aov_names():
aov = dr.ones(mi.Float, dr.width(ray))
dr.enable_grad(aov)
aovs.append(aov)
splatting_and_backward_gradient_image(
value=L * weight,
weight=1.0,
alpha=1.0
alpha=1.0,
aovs=[aov * weight for aov in aovs]
)

δL = dr.grad(L)
δaovs = dr.grad(aovs)

# Clear the dummy data splatted on the film above
film.clear()

# Launch the Monte Carlo sampling process in primal mode (1)
L, valid, state_out = self.sample(
L, valid, aovs, state_out = self.sample(
mode=dr.ADMode.Primal,
scene=scene,
sampler=sampler.clone(),
ray=ray,
depth=mi.UInt32(0),
δL=None,
δaovs=None,
state_in=None,
active=mi.Bool(True)
)

# Launch Monte Carlo sampling in backward AD mode (2)
L_2, valid_2, state_out_2 = self.sample(
L_2, valid_2, aovs_2, state_out_2 = self.sample(
mode=dr.ADMode.Backward,
scene=scene,
sampler=sampler,
ray=ray,
depth=mi.UInt32(0),
δL=δL,
δaovs=δaovs,
state_in=state_out,
active=mi.Bool(True)
)

# We don't need any of the outputs here
del L_2, valid_2, state_out, state_out_2, δL, \
ray, weight, pos, sampler
del L_2, valid_2, aovs_2, state_out, state_out_2, \
δL, δaovs, ray, weight, pos, sampler

gc.collect()

Expand Down Expand Up @@ -952,7 +981,7 @@ def render_ad(self,
ray, weight, pos = self.sample_rays(scene, sensor, sampler)

# Launch the Monte Carlo sampling process in differentiable mode
L, valid, _ = self.sample(
L, valid, aovs, _ = self.sample(
mode = mode,
scene = scene,
sampler = sampler,
Expand All @@ -973,6 +1002,7 @@ def render_ad(self,
value=L * weight,
weight=1.0,
alpha=dr.select(valid, mi.Float(1), mi.Float(0)),
aovs=[aov * weight for aov in aovs],
wavelengths=ray.wavelengths
)

Expand Down Expand Up @@ -1225,11 +1255,12 @@ def sample(self,
ray: mi.Ray3f,
depth: mi.UInt32,
δL: Optional[mi.Spectrum],
δaovs: Optional[mi.Spectrum],
state_in: Any,
active: mi.Bool,
project: bool = False,
si_shade: Optional[mi.SurfaceInteraction3f] = None
) -> Tuple[mi.Spectrum, mi.Bool, Any]:
) -> Tuple[mi.Spectrum, mi.Bool, List[mi.Float], Any]:
"""
See ADIntegrator.sample() for a description of this function's purpose.
Expand Down Expand Up @@ -1259,6 +1290,11 @@ def sample(self,
Indicates whether the rays intersected a surface, which can be used
to compute an alpha channel.
Output ``aovs`` (``Sequence[mi.Float]``):
Integrators may return one or more arbitrary output variables (AOVs).
The implementation has to guarantee that the number of returned AOVs
matches the length of self.aov_names().
Output ``seedray`` / ``state_out`` (``any``):
If ``project`` is true, the integrator returns the seed rays to be
projected as the third output. The seed rays is a python list of
Expand Down
8 changes: 4 additions & 4 deletions src/python/python/ad/integrators/direct_projective.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ def sample(self,
project: bool = False,
si_shade: Optional[mi.SurfaceInteraction3f] = None,
**kwargs # Absorbs unused arguments
) -> Tuple[mi.Spectrum, mi.Bool, Any]:
) -> Tuple[mi.Spectrum, mi.Bool, List[mi.Float], Any]:
"""
See ``PSIntegrator.sample()`` for a description of this interface and
the role of the various parameters and return values.
Expand Down Expand Up @@ -253,7 +253,7 @@ def sample(self,

guide_seed = [dr.detach(ray_seed), active_guide | mask_replace]

return L, active, guide_seed if project else None
return L, active, [], guide_seed if project else None


def sample_radiance_difference(self, scene, ss, curr_depth, sampler, active):
Expand Down Expand Up @@ -282,7 +282,7 @@ def sample_radiance_difference(self, scene, ss, curr_depth, sampler, active):

# ----------- Estimate the radiance of the background -----------
ray_bg = ss.spawn_ray()
radiance_bg, _, _ = self.sample(
radiance_bg, _, _, _ = self.sample(
dr.ADMode.Primal, scene, sampler, ray_bg, curr_depth, None, None, active, False, None)

# ----------- Estimate the radiance of the foreground -----------
Expand Down Expand Up @@ -327,7 +327,7 @@ def sample_radiance_difference(self, scene, ss, curr_depth, sampler, active):
si_fg.wi[wrong_side] = si_fg.to_local(-ss.d)

# Estimate the radiance starting from the surface interaction
radiance_fg, _, _ = self.sample(
radiance_fg, _, _, _ = self.sample(
dr.ADMode.Primal, scene, sampler, ray_bg, curr_depth, None, None, active, False, si_fg)

else:
Expand Down
4 changes: 2 additions & 2 deletions src/python/python/ad/integrators/prb.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,8 +66,7 @@ def sample(self,
state_in: Optional[mi.Spectrum],
active: mi.Bool,
**kwargs # Absorbs unused arguments
) -> Tuple[mi.Spectrum,
mi.Bool, mi.Spectrum]:
) -> Tuple[mi.Spectrum, mi.Bool, List[mi.Float], mi.Spectrum]:
"""
See ``ADIntegrator.sample()`` for a description of this interface and
the role of the various parameters and return values.
Expand Down Expand Up @@ -250,6 +249,7 @@ def sample(self,
return (
L if primal else δL, # Radiance/differential radiance
(depth != 0), # Ray validity flag for alpha blending
[], # Empty typle of AOVs
L # State for the differential phase
)

Expand Down
4 changes: 2 additions & 2 deletions src/python/python/ad/integrators/prb_basic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,7 @@ def sample(self,
state_in: Optional[mi.Spectrum],
active: mi.Bool,
**kwargs # Absorbs unused arguments
) -> Tuple[mi.Spectrum,
mi.Bool, mi.Spectrum]:
) -> Tuple[mi.Spectrum, mi.Bool, List[mi.Float], mi.Spectrum]:
"""
See ``ADIntegrator.sample()`` for a description of this interface and
the role of the various parameters and return values.
Expand Down Expand Up @@ -161,6 +160,7 @@ def sample(self,
return (
L if primal else δL, # Radiance/differential radiance
dr.neq(depth, 0), # Ray validity flag for alpha blending
[], # Empty typle of AOVs
L # State the for differential phase
)

Expand Down
7 changes: 4 additions & 3 deletions src/python/python/ad/integrators/prb_projective.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ def sample(self,
project: bool = False,
si_shade: Optional[mi.SurfaceInteraction3f] = None,
**kwargs # Absorbs unused arguments
) -> Tuple[mi.Spectrum, mi.Bool, Any]:
) -> Tuple[mi.Spectrum, mi.Bool, List[mi.Float], Any]:
"""
See ``PSIntegrator.sample()`` for a description of this interface and
the role of the various parameters and return values.
Expand Down Expand Up @@ -372,6 +372,7 @@ def sample(self,
return (
L if primal else δL, # Radiance/differential radiance
depth != 0, # Ray validity flag for alpha blending
[], # Empty tuple of AOVs.
# Seed rays, or the state for the differential phase
[dr.detach(ray_seed), active_seed] if project else L
)
Expand All @@ -389,7 +390,7 @@ def sample_radiance_difference(self, scene, ss, curr_depth, sampler, active):
ss.p + (1 + dr.max(dr.abs(ss.p))) * (ss.d * ss.offset + ss.n * mi.math.ShapeEpsilon),
ss.d
)
radiance_bg, _, _ = self.sample(
radiance_bg, _, _, _ = self.sample(
dr.ADMode.Primal, scene, sampler, ray_bg, curr_depth, None, None, active, False, None)

# ----------- Estimate the radiance of the foreground -----------
Expand Down Expand Up @@ -427,7 +428,7 @@ def sample_radiance_difference(self, scene, ss, curr_depth, sampler, active):

# Call `sample()` to estimate the radiance starting from the surface
# interaction
radiance_fg, _, _ = self.sample(
radiance_fg, _, _, _ = self.sample(
dr.ADMode.Primal, scene, sampler, ray_bg, curr_depth, None, None, active, False, si_fg)

# Compute the radiance difference
Expand Down
Loading

0 comments on commit 3012b0a

Please sign in to comment.