...
 
Commits (3)
......@@ -21,3 +21,5 @@ pip-log.txt
.*.swp
.DS_Store
doc/_build
......@@ -16,3 +16,7 @@
.wy-side-nav-search {
background-color: #B80000;
}
.version {
color:white !important;
}
......@@ -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
......
......@@ -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,
......
......@@ -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)