Bumped version

doc/conf.py

@@ -56,7 +56,7 @@ master_doc = 'index'
 
 # General information about the project.
 project = u'westpy'
-copyright = u'2018, Marco Govoni'
+copyright = u'2019, Marco Govoni'
 author = u'Marco Govoni'
 
 # The version info for the project you're documenting, acts as replacement for
@@ -64,7 +64,7 @@ author = u'Marco Govoni'
 # built documents.
 #
 # The short X.Y version.
-version = '3.1.1'
+version = '4.0.0'
 # The full version, including alpha/beta/rc tags.
 release = version
 

setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
     long_description = fh.read()
 
 setup(name='westpy',
-      version='3.1.1',
+      version='4.0.0',
       packages=find_packages(),
       description='Python analysis tools for WEST',
       long_description=long_description,

westpy/utils.py

@@ -233,3 +233,139 @@ def gaussian(x, mu, sig):
    """
    import numpy as np
    return 1./(np.sqrt(2.*np.pi)*sig)*np.exp(-np.power((x - mu)/sig, 2.)/2)
+
+
+def _putline(*args):
+    """
+    Generate a line to be written to a cube file where 
+    the first field is an int and the remaining fields are floats.
+    
+    params:
+        *args: first arg is formatted as int and remaining as floats
+    
+    returns: formatted string to be written to file with trailing newline
+    """
+    s = "{0:^ 8d}".format(args[0])
+    s += "".join("{0:< 12.6f}".format(arg) for arg in args[1:])
+    return s + "\n"
+
+def _getline(cube):
+    """
+    Read a line from cube file where first field is an int 
+    and the remaining fields are floats.
+    
+    params:
+        cube: file object of the cube file
+    
+    returns: (int, list<float>)
+    """
+    l = cube.readline().strip().split()
+    return int(l[0]), map(float, l[1:])
+
+def read_cube(fname):
+    """ 
+    Read cube file into numpy array
+    
+    params:
+        fname: filename of cube file
+        
+    returns: (data: np.array, metadata: dict)
+    """
+    import numpy as np
+    meta = {}
+    with open(fname, 'r') as cube:
+        cube.readline(); cube.readline()  # ignore comments
+        natm, meta['org'] = _getline(cube)
+        nx, meta['xvec'] = _getline(cube)
+        ny, meta['yvec'] = _getline(cube)
+        nz, meta['zvec'] = _getline(cube)
+        meta['atoms'] = [_getline(cube) for i in range(natm)]
+        data = np.zeros((nx*ny*nz))
+        idx = 0
+        for line in cube:
+            for val in line.strip().split():
+                data[idx] = float(val)
+                idx += 1
+    data = np.reshape(data, (nx, ny, nz))
+    return data, meta
+
+
+def read_imcube(rfname, ifname = ""):
+    """
+    Convenience function to read in two cube files at once, 
+    where one contains the real part and the other contains the 
+    imag part. If only one filename given, other filename is inferred.
+    
+    params:
+        rfname: filename of cube file of real part
+        ifname: optional, filename of cube file of imag part
+        
+    returns: np.array (real part + j*imag part)
+    """
+    import numpy as np
+    ifname = ifname or rfname.replace('real', 'imag')
+    _debug("reading from files", rfname, "and", ifname)
+    re, im = read_cube(rfname), read_cube(ifname)
+    fin = np.zeros(re[0].shape, dtype='complex128')
+    if re[1] != im[1]:
+        _debug("warning: meta data mismatch, real part metadata retained")
+    fin += re[0]
+    fin += 1j*im[0]
+    return fin, re[1]
+
+
+
+def write_cube(data, meta, fname):
+    """
+    Write volumetric data to cube file along
+    
+    params:
+        data: volumetric data consisting real values
+        meta: dict containing metadata with following keys
+            atoms: list of atoms in the form (mass, [position])
+            org: origin
+            xvec,yvec,zvec: lattice vector basis
+        fname: filename of cubefile (existing files overwritten)
+    
+    returns: None
+    """
+    with open(fname, "w") as cube:
+        # first two lines are comments
+        cube.write(" Cubefile created by cubetools.py\n  source: none\n")
+        natm = len(meta['atoms'])
+        nx, ny, nz = data.shape
+        cube.write(_putline(natm, *meta['org'])) # 3rd line #atoms and origin
+        cube.write(_putline(nx, *meta['xvec']))
+        cube.write(_putline(ny, *meta['yvec']))
+        cube.write(_putline(nz, *meta['zvec']))
+        for atom_mass, atom_pos in meta['atoms']:
+            cube.write(_putline(atom_mass, *atom_pos)) #skip the newline
+        for i in range(nx):
+            for j in range(ny):
+                for k in range(nz):
+                    if (i or j or k) and k%6==0:
+                        cube.write("\n")
+                    cube.write(" {0: .5E}".format(data[i,j,k]))
+
+def write_imcube(data, meta, rfname, ifname=""):
+    """
+    Convenience function to write two cube files from complex valued 
+    volumetric data, one for the real part and one for the imaginary part.
+    Data about atoms, origin and lattice vectors are kept same for both.
+    If only one filename given, other filename is inferred.
+    
+    params: 
+        data: volumetric data consisting complex values
+        meta: dict containing metadata with following keys
+            atoms: list of atoms in the form (mass, [position])
+            org: origin
+            xvec,yvec,zvec: lattice vector basis
+        rfname: filename of cube file containing real part
+        ifname: optional, filename of cube file containing imag part
+        
+    returns: None
+    """
+    ifname = ifname or rfname.replace('real', 'imag')
+    _debug("writing data to files", rfname, "and", ifname)
+    write_cube(data.real, meta, rfname)
+    write_cube(data.imag, meta, ifname)