OverviewFilesCommitsIssuesDocsReleasesPeers

ifsrenderer-linux / blender / ifs_volume.py

517 lines · python ref:
"""
Bakes a fractal into an OpenVDB volume, so Cycles renders it like any other data.

The live viewport is a screen-space overlay: no depth, no lighting, invisible to
F12. This is the other thing entirely — the chaos game accumulates into a 3D voxel
grid instead of a 2D histogram, the grid is written as `.vdb`, and Blender imports
it as a Volume object. From there it lights, occludes, motion-blurs and composites
like any other volume.

Blender ships `openvdb` in its Python (renamed from `pyopenvdb`), so no external
dependency is needed.

## Bounds

A chaos game is unbounded — a few samples fly arbitrarily far — so a bake needs a
box, and a box that is too big wastes resolution on empty space while one too
small clips the fractal. `auto_bounds` runs a cheap coarse pass first, finds the
voxels that actually got hit, and fits the real bake to that with a margin. It is
a percentile rather than the extremes, because the extremes are exactly the
handful of outliers that would blow the box back up.

## Memory

The grid is 4N x N x N uint32 — four contiguous channels per voxel, mirroring the
2D histogram's interleaving for the same coalescing reason. At N=256 that is
268 MB on the GPU; N=512 is 2.1 GB and will not fit alongside anything else on a
consumer card. 256 is the practical default.
"""

from __future__ import annotations

import os

import gpu
import numpy as np

import ifs_codegen
import ifs_document
import ifs_host
import ifs_layout
import ifs_render

HIST_SCALE = ifs_layout.DEFAULT_HIST_SCALE



def build_shader(hist_scale: float = HIST_SCALE):
    """Compiles the volume variant of the kernel."""
    library = ifs_codegen.load_library(
        os.path.join(ifs_render.REPO, "Transforms"),
        os.path.join(ifs_render.REPO, "WpfDisplay", "Library", "Includes"))
    template_path = os.path.join(
        ifs_render.REPO, "IFSEngine", "Rendering", "Shaders", "ifs_kernel.comp.shader")
    with open(template_path, "r", encoding="utf-8-sig") as handle:
        kernel = ifs_codegen.generate_kernel(handle.read(), library)
    kernel, _ = ifs_codegen.apply_portability_fixes(kernel)
    body, typedefs, report = ifs_layout.rewrite(
        kernel, hist_scale, ifs_render.STATE_W, volume=True)
    if not report.ok:
        raise RuntimeError("volume rewrite incomplete: " + "; ".join(report.missing))

    info = gpu.types.GPUShaderCreateInfo()
    info.typedef_source(typedefs)
    info.local_group_size(ifs_render.LOCAL_SIZE)
    rw = {'NO_RESTRICT', 'READ', 'WRITE'}
    info.image(0, 'R32UI', 'UINT_3D', 'volume_img', qualifiers=rw)
    info.image(1, 'RGBA32F', 'FLOAT_2D', 'state_img', qualifiers=rw)
    info.image(2, 'R32UI', 'UINT_3D', 'colour_img', qualifiers=rw)
    info.sampler(0, 'FLOAT_2D', 'alias_tables_tex')
    info.push_constant('VEC3', 'volume_min')
    info.push_constant('VEC3', 'volume_max')
    info.push_constant('INT', 'volume_size')
    info.uniform_buf(0, 'FrameParams', 'frame')
    info.uniform_buf(1, 'SettingsParams', 'settings')
    info.uniform_buf(2, 'IteratorsParams', 'iterators_buf')
    info.uniform_buf(3, 'PaletteParams', 'palette_buf')
    info.uniform_buf(4, 'RealParams', 'real_params_buf')
    info.uniform_buf(5, 'Vec3Params', 'vec3_params_buf')
    info.compute_source(body)
    return gpu.shader.create_from_info(info), library


def _document_buffers(doc):
    return {
        'settings': gpu.types.GPUUniformBuf(
            ifs_host.pack_settings(doc.camera, doc.settings)),
        'iterators_buf': gpu.types.GPUUniformBuf(
            ifs_host.pack_iterators(doc.iterators)),
        'palette_buf': gpu.types.GPUUniformBuf(
            ifs_host.pack_palette(doc.palette(doc.settings.gradient_resolution))),
        'real_params_buf': gpu.types.GPUUniformBuf(
            ifs_host.pack_real_params(ifs_host.flatten_real_params(doc.iterators))),
        'vec3_params_buf': gpu.types.GPUUniformBuf(
            ifs_host.pack_vec3_params(ifs_host.flatten_vec3_params(doc.iterators))),
    }


