drtk.pixel_grid_mesh

pixel_grid_mesh	Build a displaced pixel-cell mesh with optional cuts and feathered borders.
bernoulli_topology_samples	Sample two correlated hard topologies for the ARM estimator.
bernoulli_render_loss	Estimate expected rendered loss and its Bernoulli-logit gradient with ARM.
pixel_grid_mesh_ref	Tensor reference for pixel_grid_mesh(), supporting CPU and CUDA.

Build a mesh from an image-aligned grid, deform its cells, and learn which cells and connections to keep. The mesh builder uses hard topology; the Bernoulli helpers estimate gradients of the expected loss over sampled meshes.

See Pixel-grid topology and gradients for camera conventions, cuts, feathering, and the gradient model, and Learning pixel-grid geometry and topology for a rendered optimization example.

drtk.pixel_grid_mesh(depth_z, displacement_xyz, intrinsics, alive, join_right, join_down, diagonal=None, feather_width=None, *, return_metadata=False)

Build a displaced pixel-cell mesh with optional cuts and feathered borders.

Supports float32/float64 CUDA tensors and first-order gradients.

Parameters:

• depth_z – Positive camera Z depth [H, W] (float32 or float64).

• displacement_xyz – Camera-space translation proposals [H, W, 3]. Shared corners are least-squares averages of incident proposals; a cell’s final center need not equal its proposed center exactly.

• intrinsics – Nonsingular pinhole matrix [3, 3] with last row 0,0,1. Pixel centers have integer coordinates, corners half-integers.

• alive – Boolean cell participation [H, W].

• join_right – Boolean horizontal adjacency [H, W-1].

• join_down – Boolean vertical adjacency [H-1, W]. Isolated interior one-edge cuts are healed; cuts touching a live/dead or grid border survive. Adjacent cut edges are evaluated before any healing.

• diagonal – Optional boolean [H, W]: False selects TL–BR, True TR–BL. None selects the diagonal whose midpoint best matches the proposed cell center; ties select TL–BR.

• feather_width – Optional nonnegative widths [H, W] in source pixels. Generates boundary strips and fixed vertex opacity, not a learned opacity head. Widths are averaged within each boundary fan. Large widths can cause overlap; there is no global collision resolution.

• return_metadata – Return source-cell ownership and stable face-slot IDs.

Returns:

A fixed seven-element tuple (v, vt0, vt1, vi, vertex_opacity, face_cell_ids, face_ids). Vertices already include displacement and optional strips, in source-camera 3D, shape [V, 3]. UVs have shape [V, 2] and faces [F, 3] int32. vt0 is the source anchor; vt1 is the displaced source projection. Vertex opacity is [V] or None without feathering. The two metadata tensors are [F] int64 or None. No optional tuple slot is omitted.

All tensors must share device; floating inputs must share dtype. Inputs must be finite, depths/widths valid, and displaced vertices in front of the source camera. CUDA graph capture is unsupported.

See pixel_grid_mesh_ref() for the tensor reference implementation. Vertex/face storage order may differ; use stable face IDs when comparing geometry.

drtk.bernoulli_topology_samples(logits, *, generator=None)

Sample two correlated hard topologies for the ARM estimator.

Parameters:

• logits – Floating-point Bernoulli logits of any shape. Each element describes one independent proposal decision.

• generator – Optional random generator matching the logits’ device.

Returns:

A three-element tuple (plus, minus, uniform) with the same shape and device as logits. For p = sigmoid(logits) and shared u = uniform, plus = (u > 1 - p) and minus = (u < p) are boolean tensors. uniform has the logits’ floating dtype.

Both hard samples have Bernoulli marginals with probability p; they are correlated, not independent samples or elementwise complements. The returned samples and random numbers have no gradient to logits. Use bernoulli_render_loss() to obtain an ARM loss gradient, or use the pair and shared noise to implement a custom estimator.

drtk.bernoulli_render_loss(logits, evaluate, *, pairs=1, generator=None)

Estimate expected rendered loss and its Bernoulli-logit gradient with ARM.

Parameters:

• logits – Floating-point Bernoulli logits of any shape, typically a learnable tensor or the output of a prediction network.

• evaluate – Callable accepting a boolean tensor with the same shape as logits and returning a scalar loss tensor. It constructs a hard mesh, renders it, and compares the render with the target. Continuous parameters captured by the callable retain ordinary autograd through the sampled mesh and renderer.

• pairs – Positive number of correlated sample pairs. Each pair calls evaluate twice; more pairs reduce sampling variance at the cost of more evaluations.

• generator – Optional random generator matching the logits’ device.

Returns:

A scalar whose forward value is the mean of the 2 * pairs loss evaluations. Backward adds the ARM estimator for logits and preserves ordinary continuous-input gradients from those evaluations. Calling backward() accumulates these gradients normally.

The score term uses detached loss differences. Any topology cleanup must depend on sampled bits, not detached logits. Use deterministic evaluation, common auxiliary randomness within each pair, or auxiliary randomness independent of the sampled topology.

This helper estimates the expectation of the full loss: average losses, not rendered images. It does not compute visibility masks or supply task-specific regularizers. bernoulli_topology_samples() exposes the underlying correlated samples for custom estimators.

drtk.pixel_grid_mesh_ref(depth_z, displacement_xyz, intrinsics, alive, join_right, join_down, diagonal=None, feather_width=None, *, return_metadata=False)

Tensor reference for pixel_grid_mesh(), supporting CPU and CUDA.

See pixel_grid_mesh() for the complete argument and return-value contract.