Coverage for /builds/ase/ase/ase/io/cube.py: 89.32%
103 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3"""
4IO support for the Gaussian cube format.
6See the format specifications on:
7http://local.wasp.uwa.edu.au/~pbourke/dataformats/cube/
8"""
10import time
12import numpy as np
14from ase.atoms import Atoms
15from ase.io import read
16from ase.units import Bohr
18ATOMS = 'atoms'
19CASTEP = 'castep'
20DATA = 'data'
23def write_cube(file_obj, atoms, data=None, origin=None, comment=None):
24 """Function to write a cube file.
26 file_obj: str or file object
27 File to which output is written.
28 atoms: Atoms
29 The Atoms object specifying the atomic configuration.
30 data : 3-dim numpy array, optional (default = None)
31 Array containing volumetric data as e.g. electronic density
32 origin : 3-tuple
33 Origin of the volumetric data (units: Angstrom)
34 comment : str, optional (default = None)
35 Comment for the first line of the cube file.
36 """
38 if data is None:
39 data = np.ones((2, 2, 2))
40 data = np.asarray(data)
42 if data.dtype == complex:
43 data = np.abs(data)
45 if comment is None:
46 comment = "Cube file from ASE, written on " + time.strftime("%c")
47 else:
48 comment = comment.strip()
49 file_obj.write(comment)
51 file_obj.write("\nOUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n")
53 if origin is None:
54 origin = np.zeros(3)
55 else:
56 origin = np.asarray(origin) / Bohr
58 file_obj.write(
59 "{:5}{:12.6f}{:12.6f}{:12.6f}\n".format(
60 len(atoms), *origin))
62 for i in range(3):
63 n = data.shape[i]
64 d = atoms.cell[i] / n / Bohr
65 file_obj.write("{:5}{:12.6f}{:12.6f}{:12.6f}\n".format(n, *d))
67 positions = atoms.positions / Bohr
68 numbers = atoms.numbers
69 for Z, (x, y, z) in zip(numbers, positions):
70 file_obj.write(
71 "{:5}{:12.6f}{:12.6f}{:12.6f}{:12.6f}\n".format(
72 Z, 0.0, x, y, z)
73 )
75 data.tofile(file_obj, sep="\n", format="%e")
78def read_cube(file_obj, read_data=True, program=None, verbose=False):
79 """Read atoms and data from CUBE file.
81 file_obj : str or file
82 Location to the cube file.
83 read_data : boolean
84 If set true, the actual cube file content, i.e. an array
85 containing the electronic density (or something else )on a grid
86 and the dimensions of the corresponding voxels are read.
87 program: str
88 Use program='castep' to follow the PBC convention that first and
89 last voxel along a direction are mirror images, thus the last
90 voxel is to be removed. If program=None, the routine will try
91 to catch castep files from the comment lines.
92 verbose : bool
93 Print some more information to stdout.
95 Returns a dict with the following keys:
97 * 'atoms': Atoms object
98 * 'data' : (Nx, Ny, Nz) ndarray
99 * 'origin': (3,) ndarray, specifying the cube_data origin.
100 * 'spacing': (3, 3) ndarray, representing voxel size
101 """
103 readline = file_obj.readline
104 line = readline() # the first comment line
105 line = readline() # the second comment line
107 # The second comment line *CAN* contain information on the axes
108 # But this is by far not the case for all programs
109 axes = []
110 if "OUTER LOOP" in line.upper():
111 axes = ["XYZ".index(s[0]) for s in line.upper().split()[2::3]]
112 if not axes:
113 axes = [0, 1, 2]
115 # castep2cube files have a specific comment in the second line ...
116 if "castep2cube" in line:
117 program = CASTEP
118 if verbose:
119 print("read_cube identified program: castep")
121 # Third line contains actual system information:
122 line = readline().split()
123 num_atoms = int(line[0])
125 # num_atoms can be negative.
126 # Negative num_atoms indicates we have extra data to parse after
127 # the coordinate information.
128 has_labels = num_atoms < 0
129 num_atoms = abs(num_atoms)
131 # There is an optional last field on this line which indicates
132 # the number of values at each point. It is typically 1 (the default)
133 # in which case it can be omitted, but it may also be > 1,
134 # for example if there are multiple orbitals stored in the same cube.
135 num_val = int(line[4]) if len(line) == 5 else 1
137 # Origin around which the volumetric data is centered
138 # (at least in FHI aims):
139 origin = np.array([float(x) * Bohr for x in line[1:4:]])
141 cell = np.empty((3, 3))
142 shape = []
143 spacing = np.empty((3, 3))
145 # The upcoming three lines contain the cell information
146 for i in range(3):
147 n, x, y, z = (float(s) for s in readline().split())
148 shape.append(int(n))
150 # different PBC treatment in castep, basically the last voxel row is
151 # identical to the first one
152 if program == CASTEP:
153 n -= 1
154 cell[i] = n * Bohr * np.array([x, y, z])
155 spacing[i] = np.array([x, y, z]) * Bohr
156 pbc = [(v != 0).any() for v in cell]
158 numbers = np.empty(num_atoms, int)
159 positions = np.empty((num_atoms, 3))
160 for i in range(num_atoms):
161 line = readline().split()
162 numbers[i] = int(line[0])
163 positions[i] = [float(s) for s in line[2:]]
165 positions *= Bohr
167 atoms = Atoms(numbers=numbers, positions=positions, cell=cell, pbc=pbc)
169 # CASTEP will always have PBC, although the cube format does not
170 # contain this kind of information
171 if program == CASTEP:
172 atoms.pbc = True
174 dct = {ATOMS: atoms}
175 labels = []
177 # If we originally had a negative num_atoms, parse the extra fields now.
178 # The first field of the first line tells us how many other fields there
179 # are to parse, but we have to guess how many rows this information is
180 # split over.
181 if has_labels:
182 # Can't think of a more elegant way of doing this...
183 fields = readline().split()
184 nfields = int(fields[0])
185 labels.extend(fields[1:])
187 while len(labels) < nfields:
188 fields = readline().split()
189 labels.extend(fields)
191 labels = [int(x) for x in labels]
193 if read_data:
194 # Cube files can contain more than one density,
195 # so we need to be a little bit careful about where one ends
196 # and the next begins.
197 raw_volume = [float(s) for s in file_obj.read().split()]
198 # Split each value at each point into a separate list.
199 raw_volumes = [np.array(raw_volume[offset::num_val])
200 for offset in range(num_val)]
202 datas = []
204 # Adjust each volume in turn.
205 for data in raw_volumes:
206 data = data.reshape(shape)
207 if axes != [0, 1, 2]:
208 data = data.transpose(axes).copy()
210 if program == CASTEP:
211 # Due to the PBC applied in castep2cube, the last entry
212 # along each dimension equals the very first one.
213 data = data[:-1, :-1, :-1]
215 datas.append(data)
217 datas = np.array(datas)
219 dct[DATA] = datas[0]
220 dct["origin"] = origin
221 dct["spacing"] = spacing
222 dct["labels"] = labels
223 dct["datas"] = datas
225 return dct
228def read_cube_data(filename):
229 """Wrapper function to read not only the atoms information from a cube file
230 but also the contained volumetric data.
231 """
232 dct = read(filename, format="cube", read_data=True, full_output=True)
233 return dct[DATA], dct[ATOMS]