[Closed] [maxscript] KDTree. Can we make it faster?

I have looked and tried your code … First of all, I want to say that the code is very pro and clean. To be honest I can’t improve this with just mxs coding tricks.

lookup time is not bad… but the build time…
yes. Doing almost the same on c++ is 100 times faster. The killing issues are:

(this is my guess only yet)

#1 create mxs array
#2 create mxs struct

Replacing the arrays with value type makes it even slower because of the constant random memory access I guess.
So it looks like we’re out of options to improve it except switching to c# tree implementation.

struct KDTree
(
	root,
	points,
	tree_nodes = #(), -- point4 values [ left_node_index, right_node_index, point_index, axis ]
	
	fn AppendKDTreeNode tree_node =
	(
		potential_parent = root	
		tree_node_pt = points[ tree_node[3] ]
		active = true
		
		while active do
		(
			local axis = if potential_parent[4] == 0 then 1 else 2

			if tree_node_pt[axis] <= points[ potential_parent[3] ][axis] then
			(			
				if potential_parent[1] != 0 then
				(
					potential_parent = tree_nodes[ potential_parent[1] ]
				)
				else
				(
					tree_node[4] = if axis > 1 then 0 else 1
					tree_nodes[ tree_node[3] ] = tree_node
					potential_parent[1] = tree_node[3]						
					active = false
				)
			)
			else
			(
				if potential_parent[2] != 0 then
				(
					potential_parent = tree_nodes[ potential_parent[2] ]
				)
				else
				(
					tree_node[4] = if axis > 0 then 0 else 1
					tree_nodes[ tree_node[3] ] = tree_node
					potential_parent[2] = tree_node[3]						
					active = false
				)			
			)
		)	
	),
...

What if you change the array[4] to directly hold a value of 1 or 2 instead of a boolean?

#( undefined, undefined, points[i], true )

And then access it directly too, without storing a variable:

if tree_node_pt[axis] <= potential_parent[3][axis] then...

to:

if tree_node_pt[potential_parent[4]] <= potential_parent[3][potential_parent[4]] then...

And:

tree_node[4] = axis == 2

to:

tree_node[4] = if potential_parent[4] == 2 then 1 else 2

append tree node function might be cleaner (it’s not faster):

	fn appendTreeNode root node =
	(
		active = true
		while active do
		(
			if node[3][root[4]] <= root[3][root[4]] then 
			(
				if root[1] != undefined then root = root[1] else
				(
					root[1] = node
					active = false
				)
			)
			else
			(
				if root[2] != undefined then root = root[2] else
				(
					root[2] = node
					active = false
				)
			)
		)
		node[4] = if root[4] == 2 then 1 else 2
	),

Sorry to dredge up the thread again, but am I correct in that this would only split on the X and Y axes, and is therefore a K-D tree of only 2 dimensions?

I tried adding the third dimension by adding this method to change the split axis, rather than toggling back and forth on X and Y.

fn nextDimension dim = return 1 + (mod dim 3) as integer

I changed it in the appropriate places:

tree_node[4] = nextDimension potential_parent[4]

Also, in GetNearestNeighbor I send the axis directly:

fn GetNearestNeighbor nodeRoot target_pt axis = ...

and I increment it inside this function also, rather than toggle on state

GetNearestNeighbor nextBranch target_pt (nextDimension axis)

and, of course, use 1 for the initial call instead of a boolean:

fn FindNearestNeighbor pt = ( this.GetNearestNeighbor this.root pt 1 )	

The end of result of all this is… not much noticeable performance improvement. ^^; I did see some gain on initialization (I actually expected the effect to happen on search, not initialization) but not enough to be conclusive. To test in three dimensions instead of two, I replaced the text object with a series of nested GeoSpheres attached into a single Editable Mesh. Here is the result:

– original
KDTree build. Time: 0.081sec. Mem: 548240L Points count:7210
OK
Lookup: Time: 0.025sec. Mem: 480L
Lookup: Time: 0.024sec. Mem: 128L
Lookup: Time: 0.009sec. Mem: 128L
Lookup: Time: 0.024sec. Mem: 128L
Lookup: Time: 0.007sec. Mem: 128L
Lookup: Time: 0.024sec. Mem: 128L
Lookup: Time: 0.006sec. Mem: 128L
Lookup: Time: 0.024sec. Mem: 128L
Lookup: Time: 0.006sec. Mem: 128L
Lookup: Time: 0.024sec. Mem: 128L
Lookup: Time: 0.006sec. Mem: 128L
Lookup: Time: 0.023sec. Mem: 128L
Lookup: Time: 0.007sec. Mem: 128L
Lookup: Time: 0.023sec. Mem: 128L
Lookup: Time: 0.006sec. Mem: 128L
Lookup: Time: 0.024sec. Mem: 128L

– 3d
KDTree build. Time: 0.063sec. Mem: 548440L Points count:7210
OK
Lookup: Time: 0.021sec. Mem: 356L
Lookup: Time: 0.026sec. Mem: 128L
Lookup: Time: 0.023sec. Mem: 128L
Lookup: Time: 0.027sec. Mem: 128L
Lookup: Time: 0.013sec. Mem: 128L
Lookup: Time: 0.026sec. Mem: 128L
Lookup: Time: 0.021sec. Mem: 128L
Lookup: Time: 0.026sec. Mem: 128L
Lookup: Time: 0.022sec. Mem: 128L
Lookup: Time: 0.024sec. Mem: 128L
Lookup: Time: 0.02sec. Mem: 128L
Lookup: Time: 0.026sec. Mem: 128L
Lookup: Time: 0.021sec. Mem: 128L
Lookup: Time: 0.025sec. Mem: 128L
Lookup: Time: 0.021sec. Mem: 128L
Lookup: Time: 0.026sec. Mem: 128L
Lookup: Time: 0.02sec. Mem: 128L

I don’t know if this adds much; I’m actually not sure why the 2D version posted above works seemingly as correctly and efficiently… but just putting it out there. Here’s the whole code with my changes:

(
fn nextDimension dim = return 1 + (mod dim 3) as integer

struct KDTree
(
	root,
	
	fn AppendKDTreeNode root tree_node =
	(
		potential_parent = root
		active = true
		
		while active do
		(
			if tree_node[3][potential_parent[4]] <= potential_parent[3][potential_parent[4]] then
			(			
				if potential_parent[1] != undefined then
				(
					potential_parent = potential_parent[1]
				)
				else
				(
					potential_parent[1] = tree_node
					tree_node[4] = nextDimension potential_parent[4]
					active = false
				)
			)
			else
			(
				if potential_parent[2] != undefined then
				(
					potential_parent = potential_parent[2]
				)
				else
				(
					potential_parent[2] = tree_node
					tree_node[4] = nextDimension potential_parent[4]
					active = false
				)			
			)
		)	
	),

	fn CreateFromPoints points =
	(
		local divideByAxis = 1
		local tree = KDTree root:#( undefined, undefined, points[1], divideByAxis )		
		local tree_root = tree.root
		local append_tree_node = tree.AppendKDTreeNode
		
		for i = 2 to Points.Count do append_tree_node tree_root #( undefined, undefined, points[i], 1 )
		tree
	),
	
	fn Closest n0 n1 target_pt =
	(
		if n0 == undefined then n1 else if n1 == undefined then n0 else
		(
			if distance n0[3] target_pt < distance n1[3] target_pt then n0 else n1
		)		
	),

	fn GetNearestNeighbor nodeRoot target_pt axis = if nodeRoot != undefined do
	(	
		local nextBranch, otherBranch, best
		if target_pt[axis] < nodeRoot[3][axis] then
		(
			nextBranch  = nodeRoot[1]
			otherBranch = nodeRoot[2]
		) 
		else
		(
			nextBranch  = nodeRoot[2]
			otherBranch = nodeRoot[1]
		)
					
		temp = GetNearestNeighbor nextBranch target_pt (nextDimension axis)
		best = Closest temp nodeRoot target_pt
				
		dist = target_pt[axis] - nodeRoot[3][axis]

		if (distance target_pt best[3]) >= dist do
		(
			temp = GetNearestNeighbor otherBranch target_pt (nextDimension axis)
			best = Closest temp best target_pt
		)
		
		best		
	),		
	
	fn FindNearestNeighbor pt = ( this.GetNearestNeighbor this.root pt 1 )	
	
)


delete objects
gc()

txt = editable_mesh()
for i = 1 to 5 do (
	local geo = GeoSphere()
	geo.segs = 12
	geo.radius = i * 5
	addModifier geo (turn_to_mesh())
	maxOps.collapseNode geo true
	attach txt geo
)

maxOps.CollapseNode txt true
txt.xray = true
pts = for i=1 to txt.numverts collect getvert txt i


t1=timestamp();hf = heapfree

::kd = KDTree.CreateFromPoints pts	

format "KDTree build. Time: %sec. Mem: %  Points count:%\n" ((timestamp()-t1)/1000 as float) (hf-heapfree) pts.count


-- testing

delete helpers
::h1 = point centermarker:on cross:off wirecolor:red
::h2 = point pos:txt.center centermarker:off cross:true wirecolor:green isSelected:true

deleteAllChangeHandlers id:#xxx
when transform h2 changes id:#xxx val do
(
	t1=timestamp();hf = heapfree

	n = kd.FindNearestNeighbor ::h2.pos
	
	format "Lookup: Time: %sec. Mem: %\n" ((timestamp()-t1)/1000 as float) (hf-heapfree)
	::h1.pos = [ n[3].x, n[3].y, n[3].z ]
)

select h2
)

Here’s an interesting finding after doing some profiling work… I wrote a basic linear search:

	struct PointSearcher
(
	points = #(),
	
	fn linearSearchIndex target = (
		if points.count < 1 then return undefined
		local closestIndex = 1
		for i = 2 to points.count do (
			if (distance target points[i] < distance target points[closestIndex]) then (
				closestIndex = i
			)
		)
		return closestIndex
	),
	
	fn linearSearch target = (
		if points.count < 1 then return undefined
		return (points[linearSearchIndex target])
	)
)

and for some reason, this turns out to be much faster. I tried it against both my adjustments and the last version posted by @Serejah, and adjusted the time stamp print statement to output the average for all runs – here’s what I got, for 12,180 verts:

k-D lookup: Time: 0.0366896sec. avg
Linear lookup: Time: 0.00656522sec. avg

for 115,210 verts:

k-D lookup: Time: 0.298429sec. avg
Linear lookup: Time: 0.05544sec. avg

for 460,810 verts:

Linear lookup: Time: 0.221226sec.
k-D: (max froze while processing callbacks for maybe a second of dragging)

So, I’d call this a pretty baffling result. Maxscript is pretty mercurial but I suppose what we’re running into is the language overhead mattering more than whatever patterns we come up with? Unfortunate.

Yeah, seems like recursion performance penalty is way too big. But sometimes 2d uniform grid in enough for the task

Time: 0.207sec. Mem: 1936L – PointSearcher
pt: [0.13432,0.92957,0] – avg query

100k points Uniform 2D grid build Time: 1.269sec.
Time: 0.001sec. Mem: 1040L – avg query
pt: [0.13432,0.92957,0]

Interesting – what script are those output logs from?

Not sure if it was the latest version, but anyway. Should be enough for tests