Coverage for /builds/ase/ase/ase/io/cube.py : 88.75%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2IO support for the Gaussian cube format.
4See the format specifications on:
5http://local.wasp.uwa.edu.au/~pbourke/dataformats/cube/
6"""
9import numpy as np
10import time
11from ase.atoms import Atoms
12from ase.io import read
13from ase.units import Bohr
16def write_cube(fileobj, atoms, data=None, origin=None, comment=None):
17 """
18 Function to write a cube file.
20 fileobj: str or file object
21 File to which output is written.
22 atoms: Atoms object
23 Atoms object specifying the atomic configuration.
24 data : 3dim numpy array, optional (default = None)
25 Array containing volumetric data as e.g. electronic density
26 origin : 3-tuple
27 Origin of the volumetric data (units: Angstrom)
28 comment : str, optional (default = None)
29 Comment for the first line of the cube file.
30 """
32 if data is None:
33 data = np.ones((2, 2, 2))
34 data = np.asarray(data)
36 if data.dtype == complex:
37 data = np.abs(data)
39 if comment is None:
40 comment = "Cube file from ASE, written on " + time.strftime("%c")
41 else:
42 comment = comment.strip()
43 fileobj.write(comment)
45 fileobj.write("\nOUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n")
47 if origin is None:
48 origin = np.zeros(3)
49 else:
50 origin = np.asarray(origin) / Bohr
52 fileobj.write(
53 "{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n".format(
54 len(atoms), *origin))
56 for i in range(3):
57 n = data.shape[i]
58 d = atoms.cell[i] / n / Bohr
59 fileobj.write("{0:5}{1:12.6f}{2:12.6f}{3:12.6f}\n".format(n, *d))
61 positions = atoms.positions / Bohr
62 numbers = atoms.numbers
63 for Z, (x, y, z) in zip(numbers, positions):
64 fileobj.write(
65 "{0:5}{1:12.6f}{2:12.6f}{3:12.6f}{4:12.6f}\n".format(
66 Z, 0.0, x, y, z)
67 )
69 data.tofile(fileobj, sep="\n", format="%e")
72def read_cube(fileobj, read_data=True, program=None, verbose=False):
73 """Read atoms and data from CUBE file.
75 fileobj : str or file
76 Location to the cubefile.
77 read_data : boolean
78 If set true, the actual cube file content, i.e. an array
79 containing the electronic density (or something else )on a grid
80 and the dimensions of the corresponding voxels are read.
81 program: str
82 Use program='castep' to follow the PBC convention that first and
83 last voxel along a direction are mirror images, thus the last
84 voxel is to be removed. If program=None, the routine will try
85 to catch castep files from the comment lines.
86 verbose : bool
87 Print some more information to stdout.
89 Returns a dict with the following keys:
91 * 'atoms': Atoms object
92 * 'data' : (Nx, Ny, Nz) ndarray
93 * 'origin': (3,) ndarray, specifying the cube_data origin.
94 * 'spacing': (3, 3) ndarray, representing voxel size
95 """
97 readline = fileobj.readline
98 line = readline() # the first comment line
99 line = readline() # the second comment line
101 # The second comment line *CAN* contain information on the axes
102 # But this is by far not the case for all programs
103 axes = []
104 if "OUTER LOOP" in line.upper():
105 axes = ["XYZ".index(s[0]) for s in line.upper().split()[2::3]]
106 if not axes:
107 axes = [0, 1, 2]
109 # castep2cube files have a specific comment in the second line ...
110 if "castep2cube" in line:
111 program = "castep"
112 if verbose:
113 print("read_cube identified program: castep")
115 # Third line contains actual system information:
116 line = readline().split()
117 natoms = int(line[0])
119 # Origin around which the volumetric data is centered
120 # (at least in FHI aims):
121 origin = np.array([float(x) * Bohr for x in line[1::]])
123 cell = np.empty((3, 3))
124 shape = []
125 spacing = np.empty((3, 3))
127 # the upcoming three lines contain the cell information
128 for i in range(3):
129 n, x, y, z = [float(s) for s in readline().split()]
130 shape.append(int(n))
132 # different PBC treatment in castep, basically the last voxel row is
133 # identical to the first one
134 if program == "castep":
135 n -= 1
136 cell[i] = n * Bohr * np.array([x, y, z])
137 spacing[i] = np.array([x, y, z]) * Bohr
138 pbc = [(v != 0).any() for v in cell]
140 numbers = np.empty(natoms, int)
141 positions = np.empty((natoms, 3))
142 for i in range(natoms):
143 line = readline().split()
144 numbers[i] = int(line[0])
145 positions[i] = [float(s) for s in line[2:]]
147 positions *= Bohr
149 atoms = Atoms(numbers=numbers, positions=positions, cell=cell, pbc=pbc)
151 # CASTEP will always have PBC, although the cube format does not
152 # contain this kind of information
153 if program == "castep":
154 atoms.pbc = True
156 dct = {"atoms": atoms}
158 if read_data:
159 data = np.array([float(s)
160 for s in fileobj.read().split()]).reshape(shape)
161 if axes != [0, 1, 2]:
162 data = data.transpose(axes).copy()
164 if program == "castep":
165 # Due to the PBC applied in castep2cube, the last entry along each
166 # dimension equals the very first one.
167 data = data[:-1, :-1, :-1]
169 dct["data"] = data
170 dct["origin"] = origin
171 dct["spacing"] = spacing
173 return dct
176def read_cube_data(filename):
177 """Wrapper function to read not only the atoms information from a cube file
178 but also the contained volumetric data.
179 """
180 dct = read(filename, format="cube", read_data=True, full_output=True)
181 return dct["data"], dct["atoms"]