Source code for es_fitting

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Calculates the Wigner-Seitz projection of cube file data of periodic data.

.. autofunction:: main

Command Line Interface
----------------------

.. program:: es_fitting.py 

.. option:: filename

   The cubefile to read from. Both Bohr and Angstrom units are supported.

.. option:: --periodic

   Whether to treat input file in periodic boundary conditions. In the current implementation, this requires copying all atom positions for the 26 direct neighbouring cubes and therefore carries substantial cost in terms of memory requirements. The runtime scales O(log n) with the number of atoms.

.. option:: --leafsize

   Number of points at which brute-force nearest neighbour search is employed. Increasing this number introduces higher memory requirements, but may speed up processing. Default: 5000.

Implementation
--------------
"""

# system modules
import argparse

# third-party modules
from scipy.spatial import KDTree
import numpy as np

# custom modules
import euston.io as io
import euston.geometry as geom

parser = argparse.ArgumentParser(description='Calculates the Wigner-Seitz projection of cube file data of periodic data.')
parser.add_argument('filename', type=str, help='The cube file.')
parser.add_argument('--periodic', action='store_true', help='Treats cube data periodically.')
parser.add_argument('--leafsize', type=int, help='Number of points at which brute-force nearest neighbour search is employed.', default=5000)

[docs]def main(parser):
    """
    Main routine wrapper.

    :param argparse.ArgumentParser parser: Argument parser
    """
    args = parser.parse_args()

    print 'Reading cubefile...                 ',
    cube = io.CubeFile(args.filename)
    print 'Completed, %d atoms %d voxels.' % (cube.count_atoms(), cube.count_voxels())

    print 'Calculating cell dimensions...      ',
    h_mat = cube.get_h_matrix()
    print 'Completed.'

    print 'Adding image atoms...               ',
    coord = cube.get_coordinates()
    images = np.zeros(((len(coord)*27), 3))
    counter = 0
    for x in range(-1, 2):
        for y in range(-1, 2):
            for z in range(-1, 2):
                if x*y*z == 0:
                    continue
                shift = (h_mat.transpose()*np.array([x, y, z])).sum(axis=0)
                images[counter*cube.count_atoms():(counter+1)*cube.count_atoms(), :] = coord+shift
                counter += 1
    print 'Completed, %d atoms added.' % (cube.count_atoms()*26)

    images[-cube.count_atoms():] = coord

    if args.periodic:
        kdt = KDTree(images, leafsize=args.leafsize)
    else:
        kdt = KDTree(coord, leafsize=args.leafsize)
    vals = np.zeros(cube.count_atoms())
    for x in range(cube.get_xlen()):
        print x
        for y in range(cube.get_ylen()):
            for z in range(cube.get_zlen()):
                d, i = kdt.query(cube.get_voxel_pos(x, y, z, centered=True))
                i = i % cube.count_atoms()
                vals[i] += cube.get_val(x, y, z)

    print 'Results (atom - value)'
    vals *= cube.get_voxel_volume()
    for idx, val in enumerate(vals):
        print idx, val

if __name__ == '__main__':
    main(parser)