It comes up sometimes in my research that I need to generate coordinates along a lattice. For example, when generating an initial structure for a molecular dynamics simulation. I thought I share my code for this. The code works by recursively enumerating each dimension in a lattice and calling a given function each time it reaches a point. So, it could be used to execute any function along a lattice. Here's what it looks like for a 2D and 3D example. In the 3D example I didn't choose a cubic number of points on the lattice, so that it sort of stops in the middle.
Here's the code:
Here's the code:
#!/usr/bin/python
import sys, copy
from math import *
def print_usage():
print "Usage: build_lattice.py [dims] [box_1_size] [box_..._size] [box_dims_size] [number]"
if(len(sys.argv) < 2):
print_usage()
exit()
n_dims = int(sys.argv[1])
if(len(sys.argv) < 2 + n_dims + 1):
print_usage()
exit()
box_size = [float(sys.argv[x]) for x in range(2,2+n_dims)]
number = int(sys.argv[2 + n_dims])
volume = 1
for x in box_size:
volume = volume * x
increment = (volume / number) ** (1. / n_dims)
def prepend_emit(array, element):
array_copy = copy.copy(array)
array_copy.insert(0, element)
return array_copy
def enumerate_grid(fxn, dim, sizes, indices):
if(dim > 0):
for i in range(sizes[dim]):
enumerate_grid(fxn,
dim - 1,
sizes,
prepend_emit(indices, i))
else:
for i in range(sizes[dim]):
fxn(prepend_emit(indices, i))
def print_point(point):
print point,
def print_grid(indices):
global number
#check if to make sure we still need to make particles
if(number > 0):
number -= 1
global increment
[print_point(x * increment) for x in indices]
print ""
enumerate_grid(print_grid, n_dims - 1,
[int(ceil(x / increment)) for x in box_size],
[])