kdtree.java

Sunflow是一个照片级的渲染系统

package org.sunflow.core.accel;

import java.io.FileWriter;
import java.io.IOException;

import org.sunflow.core.AccelerationStructure;
import org.sunflow.core.IntersectionState;
import org.sunflow.core.PrimitiveList;
import org.sunflow.core.Ray;
import org.sunflow.image.Color;
import org.sunflow.math.BoundingBox;
import org.sunflow.math.Point3;
import org.sunflow.system.Memory;
import org.sunflow.system.Timer;
import org.sunflow.system.UI;
import org.sunflow.system.UI.Module;
import org.sunflow.util.IntArray;

public class KDTree implements AccelerationStructure {
    private int[] tree;
    private int[] primitives;
    private PrimitiveList primitiveList;
    private BoundingBox bounds;

    private int maxPrims;

    private static final float INTERSECT_COST = 0.5f;
    private static final float TRAVERSAL_COST = 1;
    private static final float EMPTY_BONUS = 0.2f;
    private static final int MAX_DEPTH = 64;

    private static boolean dump = false;
    private static String dumpPrefix = "kdtree";

    public KDTree() {
        this(0);
    }

    public KDTree(int maxPrims) {
        this.maxPrims = maxPrims;
    }

    private static class BuildStats {
        private int numNodes;
        private int numLeaves;
        private int sumObjects;
        private int minObjects;
        private int maxObjects;
        private int sumDepth;
        private int minDepth;
        private int maxDepth;
        private int numLeaves0;
        private int numLeaves1;
        private int numLeaves2;
        private int numLeaves3;
        private int numLeaves4;
        private int numLeaves4p;

        BuildStats() {
            numNodes = numLeaves = 0;
            sumObjects = 0;
            minObjects = Integer.MAX_VALUE;
            maxObjects = Integer.MIN_VALUE;
            sumDepth = 0;
            minDepth = Integer.MAX_VALUE;
            maxDepth = Integer.MIN_VALUE;
            numLeaves0 = 0;
            numLeaves1 = 0;
            numLeaves2 = 0;
            numLeaves3 = 0;
            numLeaves4 = 0;
            numLeaves4p = 0;
        }

        void updateInner() {
            numNodes++;
        }

        void updateLeaf(int depth, int n) {
            numLeaves++;
            minDepth = Math.min(depth, minDepth);
            maxDepth = Math.max(depth, maxDepth);
            sumDepth += depth;
            minObjects = Math.min(n, minObjects);
            maxObjects = Math.max(n, maxObjects);
            sumObjects += n;
            switch (n) {
                case 0:
                    numLeaves0++;
                    break;
                case 1:
                    numLeaves1++;
                    break;
                case 2:
                    numLeaves2++;
                    break;
                case 3:
                    numLeaves3++;
                    break;
                case 4:
                    numLeaves4++;
                    break;
                default:
                    numLeaves4p++;
                    break;
            }
        }

        void printStats() {
            UI.printDetailed(Module.ACCEL, "KDTree stats:");
            UI.printDetailed(Module.ACCEL, "  * Nodes:          %d", numNodes);
            UI.printDetailed(Module.ACCEL, "  * Leaves:         %d", numLeaves);
            UI.printDetailed(Module.ACCEL, "  * Objects: min    %d", minObjects);
            UI.printDetailed(Module.ACCEL, "             avg    %.2f", (float) sumObjects / numLeaves);
            UI.printDetailed(Module.ACCEL, "           avg(n>0) %.2f", (float) sumObjects / (numLeaves - numLeaves0));
            UI.printDetailed(Module.ACCEL, "             max    %d", maxObjects);
            UI.printDetailed(Module.ACCEL, "  * Depth:   min    %d", minDepth);
            UI.printDetailed(Module.ACCEL, "             avg    %.2f", (float) sumDepth / numLeaves);
            UI.printDetailed(Module.ACCEL, "             max    %d", maxDepth);
            UI.printDetailed(Module.ACCEL, "  * Leaves w/: N=0  %3d%%", 100 * numLeaves0 / numLeaves);
            UI.printDetailed(Module.ACCEL, "               N=1  %3d%%", 100 * numLeaves1 / numLeaves);
            UI.printDetailed(Module.ACCEL, "               N=2  %3d%%", 100 * numLeaves2 / numLeaves);
            UI.printDetailed(Module.ACCEL, "               N=3  %3d%%", 100 * numLeaves3 / numLeaves);
            UI.printDetailed(Module.ACCEL, "               N=4  %3d%%", 100 * numLeaves4 / numLeaves);
            UI.printDetailed(Module.ACCEL, "               N>4  %3d%%", 100 * numLeaves4p / numLeaves);
        }
    }

