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

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

const NUM_BUCKETS = 64;
class GSplatBudgetBalancer {
		_initBuckets() {
				if (!this._buckets) {
						this._buckets = new Array(NUM_BUCKETS);
						for(let i = 0; i < NUM_BUCKETS; i++){
								this._buckets[i] = [];
						}
				}
		}
		balance(octreeInstances, budget, globalMaxDistance) {
				this._initBuckets();
				for(let i = 0; i < NUM_BUCKETS; i++){
						this._buckets[i].length = 0;
				}
				const bucketScale = NUM_BUCKETS / Math.sqrt(globalMaxDistance);
				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;
								const bucket = Math.sqrt(nodeInfo.worldDistance) * bucketScale >>> 0;
								const bucketIdx = bucket < NUM_BUCKETS ? bucket : NUM_BUCKETS - 1;
								this._buckets[bucketIdx].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;
						}
				}
		}
		constructor(){
				this._buckets = null;
		}
}

export { GSplatBudgetBalancer };