Coverage for eminus/io/cube.py: 95.16%
62 statements
« prev ^ index » next coverage.py v7.6.4, created at 2024-11-08 12:59 +0000
« prev ^ index » next coverage.py v7.6.4, created at 2024-11-08 12:59 +0000
1# SPDX-FileCopyrightText: 2021 The eminus developers
2# SPDX-License-Identifier: Apache-2.0
3"""CUBE file handling."""
5import textwrap
6import time
8import numpy as np
10from ..data import NUMBER2SYMBOL, SYMBOL2NUMBER
11from ..logger import log
12from ..version import __version__
15def read_cube(filename):
16 """Load atom and cell data from CUBE files.
18 There is no standard for CUBE files. The following format has been used.
19 File format definition: https://h5cube-spec.readthedocs.io/en/latest/cubeformat.html
21 Args:
22 filename: CUBE input file path/name.
24 Returns:
25 Species, positions, charges, cell size, sampling, and field array.
26 """
27 if not filename.endswith((".cub", ".cube")):
28 filename += ".cube"
30 # Atomic units and a cell that starts at (0,0,0) are assumed.
31 with open(filename, encoding="utf-8") as fh:
32 lines = fh.readlines()
34 # The first and second lines can contain comments, print them if available
35 comment = f"{lines[0].strip()}\n{lines[1].strip()}"
36 if comment:
37 log.info(f'CUBE file comment: "{comment}"')
39 # Lines 4 to 6 contain the sampling per axis and the cell basis vectors with length a/s
40 s = np.empty(3, dtype=int)
41 a = np.empty((3, 3))
42 for i, line in enumerate(lines[3:6]):
43 line_split = line.strip().split()
44 s[i] = line_split[0]
45 a[i] = s[i] * np.float64(line_split[1:])
47 atom = []
48 pos = []
49 Z = []
50 # Following lines contain atom positions with the format: atom-id charge x-pos y-pos z-pos
51 _offset = 0
52 for _offset, line in enumerate(lines[6:]):
53 line_split = line.strip().split()
54 # If the first value is not a (positive) integer, we have reached the field data
55 if not line_split[0].isdigit():
56 break
57 atom.append(NUMBER2SYMBOL[int(line_split[0])])
58 Z.append(float(line_split[1]))
59 pos.append(np.float64(line_split[2:5]))
60 pos = np.asarray(pos)
62 # The rest of the data is the field data
63 # Split the strings, flatten the lists of lists, and convert to a float numpy array
64 tmp_list = [l.split() for l in lines[6 + _offset :]]
65 field_list = [item for sublist in tmp_list for item in sublist]
66 field = np.asarray(field_list, dtype=float)
67 return atom, pos, Z, a, s, field
70def write_cube(obj, filename, field, fods=None, elec_symbols=("X", "He")):
71 """Generate CUBE files from given field data.
73 There is no standard for CUBE files. The following format has been used to work with VESTA.
74 File format definition: https://h5cube-spec.readthedocs.io/en/latest/cubeformat.html
76 Args:
77 obj: Atoms or SCF object.
78 filename: CUBE output file path/name.
79 field: Real-space field data.
81 Keyword Args:
82 fods: FOD coordinates to write.
83 elec_symbols: Identifier for up and down FODs.
84 """
85 # Atomic units are assumed, so there is no need for conversion.
86 atoms = obj._atoms
88 if not filename.endswith((".cub", ".cube")):
89 filename += ".cube"
91 if "He" in atoms.atom and atoms.unrestricted:
92 log.warning(
93 'You need to modify "elec_symbols" to write helium with FODs in the spin-'
94 "polarized case."
95 )
97 # Make sure we have real-valued data in the correct order
98 if field is None:
99 log.warning('The provided field is "None".')
100 return
101 field = np.real(field)
103 with open(filename, "w", encoding="utf-8") as fp:
104 # The first two lines have to be a comment
105 # Print information about the file and program, and the file creation time
106 fp.write(f"File generated with eminus {__version__} on {time.ctime()}\n\n")
107 # Number of atoms (int), and origin of the coordinate system (float)
108 # The origin is normally at 0,0,0 but we could move our box, so take the minimum
109 if fods is None:
110 fp.write(f"{atoms.Natoms} ")
111 else:
112 fp.write(f"{atoms.Natoms + sum(len(i) for i in fods)} ")
113 fp.write("0.0 0.0 0.0\n")
114 # Number of points per axis (int), and vector defining the axis (float)
115 fp.write(
116 f"{atoms.s[0]} {atoms.a[0, 0] / atoms.s[0]:.6f} {atoms.a[0, 1] / atoms.s[0]:.6f}"
117 f" {atoms.a[0, 2] / atoms.s[0]:.6f}\n"
118 f"{atoms.s[1]} {atoms.a[1, 0] / atoms.s[1]:.6f} {atoms.a[1, 1] / atoms.s[1]:.6f}"
119 f" {atoms.a[1, 2] / atoms.s[1]:.6f}\n"
120 f"{atoms.s[2]} {atoms.a[2, 0] / atoms.s[2]:.6f} {atoms.a[2, 1] / atoms.s[2]:.6f}"
121 f" {atoms.a[2, 2] / atoms.s[2]:.6f}\n"
122 )
123 # Atomic number (int), atomic charge (float), and atom position (floats) for every atom
124 for ia in range(atoms.Natoms):
125 fp.write(
126 f"{SYMBOL2NUMBER[atoms.atom[ia]]} {atoms.Z[ia]:.3f} "
127 f"{atoms.pos[ia, 0]: .6f} {atoms.pos[ia, 1]: .6f} {atoms.pos[ia, 2]: .6f}\n"
128 )
129 if fods is not None:
130 for s in range(len(fods)):
131 for ie in fods[s]:
132 fp.write(
133 f"{SYMBOL2NUMBER[elec_symbols[s]]} 0.000 "
134 f"{ie[0]: .6f} {ie[1]: .6f} {ie[2]: .6f}\n"
135 )
136 # Field data (float) with scientific formatting
137 # We have s[0]*s[1] chunks values with a length of s[2]
138 for i in range(atoms.s[0] * atoms.s[1]):
139 # Print every round of values, so we can add empty lines between them
140 data_str = "%+1.6e " * atoms.s[2] % tuple(field[i * atoms.s[2] : (i + 1) * atoms.s[2]])
141 # Print a maximum of 6 values per row
142 # Max width for this formatting is 90, since 6*len("+1.00000e-000 ")=90
143 fp.write(f"{textwrap.fill(data_str, width=90)}\n\n")