    public static void setDumpMode(boolean dump, String prefix) {
        KDTree.dump = dump;
        KDTree.dumpPrefix = prefix;
    }

    public void build(PrimitiveList primitives) {
        UI.printDetailed(Module.ACCEL, "KDTree settings");
        UI.printDetailed(Module.ACCEL, "  * Max Leaf Size:  %d", maxPrims);
        UI.printDetailed(Module.ACCEL, "  * Max Depth:      %d", MAX_DEPTH);
        UI.printDetailed(Module.ACCEL, "  * Traversal cost: %.2f", TRAVERSAL_COST);
        UI.printDetailed(Module.ACCEL, "  * Intersect cost: %.2f", INTERSECT_COST);
        UI.printDetailed(Module.ACCEL, "  * Empty bonus:    %.2f", EMPTY_BONUS);
        UI.printDetailed(Module.ACCEL, "  * Dump leaves:    %s", dump ? "enabled" : "disabled");
        Timer total = new Timer();
        total.start();
        this.primitiveList = primitives;
        // get the object space bounds
        bounds = primitives.getWorldBounds(null);
        int nPrim = primitiveList.getNumPrimitives(), nSplits = 0;
        BuildTask task = new BuildTask(nPrim);
        Timer prepare = new Timer();
        prepare.start();
        for (int i = 0; i < nPrim; i++) {
            for (int axis = 0; axis < 3; axis++) {
                float ls = primitiveList.getPrimitiveBound(i, 2 * axis + 0);
                float rs = primitiveList.getPrimitiveBound(i, 2 * axis + 1);
                if (ls == rs) {
                    // flat in this dimension
                    task.splits[nSplits] = pack(ls, PLANAR, axis, i);
                    nSplits++;
                } else {
                    task.splits[nSplits + 0] = pack(ls, OPENED, axis, i);
                    task.splits[nSplits + 1] = pack(rs, CLOSED, axis, i);
                    nSplits += 2;
                }
            }
        }
        task.n = nSplits;
        prepare.end();
        Timer t = new Timer();
        IntArray tempTree = new IntArray();
        IntArray tempList = new IntArray();
        tempTree.add(0);
        tempTree.add(1);
        t.start();
        // sort it
        Timer sorting = new Timer();
        sorting.start();
        radix12(task.splits, task.n);
        sorting.end();
        // build the actual tree
        BuildStats stats = new BuildStats();
        buildTree(bounds.getMinimum().x, bounds.getMaximum().x, bounds.getMinimum().y, bounds.getMaximum().y, bounds.getMinimum().z, bounds.getMaximum().z, task, 1, tempTree, 0, tempList, stats);
        t.end();
        // write out final arrays
        // free some memory
        task = null;
        tree = tempTree.trim();
        tempTree = null;
        this.primitives = tempList.trim();
        tempList = null;
        total.end();
        // display some extra info
        stats.printStats();
        UI.printDetailed(Module.ACCEL, "  * Node memory:    %s", Memory.sizeof(tree));
        UI.printDetailed(Module.ACCEL, "  * Object memory:  %s", Memory.sizeof(this.primitives));
        UI.printDetailed(Module.ACCEL, "  * Prepare time:   %s", prepare);
        UI.printDetailed(Module.ACCEL, "  * Sorting time:   %s", sorting);
        UI.printDetailed(Module.ACCEL, "  * Tree creation:  %s", t);
        UI.printDetailed(Module.ACCEL, "  * Build time:     %s", total);
        if (dump) {
            try {
                UI.printInfo(Module.ACCEL, "Dumping mtls to %s.mtl ...", dumpPrefix);
                FileWriter mtlFile = new FileWriter(dumpPrefix + ".mtl");
                int maxN = stats.maxObjects;
                for (int n = 0; n <= maxN; n++) {
                    float blend = (float) n / (float) maxN;
                    Color nc;
                    if (blend < 0.25)
                        nc = Color.blend(Color.BLUE, Color.GREEN, blend / 0.25f);
                    else if (blend < 0.5)
                        nc = Color.blend(Color.GREEN, Color.YELLOW, (blend - 0.25f) / 0.25f);
                    else if (blend < 0.75)
                        nc = Color.blend(Color.YELLOW, Color.RED, (blend - 0.50f) / 0.25f);
                    else
                        nc = Color.MAGENTA;
                    mtlFile.write(String.format("newmtl mtl%d\n", n));
                    float[] rgb = nc.getRGB();
                    mtlFile.write("Ka 0.1 0.1 0.1\n");
                    mtlFile.write(String.format("Kd %.12g %.12g %.12g\n", rgb[0], rgb[1], rgb[2]));
                    mtlFile.write("illum 1\n\n");
                }
                FileWriter objFile = new FileWriter(dumpPrefix + ".obj");
                UI.printInfo(Module.ACCEL, "Dumping tree to %s.obj ...", dumpPrefix);
                dumpObj(0, 0, maxN, new BoundingBox(bounds), objFile, mtlFile);
                objFile.close();
                mtlFile.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private int dumpObj(int offset, int vertOffset, int maxN, BoundingBox bounds, FileWriter file, FileWriter mtlFile) throws IOException {
        if (offset == 0)
            file.write(String.format("mtllib %s.mtl\n", dumpPrefix));
        int nextOffset = tree[offset];
        if ((nextOffset & (3 << 30)) == (3 << 30)) {
            // leaf
            int n = tree[offset + 1];
            if (n > 0) {
                // output the current voxel to the file
                Point3 min = bounds.getMinimum();
                Point3 max = bounds.getMaximum();
                file.write(String.format("o node%d\n", offset));
                file.write(String.format("v %g %g %g\n", max.x, max.y, min.z));
                file.write(String.format("v %g %g %g\n", max.x, min.y, min.z));
                file.write(String.format("v %g %g %g\n", min.x, min.y, min.z));
                file.write(String.format("v %g %g %g\n", min.x, max.y, min.z));
                file.write(String.format("v %g %g %g\n", max.x, max.y, max.z));
                file.write(String.format("v %g %g %g\n", max.x, min.y, max.z));
                file.write(String.format("v %g %g %g\n", min.x, min.y, max.z));
                file.write(String.format("v %g %g %g\n", min.x, max.y, max.z));
                int v0 = vertOffset;
                file.write(String.format("usemtl mtl%d\n", n));
                file.write("s off\n");
                file.write(String.format("f %d %d %d %d\n", v0 + 1, v0 + 2, v0 + 3, v0 + 4));
                file.write(String.format("f %d %d %d %d\n", v0 + 5, v0 + 8, v0 + 7, v0 + 6));
                file.write(String.format("f %d %d %d %d\n", v0 + 1, v0 + 5, v0 + 6, v0 + 2));
                file.write(String.format("f %d %d %d %d\n", v0 + 2, v0 + 6, v0 + 7, v0 + 3));
                file.write(String.format("f %d %d %d %d\n", v0 + 3, v0 + 7, v0 + 8, v0 + 4));
                file.write(String.format("f %d %d %d %d\n", v0 + 5, v0 + 1, v0 + 4, v0 + 8));
                vertOffset += 8;
            }
            return vertOffset;
        } else {
            // node, recurse
            int axis = nextOffset & (3 << 30), v0;
            float split = Float.intBitsToFloat(tree[offset + 1]), min, max;