NBody in LuaJIT

If you're tired of reading about me doing NBody simulations, then you might want to skip this post. The interesting thing about this one is not the NBody part, but the LuaJIT part. I've been toying with the idea of learning Lua for awhile now. It doesn't have a lot of the convenient functionalities of Python, but the Just in Time (JIT) compiler promises to be really darn fast. As an added bonus, Lua is used in Torch, a neural net framework that I've played around with in the past. This time however, I wanted to build something in pure Lua, with no help from an external library. I thought it would be instructive to see if I could code up an array object that sort of simulates the behavior of the Numpy array. With this object in hand, it would (presumably) be a piece of cake to code up an NBody simulation. Additionally, I thought it would be fun to do a little speed comparison between my Lua implementation and a Python (numpy) implementation.




First, the array class. Lua does not have classes in the same way that Python has classes. Python has a reserved keyword class that is used to declare a class. In Lua, classes are just tables (the list/dictionary type) with functions as fields. Now, there are many ways to code up a class in Lua, but here's how I did it. Looking at this, you may think that it's a little ridiculous that I define two separate tables, one with the arithmetic functions, and another with everything else. I agree, but I couldn't get the magic of being able to add two objects to work otherwise! If you have any suggestions about how to better implement this, let me know.

	function Array2D(...)
	local args = {...}
	if #args == 0 then
	args = {{{}}}
	end
	local array2d = {
	__index = self,
	vals = args[1],
	shape = {#args[1], #args[1][1]},
	zeros = function(shape)
	local zeroArray = {}
	for i=1,shape[1] do
	zeroArray[i] = {}
	for j=1,shape[2] do
	zeroArray[i][j] = 0.0
	end
	end
	return Array2D(zeroArray)
	end,
	processArgs = function(arr1, arr2)
	local getVal1
	local getVal2
	local shape
	if type(arr1) == 'number' then
	function getVal1(i,j) return arr1 end
	function getVal2(i,j) return arr2.vals[i][j] end
	shape = arr2.shape
	elseif type(arr2) == 'number' then
	function getVal1(i,j) return arr1.vals[i][j] end
	function getVal2(i,j) return arr2 end
	shape = arr1.shape
	elseif type(arr1) == 'table' and type(arr2) == 'table' then
	function getVal1(i,j) return arr1.vals[i][j] end
	function getVal2(i,j) return arr2.vals[i][j] end
	shape = arr1.shape
	end
	return {getVal1, getVal2, shape}
	end,
	printArr = function(self)
	val_str = ""
	for i=1,self.shape[1] do
	for j=1,self.shape[2] do
	val_str = val_str .. self.vals[i][j] .. " "
	end
	if i ~= self.shape[1] then
	val_str = val_str .. '\n'
	end
	end
	print(val_str)
	end,
	}
	local arithmetic = {
	__add = function(arr1, arr2)
	local returnVal = array2d.processArgs(arr1, arr2)
	local getVal1 = returnVal[1]
	local getVal2 = returnVal[2]
	local shape = returnVal[3]
	local new = {}
	for i=1,shape[1] do
	new[i] = {}
	for j=1,shape[2] do
	new[i][j] = getVal1(i,j) + getVal2(i,j)
	end
	end
	return Array2D(new)
	end,
	__sub = function(arr1, arr2)
	local returnVal = array2d.processArgs(arr1, arr2)
	local getVal1 = returnVal[1]
	local getVal2 = returnVal[2]
	local shape = returnVal[3]
	local new = {}
	for i=1,shape[1] do
	new[i] = {}
	for j=1,shape[2] do
	new[i][j] = getVal1(i,j) - getVal2(i,j)
	end
	end
	return new
	end,
	__mul = function(arr1, arr2)
	local new = {}
	local returnVal = array2d.processArgs(arr1, arr2)
	local getVal1 = returnVal[1]
	local getVal2 = returnVal[2]
	local shape = returnVal[3]
	-- local getVal1, local getVal2, local shape = self:processArgs(arr1, arr2)
	for i=1,shape[1] do
	new[i] = {}
	for j=1,shape[2] do
	new[i][j] = getVal1(i,j) * getVal2(i,j)
	end
	end
	return Array2D(new)
	end,
	__div = function(arr1, arr2)
	local new = {}
	local returnVal = array2d.processArgs(arr1, arr2)
	local getVal1 = returnVal[1]
	local getVal2 = returnVal[2]
	local shape = returnVal[3]
	for i=1,shape[1] do
	new[i] = {}
	for j=1,shape[2] do
	new[i][j] = getVal1(i,j) / getVal2(i,j)
	end
	end
	return new
	end,
	}
	setmetatable(array2d, arithmetic)
	return array2d
	end

view raw array2d.lua hosted with ❤ by GitHub

Now on to the speed comparison. I encourage input on both implementations of this simulation. The Python version is basically a direct translation of the Lua version, so if you have any pointers, please let me know. Using the built in time command, here's the output, running each program for 20000 iterations, using a step size of 0.001. Find the code for both in this repository. The Python code was slower, and I imagine it's because the code doesn't leverage the power of the Numpy array. This is in part due to the nature of the simulation -- we have to manually iterate through every step of the simulation, and the squared part of the equations of motion have to be computed manually. Moving forward, I'm eager to see how the scientific libraries for Lua perform compared to this vanilla implementation.