Wednesday, April 24, 2013

Generating n-Dimensional Lattice Coordinates

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:
#!/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], 
               [])


No comments:

Post a Comment