[maxscript] KDTree. Can we make it faster?


#1

Here’s my attempt to build a non-recursive kd-tree in plain maxscript.
It would be great if someone could show how to improve this code
Results:

KDTree build. Time: 0.322sec. Mem: 660280L Points count:9167
Lookup time up to 0.08 sec

(
	fn AppendKDTreeNode root tree_node =
	(
		potential_parent = root	
		tree_node_pt = tree_node[3]
		active = true
		
		while active do
		(
			local axis = if potential_parent[4] then 1 else 2

			if tree_node_pt[axis] <= potential_parent[3][axis] then
			(			
				if potential_parent[1] != undefined then
				(
					potential_parent = potential_parent[1]
				)
				else
				(
					potential_parent[1] = tree_node
					tree_node[4] = axis == 2
					active = false
				)
			)
			else
			(
				if potential_parent[2] != undefined then
				(
					potential_parent = potential_parent[2]
				)
				else
				(
					potential_parent[2] = tree_node
					tree_node[4] = axis == 2
					active = false
				)			
			)
		)	
	)

	struct KDTree
	(
		root,
		
		fn CreateFromPoints points =
		(
			local divide_by_x_axis = true
			local tree = KDTree root:#( undefined, undefined, points[1], divide_by_x_axis )		
			local tree_root = tree.root		
			for i = 2 to Points.Count do AppendKDTreeNode tree_root #( undefined, undefined, points[i], true )
			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 state = if nodeRoot != undefined do
		(	
			local nextBranch, otherBranch, best

			local axis = if state then 1 else 2
			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 (not state)
			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 (not state)
				best = Closest temp best target_pt
			)
			
			best		
		),		
		
		fn FindNearestNeighbor pt = ( this.GetNearestNeighbor this.root pt true )	
		
	)

	
	delete objects
	gc()

txt = Text text:"MAXS\r\ncript"
	addModifier txt (subdivide())
	txt.modifiers[1].size = (distance txt.max txt.min) / 250.0

	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, 0 ]
	)

)

#2

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


#3

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 :confused: 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
				)			
			)
		)	
	),
...

#4

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

#5

thanks!
it helps indeed, I also removed tree_node_pt since it doesn’t affect the performance at all

KDTree build. Time: 0.226sec. Mem: 663336L Points count:9167

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] = if potential_parent[4] == 2 then 1 else 2
					active = false
				)
			)
			else
			(
				if potential_parent[2] != undefined then
				(
					potential_parent = potential_parent[2]
				)
				else
				(
					potential_parent[2] = tree_node
					tree_node[4] = if potential_parent[4] == 2 then 1 else 2
					active = false
				)			
			)
		)	
	),

	fn CreateFromPoints points =
	(
		local divide_by_x_axis = 1
		local tree = KDTree root:#( undefined, undefined, points[1], divide_by_x_axis )		
		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 state = if nodeRoot != undefined do
	(	
		local nextBranch, otherBranch, best

		local axis = if state then 1 else 2
		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 (not state)
		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 (not state)
			best = Closest temp best target_pt
		)
		
		best		
	),		
	
	fn FindNearestNeighbor pt = ( this.GetNearestNeighbor this.root pt true )	
	
)

#6

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
	),