From b724231ecd449a396bc22eeb93a00b111b28024d Mon Sep 17 00:00:00 2001 From: oarriaga Date: Sun, 2 Aug 2026 20:35:14 +0200 Subject: [PATCH 1/5] refactor: split the renderer into a package renderer.py becomes paz/graphics/renderer/, one module per stage of the pipeline: rays (the bounce loop), tiling, intersect, shade, shadow, optics, material. The loop body now reads intersect -> shade -> advance instead of burying the scatter step inside update_state. Avoids "trace", which in a JAX-first library already means jit tracing: trace_bounces is rays.render, trace_chunks is rays.render_chunks. The four shape-rendering examples patched paz.graphics.renderer.compute_soft_occlusion, which the split would have silently turned into a no-op; they now patch renderer.shadow directly. The tiling module is not called "tiles" because render() takes a public tiles= keyword that examples pass by name. No behaviour change. 106 tests pass in renderer_test.py and mesh_test.py including all snapshots; all 16 benchmark gradients are bit-identical to the pre-split values. --- .../phong_realism_levels.py | 6 +- .../realistic_primitives.py | 6 +- .../visualize_scene_gradients.py | 6 +- paz/graphics/renderer.py | 606 ------------------ paz/graphics/renderer/__init__.py | 48 ++ paz/graphics/renderer/intersect.py | 92 +++ paz/graphics/renderer/material.py | 25 + paz/graphics/renderer/optics.py | 77 +++ paz/graphics/renderer/rays.py | 146 +++++ paz/graphics/renderer/shade.py | 134 ++++ paz/graphics/renderer/shadow.py | 95 +++ paz/graphics/renderer/tiling.py | 34 + paz/graphics/renderer_test.py | 38 +- paz/graphics/types.py | 2 + 14 files changed, 682 insertions(+), 633 deletions(-) delete mode 100644 paz/graphics/renderer.py create mode 100644 paz/graphics/renderer/__init__.py create mode 100644 paz/graphics/renderer/intersect.py create mode 100644 paz/graphics/renderer/material.py create mode 100644 paz/graphics/renderer/optics.py create mode 100644 paz/graphics/renderer/rays.py create mode 100644 paz/graphics/renderer/shade.py create mode 100644 paz/graphics/renderer/shadow.py create mode 100644 paz/graphics/renderer/tiling.py diff --git a/examples/differentiable_shape_rendering/phong_realism_levels.py b/examples/differentiable_shape_rendering/phong_realism_levels.py index 41b08d4be..125e65c13 100644 --- a/examples/differentiable_shape_rendering/phong_realism_levels.py +++ b/examples/differentiable_shape_rendering/phong_realism_levels.py @@ -1,10 +1,10 @@ import jax import jax.numpy as jp import paz -import paz.graphics.renderer as paz_renderer +from paz.graphics.renderer import shadow -soft_occlusion = paz_renderer.compute_soft_occlusion -paz_renderer.compute_soft_occlusion = paz.partial(soft_occlusion, slope=1.0) +soft_occlusion = shadow.compute_soft_occlusion +shadow.compute_soft_occlusion = paz.partial(soft_occlusion, slope=1.0) H, W = 480, 640 diff --git a/examples/differentiable_shape_rendering/realistic_primitives.py b/examples/differentiable_shape_rendering/realistic_primitives.py index 42d1223fe..85eca7f23 100644 --- a/examples/differentiable_shape_rendering/realistic_primitives.py +++ b/examples/differentiable_shape_rendering/realistic_primitives.py @@ -1,10 +1,10 @@ import jax import jax.numpy as jp import paz -import paz.graphics.renderer as paz_renderer +from paz.graphics.renderer import shadow -soft_occlusion = paz_renderer.compute_soft_occlusion -paz_renderer.compute_soft_occlusion = paz.partial(soft_occlusion, slope=1.0) +soft_occlusion = shadow.compute_soft_occlusion +shadow.compute_soft_occlusion = paz.partial(soft_occlusion, slope=1.0) BLUE = jp.array([0.324, 0.692, 0.863]) # MyBlue GREEN = jp.array([154 / 255, 213 / 255, 135 / 255]) # YlGnI diff --git a/examples/differentiable_shape_rendering/visualize_scene_gradients.py b/examples/differentiable_shape_rendering/visualize_scene_gradients.py index 0b5db7024..f89314dab 100644 --- a/examples/differentiable_shape_rendering/visualize_scene_gradients.py +++ b/examples/differentiable_shape_rendering/visualize_scene_gradients.py @@ -8,7 +8,7 @@ import matplotlib.pyplot as plt import numpy as np import paz -import paz.graphics.renderer as paz_renderer +from paz.graphics.renderer import shadow import paz.utils.plot as plot from mpl_toolkits.axes_grid1 import make_axes_locatable @@ -23,7 +23,7 @@ CURVATURE_RATIO = 1e-3 LUMINANCE_WEIGHTS = jp.array([0.2126, 0.7152, 0.0722]) SIGNED_CMAP = "RdBu_r" -SOFT_OCCLUSION = paz_renderer.compute_soft_occlusion +SOFT_OCCLUSION = shadow.compute_soft_occlusion CAMERA_ARGS = ( jp.array([0.0, 2.0, 2.0]), jp.array([0.0, 0.0, 0.0]), @@ -77,7 +77,7 @@ def build_scene(shape_transform=jp.eye(4)): def configure_soft_occlusion(): - paz_renderer.compute_soft_occlusion = paz.partial(SOFT_OCCLUSION, slope=1.0) + shadow.compute_soft_occlusion = paz.partial(SOFT_OCCLUSION, slope=1.0) def compute_autodiff_gradient(function, args, basis): diff --git a/paz/graphics/renderer.py b/paz/graphics/renderer.py deleted file mode 100644 index d6dd9adc1..000000000 --- a/paz/graphics/renderer.py +++ /dev/null @@ -1,606 +0,0 @@ -from collections import namedtuple - -import jax -import jax.numpy as jp - -import paz - -from paz.graphics.composite import ( - compute_scene_hit_mask, - find_closest_intersection_args, - postprocess, - take_closest, -) - -# TODO retune: measured best under 6k faces, but 2048 is 1.75x faster -# at 82k, so this is wrong for the room-scale meshes ahead. -FACE_CHUNK_SIZE = 128 -SHADOW_ORIGIN_EPSILON = 1e-5 -SHADOW_SELF_HIT_EPSILON = 1e-5 -BOUNCE_ORIGIN_EPSILON = 1e-2 - -TRIANGLE_HIT_NAMES = "hit_mask depth points normals eyes albedo primitive" -TriangleHit = namedtuple("TriangleHit", TRIANGLE_HIT_NAMES.split()) - -STATE_NAMES = "color depth hit_mask throughput active_mask " -STATE_NAMES += "refractive_index rays" -RenderState = namedtuple("RenderState", STATE_NAMES.split()) - -MATERIAL_NAMES = "reflectivities transparencies refractive_indices" -HitMaterial = namedtuple("HitMaterial", MATERIAL_NAMES.split()) - -Surfaces = namedtuple("Surfaces", ["points", "normals", "eyes"]) - - -def render( - shape, y_FOV, pose, scene, mask, lights, tiles, chunk_size, - shadows=False, shadow_mask=None, num_bounces=1, - face_chunk_size=FACE_CHUNK_SIZE, -): - compiled = paz.graphics.scene.compile(scene, lights, mask, shadow_mask) - trace_args = compiled, shadows, num_bounces, face_chunk_size, chunk_size - trace = paz.lock(trace_chunks, *trace_args) - tile_step = paz.lock(render_tile_step, shape, y_FOV, pose, tiles, trace) - images, depths = scan_tiles(shape, tiles, tile_step) - image = assemble_tiles(shape, tiles, images) - depth = assemble_tiles(shape, tiles, depths)[..., 0] - return image, depth - - -def render_masks( - shape, y_FOV, pose, scene, lights, depth, tiles, chunk_size, - num_objects=None, shadows=False, shadow_mask=None, num_bounces=1, - face_chunk_size=FACE_CHUNK_SIZE, -): - if num_objects is None: - num_objects = len(scene.nodes) - min_depth, max_depth = depth - masks = [] - for object_arg in range(num_objects): - mask = build_object_mask(len(scene.nodes), object_arg) - args = shape, y_FOV, pose, scene, mask, lights, tiles, chunk_size - args += shadows, shadow_mask, num_bounces, face_chunk_size - _, depth_image = render(*args) - soft = paz.depth.to_soft_mask(depth_image, min_depth, max_depth) - masks.append(jp.expand_dims(soft, axis=-1)) - return jp.stack(masks) - - -def build_object_mask(num_nodes, object_arg): - return jp.zeros((num_nodes,), dtype=bool).at[object_arg].set(True) - - -def scan_tiles(shape, tiles, tile_step): - H, W = shape - H_tiles, W_tiles = tiles - paz.graphics.mesh.assert_exact_tile_side(H, H_tiles) - paz.graphics.mesh.assert_exact_tile_side(W, W_tiles) - coordinates = paz.graphics.mesh.make_tile_coordinates(H_tiles, W_tiles) - return jax.lax.scan(tile_step, None, coordinates)[1] - - -def render_tile_step(carry, tile_arg, shape, y_FOV, pose, tiles, trace): - H, W = shape - H_tiles, W_tiles = tiles - camera_to_world = jp.linalg.inv(pose) - tile_args = H, W, H_tiles, W_tiles, y_FOV, camera_to_world - rays = paz.graphics.mesh.build_tile_rays(*tile_args, tile_arg) - hit_mask, depth, color = trace(rays) - tile_H, tile_W = H // H_tiles, W // W_tiles - post_args = hit_mask, depth, color, pose, rays, tile_H, tile_W - image, depth = postprocess(*post_args) - return carry, (image, jp.expand_dims(depth, -1)) - - -def assemble_tiles(shape, tiles, images): - H, W = shape - H_tiles, W_tiles = tiles - return paz.graphics.mesh.assemble(H, W, H_tiles, W_tiles, images) - - -def trace_chunks(rays, compiled, shadows, num_bounces, face_chunk, chunk_size): - ray_chunks = split_ray_chunks(rays, chunk_size) - step_args = compiled, shadows, num_bounces, face_chunk - trace_step = paz.lock(trace_chunk_step, *step_args) - hit_mask, depth, color = jax.lax.scan(trace_step, None, ray_chunks)[1] - return flatten_chunk_results(hit_mask, depth, color, len(rays[0])) - - -def trace_chunk_step(carry, rays, compiled, shadows, num_bounces, face_chunk): - args = rays, compiled, shadows, num_bounces, face_chunk - return carry, trace_bounces(*args) - - -def split_ray_chunks(rays, chunk_size): - origins, directions = rays - origins = pad_to_chunks(origins, chunk_size) - directions = pad_to_chunks(directions, chunk_size) - num_chunks = origins.shape[0] // chunk_size - shape = num_chunks, chunk_size, 3 - return origins.reshape(shape), directions.reshape(shape) - - -def pad_to_chunks(array, chunk_size): - remainder = array.shape[0] % chunk_size - if remainder == 0: - padded = array - else: - padding = jp.repeat(array[-1:], chunk_size - remainder, axis=0) - padded = jp.concatenate([array, padding], axis=0) - return padded - - -def flatten_chunk_results(hit_mask, depth, color, num_rays): - hit_mask = flatten_chunk_array(hit_mask, num_rays) - depth = flatten_chunk_array(depth, num_rays) - color = flatten_chunk_array(color, num_rays) - return hit_mask, depth, color - - -def flatten_chunk_array(array, num_rays): - shape = (-1,) + array.shape[2:] - return array.reshape(shape)[:num_rays] - - -def trace_bounces(rays, compiled, shadows, bounces, face_chunk): - state = initialize_state(rays) - bounce = paz.lock(bounce_step, compiled, shadows, face_chunk) - for step_arg in range(bounces): - state = bounce(state, step_arg) - return state.hit_mask, state.depth, state.color - - -def initialize_state(rays): - num_rays = rays[0].shape[0] - color = jp.zeros((num_rays, 3)) - depth = jp.full((num_rays,), paz.graphics.FARAWAY) - hit_mask = jp.zeros((num_rays,), dtype=bool) - throughput = jp.ones((num_rays, 3)) - active_mask = jp.ones((num_rays,), dtype=bool) - refractive_index = jp.ones((num_rays,)) - fields = color, depth, hit_mask, throughput, active_mask - return RenderState(*fields, refractive_index, rays) - - -def bounce_step(state, bounce, compiled, shadows, face_chunk): - triangle_hit = compute_triangle_hit(compiled, state.rays, face_chunk) - hit_masks, depths, surfaces = intersect(compiled, state.rays, triangle_hit) - closest = gather_closest(hit_masks, depths, surfaces) - state = update_first_hit(state, closest, bounce) - state = update_active_mask(state, closest) - color_args = compiled, closest, surfaces, shadows, triangle_hit - colors = compute_hit_colors(*color_args) - return update_state(state, compiled, closest, colors) - - -def update_first_hit(state, closest, bounce): - if bounce == 0: - depth, hit_mask = closest.depth, closest.hit_mask - else: - depth, hit_mask = state.depth, state.hit_mask - return state._replace(depth=depth, hit_mask=hit_mask) - - -def update_active_mask(state, closest): - active_mask = state.active_mask & closest.hit_mask - return state._replace(active_mask=active_mask) - - -def compute_triangle_hit(compiled, rays, face_chunk): - if compiled.triangles is None: - triangle_hit = None - else: - triangle_hit = build_triangle_hit(compiled, rays, face_chunk) - return triangle_hit - - -def build_triangle_hit(compiled, rays, face_chunk): - triangles = compiled.triangles - args = triangles, rays, face_chunk - result = paz.graphics.mesh.intersect_triangles(*args) - hit_mask, depth, points, normals, eyes, face_index, u, v = result - primitive = triangles.primitive_index[face_index] - hit_mask = jp.logical_and(hit_mask, compiled.triangle_mask[primitive]) - depth = jp.where(hit_mask, depth, paz.graphics.FARAWAY) - albedo_args = triangles, face_index, u, v - albedo = paz.graphics.albedo.compute_triangle_albedo(*albedo_args) - args = hit_mask, depth, points, normals, eyes, albedo, primitive - return TriangleHit(*args) - - -def intersect(compiled, rays, triangle_hit): - rows = [] - if len(compiled.shapes) > 0: - rows.append(intersect_shapes(compiled.shapes, rays, compiled.mask)) - if triangle_hit is not None: - rows.append(build_triangle_row(triangle_hit)) - joined = tuple(jp.concatenate(fields, axis=0) for fields in zip(*rows)) - hit_masks, depths, points, normals, eyes = joined - return hit_masks, depths, Surfaces(points, normals, eyes) - - -def intersect_shapes(shapes, rays, mask): - merged = paz.graphics.shapes.field_merge(shapes, ["transform", "type"]) - intersect_fun = paz.lock(paz.graphics.shapes.intersect, *rays) - hit_masks, depths, points, normals, eyes = jax.vmap(intersect_fun)(merged) - hit_masks = jp.where(jp.expand_dims(mask, 1), hit_masks, False) - return hit_masks, depths, points, normals, eyes - - -def build_triangle_row(triangle_hit): - depth = jp.expand_dims(triangle_hit.depth, -1) - fields = triangle_hit.hit_mask, depth, triangle_hit.points - fields += triangle_hit.normals, triangle_hit.eyes - return tuple(jp.expand_dims(field, 0) for field in fields) - - -def gather_closest(hit_masks, depths, surfaces): - indices = find_closest_intersection_args(hit_masks, depths) - fields = hit_masks, depths, *surfaces - closest = tuple(take_closest(field, indices) for field in fields) - return paz.graphics.Hit(*closest, indices) - - -def compute_hit_colors(compiled, closest, surfaces, shadows, triangle_hit): - if len(compiled.shapes) == 0: - colors = compute_triangle_colors(compiled, triangle_hit) - elif triangle_hit is None: - colors = compute_shape_colors(compiled, closest, surfaces, shadows) - else: - args = compiled, closest, surfaces, shadows, triangle_hit - colors = blend_hit_colors(*args) - return colors - - -def blend_hit_colors(compiled, closest, surfaces, shadows, triangle_hit): - # TODO a mixed scene shades both paths for every ray and throws one - # away. Joining them before selection needs color_with_shadows to - # return rows instead of selecting inside its per-light scan. - shape_colors = compute_shape_colors(compiled, closest, surfaces, shadows) - triangle_colors = compute_triangle_colors(compiled, triangle_hit) - is_triangle = closest.primitive_index == len(compiled.shapes) - is_triangle = jp.expand_dims(is_triangle, -1) - return jp.where(is_triangle, triangle_colors, shape_colors) - - -def compute_shape_colors(compiled, closest, surfaces, shadows): - num_shapes = len(compiled.shapes) - indices = jp.minimum(closest.primitive_index, num_shapes - 1) - surfaces = slice_surfaces(surfaces, 0, num_shapes) - if shadows: - colors = color_with_shadows(compiled, closest, surfaces, indices) - else: - colors = color_without_shadow(compiled, surfaces, indices) - return colors - - -def compute_triangle_colors(compiled, triangle_hit): - materials = compiled.triangles.materials - material = gather_triangle_material(materials, triangle_hit.primitive) - shader = select_shader(materials) - colors = jp.zeros_like(triangle_hit.albedo) - for light in compiled.lights: - args = triangle_hit.albedo, material, triangle_hit.points - args += triangle_hit.normals, triangle_hit.eyes, light - colors = colors + shader.compute_colors(*args) - return colors - - -def gather_triangle_material(materials, primitive): - material = jax.tree.map(lambda field: field[primitive], materials) - return jax.tree.map(expand_scalar_field, material) - - -def expand_scalar_field(field): - if field.ndim == 1: - field = jp.expand_dims(field, -1) - return field - - -def select_shader(material): - if isinstance(material, paz.graphics.CookTorranceMaterial): - shader = paz.graphics.cook_torrance - else: - shader = paz.graphics.phong - return shader - - -def color_without_shadow(compiled, surfaces, indices): - colors = [] - lights = paz.graphics.shapes.merge(*compiled.lights) - for group, start_arg, final_arg in iterate_shape_groups(compiled.shapes): - group_surfaces = slice_surfaces(surfaces, start_arg, final_arg) - albedo = compute_group_albedo(group, group_surfaces.points) - shader = select_shader(group.material) - axes = 0, 0, 0, 0, 0, None - color_per_light = jax.vmap(shader.compute_colors, axes) - color = jax.vmap(color_per_light, (None, None, None, None, None, 0)) - args = albedo, group.material, *group_surfaces, lights - colors.append(jp.sum(color(*args), axis=0)) - return take_closest(jp.concatenate(colors, axis=0), indices) - - -def iterate_shape_groups(shapes): - start_arg = 0 - for group in paz.graphics.shapes.group_by_pattern_size(shapes).values(): - final_arg = start_arg + len(group) - yield paz.graphics.shapes.merge(*group), start_arg, final_arg - start_arg = final_arg - - -def slice_surfaces(surfaces, start_arg, final_arg): - points = surfaces.points[start_arg:final_arg] - normals = surfaces.normals[start_arg:final_arg] - eyes = surfaces.eyes[start_arg:final_arg] - return Surfaces(points, normals, eyes) - - -def compute_group_albedo(group, points): - compute_albedo = paz.graphics.albedo.compute_shape_albedo - return jax.vmap(compute_albedo)(group, group.material, points) - - -def color_with_shadows(compiled, closest, surfaces, indices): - colors = jp.zeros((len(surfaces.points[0]), 3)) - lights = paz.graphics.shapes.merge(*compiled.lights) - body = paz.lock(scan_light_step, compiled, closest, surfaces, indices) - return jax.lax.scan(body, colors, lights)[0] - - -def scan_light_step(colors, light, compiled, closest, surfaces, indices): - args = compiled, closest, surfaces, indices, light - return colors + compute_light_colors(*args), None - - -def compute_light_colors(compiled, closest, surfaces, indices, light): - directions, distance = compute_light_directions(light, closest.point) - occlusion_args = compiled, closest, indices, directions, distance - is_shadow = compute_light_occlusion(*occlusion_args) - color_args = compiled.shapes, light, surfaces, is_shadow - return take_closest(compute_shadowed_colors(*color_args), indices) - - -def compute_light_directions(light, points): - vector = light.position - points - norm = paz.algebra.compute_norms(vector, 1) - return vector / norm, jp.squeeze(norm, axis=1) - - -def compute_light_occlusion(compiled, closest, indices, directions, distance): - origins = compute_shadow_ray_origins(closest.point, closest.normal) - shadow_args = compiled.shapes, origins, directions - hit_masks, depths, _, _, _, casters = intersect_shadow_groups(*shadow_args) - masks = resolve_shadow_masks(compiled, hit_masks) - depth_args = masks, depths, casters, indices, closest.normal, directions - masks, depths = select_shadow_depths(*depth_args) - return compute_soft_occlusion(masks, depths, distance) - - -def compute_shadow_ray_origins(points, normals): - over_point, _ = compute_surface_points(points, normals) - return over_point - - -def compute_surface_points(point, normal): - over_point = point + normal * SHADOW_ORIGIN_EPSILON - under_point = point - normal * SHADOW_ORIGIN_EPSILON - return over_point, under_point - - -def intersect_shadow_groups(shapes, origins, directions): - intersect_all = paz.graphics.shapes.intersect_all - intersect_group = jax.vmap(paz.lock(intersect_all, origins, directions)) - rows = [] - for group, start_arg, final_arg in iterate_shape_groups(shapes): - indices = jp.arange(start_arg, final_arg) - rows.append((*intersect_group(group), indices)) - return tuple(jp.concatenate(fields, axis=0) for fields in zip(*rows)) - - -def resolve_shadow_masks(compiled, hit_masks): - transparencies = compute_transparencies(compiled.shapes) - masks = jp.where(jp.expand_dims(compiled.mask, 1), hit_masks, False) - masks = hide_transparent_shapes(masks, transparencies > 0.0) - if compiled.shadow_mask is None: - resolved = masks - else: - resolved = hide_non_casting_shapes(masks, compiled.shadow_mask) - return resolved - - -def compute_transparencies(shapes): - return jp.array([shape.material.transparency for shape in shapes]) - - -def hide_transparent_shapes(shadow_masks, is_transparent): - return jp.where(jp.expand_dims(is_transparent, 1), False, shadow_masks) - - -def hide_non_casting_shapes(shadow_masks, shadow_mask): - return jp.where(jp.expand_dims(shadow_mask, 1), shadow_masks, False) - - -def select_shadow_depths( - hit_masks, depths, casters, receivers, normals, directions -): - same_shape = casters[:, None] == receivers[None, :] - front_args = same_shape, normals, directions - front_side_hits = compute_front_side_shadow_mask(*front_args) - root_args = hit_masks, depths, same_shape, front_side_hits - valid_roots = compute_valid_roots(*root_args) - depths = jp.where(valid_roots, depths, paz.graphics.FARAWAY) - return jp.any(valid_roots, axis=1), jp.min(depths, axis=1) - - -def compute_front_side_shadow_mask(same_shape, normals, directions): - front_side = paz.algebra.dot(normals, directions) >= 0.0 - return jp.logical_and(same_shape, front_side[None, :]) - - -def compute_valid_roots(hit_masks, depths, same_shape, front_side_hits): - thresholds = compute_shadow_depth_thresholds(same_shape) - valid_roots = depths > thresholds[:, None, :] - valid_roots = jp.logical_and(valid_roots, depths < paz.graphics.FARAWAY) - valid_roots = jp.logical_and(jp.expand_dims(hit_masks, 1), valid_roots) - return jp.logical_and(valid_roots, ~front_side_hits[:, None, :]) - - -def compute_shadow_depth_thresholds(same_shape): - return jp.where(same_shape, SHADOW_SELF_HIT_EPSILON, paz.graphics.EPSILON) - - -def compute_soft_occlusion(hit_masks, depths, light_lengths, slope=0.01): - closest_depths = jp.where(hit_masks, depths, paz.graphics.FARAWAY) - closest_depths = jp.min(closest_depths, axis=0) - scene_hit_mask = compute_scene_hit_mask(hit_masks) - blockers = closest_depths <= light_lengths - blocker_mask = jp.logical_and(scene_hit_mask, blockers) - difference = light_lengths - closest_depths - occlusion = jax.nn.sigmoid(slope * difference) - return jp.where(blocker_mask, occlusion, 0.0) - - -def compute_shadowed_colors(shapes, light, surfaces, is_shadow): - colors = [] - for group, start_arg, final_arg in iterate_shape_groups(shapes): - group_surfaces = slice_surfaces(surfaces, start_arg, final_arg) - albedo = compute_group_albedo(group, group_surfaces.points) - shader = select_shader(group.material) - axes = 0, 0, 0, 0, 0, None, None - color = jax.vmap(shader.compute_colors_with_shadow, axes) - args = albedo, group.material, *group_surfaces, light, is_shadow - colors.append(color(*args)) - return jp.concatenate(colors, axis=0) - - -def update_state(state, compiled, closest, hit_colors): - material = compute_material_properties(compiled, closest.primitive_index) - color_args = state.color, state.throughput, state.active_mask, hit_colors - color_args += material.reflectivities, material.transparencies - color = accumulate_color(*color_args) - new_rays, n_2, reflectance = compute_bounce(state, closest, material) - args = new_rays, n_2, material, reflectance - return apply_bounce_update(state._replace(color=color), *args) - - -def compute_material_properties(compiled, hit_shape_args): - values = collect_material_values(compiled) - gathered = tuple(jp.array(row)[hit_shape_args] for row in values) - return HitMaterial(*gathered) - - -def collect_material_values(compiled): - reflectivities, transparencies, refractive_indices = [], [], [] - for shape in compiled.shapes: - reflectivities.append(shape.material.reflective) - transparencies.append(shape.material.transparency) - refractive_indices.append(shape.material.refractive_index) - if compiled.triangles is not None: - reflectivities.append(0.0) - transparencies.append(0.0) - refractive_indices.append(1.0) - return reflectivities, transparencies, refractive_indices - - -def accumulate_color( - colors, throughput, active_mask, hit_colors, reflectivities, transparencies -): - weights = jp.maximum(1.0 - reflectivities - transparencies, 0.0) - weights = jp.expand_dims(weights, -1) - active_mask = jp.expand_dims(active_mask, -1) - return colors + (throughput * active_mask * weights * hit_colors) - - -def compute_bounce(state, closest, material): - args = state.rays[1], closest.normal, state.refractive_index - terms = compute_refraction_terms(*args, material.refractive_indices) - normal, eye, n_1, n_2, n_ratio = terms - reflectance = compute_reflectance(normal, eye, n_1, n_2) - ray_args = normal, eye, n_ratio, closest.point, material.transparencies - return compute_new_rays(*ray_args), n_2, reflectance - - -def compute_refraction_terms(directions, normal, n_1, refractive_indices): - eye = -directions - normal, is_inside = flip_normal_if_inside(eye, normal) - # TODO why 1.0 hardcoded - n_2 = jp.where(is_inside, 1.0, refractive_indices) - return normal, eye, n_1, n_2, n_1 / n_2 - - -def flip_normal_if_inside(eye, normal): - is_inside = jp.sum(normal * eye, axis=-1) < 0.0 - return jp.where(jp.expand_dims(is_inside, -1), -normal, normal), is_inside - - -def compute_reflectance(normal, eye, n_1, n_2): - cosines = compute_transmission_cosines(eye, normal, n_1 / n_2) - cos_incident, sin_transmit_squared, cos_transmit = cosines - cos = jp.where(n_1 > n_2, cos_transmit, cos_incident) - base_reflectance = ((n_1 - n_2) / (n_1 + n_2)) ** 2 - grazing = (1.0 - cos) ** 5 - reflectance = base_reflectance + (1.0 - base_reflectance) * grazing - return jp.where(sin_transmit_squared > 1.0, 1.0, reflectance) - - -def compute_transmission_cosines(eye, normal, n_ratio): - cos_incident = jp.sum(eye * normal, axis=-1) - sin_transmit_squared = (n_ratio**2) * (1.0 - (cos_incident**2)) - cos_transmit = jp.sqrt(jp.maximum(0.0, 1.0 - sin_transmit_squared)) - return cos_incident, sin_transmit_squared, cos_transmit - - -def compute_new_rays(normal, eye, n_ratio, point, transparencies): - is_transparent = transparencies > 0.0 - do_reflect = jp.expand_dims(~is_transparent, -1) - reflection = paz.graphics.geometry.reflect(-eye, normal) - refraction = compute_refractive_direction(eye, normal, n_ratio) - direction = jp.where(do_reflect, reflection, refraction) - lower_point, upper_point = displace_by_normal(point, normal) - origin = jp.where(do_reflect, upper_point, lower_point) - return origin, paz.algebra.normalize(direction) - - -def compute_refractive_direction(eye, normal, n_ratio): - args = eye, normal, n_ratio - cos_incident, _, cos_transmit = compute_transmission_cosines(*args) - inside_vector = -eye * jp.expand_dims(n_ratio, -1) - up_weight = n_ratio * cos_incident - cos_transmit - return jp.expand_dims(up_weight, -1) * normal + inside_vector - - -def displace_by_normal(point, normal): - upper_point = point + normal * BOUNCE_ORIGIN_EPSILON - lower_point = point - normal * BOUNCE_ORIGIN_EPSILON - return lower_point, upper_point - - -def apply_bounce_update(state, new_rays, n_2, material, reflectance): - is_transparent = material.transparencies > 0.0 - is_reflective = material.reflectivities > 0.0 - factor = compute_bounce_factor(material, reflectance) - factor = jp.where(is_transparent & (reflectance >= 1.0), 1.0, factor) - throughput = state.throughput * jp.expand_dims(factor, -1) - active_mask = state.active_mask & (is_transparent | is_reflective) - index_args = state, n_2, is_transparent, reflectance - refractive_index = update_refractive_index(*index_args) - args = throughput, active_mask, refractive_index, new_rays - return replace_bounce_state(state, *args) - - -def compute_bounce_factor(material, reflectance): - is_transparent = material.transparencies > 0.0 - is_reflective = material.reflectivities > 0.0 - transparent_factor = material.transparencies * (1.0 - reflectance) - reflective_factor = jp.where(is_reflective, material.reflectivities, 0.0) - return jp.where(is_transparent, transparent_factor, reflective_factor) - - -def update_refractive_index(state, n_2, is_transparent, reflectance): - update_mask = is_transparent & (reflectance < 1.0) - return jp.where(update_mask, n_2, state.refractive_index) - - -def replace_bounce_state(state, throughput, active_mask, index, rays): - state = state._replace(throughput=throughput, active_mask=active_mask) - return state._replace(refractive_index=index, rays=rays) diff --git a/paz/graphics/renderer/__init__.py b/paz/graphics/renderer/__init__.py new file mode 100644 index 000000000..65efde72d --- /dev/null +++ b/paz/graphics/renderer/__init__.py @@ -0,0 +1,48 @@ +import jax.numpy as jp + +import paz + +from paz.graphics.renderer import rays, tiling + +# TODO retune: measured best under 6k faces, but 2048 is 1.75x faster +# at 82k, so this is wrong for the room-scale meshes ahead. +FACE_CHUNK_SIZE = 128 + + +def render( + shape, y_FOV, pose, scene, mask, lights, tiles, chunk_size, + shadows=False, shadow_mask=None, num_bounces=1, + face_chunk_size=FACE_CHUNK_SIZE, +): + compiled = paz.graphics.scene.compile(scene, lights, mask, shadow_mask) + ray_args = compiled, shadows, num_bounces, face_chunk_size, chunk_size + render_rays = paz.lock(rays.render_chunks, *ray_args) + step_args = shape, y_FOV, pose, tiles, render_rays + tile_step = paz.lock(tiling.render_step, *step_args) + images, depths = tiling.scan(shape, tiles, tile_step) + image = tiling.assemble(shape, tiles, images) + depth = tiling.assemble(shape, tiles, depths)[..., 0] + return image, depth + + +def render_masks( + shape, y_FOV, pose, scene, lights, depth, tiles, chunk_size, + num_objects=None, shadows=False, shadow_mask=None, num_bounces=1, + face_chunk_size=FACE_CHUNK_SIZE, +): + if num_objects is None: + num_objects = len(scene.nodes) + min_depth, max_depth = depth + masks = [] + for object_arg in range(num_objects): + mask = build_object_mask(len(scene.nodes), object_arg) + args = shape, y_FOV, pose, scene, mask, lights, tiles, chunk_size + args += shadows, shadow_mask, num_bounces, face_chunk_size + _, depth_image = render(*args) + soft = paz.depth.to_soft_mask(depth_image, min_depth, max_depth) + masks.append(jp.expand_dims(soft, axis=-1)) + return jp.stack(masks) + + +def build_object_mask(num_nodes, object_arg): + return jp.zeros((num_nodes,), dtype=bool).at[object_arg].set(True) diff --git a/paz/graphics/renderer/intersect.py b/paz/graphics/renderer/intersect.py new file mode 100644 index 000000000..333b11287 --- /dev/null +++ b/paz/graphics/renderer/intersect.py @@ -0,0 +1,92 @@ +from collections import namedtuple + +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.composite import find_closest_intersection_args, take_closest +from paz.graphics.types import Surfaces + +TRIANGLE_HIT_NAMES = "hit_mask depth points normals eyes albedo primitive" +TriangleHit = namedtuple("TriangleHit", TRIANGLE_HIT_NAMES.split()) + + +def build_candidates(compiled, rays, triangle_hit): + rows = [] + if len(compiled.shapes) > 0: + rows.append(intersect_shapes(compiled.shapes, rays, compiled.mask)) + if triangle_hit is not None: + rows.append(build_triangle_row(triangle_hit)) + joined = tuple(jp.concatenate(fields, axis=0) for fields in zip(*rows)) + hit_masks, depths, points, normals, eyes = joined + return hit_masks, depths, Surfaces(points, normals, eyes) + + +def intersect_shapes(shapes, rays, mask): + merged = paz.graphics.shapes.field_merge(shapes, ["transform", "type"]) + intersect_fun = paz.lock(paz.graphics.shapes.intersect, *rays) + hit_masks, depths, points, normals, eyes = jax.vmap(intersect_fun)(merged) + hit_masks = jp.where(jp.expand_dims(mask, 1), hit_masks, False) + return hit_masks, depths, points, normals, eyes + + +def build_triangle_row(triangle_hit): + depth = jp.expand_dims(triangle_hit.depth, -1) + fields = triangle_hit.hit_mask, depth, triangle_hit.points + fields += triangle_hit.normals, triangle_hit.eyes + return tuple(jp.expand_dims(field, 0) for field in fields) + + +def find_closest(hit_masks, depths, surfaces): + indices = find_closest_intersection_args(hit_masks, depths) + fields = hit_masks, depths, *surfaces + closest = tuple(take_closest(field, indices) for field in fields) + return paz.graphics.Hit(*closest, indices) + + +def compute_triangle_hit(compiled, rays, face_chunk): + if compiled.triangles is None: + triangle_hit = None + else: + triangle_hit = build_triangle_hit(compiled, rays, face_chunk) + return triangle_hit + + +def build_triangle_hit(compiled, rays, face_chunk): + triangles = compiled.triangles + args = triangles, rays, face_chunk + result = paz.graphics.mesh.intersect_triangles(*args) + hit_mask, depth, points, normals, eyes, face_index, u, v = result + primitive = triangles.primitive_index[face_index] + hit_mask = jp.logical_and(hit_mask, compiled.triangle_mask[primitive]) + depth = jp.where(hit_mask, depth, paz.graphics.FARAWAY) + albedo_args = triangles, face_index, u, v + albedo = paz.graphics.albedo.compute_triangle_albedo(*albedo_args) + args = hit_mask, depth, points, normals, eyes, albedo, primitive + return TriangleHit(*args) + + +def intersect_shadow_groups(shapes, origins, directions): + intersect_all = paz.graphics.shapes.intersect_all + intersect_group = jax.vmap(paz.lock(intersect_all, origins, directions)) + rows = [] + for group, start_arg, final_arg in iterate_shape_groups(shapes): + indices = jp.arange(start_arg, final_arg) + rows.append((*intersect_group(group), indices)) + return tuple(jp.concatenate(fields, axis=0) for fields in zip(*rows)) + + +def iterate_shape_groups(shapes): + start_arg = 0 + for group in paz.graphics.shapes.group_by_pattern_size(shapes).values(): + final_arg = start_arg + len(group) + yield paz.graphics.shapes.merge(*group), start_arg, final_arg + start_arg = final_arg + + +def slice_surfaces(surfaces, start_arg, final_arg): + points = surfaces.points[start_arg:final_arg] + normals = surfaces.normals[start_arg:final_arg] + eyes = surfaces.eyes[start_arg:final_arg] + return Surfaces(points, normals, eyes) diff --git a/paz/graphics/renderer/material.py b/paz/graphics/renderer/material.py new file mode 100644 index 000000000..a9443d008 --- /dev/null +++ b/paz/graphics/renderer/material.py @@ -0,0 +1,25 @@ +from collections import namedtuple + +import jax.numpy as jp + +MATERIAL_NAMES = "reflectivities transparencies refractive_indices" +HitMaterial = namedtuple("HitMaterial", MATERIAL_NAMES.split()) + + +def compute_material_properties(compiled, hit_shape_args): + values = collect_material_values(compiled) + gathered = tuple(jp.array(row)[hit_shape_args] for row in values) + return HitMaterial(*gathered) + + +def collect_material_values(compiled): + reflectivities, transparencies, refractive_indices = [], [], [] + for shape in compiled.shapes: + reflectivities.append(shape.material.reflective) + transparencies.append(shape.material.transparency) + refractive_indices.append(shape.material.refractive_index) + if compiled.triangles is not None: + reflectivities.append(0.0) + transparencies.append(0.0) + refractive_indices.append(1.0) + return reflectivities, transparencies, refractive_indices diff --git a/paz/graphics/renderer/optics.py b/paz/graphics/renderer/optics.py new file mode 100644 index 000000000..b4ccf90ec --- /dev/null +++ b/paz/graphics/renderer/optics.py @@ -0,0 +1,77 @@ +import jax.numpy as jp + +import paz + +BOUNCE_ORIGIN_EPSILON = 1e-2 + + +def compute_bounce(directions, closest, refractive_index, material): + args = directions, closest.normal, refractive_index + terms = compute_refraction_terms(*args, material.refractive_indices) + normal, eye, n_1, n_2, n_ratio = terms + reflectance = compute_reflectance(normal, eye, n_1, n_2) + ray_args = normal, eye, n_ratio, closest.point, material.transparencies + return compute_new_rays(*ray_args), n_2, reflectance + + +def compute_refraction_terms(directions, normal, n_1, refractive_indices): + eye = -directions + normal, is_inside = flip_normal_if_inside(eye, normal) + # TODO why 1.0 hardcoded + n_2 = jp.where(is_inside, 1.0, refractive_indices) + return normal, eye, n_1, n_2, n_1 / n_2 + + +def flip_normal_if_inside(eye, normal): + is_inside = jp.sum(normal * eye, axis=-1) < 0.0 + return jp.where(jp.expand_dims(is_inside, -1), -normal, normal), is_inside + + +def compute_reflectance(normal, eye, n_1, n_2): + cosines = compute_transmission_cosines(eye, normal, n_1 / n_2) + cos_incident, sin_transmit_squared, cos_transmit = cosines + cos = jp.where(n_1 > n_2, cos_transmit, cos_incident) + base_reflectance = ((n_1 - n_2) / (n_1 + n_2)) ** 2 + grazing = (1.0 - cos) ** 5 + reflectance = base_reflectance + (1.0 - base_reflectance) * grazing + return jp.where(sin_transmit_squared > 1.0, 1.0, reflectance) + + +def compute_transmission_cosines(eye, normal, n_ratio): + cos_incident = jp.sum(eye * normal, axis=-1) + sin_transmit_squared = (n_ratio**2) * (1.0 - (cos_incident**2)) + cos_transmit = jp.sqrt(jp.maximum(0.0, 1.0 - sin_transmit_squared)) + return cos_incident, sin_transmit_squared, cos_transmit + + +def compute_new_rays(normal, eye, n_ratio, point, transparencies): + is_transparent = transparencies > 0.0 + do_reflect = jp.expand_dims(~is_transparent, -1) + reflection = paz.graphics.geometry.reflect(-eye, normal) + refraction = compute_refractive_direction(eye, normal, n_ratio) + direction = jp.where(do_reflect, reflection, refraction) + lower_point, upper_point = displace_by_normal(point, normal) + origin = jp.where(do_reflect, upper_point, lower_point) + return origin, paz.algebra.normalize(direction) + + +def compute_refractive_direction(eye, normal, n_ratio): + args = eye, normal, n_ratio + cos_incident, _, cos_transmit = compute_transmission_cosines(*args) + inside_vector = -eye * jp.expand_dims(n_ratio, -1) + up_weight = n_ratio * cos_incident - cos_transmit + return jp.expand_dims(up_weight, -1) * normal + inside_vector + + +def displace_by_normal(point, normal): + upper_point = point + normal * BOUNCE_ORIGIN_EPSILON + lower_point = point - normal * BOUNCE_ORIGIN_EPSILON + return lower_point, upper_point + + +def compute_bounce_factor(material, reflectance): + is_transparent = material.transparencies > 0.0 + is_reflective = material.reflectivities > 0.0 + transparent_factor = material.transparencies * (1.0 - reflectance) + reflective_factor = jp.where(is_reflective, material.reflectivities, 0.0) + return jp.where(is_transparent, transparent_factor, reflective_factor) diff --git a/paz/graphics/renderer/rays.py b/paz/graphics/renderer/rays.py new file mode 100644 index 000000000..a697a9dcf --- /dev/null +++ b/paz/graphics/renderer/rays.py @@ -0,0 +1,146 @@ +from collections import namedtuple + +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.renderer import intersect, material, optics, shade + +STATE_NAMES = "color depth hit_mask throughput active_mask " +STATE_NAMES += "refractive_index rays" +RenderState = namedtuple("RenderState", STATE_NAMES.split()) + + +def render(rays, compiled, shadows, num_bounces, face_chunk): + state = initialize_state(rays) + step = paz.lock(bounce, compiled, shadows, face_chunk) + for step_arg in range(num_bounces): + state = step(state, step_arg) + return state.hit_mask, state.depth, state.color + + +def bounce(state, step_arg, compiled, shadows, face_chunk): + hit_args = compiled, state.rays, face_chunk + triangle_hit = intersect.compute_triangle_hit(*hit_args) + candidates = intersect.build_candidates(compiled, state.rays, triangle_hit) + hit_masks, depths, surfaces = candidates + closest = intersect.find_closest(hit_masks, depths, surfaces) + state = update_first_hit(state, closest, step_arg) + state = update_active_mask(state, closest) + color_args = compiled, closest, surfaces, shadows, triangle_hit + colors = shade.compute_hit_colors(*color_args) + return advance(state, compiled, closest, colors) + + +def advance(state, compiled, closest, hit_colors): + index = closest.primitive_index + properties = material.compute_material_properties(compiled, index) + color_args = state.color, state.throughput, state.active_mask, hit_colors + color_args += properties.reflectivities, properties.transparencies + color = accumulate_color(*color_args) + bounce_args = state.rays[1], closest, state.refractive_index, properties + new_rays, n_2, reflectance = optics.compute_bounce(*bounce_args) + args = new_rays, n_2, properties, reflectance + return apply_bounce_update(state._replace(color=color), *args) + + +def render_chunks(rays, compiled, shadows, num_bounces, face_chunk, chunk): + ray_chunks = split_ray_chunks(rays, chunk) + step_args = compiled, shadows, num_bounces, face_chunk + chunk_step = paz.lock(render_chunk_step, *step_args) + hit_mask, depth, color = jax.lax.scan(chunk_step, None, ray_chunks)[1] + return flatten_chunk_results(hit_mask, depth, color, len(rays[0])) + + +def render_chunk_step(carry, rays, compiled, shadows, num_bounces, face_chunk): + args = rays, compiled, shadows, num_bounces, face_chunk + return carry, render(*args) + + +def initialize_state(rays): + num_rays = rays[0].shape[0] + color = jp.zeros((num_rays, 3)) + depth = jp.full((num_rays,), paz.graphics.FARAWAY) + hit_mask = jp.zeros((num_rays,), dtype=bool) + throughput = jp.ones((num_rays, 3)) + active_mask = jp.ones((num_rays,), dtype=bool) + refractive_index = jp.ones((num_rays,)) + fields = color, depth, hit_mask, throughput, active_mask + return RenderState(*fields, refractive_index, rays) + + +def update_first_hit(state, closest, step_arg): + if step_arg == 0: + depth, hit_mask = closest.depth, closest.hit_mask + else: + depth, hit_mask = state.depth, state.hit_mask + return state._replace(depth=depth, hit_mask=hit_mask) + + +def update_active_mask(state, closest): + active_mask = state.active_mask & closest.hit_mask + return state._replace(active_mask=active_mask) + + +def accumulate_color( + colors, throughput, active_mask, hit_colors, reflectivities, transparencies +): + weights = jp.maximum(1.0 - reflectivities - transparencies, 0.0) + weights = jp.expand_dims(weights, -1) + active_mask = jp.expand_dims(active_mask, -1) + return colors + (throughput * active_mask * weights * hit_colors) + + +def apply_bounce_update(state, new_rays, n_2, properties, reflectance): + is_transparent = properties.transparencies > 0.0 + is_reflective = properties.reflectivities > 0.0 + factor = optics.compute_bounce_factor(properties, reflectance) + factor = jp.where(is_transparent & (reflectance >= 1.0), 1.0, factor) + throughput = state.throughput * jp.expand_dims(factor, -1) + active_mask = state.active_mask & (is_transparent | is_reflective) + index_args = state, n_2, is_transparent, reflectance + refractive_index = update_refractive_index(*index_args) + args = throughput, active_mask, refractive_index, new_rays + return replace_bounce_state(state, *args) + + +def update_refractive_index(state, n_2, is_transparent, reflectance): + update_mask = is_transparent & (reflectance < 1.0) + return jp.where(update_mask, n_2, state.refractive_index) + + +def replace_bounce_state(state, throughput, active_mask, index, new_rays): + state = state._replace(throughput=throughput, active_mask=active_mask) + return state._replace(refractive_index=index, rays=new_rays) + + +def split_ray_chunks(rays, chunk_size): + origins, directions = rays + origins = pad_to_chunks(origins, chunk_size) + directions = pad_to_chunks(directions, chunk_size) + num_chunks = origins.shape[0] // chunk_size + shape = num_chunks, chunk_size, 3 + return origins.reshape(shape), directions.reshape(shape) + + +def pad_to_chunks(array, chunk_size): + remainder = array.shape[0] % chunk_size + if remainder == 0: + padded = array + else: + padding = jp.repeat(array[-1:], chunk_size - remainder, axis=0) + padded = jp.concatenate([array, padding], axis=0) + return padded + + +def flatten_chunk_results(hit_mask, depth, color, num_rays): + hit_mask = flatten_chunk_array(hit_mask, num_rays) + depth = flatten_chunk_array(depth, num_rays) + color = flatten_chunk_array(color, num_rays) + return hit_mask, depth, color + + +def flatten_chunk_array(array, num_rays): + shape = (-1,) + array.shape[2:] + return array.reshape(shape)[:num_rays] diff --git a/paz/graphics/renderer/shade.py b/paz/graphics/renderer/shade.py new file mode 100644 index 000000000..b661dfcce --- /dev/null +++ b/paz/graphics/renderer/shade.py @@ -0,0 +1,134 @@ +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.composite import take_closest +from paz.graphics.renderer import shadow +from paz.graphics.renderer.intersect import ( + iterate_shape_groups, + slice_surfaces, +) + + +def compute_hit_colors(compiled, closest, surfaces, shadows, triangle_hit): + if len(compiled.shapes) == 0: + colors = compute_triangle_colors(compiled, triangle_hit) + elif triangle_hit is None: + colors = compute_shape_colors(compiled, closest, surfaces, shadows) + else: + args = compiled, closest, surfaces, shadows, triangle_hit + colors = blend_hit_colors(*args) + return colors + + +def blend_hit_colors(compiled, closest, surfaces, shadows, triangle_hit): + # TODO a mixed scene shades both paths for every ray and throws one + # away. Joining them before selection needs color_with_shadows to + # return rows instead of selecting inside its per-light scan. + shape_colors = compute_shape_colors(compiled, closest, surfaces, shadows) + triangle_colors = compute_triangle_colors(compiled, triangle_hit) + is_triangle = closest.primitive_index == len(compiled.shapes) + is_triangle = jp.expand_dims(is_triangle, -1) + return jp.where(is_triangle, triangle_colors, shape_colors) + + +def compute_shape_colors(compiled, closest, surfaces, shadows): + num_shapes = len(compiled.shapes) + indices = jp.minimum(closest.primitive_index, num_shapes - 1) + surfaces = slice_surfaces(surfaces, 0, num_shapes) + if shadows: + colors = color_with_shadows(compiled, closest, surfaces, indices) + else: + colors = color_without_shadow(compiled, surfaces, indices) + return colors + + +def compute_triangle_colors(compiled, triangle_hit): + materials = compiled.triangles.materials + material = gather_triangle_material(materials, triangle_hit.primitive) + shader = select_shader(materials) + colors = jp.zeros_like(triangle_hit.albedo) + for light in compiled.lights: + args = triangle_hit.albedo, material, triangle_hit.points + args += triangle_hit.normals, triangle_hit.eyes, light + colors = colors + shader.compute_colors(*args) + return colors + + +def gather_triangle_material(materials, primitive): + material = jax.tree.map(lambda field: field[primitive], materials) + return jax.tree.map(expand_scalar_field, material) + + +def expand_scalar_field(field): + if field.ndim == 1: + field = jp.expand_dims(field, -1) + return field + + +def select_shader(material): + if isinstance(material, paz.graphics.CookTorranceMaterial): + shader = paz.graphics.cook_torrance + else: + shader = paz.graphics.phong + return shader + + +def color_without_shadow(compiled, surfaces, indices): + colors = [] + lights = paz.graphics.shapes.merge(*compiled.lights) + for group, start_arg, final_arg in iterate_shape_groups(compiled.shapes): + group_surfaces = slice_surfaces(surfaces, start_arg, final_arg) + albedo = compute_group_albedo(group, group_surfaces.points) + shader = select_shader(group.material) + axes = 0, 0, 0, 0, 0, None + color_per_light = jax.vmap(shader.compute_colors, axes) + color = jax.vmap(color_per_light, (None, None, None, None, None, 0)) + args = albedo, group.material, *group_surfaces, lights + colors.append(jp.sum(color(*args), axis=0)) + return take_closest(jp.concatenate(colors, axis=0), indices) + + +def compute_group_albedo(group, points): + compute_albedo = paz.graphics.albedo.compute_shape_albedo + return jax.vmap(compute_albedo)(group, group.material, points) + + +def color_with_shadows(compiled, closest, surfaces, indices): + colors = jp.zeros((len(surfaces.points[0]), 3)) + lights = paz.graphics.shapes.merge(*compiled.lights) + body = paz.lock(scan_light_step, compiled, closest, surfaces, indices) + return jax.lax.scan(body, colors, lights)[0] + + +def scan_light_step(colors, light, compiled, closest, surfaces, indices): + args = compiled, closest, surfaces, indices, light + return colors + compute_light_colors(*args), None + + +def compute_light_colors(compiled, closest, surfaces, indices, light): + directions, distance = compute_light_directions(light, closest.point) + occlusion_args = compiled, closest, indices, directions, distance + is_shadow = shadow.compute_occlusion(*occlusion_args) + color_args = compiled.shapes, light, surfaces, is_shadow + return take_closest(compute_shadowed_colors(*color_args), indices) + + +def compute_light_directions(light, points): + vector = light.position - points + norm = paz.algebra.compute_norms(vector, 1) + return vector / norm, jp.squeeze(norm, axis=1) + + +def compute_shadowed_colors(shapes, light, surfaces, is_shadow): + colors = [] + for group, start_arg, final_arg in iterate_shape_groups(shapes): + group_surfaces = slice_surfaces(surfaces, start_arg, final_arg) + albedo = compute_group_albedo(group, group_surfaces.points) + shader = select_shader(group.material) + axes = 0, 0, 0, 0, 0, None, None + color = jax.vmap(shader.compute_colors_with_shadow, axes) + args = albedo, group.material, *group_surfaces, light, is_shadow + colors.append(color(*args)) + return jp.concatenate(colors, axis=0) diff --git a/paz/graphics/renderer/shadow.py b/paz/graphics/renderer/shadow.py new file mode 100644 index 000000000..bc9a80df3 --- /dev/null +++ b/paz/graphics/renderer/shadow.py @@ -0,0 +1,95 @@ +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.composite import compute_scene_hit_mask +from paz.graphics.renderer.intersect import intersect_shadow_groups + +SHADOW_ORIGIN_EPSILON = 1e-5 +SHADOW_SELF_HIT_EPSILON = 1e-5 + + +def compute_occlusion(compiled, closest, indices, directions, distance): + origins = compute_shadow_ray_origins(closest.point, closest.normal) + shadow_args = compiled.shapes, origins, directions + intersections = intersect_shadow_groups(*shadow_args) + hit_masks, depths, _, _, _, casters = intersections + masks = resolve_shadow_masks(compiled, hit_masks) + depth_args = masks, depths, casters, indices, closest.normal, directions + masks, depths = select_shadow_depths(*depth_args) + return compute_soft_occlusion(masks, depths, distance) + + +def compute_shadow_ray_origins(points, normals): + over_point, _ = compute_surface_points(points, normals) + return over_point + + +def compute_surface_points(point, normal): + over_point = point + normal * SHADOW_ORIGIN_EPSILON + under_point = point - normal * SHADOW_ORIGIN_EPSILON + return over_point, under_point + + +def resolve_shadow_masks(compiled, hit_masks): + transparencies = compute_transparencies(compiled.shapes) + masks = jp.where(jp.expand_dims(compiled.mask, 1), hit_masks, False) + masks = hide_transparent_shapes(masks, transparencies > 0.0) + if compiled.shadow_mask is None: + resolved = masks + else: + resolved = hide_non_casting_shapes(masks, compiled.shadow_mask) + return resolved + + +def compute_transparencies(shapes): + return jp.array([shape.material.transparency for shape in shapes]) + + +def hide_transparent_shapes(shadow_masks, is_transparent): + return jp.where(jp.expand_dims(is_transparent, 1), False, shadow_masks) + + +def hide_non_casting_shapes(shadow_masks, shadow_mask): + return jp.where(jp.expand_dims(shadow_mask, 1), shadow_masks, False) + + +def select_shadow_depths( + hit_masks, depths, casters, receivers, normals, directions +): + same_shape = casters[:, None] == receivers[None, :] + front_args = same_shape, normals, directions + front_side_hits = compute_front_side_shadow_mask(*front_args) + root_args = hit_masks, depths, same_shape, front_side_hits + valid_roots = compute_valid_roots(*root_args) + depths = jp.where(valid_roots, depths, paz.graphics.FARAWAY) + return jp.any(valid_roots, axis=1), jp.min(depths, axis=1) + + +def compute_front_side_shadow_mask(same_shape, normals, directions): + front_side = paz.algebra.dot(normals, directions) >= 0.0 + return jp.logical_and(same_shape, front_side[None, :]) + + +def compute_valid_roots(hit_masks, depths, same_shape, front_side_hits): + thresholds = compute_shadow_depth_thresholds(same_shape) + valid_roots = depths > thresholds[:, None, :] + valid_roots = jp.logical_and(valid_roots, depths < paz.graphics.FARAWAY) + valid_roots = jp.logical_and(jp.expand_dims(hit_masks, 1), valid_roots) + return jp.logical_and(valid_roots, ~front_side_hits[:, None, :]) + + +def compute_shadow_depth_thresholds(same_shape): + return jp.where(same_shape, SHADOW_SELF_HIT_EPSILON, paz.graphics.EPSILON) + + +def compute_soft_occlusion(hit_masks, depths, light_lengths, slope=0.01): + closest_depths = jp.where(hit_masks, depths, paz.graphics.FARAWAY) + closest_depths = jp.min(closest_depths, axis=0) + scene_hit_mask = compute_scene_hit_mask(hit_masks) + blockers = closest_depths <= light_lengths + blocker_mask = jp.logical_and(scene_hit_mask, blockers) + difference = light_lengths - closest_depths + occlusion = jax.nn.sigmoid(slope * difference) + return jp.where(blocker_mask, occlusion, 0.0) diff --git a/paz/graphics/renderer/tiling.py b/paz/graphics/renderer/tiling.py new file mode 100644 index 000000000..54ed1b127 --- /dev/null +++ b/paz/graphics/renderer/tiling.py @@ -0,0 +1,34 @@ +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.composite import postprocess + + +def scan(shape, tiles, tile_step): + H, W = shape + H_tiles, W_tiles = tiles + paz.graphics.mesh.assert_exact_tile_side(H, H_tiles) + paz.graphics.mesh.assert_exact_tile_side(W, W_tiles) + coordinates = paz.graphics.mesh.make_tile_coordinates(H_tiles, W_tiles) + return jax.lax.scan(tile_step, None, coordinates)[1] + + +def render_step(carry, tile_arg, shape, y_FOV, pose, tiles, render_rays): + H, W = shape + H_tiles, W_tiles = tiles + camera_to_world = jp.linalg.inv(pose) + tile_args = H, W, H_tiles, W_tiles, y_FOV, camera_to_world + rays = paz.graphics.mesh.build_tile_rays(*tile_args, tile_arg) + hit_mask, depth, color = render_rays(rays) + tile_H, tile_W = H // H_tiles, W // W_tiles + post_args = hit_mask, depth, color, pose, rays, tile_H, tile_W + image, depth = postprocess(*post_args) + return carry, (image, jp.expand_dims(depth, -1)) + + +def assemble(shape, tiles, images): + H, W = shape + H_tiles, W_tiles = tiles + return paz.graphics.mesh.assemble(H, W, H_tiles, W_tiles, images) diff --git a/paz/graphics/renderer_test.py b/paz/graphics/renderer_test.py index 35422341c..9472d27c2 100644 --- a/paz/graphics/renderer_test.py +++ b/paz/graphics/renderer_test.py @@ -115,15 +115,15 @@ def compute_selected_shadow_depths(camera_pose, image_shape=(120, 160)): rays = paz.graphics.camera.build_rays(image_shape, jp.pi / 3.0, camera_pose) compiled = paz.graphics.scene.compile(scene, lights, None) shapes, mask, lights = compiled.shapes, compiled.mask, compiled.lights - intersections = renderer.intersect(compiled, rays, None) - closest = renderer.gather_closest(*intersections) + intersections = renderer.intersect.build_candidates(compiled, rays, None) + closest = renderer.intersect.find_closest(*intersections) vector = lights[0].position - closest.point distance = jp.squeeze(paz.algebra.compute_norms(vector, 1), axis=1) light_directions = vector / jp.expand_dims(distance, 1) - shadow_ray_origins = renderer.compute_shadow_ray_origins( + shadow_ray_origins = renderer.shadow.compute_shadow_ray_origins( closest.point, closest.normal ) - intersections = renderer.intersect_shadow_groups( + intersections = renderer.intersect.intersect_shadow_groups( shapes, shadow_ray_origins, light_directions ) hit_masks, depths, _, _, _, shape_indices = intersections @@ -132,7 +132,7 @@ def compute_selected_shadow_depths(camera_pose, image_shape=(120, 160)): shadow_masks = jp.where( jp.expand_dims(transparencies > 0.0, 1), False, shadow_masks ) - shadow_masks, depths = renderer.select_shadow_depths( + shadow_masks, depths = renderer.shadow.select_shadow_depths( shadow_masks, depths, shape_indices, @@ -169,7 +169,7 @@ def test_compute_soft_occlusion(): ) rows = [[True, True, True, True], [False, False, True, False]] hit_masks = jp.array(rows) - result = renderer.compute_soft_occlusion( + result = renderer.shadow.compute_soft_occlusion( hit_masks, depths, light_lengths, slope=10.0 ) assert float(result[1]) > 0.9 @@ -183,7 +183,7 @@ def test_compute_new_rays_reflection_is_normalized(): eye = jp.array([[0.0, 1.0, -1.0]]) point = jp.array([[0.0, 0.0, 0.0]]) transparencies = jp.array([0.0]) - _, direction = renderer.compute_new_rays( + _, direction = renderer.optics.compute_new_rays( normal, eye, jp.array([1.0]), point, transparencies ) norm = jp.linalg.norm(direction, axis=-1) @@ -195,7 +195,7 @@ def test_compute_new_rays_refraction_is_normalized(): eye = jp.array([[0.0, 0.0, -1.0]]) point = jp.array([[0.0, 0.0, 0.0]]) transparencies = jp.array([1.0]) - _, direction = renderer.compute_new_rays( + _, direction = renderer.optics.compute_new_rays( normal, eye, jp.array([1.0 / 1.5]), point, transparencies ) norm = jp.linalg.norm(direction, axis=-1) @@ -205,8 +205,9 @@ def test_compute_new_rays_refraction_is_normalized(): def test_compute_surface_points_offset_hit(): point = jp.array([[0.0, 0.0, 0.0]]) normal = jp.array([[0.0, 0.0, -1.0]]) - over_point, under_point = renderer.compute_surface_points(point, normal) - assert over_point[0, 2] < -(renderer.SHADOW_ORIGIN_EPSILON / 2.0) + surface_points = renderer.shadow.compute_surface_points(point, normal) + over_point, under_point = surface_points + assert over_point[0, 2] < -(renderer.shadow.SHADOW_ORIGIN_EPSILON / 2.0) assert point[0, 2] > over_point[0, 2] assert under_point[0, 2] > 0.0 @@ -223,7 +224,7 @@ def test_select_shadow_depths_discard_front_side_same_shape_hits(): receiver_indices = jp.array([0]) receiver_normals = jp.array([[0.0, 1.0, 0.0]]) light_directions = jp.array([[0.0, 1.0, 0.0]]) - hit_masks, depths = renderer.select_shadow_depths( + hit_masks, depths = renderer.shadow.select_shadow_depths( hit_masks, depths, shape_indices, @@ -249,7 +250,7 @@ def test_select_shadow_depths_keep_back_side_second_root(): receiver_indices = jp.array([0]) receiver_normals = jp.array([[0.0, 1.0, 0.0]]) light_directions = jp.array([[0.0, -1.0, 0.0]]) - hit_masks, depths = renderer.select_shadow_depths( + hit_masks, depths = renderer.shadow.select_shadow_depths( hit_masks, depths, shape_indices, @@ -258,7 +259,7 @@ def test_select_shadow_depths_keep_back_side_second_root(): light_directions, ) args = hit_masks, depths, jp.array([0.01]) - result = renderer.compute_soft_occlusion(*args) + result = renderer.shadow.compute_soft_occlusion(*args) assert bool(hit_masks[0, 0]) assert float(depths[0, 0]) == pytest.approx(0.2) assert bool(hit_masks[1, 0]) @@ -295,7 +296,7 @@ def test_select_colors(): def test_initialize_render_state(): num_rays = 10 rays = (jp.zeros((num_rays, 3)), jp.ones((num_rays, 3))) - state = renderer.initialize_state(rays) + state = renderer.rays.initialize_state(rays) assert state.color.shape == (num_rays, 3) assert state.throughput.shape == (num_rays, 3) assert jp.all(state.refractive_index == 1.0) @@ -318,7 +319,7 @@ def test_compute_material_properties(): scene = Scene([shape1, shape2]) compiled = paz.graphics.scene.compile(scene, [], None) indices = jp.array([0, 1]) - material = renderer.compute_material_properties(compiled, indices) + material = renderer.material.compute_material_properties(compiled, indices) assert material.reflectivities[0] == 0.5 assert material.transparencies[1] == 0.8 assert material.refractive_indices[0] == 1.0 @@ -332,7 +333,7 @@ def test_accumulate_color(): intersected_colors = jp.ones((num_rays, 3)) reflectivities = jp.zeros((num_rays,)) transparencies = jp.zeros((num_rays,)) - result = renderer.accumulate_color( + result = renderer.rays.accumulate_color( colors, throughput, active_mask, @@ -349,7 +350,7 @@ def test_compute_shadow_ray_origins_avoid_lit_side_self_hit(): normals = jp.array([[0.0, 1.0, 0.0]]) light_position = jp.array([[0.0, 3.0, -3.0]]) directions = paz.algebra.normalize(light_position - points) - origins = renderer.compute_shadow_ray_origins(points, normals) + origins = renderer.shadow.compute_shadow_ray_origins(points, normals) hit_mask, _, _ = intersect_canonical_sphere(origins, directions) assert not bool(hit_mask[0]) @@ -721,8 +722,9 @@ def test_saved_pose_sphere_self_shadow_keeps_later_roots(): ) args = depths, shape_indices, receiver_indices, 0 sphere_depths = take_shape_depths(*args) + epsilon = renderer.shadow.SHADOW_SELF_HIT_EPSILON assert int(jp.sum(sphere_depths < 1e-2)) > 0 - assert float(jp.min(sphere_depths)) > renderer.SHADOW_SELF_HIT_EPSILON + assert float(jp.min(sphere_depths)) > epsilon def test_saved_pose_floor_self_hits_stay_filtered(): diff --git a/paz/graphics/types.py b/paz/graphics/types.py index e7fc4da51..7e548db44 100644 --- a/paz/graphics/types.py +++ b/paz/graphics/types.py @@ -17,6 +17,8 @@ HIT_NAMES = "hit_mask depth point normal eye primitive_index" Hit = namedtuple("Hit", HIT_NAMES.split()) +Surfaces = namedtuple("Surfaces", ["points", "normals", "eyes"]) + MESH_NAMES = "vertices vertex_colors transform material faces edges " MESH_NAMES += "pattern vertex_uvs" Mesh = namedtuple("Mesh", MESH_NAMES.split(), defaults=(None, None)) From 4dc896d1da12e763d6fc83a062728963d72b7672 Mon Sep 17 00:00:00 2001 From: oarriaga Date: Sun, 2 Aug 2026 20:44:55 +0200 Subject: [PATCH 2/5] fix: honour reflective, transparent and refractive mesh materials scene.stack_materials already compiles reflective, transparency and refractive_index for every mesh, and both Material and CookTorranceMaterial declare all three. The renderer threw them away and substituted a hardcoded (0.0, 0.0, 1.0) row for all triangles, so a mirror or glass mesh rendered as matte and opaque. Gather the real values per primitive via triangle_hit.primitive and select between shape and triangle values on the closest hit. Meshes that leave these fields at their defaults render identically, so existing scenes are unaffected. Two new mesh_test cases cover a reflective and a transparent mesh; both fail before this change. --- paz/graphics/mesh_test.py | 32 +++++++++++++++++++++++++++++++ paz/graphics/renderer/material.py | 32 ++++++++++++++++++++++++++++--- paz/graphics/renderer/rays.py | 8 ++++---- paz/graphics/renderer_test.py | 18 +++++++++++++++-- 4 files changed, 81 insertions(+), 9 deletions(-) diff --git a/paz/graphics/mesh_test.py b/paz/graphics/mesh_test.py index d2ebbccf2..ac7596fac 100644 --- a/paz/graphics/mesh_test.py +++ b/paz/graphics/mesh_test.py @@ -854,6 +854,38 @@ def test_scene_mixes_meshes_and_shapes(): assert jp.any(depth > 0) +def build_cube_mesh_with_material(material): + vertices, faces, edges = build_cube(1.0) + colors = build_vertex_colors(vertices, [0.7, 0.3, 0.1]) + transform = SE3.translation(jp.array([-0.45, 0.0, 0.0])) + args = vertices, colors, transform, material, faces, edges + return Mesh(*args) + + +def render_mesh_with_material(material): + shape, y_FOV, pose, _, _, lights = make_multi_mesh_scene() + mesh = build_cube_mesh_with_material(material) + scene = paz.graphics.Scene([mesh, build_sphere_mesh()]) + args = shape, y_FOV, pose, scene, None, lights, (1, 1), 1024 + return paz.graphics.render(*args, False, None, 2) + + +def test_mesh_reflective_material_changes_render(): + matte = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100, 0.0) + mirror = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100, 0.9) + matte_image, _ = render_mesh_with_material(matte) + mirror_image, _ = render_mesh_with_material(mirror) + assert compute_max_abs_difference(matte_image, mirror_image) > 1e-3 + + +def test_mesh_transparent_material_changes_render(): + opaque = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100, 0.0, 0.0) + glass = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100, 0.0, 0.8, 1.5) + opaque_image, _ = render_mesh_with_material(opaque) + glass_image, _ = render_mesh_with_material(glass) + assert compute_max_abs_difference(opaque_image, glass_image) > 1e-3 + + def test_scene_rejects_meshes_with_mixed_pattern_sizes(): plain = build_cube_mesh() textured = build_textured_quad_mesh() diff --git a/paz/graphics/renderer/material.py b/paz/graphics/renderer/material.py index a9443d008..c426a63e0 100644 --- a/paz/graphics/renderer/material.py +++ b/paz/graphics/renderer/material.py @@ -6,13 +6,25 @@ HitMaterial = namedtuple("HitMaterial", MATERIAL_NAMES.split()) -def compute_material_properties(compiled, hit_shape_args): - values = collect_material_values(compiled) +def compute_material_properties(compiled, closest, triangle_hit): + shape_material = gather_shape_material(compiled, closest.primitive_index) + if triangle_hit is None: + resolved = shape_material + else: + materials = compiled.triangles.materials + triangle = gather_primitive_material(materials, triangle_hit.primitive) + is_triangle = closest.primitive_index == len(compiled.shapes) + resolved = select_material(is_triangle, triangle, shape_material) + return resolved + + +def gather_shape_material(compiled, hit_shape_args): + values = collect_shape_values(compiled) gathered = tuple(jp.array(row)[hit_shape_args] for row in values) return HitMaterial(*gathered) -def collect_material_values(compiled): +def collect_shape_values(compiled): reflectivities, transparencies, refractive_indices = [], [], [] for shape in compiled.shapes: reflectivities.append(shape.material.reflective) @@ -23,3 +35,17 @@ def collect_material_values(compiled): transparencies.append(0.0) refractive_indices.append(1.0) return reflectivities, transparencies, refractive_indices + + +def gather_primitive_material(materials, primitive): + reflectivities = materials.reflective[primitive] + transparencies = materials.transparency[primitive] + refractive_indices = materials.refractive_index[primitive] + return HitMaterial(reflectivities, transparencies, refractive_indices) + + +def select_material(is_triangle, triangle, shape): + fields = [] + for triangle_field, shape_field in zip(triangle, shape): + fields.append(jp.where(is_triangle, triangle_field, shape_field)) + return HitMaterial(*fields) diff --git a/paz/graphics/renderer/rays.py b/paz/graphics/renderer/rays.py index a697a9dcf..800a01620 100644 --- a/paz/graphics/renderer/rays.py +++ b/paz/graphics/renderer/rays.py @@ -30,12 +30,12 @@ def bounce(state, step_arg, compiled, shadows, face_chunk): state = update_active_mask(state, closest) color_args = compiled, closest, surfaces, shadows, triangle_hit colors = shade.compute_hit_colors(*color_args) - return advance(state, compiled, closest, colors) + return advance(state, compiled, closest, colors, triangle_hit) -def advance(state, compiled, closest, hit_colors): - index = closest.primitive_index - properties = material.compute_material_properties(compiled, index) +def advance(state, compiled, closest, hit_colors, triangle_hit): + material_args = compiled, closest, triangle_hit + properties = material.compute_material_properties(*material_args) color_args = state.color, state.throughput, state.active_mask, hit_colors color_args += properties.reflectivities, properties.transparencies color = accumulate_color(*color_args) diff --git a/paz/graphics/renderer_test.py b/paz/graphics/renderer_test.py index 9472d27c2..9e36923a3 100644 --- a/paz/graphics/renderer_test.py +++ b/paz/graphics/renderer_test.py @@ -311,7 +311,7 @@ def test_find_closest_intersection_args(): assert jp.array_equal(indices, jp.array([0, 1])) -def test_compute_material_properties(): +def test_gather_shape_material(): mat1 = Material(reflective=0.5) mat2 = Material(transparency=0.8) shape1 = Sphere(material=mat1) @@ -319,12 +319,26 @@ def test_compute_material_properties(): scene = Scene([shape1, shape2]) compiled = paz.graphics.scene.compile(scene, [], None) indices = jp.array([0, 1]) - material = renderer.material.compute_material_properties(compiled, indices) + material = renderer.material.gather_shape_material(compiled, indices) assert material.reflectivities[0] == 0.5 assert material.transparencies[1] == 0.8 assert material.refractive_indices[0] == 1.0 +def test_gather_primitive_material_reads_mesh_fields(): + materials = Material( + reflective=jp.array([0.0, 0.6]), + transparency=jp.array([0.0, 0.3]), + refractive_index=jp.array([1.0, 1.5]), + ) + primitive = jp.array([1, 0, 1]) + gather = renderer.material.gather_primitive_material + material = gather(materials, primitive) + expected_indices = jp.array([1.5, 1.0, 1.5]) + assert jp.array_equal(material.reflectivities, jp.array([0.6, 0.0, 0.6])) + assert jp.array_equal(material.refractive_indices, expected_indices) + + def test_accumulate_color(): num_rays = 2 colors = jp.zeros((num_rays, 3)) From cfe13ef232678e774525b050acd01c98e241ed78 Mon Sep 17 00:00:00 2001 From: oarriaga Date: Sun, 2 Aug 2026 21:02:21 +0200 Subject: [PATCH 3/5] feat: let meshes cast shadows Shadow rays only ever intersected shapes, so a mesh lit from above cast nothing onto the floor beneath it. Shadow rays now also traverse the triangles and join the shape blockers as one more caster row before the soft-occlusion reduction. select_meshes never received shadow_mask, so the public "this node does not cast" control had no mesh equivalent. CompiledScene gains triangle_shadow_mask and scene.compile fills it, which keeps that control working for meshes instead of silently ignoring it. The shape path rejects a shadow ray re-hitting its own surface by shape identity. Triangles have no such index, so they use a distance threshold instead; meshes do not receive shadows yet, so today the receiver is always a shape and the threshold only guards shape receivers close to a mesh. Two mesh_test cases cover casting and the mask suppressing it; both fail before this change. --- paz/graphics/mesh_test.py | 34 ++++++++++++++++++++++++++ paz/graphics/renderer/rays.py | 2 +- paz/graphics/renderer/shade.py | 38 +++++++++++++++++++---------- paz/graphics/renderer/shadow.py | 42 ++++++++++++++++++++++++++++++--- paz/graphics/scene.py | 15 ++++++++---- paz/graphics/types.py | 3 ++- 6 files changed, 112 insertions(+), 22 deletions(-) diff --git a/paz/graphics/mesh_test.py b/paz/graphics/mesh_test.py index ac7596fac..faf5f9f37 100644 --- a/paz/graphics/mesh_test.py +++ b/paz/graphics/mesh_test.py @@ -886,6 +886,40 @@ def test_mesh_transparent_material_changes_render(): assert compute_max_abs_difference(opaque_image, glass_image) > 1e-3 +def build_mesh_over_floor_scene(): + vertices, faces, edges = build_cube(1.0) + colors = build_vertex_colors(vertices, [0.7, 0.3, 0.1]) + material = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100) + transform = SE3.translation(jp.array([0.0, 1.0, 0.0])) + args = vertices, colors, transform, material, faces, edges + floor = paz.graphics.Plane(SE3.translation(jp.array([0.0, -0.5, 0.0]))) + return paz.graphics.Scene([Mesh(*args), floor]) + + +def render_mesh_over_floor(shadows, shadow_mask=None): + camera_pose = SE3.view_transform( + jp.array([0.0, 2.0, -4.0]), jp.zeros(3), jp.array([0.0, 1.0, 0.0]) + ) + lights = [PointLight(jp.ones(3), jp.array([0.0, 5.0, -1.0]))] + scene = build_mesh_over_floor_scene() + args = (32, 32), jp.pi / 3.0, camera_pose, scene, None, lights + return paz.graphics.render(*args, (1, 1), 1024, shadows, shadow_mask) + + +def test_mesh_casts_shadow_on_shape(): + lit, _ = render_mesh_over_floor(False) + shadowed, _ = render_mesh_over_floor(True) + assert compute_max_abs_difference(lit, shadowed) > 1e-2 + assert jp.all(shadowed <= lit + 1e-4) + + +def test_mesh_shadow_mask_stops_mesh_casting(): + casting, _ = render_mesh_over_floor(True) + blocked, _ = render_mesh_over_floor(True, jp.array([False, True])) + assert compute_max_abs_difference(casting, blocked) > 1e-2 + assert jp.all(casting <= blocked + 1e-4) + + def test_scene_rejects_meshes_with_mixed_pattern_sizes(): plain = build_cube_mesh() textured = build_textured_quad_mesh() diff --git a/paz/graphics/renderer/rays.py b/paz/graphics/renderer/rays.py index 800a01620..136ef6e33 100644 --- a/paz/graphics/renderer/rays.py +++ b/paz/graphics/renderer/rays.py @@ -29,7 +29,7 @@ def bounce(state, step_arg, compiled, shadows, face_chunk): state = update_first_hit(state, closest, step_arg) state = update_active_mask(state, closest) color_args = compiled, closest, surfaces, shadows, triangle_hit - colors = shade.compute_hit_colors(*color_args) + colors = shade.compute_hit_colors(*color_args, face_chunk) return advance(state, compiled, closest, colors, triangle_hit) diff --git a/paz/graphics/renderer/shade.py b/paz/graphics/renderer/shade.py index b661dfcce..c746ffe72 100644 --- a/paz/graphics/renderer/shade.py +++ b/paz/graphics/renderer/shade.py @@ -11,34 +11,41 @@ ) -def compute_hit_colors(compiled, closest, surfaces, shadows, triangle_hit): +def compute_hit_colors( + compiled, closest, surfaces, shadows, triangle_hit, face_chunk +): if len(compiled.shapes) == 0: colors = compute_triangle_colors(compiled, triangle_hit) elif triangle_hit is None: - colors = compute_shape_colors(compiled, closest, surfaces, shadows) + shape_args = compiled, closest, surfaces, shadows, face_chunk + colors = compute_shape_colors(*shape_args) else: args = compiled, closest, surfaces, shadows, triangle_hit - colors = blend_hit_colors(*args) + colors = blend_hit_colors(*args, face_chunk) return colors -def blend_hit_colors(compiled, closest, surfaces, shadows, triangle_hit): +def blend_hit_colors( + compiled, closest, surfaces, shadows, triangle_hit, face_chunk +): # TODO a mixed scene shades both paths for every ray and throws one # away. Joining them before selection needs color_with_shadows to # return rows instead of selecting inside its per-light scan. - shape_colors = compute_shape_colors(compiled, closest, surfaces, shadows) + shape_args = compiled, closest, surfaces, shadows, face_chunk + shape_colors = compute_shape_colors(*shape_args) triangle_colors = compute_triangle_colors(compiled, triangle_hit) is_triangle = closest.primitive_index == len(compiled.shapes) is_triangle = jp.expand_dims(is_triangle, -1) return jp.where(is_triangle, triangle_colors, shape_colors) -def compute_shape_colors(compiled, closest, surfaces, shadows): +def compute_shape_colors(compiled, closest, surfaces, shadows, face_chunk): num_shapes = len(compiled.shapes) indices = jp.minimum(closest.primitive_index, num_shapes - 1) surfaces = slice_surfaces(surfaces, 0, num_shapes) if shadows: - colors = color_with_shadows(compiled, closest, surfaces, indices) + args = compiled, closest, surfaces, indices, face_chunk + colors = color_with_shadows(*args) else: colors = color_without_shadow(compiled, surfaces, indices) return colors @@ -95,22 +102,27 @@ def compute_group_albedo(group, points): return jax.vmap(compute_albedo)(group, group.material, points) -def color_with_shadows(compiled, closest, surfaces, indices): +def color_with_shadows(compiled, closest, surfaces, indices, face_chunk): colors = jp.zeros((len(surfaces.points[0]), 3)) lights = paz.graphics.shapes.merge(*compiled.lights) - body = paz.lock(scan_light_step, compiled, closest, surfaces, indices) + step_args = compiled, closest, surfaces, indices, face_chunk + body = paz.lock(scan_light_step, *step_args) return jax.lax.scan(body, colors, lights)[0] -def scan_light_step(colors, light, compiled, closest, surfaces, indices): - args = compiled, closest, surfaces, indices, light +def scan_light_step( + colors, light, compiled, closest, surfaces, indices, face_chunk +): + args = compiled, closest, surfaces, indices, light, face_chunk return colors + compute_light_colors(*args), None -def compute_light_colors(compiled, closest, surfaces, indices, light): +def compute_light_colors( + compiled, closest, surfaces, indices, light, face_chunk +): directions, distance = compute_light_directions(light, closest.point) occlusion_args = compiled, closest, indices, directions, distance - is_shadow = shadow.compute_occlusion(*occlusion_args) + is_shadow = shadow.compute_occlusion(*occlusion_args, face_chunk) color_args = compiled.shapes, light, surfaces, is_shadow return take_closest(compute_shadowed_colors(*color_args), indices) diff --git a/paz/graphics/renderer/shadow.py b/paz/graphics/renderer/shadow.py index bc9a80df3..78297dc5d 100644 --- a/paz/graphics/renderer/shadow.py +++ b/paz/graphics/renderer/shadow.py @@ -8,17 +8,53 @@ SHADOW_ORIGIN_EPSILON = 1e-5 SHADOW_SELF_HIT_EPSILON = 1e-5 +# A shadow ray leaving a mesh re-hits its own triangle at a tiny depth. +# The shape path rejects that by shape identity; triangles have no such +# index, so they need a distance large enough to clear the surface. +TRIANGLE_SELF_HIT_EPSILON = 1e-3 -def compute_occlusion(compiled, closest, indices, directions, distance): +def compute_occlusion(compiled, closest, indices, directions, distance, + face_chunk): origins = compute_shadow_ray_origins(closest.point, closest.normal) + shape_args = compiled, closest, indices, origins, directions + masks, depths = compute_shape_blockers(*shape_args) + if compiled.triangles is not None: + blocker_args = compiled, origins, directions, face_chunk + mask, depth = compute_triangle_blockers(*blocker_args) + masks = jp.concatenate([masks, jp.expand_dims(mask, 0)], axis=0) + depths = jp.concatenate([depths, jp.expand_dims(depth, 0)], axis=0) + return compute_soft_occlusion(masks, depths, distance) + + +def compute_shape_blockers(compiled, closest, indices, origins, directions): shadow_args = compiled.shapes, origins, directions intersections = intersect_shadow_groups(*shadow_args) hit_masks, depths, _, _, _, casters = intersections masks = resolve_shadow_masks(compiled, hit_masks) depth_args = masks, depths, casters, indices, closest.normal, directions - masks, depths = select_shadow_depths(*depth_args) - return compute_soft_occlusion(masks, depths, distance) + return select_shadow_depths(*depth_args) + + +def compute_triangle_blockers(compiled, origins, directions, face_chunk): + triangles = compiled.triangles + args = triangles.vertices, triangles.faces, (origins, directions) + result = paz.graphics.mesh.intersect_chunked(*args, face_chunk) + hit_mask, depth, _, _, face_index = result + primitive = triangles.primitive_index[face_index] + hit_mask = jp.logical_and(hit_mask, compiled.triangle_mask[primitive]) + hit_mask = jp.logical_and(hit_mask, depth > TRIANGLE_SELF_HIT_EPSILON) + hit_mask = hide_non_casting_triangles(compiled, hit_mask, primitive) + return hit_mask, jp.where(hit_mask, depth, paz.graphics.FARAWAY) + + +def hide_non_casting_triangles(compiled, hit_mask, primitive): + if compiled.triangle_shadow_mask is None: + casting = hit_mask + else: + casting = compiled.triangle_shadow_mask[primitive] + casting = jp.logical_and(hit_mask, casting) + return casting def compute_shadow_ray_origins(points, normals): diff --git a/paz/graphics/scene.py b/paz/graphics/scene.py index 167f5d81a..8c76763c4 100644 --- a/paz/graphics/scene.py +++ b/paz/graphics/scene.py @@ -189,11 +189,13 @@ def compile(scene, lights, mask, shadow_mask=None): if shadow_mask is not None: shadow_mask = prepare_mask(shadow_mask, len(flat_scene), scene) + mesh_args = select_meshes(flat_scene, mask, shadow_mask) + meshes, triangle_mask, triangle_shadow_mask = mesh_args shape_args = select_shapes(flat_scene, mask, shadow_mask) shapes, shape_mask, shadow_mask = sort_by_group(*shape_args) - meshes, triangle_mask = select_meshes(flat_scene, mask) args = shapes, build_triangles(meshes), lights, shape_mask - return paz.graphics.CompiledScene(*args, shadow_mask, triangle_mask) + args += shadow_mask, triangle_mask, triangle_shadow_mask + return paz.graphics.CompiledScene(*args) def select_shapes(flat_scene, mask, shadow_mask): @@ -205,10 +207,15 @@ def select_shapes(flat_scene, mask, shadow_mask): return shapes, mask[indices], shadow_mask -def select_meshes(flat_scene, mask): +def select_meshes(flat_scene, mask, shadow_mask): args = [arg for arg, node in enumerate(flat_scene) if not is_shape(node)] + indices = jp.array(args, dtype=jp.int32) meshes = [flat_scene[arg] for arg in args] - return meshes, mask[jp.array(args, dtype=jp.int32)] + if shadow_mask is None: + mesh_shadow_mask = None + else: + mesh_shadow_mask = shadow_mask[indices] + return meshes, mask[indices], mesh_shadow_mask def is_shape(node): diff --git a/paz/graphics/types.py b/paz/graphics/types.py index 7e548db44..bb2ab18e3 100644 --- a/paz/graphics/types.py +++ b/paz/graphics/types.py @@ -27,7 +27,8 @@ TRIANGLE_NAMES += "materials patterns" Triangles = namedtuple("Triangles", TRIANGLE_NAMES.split()) -COMPILED_NAMES = "shapes triangles lights mask shadow_mask triangle_mask" +COMPILED_NAMES = "shapes triangles lights mask shadow_mask triangle_mask " +COMPILED_NAMES += "triangle_shadow_mask" CompiledScene = namedtuple("CompiledScene", COMPILED_NAMES.split()) From be860b39580ef50dd5bd200bc59874d58b811133 Mon Sep 17 00:00:00 2001 From: oarriaga Date: Sun, 2 Aug 2026 21:20:22 +0200 Subject: [PATCH 4/5] refactor: reuse graphics EPSILON for triangle shadow self-hits The invented TRIANGLE_SELF_HIT_EPSILON was unexercised: sweeping it over five orders of magnitude across three scene layouts produced byte-identical renders, because with only shape receivers a shadow ray can reach a triangle at near-zero depth only if the shape is coincident with the mesh. Use the EPSILON the shape path already applies to non-self hits instead of a new magic number. --- paz/graphics/renderer/shadow.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/paz/graphics/renderer/shadow.py b/paz/graphics/renderer/shadow.py index 78297dc5d..b4ed7a51b 100644 --- a/paz/graphics/renderer/shadow.py +++ b/paz/graphics/renderer/shadow.py @@ -8,10 +8,6 @@ SHADOW_ORIGIN_EPSILON = 1e-5 SHADOW_SELF_HIT_EPSILON = 1e-5 -# A shadow ray leaving a mesh re-hits its own triangle at a tiny depth. -# The shape path rejects that by shape identity; triangles have no such -# index, so they need a distance large enough to clear the surface. -TRIANGLE_SELF_HIT_EPSILON = 1e-3 def compute_occlusion(compiled, closest, indices, directions, distance, @@ -43,7 +39,7 @@ def compute_triangle_blockers(compiled, origins, directions, face_chunk): hit_mask, depth, _, _, face_index = result primitive = triangles.primitive_index[face_index] hit_mask = jp.logical_and(hit_mask, compiled.triangle_mask[primitive]) - hit_mask = jp.logical_and(hit_mask, depth > TRIANGLE_SELF_HIT_EPSILON) + hit_mask = jp.logical_and(hit_mask, depth > paz.graphics.EPSILON) hit_mask = hide_non_casting_triangles(compiled, hit_mask, primitive) return hit_mask, jp.where(hit_mask, depth, paz.graphics.FARAWAY) From c32d25c842af57772396e061697bcbcb3b52a826 Mon Sep 17 00:00:00 2001 From: oarriaga Date: Sun, 2 Aug 2026 22:15:42 +0200 Subject: [PATCH 5/5] feat: let meshes receive shadows compute_triangle_colors shaded with compute_colors and no occlusion, so a mesh was never darkened by anything and a mesh-only scene had no shadows at all regardless of the shadows flag. Triangles now take the same per-light occlusion the shape path uses. compute_occlusion takes a Receiver (points, normals, shape indices) instead of a Hit, so triangle hit points can be fed in; a triangle receiver uses NO_SHAPE so the shape self-hit rejection never matches it. Shape blockers are now skipped when a scene has no shapes, which a mesh-only scene with shadows would previously have crashed on. Reinstates TRIANGLE_SELF_HIT_EPSILON, this time with evidence: on a sphere mesh self-shadowing at its terminator, values at or below 1e-5 leak self-hits, the effect saturates at 1e-3, and 1e-3 to 5e-2 are flat. The spurious self-shadow it removes is worth up to 9 grey levels. --- paz/graphics/mesh_test.py | 29 +++++++++++++++++ paz/graphics/renderer/shade.py | 58 +++++++++++++++++++++++---------- paz/graphics/renderer/shadow.py | 52 ++++++++++++++++++++++------- 3 files changed, 110 insertions(+), 29 deletions(-) diff --git a/paz/graphics/mesh_test.py b/paz/graphics/mesh_test.py index faf5f9f37..50bfa9a80 100644 --- a/paz/graphics/mesh_test.py +++ b/paz/graphics/mesh_test.py @@ -920,6 +920,35 @@ def test_mesh_shadow_mask_stops_mesh_casting(): assert jp.all(casting <= blocked + 1e-4) +def build_mesh_only_shadow_scene(): + vertices, faces, edges = build_cube(1.0) + colors = build_vertex_colors(vertices, [0.7, 0.4, 0.2]) + material = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100) + blocker_pose = SE3.translation(jp.array([0.0, 1.4, 0.0])) + blocker = Mesh(vertices, colors, blocker_pose, material, faces, edges) + slab_pose = SE3.translation(jp.array([0.0, -0.6, 0.0])) + slab_pose = slab_pose @ SE3.scaling(jp.array([6.0, 0.2, 6.0])) + slab = Mesh(vertices, colors, slab_pose, material, faces, edges) + return paz.graphics.Scene([blocker, slab]) + + +def render_mesh_only_shadow(shadows): + camera_pose = SE3.view_transform( + jp.array([0.0, 2.4, -5.0]), jp.zeros(3), jp.array([0.0, 1.0, 0.0]) + ) + lights = [PointLight(jp.ones(3), jp.array([0.0, 6.0, -1.5]))] + scene = build_mesh_only_shadow_scene() + args = (32, 32), jp.pi / 3.0, camera_pose, scene, None, lights + return paz.graphics.render(*args, (1, 1), 1024, shadows) + + +def test_mesh_receives_shadow_from_mesh(): + lit, _ = render_mesh_only_shadow(False) + shadowed, _ = render_mesh_only_shadow(True) + assert compute_max_abs_difference(lit, shadowed) > 1e-2 + assert jp.all(shadowed <= lit + 1e-4) + + def test_scene_rejects_meshes_with_mixed_pattern_sizes(): plain = build_cube_mesh() textured = build_textured_quad_mesh() diff --git a/paz/graphics/renderer/shade.py b/paz/graphics/renderer/shade.py index c746ffe72..f266b6177 100644 --- a/paz/graphics/renderer/shade.py +++ b/paz/graphics/renderer/shade.py @@ -15,7 +15,8 @@ def compute_hit_colors( compiled, closest, surfaces, shadows, triangle_hit, face_chunk ): if len(compiled.shapes) == 0: - colors = compute_triangle_colors(compiled, triangle_hit) + triangle_args = compiled, triangle_hit, shadows, face_chunk + colors = compute_triangle_colors(*triangle_args) elif triangle_hit is None: shape_args = compiled, closest, surfaces, shadows, face_chunk colors = compute_shape_colors(*shape_args) @@ -33,7 +34,8 @@ def blend_hit_colors( # return rows instead of selecting inside its per-light scan. shape_args = compiled, closest, surfaces, shadows, face_chunk shape_colors = compute_shape_colors(*shape_args) - triangle_colors = compute_triangle_colors(compiled, triangle_hit) + triangle_args = compiled, triangle_hit, shadows, face_chunk + triangle_colors = compute_triangle_colors(*triangle_args) is_triangle = closest.primitive_index == len(compiled.shapes) is_triangle = jp.expand_dims(is_triangle, -1) return jp.where(is_triangle, triangle_colors, shape_colors) @@ -51,18 +53,47 @@ def compute_shape_colors(compiled, closest, surfaces, shadows, face_chunk): return colors -def compute_triangle_colors(compiled, triangle_hit): - materials = compiled.triangles.materials - material = gather_triangle_material(materials, triangle_hit.primitive) - shader = select_shader(materials) +def compute_triangle_colors(compiled, triangle_hit, shadows, face_chunk): + if shadows: + args = compiled, triangle_hit, face_chunk + colors = color_triangles_with_shadows(*args) + else: + colors = color_triangles(compiled, triangle_hit) + return colors + + +def color_triangles(compiled, triangle_hit): + shader, material = select_triangle_shader(compiled, triangle_hit) colors = jp.zeros_like(triangle_hit.albedo) for light in compiled.lights: - args = triangle_hit.albedo, material, triangle_hit.points - args += triangle_hit.normals, triangle_hit.eyes, light + args = build_triangle_shader_args(triangle_hit, material, light) colors = colors + shader.compute_colors(*args) return colors +def color_triangles_with_shadows(compiled, triangle_hit, face_chunk): + shader, material = select_triangle_shader(compiled, triangle_hit) + receiver = shadow.build_triangle_receiver(triangle_hit) + colors = jp.zeros_like(triangle_hit.albedo) + for light in compiled.lights: + occlusion_args = compiled, receiver, light, face_chunk + is_shadow = shadow.compute_occlusion(*occlusion_args) + args = build_triangle_shader_args(triangle_hit, material, light) + colors = colors + shader.compute_colors_with_shadow(*args, is_shadow) + return colors + + +def select_triangle_shader(compiled, triangle_hit): + materials = compiled.triangles.materials + material = gather_triangle_material(materials, triangle_hit.primitive) + return select_shader(materials), material + + +def build_triangle_shader_args(triangle_hit, material, light): + args = triangle_hit.albedo, material, triangle_hit.points + return args + (triangle_hit.normals, triangle_hit.eyes, light) + + def gather_triangle_material(materials, primitive): material = jax.tree.map(lambda field: field[primitive], materials) return jax.tree.map(expand_scalar_field, material) @@ -120,19 +151,12 @@ def scan_light_step( def compute_light_colors( compiled, closest, surfaces, indices, light, face_chunk ): - directions, distance = compute_light_directions(light, closest.point) - occlusion_args = compiled, closest, indices, directions, distance - is_shadow = shadow.compute_occlusion(*occlusion_args, face_chunk) + receiver = shadow.build_shape_receiver(closest, indices) + is_shadow = shadow.compute_occlusion(compiled, receiver, light, face_chunk) color_args = compiled.shapes, light, surfaces, is_shadow return take_closest(compute_shadowed_colors(*color_args), indices) -def compute_light_directions(light, points): - vector = light.position - points - norm = paz.algebra.compute_norms(vector, 1) - return vector / norm, jp.squeeze(norm, axis=1) - - def compute_shadowed_colors(shapes, light, surfaces, is_shadow): colors = [] for group, start_arg, final_arg in iterate_shape_groups(shapes): diff --git a/paz/graphics/renderer/shadow.py b/paz/graphics/renderer/shadow.py index b4ed7a51b..4e0fe24c5 100644 --- a/paz/graphics/renderer/shadow.py +++ b/paz/graphics/renderer/shadow.py @@ -1,3 +1,5 @@ +from collections import namedtuple + import jax import jax.numpy as jp @@ -8,27 +10,52 @@ SHADOW_ORIGIN_EPSILON = 1e-5 SHADOW_SELF_HIT_EPSILON = 1e-5 +# A shadow ray leaving a mesh surface can re-hit its own triangle at a +# grazing angle. Shapes reject that by identity; triangles have no such +# index, so they need a distance that clears float error at scene scale. +TRIANGLE_SELF_HIT_EPSILON = 1e-3 +NO_SHAPE = -1 + +Receiver = namedtuple("Receiver", ["points", "normals", "indices"]) + + +def build_shape_receiver(closest, indices): + return Receiver(closest.point, closest.normal, indices) -def compute_occlusion(compiled, closest, indices, directions, distance, - face_chunk): - origins = compute_shadow_ray_origins(closest.point, closest.normal) - shape_args = compiled, closest, indices, origins, directions - masks, depths = compute_shape_blockers(*shape_args) +def build_triangle_receiver(triangle_hit): + indices = jp.full(len(triangle_hit.points), NO_SHAPE) + return Receiver(triangle_hit.points, triangle_hit.normals, indices) + + +def compute_occlusion(compiled, receiver, light, face_chunk): + directions, distance = compute_light_directions(light, receiver.points) + origins = compute_shadow_ray_origins(receiver.points, receiver.normals) + rows = [] + if len(compiled.shapes) > 0: + shape_args = compiled, receiver, origins, directions + rows.append(compute_shape_blockers(*shape_args)) if compiled.triangles is not None: blocker_args = compiled, origins, directions, face_chunk - mask, depth = compute_triangle_blockers(*blocker_args) - masks = jp.concatenate([masks, jp.expand_dims(mask, 0)], axis=0) - depths = jp.concatenate([depths, jp.expand_dims(depth, 0)], axis=0) + rows.append(compute_triangle_blockers(*blocker_args)) + masks = jp.concatenate([row[0] for row in rows], axis=0) + depths = jp.concatenate([row[1] for row in rows], axis=0) return compute_soft_occlusion(masks, depths, distance) -def compute_shape_blockers(compiled, closest, indices, origins, directions): +def compute_light_directions(light, points): + vector = light.position - points + norm = paz.algebra.compute_norms(vector, 1) + return vector / norm, jp.squeeze(norm, axis=1) + + +def compute_shape_blockers(compiled, receiver, origins, directions): shadow_args = compiled.shapes, origins, directions intersections = intersect_shadow_groups(*shadow_args) hit_masks, depths, _, _, _, casters = intersections masks = resolve_shadow_masks(compiled, hit_masks) - depth_args = masks, depths, casters, indices, closest.normal, directions + depth_args = masks, depths, casters, receiver.indices + depth_args += receiver.normals, directions return select_shadow_depths(*depth_args) @@ -39,9 +66,10 @@ def compute_triangle_blockers(compiled, origins, directions, face_chunk): hit_mask, depth, _, _, face_index = result primitive = triangles.primitive_index[face_index] hit_mask = jp.logical_and(hit_mask, compiled.triangle_mask[primitive]) - hit_mask = jp.logical_and(hit_mask, depth > paz.graphics.EPSILON) + hit_mask = jp.logical_and(hit_mask, depth > TRIANGLE_SELF_HIT_EPSILON) hit_mask = hide_non_casting_triangles(compiled, hit_mask, primitive) - return hit_mask, jp.where(hit_mask, depth, paz.graphics.FARAWAY) + depth = jp.where(hit_mask, depth, paz.graphics.FARAWAY) + return jp.expand_dims(hit_mask, 0), jp.expand_dims(depth, 0) def hide_non_casting_triangles(compiled, hit_mask, primitive):