playcanvas/build/playcanvas.prf/src/scene/gsplat-unified/gsplat-budget-balancer.js

Open-source WebGL/WebGPU 3D engine for the web

import { NUM_BUCKETS } from "./constants.js";
class GSplatBudgetBalancer {
	_buckets = null;
	_initBuckets() {
		if (!this._buckets) {
			this._buckets = new Array(NUM_BUCKETS);
			for (let i = 0; i < NUM_BUCKETS; i++) {
				this._buckets[i] = [];
			}
		}
	}
	balance(octreeInstances, budget) {
		this._initBuckets();
		for (let i = 0; i < NUM_BUCKETS; i++) {
			this._buckets[i].length = 0;
		}
		let totalOptimalSplats = 0;
		for (const [, inst] of octreeInstances) {
			const nodes = inst.octree.nodes;
			const nodeInfos = inst.nodeInfos;
			for (let nodeIndex = 0, len = nodes.length; nodeIndex < len; nodeIndex++) {
				const nodeInfo = nodeInfos[nodeIndex];
				const optimalLod = nodeInfo.optimalLod;
				if (optimalLod < 0) continue;
				const lods = nodes[nodeIndex].lods;
				nodeInfo.lods = lods;
				this._buckets[nodeInfo.budgetBucket].push(nodeInfo);
				totalOptimalSplats += lods[optimalLod].count;
			}
		}
		let currentSplats = totalOptimalSplats;
		if (currentSplats === budget) {
			return;
		}
		const isOverBudget = currentSplats > budget;
		let done = false;
		while (!done && (isOverBudget ? currentSplats > budget : currentSplats < budget)) {
			let modified = false;
			if (isOverBudget) {
				for (let b = NUM_BUCKETS - 1; b >= 0 && !done; b--) {
					const bucket = this._buckets[b];
					for (let i = 0, len = bucket.length; i < len; i++) {
						const nodeInfo = bucket[i];
						if (nodeInfo.optimalLod < nodeInfo.inst.rangeMax) {
							const lods = nodeInfo.lods;
							const optimalLod = nodeInfo.optimalLod;
							currentSplats -= lods[optimalLod].count - lods[optimalLod + 1].count;
							nodeInfo.optimalLod = optimalLod + 1;
							modified = true;
							if (currentSplats <= budget) {
								done = true;
								break;
							}
						}
					}
				}
			} else {
				for (let b = 0; b < NUM_BUCKETS && !done; b++) {
					const bucket = this._buckets[b];
					for (let i = 0, len = bucket.length; i < len; i++) {
						const nodeInfo = bucket[i];
						if (nodeInfo.optimalLod > nodeInfo.inst.rangeMin) {
							const lods = nodeInfo.lods;
							const optimalLod = nodeInfo.optimalLod;
							const splatsAdded = lods[optimalLod - 1].count - lods[optimalLod].count;
							if (currentSplats + splatsAdded <= budget) {
								nodeInfo.optimalLod = optimalLod - 1;
								currentSplats += splatsAdded;
								modified = true;
								if (currentSplats >= budget) {
									done = true;
									break;
								}
							} else {
								done = true;
								break;
							}
						}
					}
				}
			}
			if (!modified) {
				break;
			}
		}
	}
}
export {
	GSplatBudgetBalancer
};