def accumulate(shader, doc, ubos, alias_tex, size, bounds, dispatches,
               invocation_iters=500):
    """
    Runs the chaos game into a fresh density grid plus a HALF-resolution colour
    grid, and reads both back.

    Returns `(density, colour)` — density (z, y, x) float32, colour
    (z/2, y/2, x/2, 3) float32 and still premultiplied by opacity. Splitting them
    costs 5.5 bytes per voxel against the old interleaved 16.
    """
    lo, hi = bounds
    half = max(1, size // 2)
    volume = gpu.types.GPUTexture((size, size, size), format='R32UI')
    volume.clear(format='UINT', value=(0, 0, 0, 0))
    colour_tex = gpu.types.GPUTexture((half * 3, half, half), format='R32UI')
    colour_tex.clear(format='UINT', value=(0, 0, 0, 0))
    rows = (ifs_render.INVOCATIONS + ifs_render.STATE_W - 1) // ifs_render.STATE_W
    state = gpu.types.GPUTexture((ifs_render.STATE_W * 2, rows), format='RGBA32F')
    state.clear(format='FLOAT', value=(0.0, 0.0, 0.0, 0.0))

    for index in range(dispatches):
        frame = gpu.types.GPUUniformBuf(ifs_host.pack_frame(
            seed=0x9E3779B9, width=size, height=size, dispatch_cnt=index,
            reset_points_state=1 if index == 0 else 0,
            invocation_iters=invocation_iters))
        shader.bind()
        shader.image('volume_img', volume)
        shader.image('colour_img', colour_tex)
        shader.image('state_img', state)
        shader.uniform_sampler('alias_tables_tex', alias_tex)
        shader.uniform_float('volume_min', lo)
        shader.uniform_float('volume_max', hi)
        shader.uniform_int('volume_size', (size,))
        shader.uniform_block('frame', frame)
        for name, ubo in ubos.items():
            shader.uniform_block(name, ubo)
        gpu.compute.dispatch(shader, ifs_render.WORKGROUPS, 1, 1)

    # ONE allocation each, not two. `astype(float32) / HIST_SCALE` materialises the
    # float copy AND the quotient; np.divide with an explicit dtype does both in a
    # single pass, which is most of a 2.4x peak-RSS overshoot.
    raw = np.array(volume.read(), dtype=np.uint32)
    density = np.divide(raw.reshape(size, size, size), HIST_SCALE, dtype=np.float32)
    del raw
    raw_colour = np.array(colour_tex.read(), dtype=np.uint32)
    colour = np.divide(raw_colour.reshape(half, half, half, 3), HIST_SCALE,
                       dtype=np.float32)
    del raw_colour
    return density, colour


def auto_bounds(shader, doc, ubos, alias_tex, probe_size=64, dispatches=8,
                percentile=99.5, margin=0.15, start_extent=4.0, max_extent=256.0,
                refine_passes=2, collapse_voxels=8):
    """
    Fits the bake box to where the fractal actually is.

    Two things this has to get right, and the first is easy to miss:

    * **The probe box must not clip.** Fitting inside a box the fractal overflows
      silently returns the box, not the fractal - the first version probed a fixed
      +/-8 and reported x = +/-10 for a gasket that simply ran off the edge. The
      probe now starts small and DOUBLES while a meaningful share of occupied
      voxels still touch its boundary.
    * **Fit to a percentile, not the extremes.** A chaos game throws a handful of
      samples arbitrarily far; fitting to those spends the whole grid on empty
      space. The outliers are excluded on purpose, so a bake does clip its faintest
      tails - that is the trade, and it is why `bounds` can be given explicitly.
    """
    extent = start_extent
    while True:
        box = ((-extent, -extent, -extent), (extent, extent, extent))
        density, _coarse_colour = accumulate(
            shader, doc, ubos, alias_tex, probe_size, box, dispatches)
        hit = np.argwhere(density > 0)
        if hit.size == 0:
            if extent >= max_extent:
                raise RuntimeError(
                    f"nothing found within +/-{extent:g}; the document may render "
                    "nothing, or need explicit bounds")
            extent *= 2.0
            continue

        edge = ((hit == 0) | (hit == probe_size - 1)).any(axis=1).mean()

        # Grow while the BODY is clipped, not while stragglers touch the edge.
        #
        # The old test grew whenever >1% of occupied voxels sat on the boundary. A
        # chaos game throws samples arbitrarily far, so for a heavy-tailed document
        # that fraction never falls: growing the box just captures more outliers.
        # light_rods_sample doubled all the way to the max_extent of 256 and still
        # reported 4.3% on the boundary, so it fitted a 655-unit box - roughly 6.8
        # units per voxel at size 96, smearing the fractal into a featureless blob.
        #
        # The percentile fit below already excludes those outliers deliberately; the
        # growth decision simply was not using it. Now it asks the same question the
        # fit does - is the 99% body inside the probe? - so the loop stops once the
        # part we are actually going to keep fits.
        body_lo, body_hi = _mass_bounds(density, probe_size, percentile)
        clipped = bool((body_lo <= 0.5).any() or (body_hi >= probe_size - 1.5).any())
        if not clipped or extent >= max_extent:
            break
        extent *= 2.0
    box = ((-extent, -extent, -extent), (extent, extent, extent))
    lo_w, hi_w = _fit(density, box, probe_size, percentile, margin)

    # Re-probe any axis that collapsed to a handful of probe voxels.
    #
    # The fit can only resolve an axis to the PROBE's voxel size, so a flat fractal -
    # spherical_gasket_base is planar, 2 voxels thick out of 64 - gets a box barely
    # thicker than one probe voxel AND offset by up to one, because both edges are
    # rounding artifacts. That put the plane at ~88% of the box height instead of the
    # middle, so the bake spent its resolution on empty space and pressed the fractal
    # against a face. Re-probing inside the fitted box gives that axis a voxel size
    # smaller by the ratio of the two spans, which is where the precision comes from.
    #
    # A truly planar attractor would refine forever, each pass halving the thickness
    # toward a zero-size voxel, so this stops when an axis is thinner than one voxel
    # of the LARGEST axis: past that it is planar in any sense the bake can represent.
    for _ in range(refine_passes):
        # `axis` is a world axis and the array is [x][y][z], so they coincide; this
        # used to be `2 - axis` to undo the old [z][y][x] readback.
        occupied = [int((density.sum(axis=tuple(a for a in range(3)
                        if a != axis)) > 0).sum()) for axis in range(3)]
        span = hi_w - lo_w
        floor = float(span.max()) / probe_size
        collapsed = [a for a in range(3)
                     if occupied[a] < collapse_voxels and span[a] > floor]
        if not collapsed:
            break

        # Pad the collapsed axes by a probe voxel of the box we just used: the fitted
        # edges are themselves quantised, so re-probing exactly the fit would clip
        # whatever sits in the rounding.
        pad = np.zeros(3)
        for a in collapsed:
            pad[a] = (box[1][a] - box[0][a]) / probe_size
        probe_lo = tuple(float(lo_w[a] - pad[a]) for a in range(3))
        probe_hi = tuple(float(hi_w[a] + pad[a]) for a in range(3))

        finer_density, _finer_colour = accumulate(
            shader, doc, ubos, alias_tex, probe_size, (probe_lo, probe_hi), dispatches)
        if not (finer_density > 0).any():
            break          # the tighter box found nothing; keep what we had

        box = (probe_lo, probe_hi)
        density = finer_density
        lo_w, hi_w = _fit(density, box, probe_size, percentile, margin)

    return (tuple(float(v) for v in lo_w), tuple(float(v) for v in hi_w),
            {"probe_extent": extent, "edge_fraction": float(edge),
             "clipped": bool(clipped)})


def _mass_bounds(density, probe_size, percentile):
    """
    Index bounds holding the central `percentile` of the fractal's MASS, per world axis.

    The fit used to percentile over `np.argwhere(density > 0)` - the INDICES of occupied
    voxels - which counts a voxel holding one sample the same as one holding a million.
    A chaos game's halo is enormous in voxel count and negligible in mass, so that
    statistic tracked the halo rather than the fractal. light_rods_sample grew its probe
    to the 256 cap and still reported the body clipped, fitting a 655-unit box.

    Summing the marginal along each axis and cutting the cumulative distribution asks
    where the mass is instead, which is what the bake should spend its resolution on.
    """
    lo_idx, hi_idx = [], []
    # The array is [x][y][z] now (the kernel stores transposed), so world axis and
    # array axis coincide — this used to iterate (2, 1, 0) to undo [z][y][x].
    for array_axis in (0, 1, 2):
        others = tuple(a for a in range(3) if a != array_axis)
        marginal = density.sum(axis=others)
        total = float(marginal.sum())
        if total <= 0:
            lo_idx.append(0.0)
            hi_idx.append(float(probe_size - 1))
            continue
        cumulative = np.cumsum(marginal) / total
        tail = (100.0 - percentile) / 100.0
        lo_idx.append(float(np.searchsorted(cumulative, tail)))
        hi_idx.append(float(np.searchsorted(cumulative, 1.0 - tail)))
    return np.array(lo_idx), np.array(hi_idx)


def _fit(density, box, probe_size, percentile, margin):
    """
    Percentile-fits world bounds to probe hits, for an ARBITRARY probe box.

    `hit` columns are (z, y, x) - the readback order - so they are reversed here into
    world xyz. This was inline against a symmetric +/-extent cube; it takes the box as
    an argument so a refinement pass can re-fit inside a tighter one.
    """
    lo, hi = np.array(box[0], dtype=float), np.array(box[1], dtype=float)
    # Already in world-axis order (x, y, z) - _mass_bounds does the reordering.
    lo_idx, hi_idx = _mass_bounds(density, probe_size, percentile)
    lo_n = lo_idx / probe_size
    hi_n = hi_idx / probe_size

    lo_w = lo + lo_n * (hi - lo)
    hi_w = lo + hi_n * (hi - lo)
    span = np.maximum(hi_w - lo_w, 1e-3)
    return lo_w - span * margin, hi_w + span * margin


def write_vdb(density_grid: np.ndarray, colour_grid: np.ndarray, bounds,
              path: str, name: str = "density", gamma: float = 1.0):
    """
    Writes density and colour grids to a `.vdb`.

    Density is normalised to a 0..1 peak: absolute sample counts depend on how long
    the bake ran, and a volume whose density scales with render time is unusable.
    """
    import openvdb

    size = density_grid.shape[0]
    lo, hi = bounds

    # Per-axis scale AND a translation to `lo`. Both matter and both were missing:
    #
    # * `auto_bounds` fits each axis independently (per-axis percentile + margin), so
    #   non-cubic bounds are the NORMAL case. One voxel size taken from the x extent
    #   squashed y and z by whatever the aspect happened to be.
    # * A transform carrying only a voxel size leaves the grid at world
    #   (0 .. size*voxel), because copyFromArray writes from index origin. The bake
    #   landed nowhere near the bounds it was fitted to.
    #
    # openvdb uses the ROW-vector convention, so translation goes in the last ROW -
    # verified against indexToWorld rather than assumed.
    #
    # The translation is `lo + voxel/2`, not `lo`, because a voxel INDEX addresses
    # the CENTRE of its cell rather than the cell's near corner. Measured, not
    # assumed: a grid with one lit voxel at index (0,0,0) and voxel size 1 imports
    # into Blender with a bounding box starting at -0.5. `accumulate_voxel` bins a
    # sample by `ivec3(t * volume_size)`, so the mass written into voxel i came from
    # world [lo + i*s, lo + (i+1)*s) - centred half a voxel PAST lo + i*s. Placing
    # the grid at `lo` therefore displaced the whole bake by half a voxel toward
    # -x, -y, -z. Small (1/512 of the box at size 256) and invisible against a
    # fractal, which is exactly why it survived: test_volume_axes checks a marker
    # against the same convention it was written with, so it agreed with itself.
    # test_bake_orientation measures the baked centroid against an analytically
    # known one instead, where it showed up as a clean -0.5 on all three axes.
    def _transform(scale=1.0):
        """A fresh Transform per grid; sharing one between grids is asking for it."""
        sx, sy, sz = ((hi[a] - lo[a]) / size * scale for a in range(3))
        return openvdb.createLinearTransform(matrix=[
            [sx, 0.0, 0.0, 0.0],
            [0.0, sy, 0.0, 0.0],
            [0.0, 0.0, sz, 0.0],
            [lo[0] + sx / 2.0, lo[1] + sy / 2.0, lo[2] + sz / 2.0, 1.0],
        ])

    # TRANSPOSE, and it is load-bearing. Two conventions meet here and disagree:
    # `GPUTexture.read()` hands back a 3D texture as (z, y, x*4), so `grid` is
    # indexed [z][y][x][channel] - while openvdb's `copyFromArray` reads array index
    # (i, j, k) as world (x, y, z). Passing the array straight through swaps x and z,
    # and a transposed fractal is still a plausible fractal, so nothing downstream
    # notices; it shows only as a bake whose orientation disagrees with the still
    # render. `auto_bounds` above already does this reorder (its `zyx_*` locals), so
    # before this the same file computed correct bounds around transposed data.
    # GAMMA, off by default. A chaos game's density histogram is
    # violently long-tailed: measured on nebula_sample the median occupied voxel sits at
    # 2e-5 of the peak and the 90th percentile at 3e-4, so dividing by the max alone
    # crushes the entire fractal to zero and leaves a couple of voxels at 1.0. There is
    # gamma=4 lifts that median to 0.063 and the 90th to 0.136. It is OFF by default
    # because it does not fix what it looks like it should: the featureless blob was
    # Principled Volume's Emission Strength not scaling with density, so every voxel
    # emitted equally regardless of the histogram. Lifting the faint bulk toward uniform
    # actually moves AWAY from the sparse, mostly-black look the still render has. Kept
    # as a knob for deliberately hazier bakes, not as a fix. (A log curve is no use here
    # either: log(1+x) ~ x for small x, so it compresses the top rather than lifting the
    # bottom.)
    #
    # Applied to density ONLY, and only after colour has been divided by the ORIGINAL
    # density below - colour is accumulated weighted by density, so dividing it by a
    # gamma-lifted density collapses it to black. That mistake renders an empty volume.
    density = density_grid
    peak = float(density.max())
    if peak <= 0:
        raise RuntimeError("nothing was accumulated - the bounds may exclude the fractal")
    # NO transpose. The kernel now stores voxel (x,y,z) at texel (z,y,x), so the
    # readback already arrives as [x][y][z] — exactly what copyFromArray wants.
    # Doing it here instead cost 17.6s at N=512, ~70% of the whole write, because
    # a (N,N,N) transpose is a full strided copy with cache-hostile access.
    #
    # Still a COPY, not in place: `density_grid` is read again below for the colour
    # blocks and again by bake_document for its occupancy and exposure statistics,
    # so normalising the caller's array would silently double-normalise. A
    # contiguous divide is ~0.2s; the strided copy was the expensive part.
    density = density_grid / peak
    if gamma and gamma != 1.0:
        np.power(density, 1.0 / gamma, out=density)

    grids = []
    dgrid = openvdb.FloatGrid()
    dgrid.copyFromArray(density)
    dgrid.name = name
    dgrid.transform = _transform()
    grids.append(dgrid)

    # Colour is accumulated weighted by density, so divide it back out.
    # Colour is HALF resolution and still premultiplied by opacity, so dividing it
    # back out needs the density of the SAME colour voxel — the sum of its 2x2x2
    # density block, not a single voxel. Getting that wrong tints the fractal in a
    # way that looks like a palette bug rather than an indexing one.
    half = colour_grid.shape[0]
    usable = half * 2
    block = density_grid[:usable, :usable, :usable].reshape(
        half, 2, half, 2, half, 2).sum(axis=(1, 3, 5))
    safe = np.maximum(block[:, :, :, None], 1e-9)

    # No transpose here either, for the same reason — and the two grids must agree,
    # because a colour grid oriented differently from density tints the wrong parts
    # of the fractal, which is far harder to spot than a whole-volume rotation.
    colour = np.clip(colour_grid / safe, 0.0, 1.0)

    try:
        cgrid = openvdb.Vec3SGrid()
        cgrid.copyFromArray(colour)
        cgrid.name = "color"
        # Twice the voxel size, because it is half the resolution. VDB grids carry
        # independent transforms and Blender honours them, so no upsampling.
        cgrid.transform = _transform(scale=2.0)
        grids.append(cgrid)
    except Exception:
        # A density-only volume is still useful; do not lose the bake over colour.
        pass

    openvdb.write(path, grids=grids)
    return path, peak


def bake_document(doc, out_path, size=256, dispatches=200, bounds=None,
                  invocation_iters=500, shader=None, gamma=1.0):
    """
    Bakes a document that is ALREADY loaded, which is what makes the bake follow what
    you are looking at.

    The node graph and the panel sliders both write into the live viewport session's
    document; `bake` below reloads a fresh copy from disk, so everything edited inside
    Blender was silently discarded and a bake could only ever reproduce whatever the
    .ifsjson already contained. Taking the document as an argument is the whole fix -
    the operator hands over the session's own object.

    The document's camera cannot skew the result: volume mode replaces the projection
    block outright (ifs_layout._PLOT_VOLUME), so a live session's viewport camera is
    simply unused here.
    """
    gpu.init()
    if shader is None:
        shader, _library = build_shader()
    ubos = _document_buffers(doc)

    alias = ifs_host.build_alias_tables(doc.iterators)
    n = len(doc.iterators)
    alias_tex = gpu.types.GPUTexture(
        (n, n), format='RGBA32F',
        data=gpu.types.Buffer('FLOAT', len(alias), alias))

    fit = {}
    if bounds is None:
        lo, hi, fit = auto_bounds(shader, doc, ubos, alias_tex)
        bounds = (lo, hi)

    grid, colour_grid = accumulate(shader, doc, ubos, alias_tex, size, bounds,
                                   dispatches, invocation_iters)
    path, peak = write_vdb(grid, colour_grid, bounds, out_path, gamma=gamma)
    occupied = int((grid > 0).sum())

    # How bright to drive emission, DERIVED rather than guessed. A fixed number cannot
    # work: measured across two shipped templates the required scale differs by ~200x
    # (nebula wants ~1100, the gasket ~5), because peak-normalised density depends
    # entirely on how long-tailed that particular fractal's histogram is. A constant
    # tuned on one document blows the other to white, which is exactly what happened.
    #
    # The statistic is what the renderer actually integrates: density summed ALONG A
    # RAY. Sum each column, average over the columns that contain anything, and scale
    # so a typical ray through the fractal integrates to ~1. Not a heuristic - it is
    # the quantity the volume integrator accumulates.
    # Averaged over ALL THREE axes, not one. Summing a single axis measures the path
    # a ray takes along that axis only, and for a PLANAR fractal - the gasket is two
    # probe voxels thick - the fixed axis is the thin direction, so the statistic
    # reports the shortest possible path and derives an exposure far too dark. The
    # camera can look from anywhere, so the metric has to be view-independent.
    normalised = grid / peak if peak > 0 else grid
    per_axis = []
    for axis in range(3):
        columns = normalised.sum(axis=axis)
        lit_columns = columns[columns > 0]
        if lit_columns.size:
            per_axis.append(float(lit_columns.mean()))
    mean_column = float(np.mean(per_axis)) if per_axis else 0.0
    emission_scale = (1.0 / mean_column) if mean_column > 0 else 1.0
    return path, {
        **fit,
        "bounds": bounds,
        "size": size,
        "occupied_voxels": occupied,
        "total_voxels": size ** 3,
        "peak_density": peak,
        "emission_scale": emission_scale,
        "mean_column_density": mean_column,
    }


def bake(document_path, out_path, size=256, dispatches=200, bounds=None,
         invocation_iters=500, gamma=1.0):
    """
    Loads a document from disk and bakes it. Returns (path, stats).

    Kept as-is for callers that genuinely mean "bake this file" - the CLI, and the
    operator when no viewport session is running. Anything that wants what is on
    screen should call `bake_document` with the session's document instead.
    """
    gpu.init()
    shader, library = build_shader()
    doc = ifs_document.load(document_path, library)
    return bake_document(doc, out_path, size, dispatches, bounds,
                         invocation_iters, shader=shader, gamma=gamma)