Test program for 'minimum image convention' routines for triclinic cells

mictest.py

#!/usr/bin/env python
'''Test program for 'minimum image convention' routines for triclinic cells

   This program is public domain.
   
   Author(s): Toon Verstraelen
   
   The minimum image convention (MIC) is a standard method in molecular
   simulations of periodic systems. The convention states that one can
   approximate the interaction between atom i and all periodic images of atom j
   by considering only the pair (i,j) whose separation is the smallest of all
   possible periodic images of j.
   
   The problem is that most algorithms to find (efficiently) the minimum image
   pair (i,j) only work well for orthorombic cells. They become somewhat
   unreliable for triclinic cells, especially when the cell is very skewed and
   the length of the relative vectors is of the order of the distance between
   two principal crystal planes.

   The reference implementation is a brute force approach that loops over
   all near images and takes the shortest relative vector. It is too slow to be
   useful in serious simulations. Two 'faster' implementations (naive and
   lammps) are also added for comparison. Both do not work for all possible
   relative vectors. This program compares the brute force implementation with
   one of the fast implementations, using randomly generated relative vectors,
   until 100 failures are encountered.
   
   The top of this file contains all control parameters: unit cell & choice of
   the fast mic routine.
   
   If you encounter mistakes or know better MIC routines, please contact me
   through [email protected]. This program was written as part of a
   question posted at scicomp.stackexchange.com:
   
   http://scicomp.stackexchange.com/questions/3107/minimum-image-convention-for-triclinic-unit-cell
'''

import numpy as np


################################################################################
### CONTROL PARAMETERS
################################################################################

# real-space cell vectors
rvecs = np.array([
    [5.0, 0.0, 0.0], # cell vector a
    [1.0, 5.0, 0.0], # cell vector b
    [1.0, 1.0, 5.0], # cell vector c
])

#fast_name = 'naive'
fast_name = 'lammps'


################################################################################
### DERIVED PARAMETERS
################################################################################


# reciprocal cell vectors as rows, without factor 2*pi that is common in
# plane wave codes
gvecs = np.linalg.inv(rvecs).T

if fast_name == 'lammps':
    # compute the rectangular bounding box around the unit cell
    box_low = np.array([np.PINF, np.PINF, np.PINF])
    box_high = np.array([np.NINF, np.NINF, np.NINF])
    for i0 in -0.5, 0.5:
        for i1 in -0.5, 0.5:
            for i2 in -0.5, 0.5:
                corner = np.dot(rvecs.T, [i0, i1, i2])
                mask = box_low > corner
                box_low[mask] = corner[mask]
                mask = box_high < corner
                box_high[mask] = corner[mask]


    # check conditions from lammps domain code, for which lammps_mic should work
    assert rvecs[0,1] == 0.0
    assert rvecs[0,2] == 0.0
    assert rvecs[1,2] == 0.0
    assert abs(rvecs[1,0]/(box_high[0]-box_low[0])) < 0.5
    assert abs(rvecs[2,0]/(box_high[0]-box_low[0])) < 0.5
    assert abs(rvecs[2,1]/(box_high[1]-box_low[1])) < 0.5


################################################################################
### MIC ROUTINES
################################################################################


def brute_mic(r):
    # Basic idea = triple loop over all nearby image, pick shortest relative vector.
    # This is horribly slow, and for very skewed cells, n has to become very large.
    r = naive_mic(r)
    result = None
    n = 2
    for i0 in xrange(-n, n+1):
        for i1 in xrange(-n, n+1):
            for i2 in xrange(-n, n+1):
                rp = r+np.dot(rvecs.T, [i0,i1,i2])
                d = np.linalg.norm(rp)
                if (result is None) or (result[1] > d):
                    result = (rp, d)
    return result[0]


def naive_mic(r):
    # transform to fractional coordinates, apply mic for cubic and transform
    # back. This is what I came up with myself.
    f = np.dot(gvecs, r)
    f -= np.floor(f+0.5)
    return np.dot(rvecs.T, f)


def lammps_mic(r):
    # the one implemented in lammps, unless I'm mistaken.
    rp = r.copy()
    for i in 2, 1, 0:
        if abs(rp[i]) > rvecs[i,i]/2:
            if rp[i] < 0:
                rp += rvecs[i]
            else:
                rp -= rvecs[i]
    return rp    

# select mic routine
if fast_name == 'naive':
    fast_mic = naive_mic
elif fast_name == 'lammps':
    fast_mic = lammps_mic
else:
    raise NotImplementedError


################################################################################
### AUXILIARY STUFF
################################################################################


def print_vec(label, vec):
    d = np.linalg.norm(vec)
    print '%s: (%10.5f, %10.5f, %10.5f)  d = %10.5f' % (
        label.rjust(20), vec[0], vec[1], vec[2], d
    )
    return d


################################################################################
### TEST LOOP
################################################################################


# compare mic routines using random relative vectors
counter = 0
fail_counter = 0
eps = 1e-5 # safe margin for roundoff errors
while fail_counter < 100:
    counter += 1

    f_orig = np.random.uniform(-1, 1, 3) # fractional
    r_orig = np.dot(rvecs.T, f_orig)     # real space
    r_brute = brute_mic(r_orig)
    r_fast = lammps_mic(r_orig)
    if abs(r_brute-r_fast).max() > eps:
        fail_counter += 1
        print 'Failure %i/%i' % (fail_counter, counter)
        results = []
        for label, r in ('original', r_orig), ('brute', r_brute), (fast_name, r_fast):
            d = print_vec(label, r)
            results.append((d, label))
        results.sort()
        print 'Shortest:', [label for d, label in results if d < results[0][0]+eps]
        print