Coverage for ase / io / cube.py: 89.32%
103 statements
« prev ^ index » next coverage.py v7.13.3, created at 2026-02-04 10:20 +0000
1"""
2IO support for the Gaussian cube format.
4See the format specifications on: https://paulbourke.net/dataformats/cube/
5"""
7import time
9import numpy as np
11from ase.atoms import Atoms
12from ase.io import read
13from ase.units import Bohr
15ATOMS = 'atoms'
16CASTEP = 'castep'
17DATA = 'data'
20def write_cube(file_obj, atoms, data=None, origin=None, comment=None):
21 """Function to write a cube file.
23 file_obj: str or file object
24 File to which output is written.
25 atoms: Atoms
26 The Atoms object specifying the atomic configuration.
27 data : 3-dim numpy array, optional (default = None)
28 Array containing volumetric data as e.g. electronic density
29 origin : 3-tuple
30 Origin of the volumetric data (units: Angstrom)
31 comment : str, optional (default = None)
32 Comment for the first line of the cube file.
33 """
35 if data is None:
36 data = np.ones((2, 2, 2))
37 data = np.asarray(data)
39 if data.dtype == complex:
40 data = np.abs(data)
42 if comment is None:
43 comment = 'Cube file from ASE, written on ' + time.strftime('%c')
44 else:
45 comment = comment.strip()
46 file_obj.write(comment)
48 file_obj.write('\nOUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n')
50 if origin is None:
51 origin = np.zeros(3)
52 else:
53 origin = np.asarray(origin) / Bohr
55 file_obj.write('{:5}{:12.6f}{:12.6f}{:12.6f}\n'.format(len(atoms), *origin))
57 for i in range(3):
58 n = data.shape[i]
59 d = atoms.cell[i] / n / Bohr
60 file_obj.write('{:5}{:12.6f}{:12.6f}{:12.6f}\n'.format(n, *d))
62 positions = atoms.positions / Bohr
63 numbers = atoms.numbers
64 for Z, (x, y, z) in zip(numbers, positions):
65 file_obj.write(
66 '{:5}{:12.6f}{:12.6f}{:12.6f}{:12.6f}\n'.format(Z, 0.0, x, y, z)
67 )
69 data.tofile(file_obj, sep='\n', format='%e')
72def read_cube(file_obj, read_data=True, program=None, verbose=False):
73 """Read atoms and data from CUBE file.
75 file_obj : str or file
76 Location to the cube file.
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 = file_obj.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 num_atoms = int(line[0])
119 # num_atoms can be negative.
120 # Negative num_atoms indicates we have extra data to parse after
121 # the coordinate information.
122 has_labels = num_atoms < 0
123 num_atoms = abs(num_atoms)
125 # There is an optional last field on this line which indicates
126 # the number of values at each point. It is typically 1 (the default)
127 # in which case it can be omitted, but it may also be > 1,
128 # for example if there are multiple orbitals stored in the same cube.
129 num_val = int(line[4]) if len(line) == 5 else 1
131 # Origin around which the volumetric data is centered
132 # (at least in FHI aims):
133 origin = np.array([float(x) * Bohr for x in line[1:4:]])
135 cell = np.empty((3, 3))
136 shape = []
137 spacing = np.empty((3, 3))
139 # The upcoming three lines contain the cell information
140 for i in range(3):
141 n, x, y, z = (float(s) for s in readline().split())
142 shape.append(int(n))
144 # different PBC treatment in castep, basically the last voxel row is
145 # identical to the first one
146 if program == CASTEP:
147 n -= 1
148 cell[i] = n * Bohr * np.array([x, y, z])
149 spacing[i] = np.array([x, y, z]) * Bohr
150 pbc = [(v != 0).any() for v in cell]
152 numbers = np.empty(num_atoms, int)
153 positions = np.empty((num_atoms, 3))
154 for i in range(num_atoms):
155 line = readline().split()
156 numbers[i] = int(line[0])
157 positions[i] = [float(s) for s in line[2:]]
159 positions *= Bohr
161 atoms = Atoms(numbers=numbers, positions=positions, cell=cell, pbc=pbc)
163 # CASTEP will always have PBC, although the cube format does not
164 # contain this kind of information
165 if program == CASTEP:
166 atoms.pbc = True
168 dct = {ATOMS: atoms}
169 labels = []
171 # If we originally had a negative num_atoms, parse the extra fields now.
172 # The first field of the first line tells us how many other fields there
173 # are to parse, but we have to guess how many rows this information is
174 # split over.
175 if has_labels:
176 # Can't think of a more elegant way of doing this...
177 fields = readline().split()
178 nfields = int(fields[0])
179 labels.extend(fields[1:])
181 while len(labels) < nfields:
182 fields = readline().split()
183 labels.extend(fields)
185 labels = [int(x) for x in labels]
187 if read_data:
188 # Cube files can contain more than one density,
189 # so we need to be a little bit careful about where one ends
190 # and the next begins.
191 raw_volume = [float(s) for s in file_obj.read().split()]
192 # Split each value at each point into a separate list.
193 raw_volumes = [
194 np.array(raw_volume[offset::num_val]) for offset in range(num_val)
195 ]
197 datas = []
199 # Adjust each volume in turn.
200 for data in raw_volumes:
201 data = data.reshape(shape)
202 if axes != [0, 1, 2]:
203 data = data.transpose(axes).copy()
205 if program == CASTEP:
206 # Due to the PBC applied in castep2cube, the last entry
207 # along each dimension equals the very first one.
208 data = data[:-1, :-1, :-1]
210 datas.append(data)
212 datas = np.array(datas)
214 dct[DATA] = datas[0]
215 dct['origin'] = origin
216 dct['spacing'] = spacing
217 dct['labels'] = labels
218 dct['datas'] = datas
220 return dct
223def read_cube_data(filename):
224 """Wrapper function to read not only the atoms information from a cube file
225 but also the contained volumetric data.
226 """
227 dct = read(filename, format='cube', read_data=True, full_output=True)
228 return dct[DATA], dct[ATOMS]