"""
Draws the emitted points as gaussian splats.
The viewport's normal path accumulates samples into a 2D histogram and tonemaps
it. This is the other way of showing the same points: give each one a camera-
facing quad and a gaussian falloff, and add them together. It renders the
PointCloud we already emit, at float precision, with no grid to be finer than.
## Why this is so much smaller than a general 3DGS renderer
A published gaussian-splat renderer spends most of its code on two things we do
not need. See `NOTES-splat-rendering.md` for the reading this came from.
* **The covariance projection.** A trained splat is an oriented, anisotropic
ellipsoid, so the renderer packs a 3D covariance, builds the Jacobian of the
perspective projection per splat, computes `cov2d = T' Vrk T`, eigen-decomposes
the 2x2 result and emits a quad along its major/minor axes. Ours are
ISOTROPIC — `to_pointcloud` writes the identity quaternion because the chaos
game gives sample positions, not local anisotropy — so `cov3d = r^2 I`, the
projection of a sphere is a circle, and all of that collapses to offsetting the
quad corner in CAMERA space by the radius. Perspective then shrinks distant
splats for free.
* **The depth sort.** Splat renderers sort back-to-front every frame because
alpha compositing is order-dependent; that sort is the expensive and fiddly
part. A fractal flame is EMISSION, contributions add, and addition commutes —
so there is nothing to sort. This is the same reason the viewport's tonemap
already blends `ADDITIVE_PREMULT`.
The cost of skipping the sort is that splats do not occlude EACH OTHER. For
emission that is correct rather than an approximation. It does mean this pass
suits the EMISSION shading mode; FLAT and METAL are alpha modes and would need
the sort back before they could use it.
## Getting the splats to the shader
Blender's GPU module has no SSBOs, so the per-splat data rides in two RGBA32F
textures indexed by `texelFetch` — the same substitute `ifs_points` uses for the
points buffer, and the same one the WebGL reference uses. The vertex buffer holds
only a quad corner and a splat index; four vertices per splat, six indices.
Position texture: `xyz` = centre, `w` = radius (sigma).
Colour texture: `rgb` = colour, `w` = opacity.
That is exactly the layout `ifs_points.emit` already returns, so no repacking.
"""
from __future__ import annotations
import gpu
import numpy as np
from gpu_extras.batch import batch_for_shader
# One texel per splat, same reasoning as ifs_points.POINTS_WIDTH: comfortably
# inside every GL_MAX_TEXTURE_SIZE worth worrying about.
SPLAT_WIDTH = 4096
# The quad spans +/-2 sigma; the fragment shader discards beyond that, so the
# corners are cut and a splat is round rather than square. exp(-4) = 0.018, i.e.
# under 2% of peak, which is where the falloff stops being worth a fragment.
CUTOFF_SIGMA = 2.0
VERTEX_SHADER = """
void main()
{
ivec2 texel = ivec2(splat_index % splat_width, splat_index / splat_width);
vec4 centre = texelFetch(splat_pos, texel, 0);
vec4 tint = texelFetch(splat_col, texel, 0);
/* Offset in CAMERA space, not clip space. A sphere projects to a circle, so
an isotropic splat needs no covariance maths at all — and doing it here
rather than in NDC means perspective shrinks distant splats for free. */
vec4 cam = view * vec4(centre.xyz, 1.0);
cam.xy += corner * centre.w;
vColor = tint.rgb;
vAlpha = tint.w;
vOffset = corner;
gl_Position = projection * cam;
}
"""
FRAGMENT_SHADER = """
void main()
{
/* vOffset is in units of sigma, so this is -(d/sigma)^2 directly. */
float a = -dot(vOffset, vOffset);
if (a < -CUTOFF_SQ) discard; /* round, not square */
/* Nearest-splat depth, in exactly the encoding the kernel uses so the shared
tonemap can read this buffer without knowing which pass filled it:
floatBitsToUint of window-space z, which is monotonic over [0,1] so an
atomic MIN is an exact minimum. We draw with Blender's own window matrix,
so gl_FragCoord.z IS that value — no re-derivation, no near/far guesswork.
Written for every surviving fragment rather than at the centre, for the
reason the density-estimation filter had to: a fragment left at the far
value silently refuses to be occluded. */
imageAtomicMin(depth_img, ivec2(gl_FragCoord.xy), floatBitsToUint(gl_FragCoord.z));
float weight = exp(a) * vAlpha;
/* Premultiplied and ADDITIVE: this lands in the float accumulation buffer that
the histogram path also fills, so the shared maxpass + tonemap turn it into
the image. Values above 1 are wanted, not clamped — overlapping splats are
how a flame gets its core, and the tonemap is what compresses that back. */
fragColor = vec4(vColor * weight, weight);
}
"""
def build_shader():
"""Compiles the splat pass."""
info = gpu.types.GPUShaderCreateInfo()
interface = gpu.types.GPUStageInterfaceInfo("ifs_splat_iface")
interface.smooth('VEC3', 'vColor')
interface.smooth('FLOAT', 'vAlpha')
interface.smooth('VEC2', 'vOffset')
info.vertex_in(0, 'VEC2', 'corner')
info.vertex_in(1, 'INT', 'splat_index')
info.vertex_out(interface)
info.fragment_out(0, 'VEC4', 'fragColor')
info.sampler(0, 'FLOAT_2D', 'splat_pos')
info.sampler(1, 'FLOAT_2D', 'splat_col')
info.image(0, 'R32UI', 'UINT_2D', 'depth_img',
qualifiers={'NO_RESTRICT', 'READ', 'WRITE'})
info.push_constant('MAT4', 'view')
info.push_constant('MAT4', 'projection')
info.push_constant('INT', 'splat_width')
info.define('CUTOFF_SQ', str(CUTOFF_SIGMA * CUTOFF_SIGMA))
info.vertex_source(VERTEX_SHADER)
info.fragment_source(FRAGMENT_SHADER)
return gpu.shader.create_from_info(info)
class SplatCloud:
"""The GPU-side state for one cloud: two data textures and a quad batch."""
def __init__(self, batch, positions_tex, colors_tex, width, count):
self.batch = batch
self.positions_tex = positions_tex
self.colors_tex = colors_tex
self.width = width
self.count = count
def _to_texture(rows, width):
"""Packs an (N, 4) float array into an RGBA32F texture of `width` texels."""
count = len(rows)
height = max(1, (count + width - 1) // width)
padded = np.zeros((height * width, 4), dtype=np.float32)
padded[:count] = rows
flat = np.ascontiguousarray(padded).ravel()
return gpu.types.GPUTexture(
(width, height), format='RGBA32F',
data=gpu.types.Buffer('FLOAT', len(flat), flat))
def upload(positions, colors, radius, width=SPLAT_WIDTH):
"""
Builds a `SplatCloud` from the arrays `ifs_points.emit_document` returns.
`radius` is the sigma of every splat — one value, because the cloud is
isotropic and `estimate_spacing` derives a single spacing for it. Per-point
opacity rides in `colors[:, 3]`, which is the w the kernel emits and which
`to_pointcloud` now also writes as the `opacity` attribute.
"""
positions = np.asarray(positions, dtype=np.float32)
colors = np.asarray(colors, dtype=np.float32)
count = len(positions)
if count == 0:
raise ValueError('cannot build a splat cloud from zero points')
pos_rows = np.zeros((count, 4), dtype=np.float32)
pos_rows[:, :3] = positions[:, :3]
pos_rows[:, 3] = radius
col_rows = np.ones((count, 4), dtype=np.float32)
col_rows[:, :colors.shape[1]] = colors[:, :4]
if colors.shape[1] < 4:
col_rows[:, 3] = 1.0
# Four corners per splat, spanning +/-CUTOFF_SIGMA, and two triangles.
quad = np.array([(-1.0, -1.0), (1.0, -1.0), (1.0, 1.0), (-1.0, 1.0)],
dtype=np.float32) * CUTOFF_SIGMA
corners = np.tile(quad, (count, 1))
indices = np.repeat(np.arange(count, dtype=np.int32), 4)
tris = (np.array([0, 1, 2, 0, 2, 3], dtype=np.int32)[None, :]
+ (np.arange(count, dtype=np.int32) * 4)[:, None]).reshape(-1, 3)
shader_format = gpu.types.GPUVertFormat()
shader_format.attr_add(id="corner", comp_type='F32', len=2, fetch_mode='FLOAT')
shader_format.attr_add(id="splat_index", comp_type='I32', len=1, fetch_mode='INT')
vbo = gpu.types.GPUVertBuf(shader_format, len=count * 4)
# numpy STRAIGHT in, no .tolist(). Measured at 1M splats: the two list
# conversions were 739 ms + 445 ms of a 1952 ms upload — 61% of the total,
# spent building four million Python tuples that the API then unpacks again.
# Both attr_fill and GPUIndexBuf accept arrays directly.
vbo.attr_fill(id="corner", data=corners)
vbo.attr_fill(id="splat_index", data=indices)
ibo = gpu.types.GPUIndexBuf(type='TRIS', seq=tris)
batch = gpu.types.GPUBatch(type='TRIS', buf=vbo, elem=ibo)
return SplatCloud(batch, _to_texture(pos_rows, width), _to_texture(col_rows, width),
width, count)
def draw(cloud, shader, view_matrix, projection_matrix, depth_img,
additive=True):
"""
Accumulates the cloud into whatever framebuffer is currently bound.
This does NOT draw to the screen. It is a density pass: the caller binds a
float colour attachment (the viewport's `accum`), and the tonemap turns the
result into an image afterwards. That indirection is the whole point — the
fractal's look comes from a curve applied to ACCUMULATED density, not from
what one sample looks like, so drawing splats straight to the viewport
produced fuzzy beads that no panel control could fix.
`depth_img` is the R32UI nearest-depth buffer the tonemap reads, written here
by the fragment shader in the kernel's own encoding.
"""
shader.bind()
shader.uniform_sampler('splat_pos', cloud.positions_tex)
shader.uniform_sampler('splat_col', cloud.colors_tex)
shader.image('depth_img', depth_img)
shader.uniform_int('splat_width', (cloud.width,))
shader.uniform_float('view', view_matrix)
shader.uniform_float('projection', projection_matrix)
# No depth TEST here: this pass has no scene depth to test against, and
# occlusion happens later when the tonemap blit emits gl_FragDepth from
# depth_img and is tested against the real viewport buffer.
gpu.state.blend_set('ADDITIVE_PREMULT' if additive else 'ALPHA')
gpu.state.depth_mask_set(False)
cloud.batch.draw(shader)
gpu.state.blend_set('NONE')