Coverage for gpaw/utilities/tools.py: 76%
82 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1import numpy as np
3from ase.units import Hartree, Bohr
6def split_formula(formula):
7 """Count elements in a chemical formula.
9 E.g. split_formula('C2H3Mg') -> ['C', 'C', 'H', 'H', 'H', 'Mg']
10 """
11 res = []
12 for c in formula:
13 if c.isupper():
14 res.append(c)
15 elif c.islower():
16 res[-1] += c
17 else:
18 res.extend([res[-1]] * (eval(c) - 1))
19 return res
22def construct_reciprocal(gd, q_c=None):
23 """Construct the reciprocal lattice from ``GridDescriptor`` instance.
25 The generated reciprocal lattice has lattice vectors corresponding to the
26 real-space lattice defined in the input grid. Note that it is the squared
27 length of the reciprocal lattice vectors that are returned.
29 The ordering of the reciprocal lattice agrees with the one typically used
30 in fft algorithms, i.e. positive k-values followed by negative.
32 Note that the G=(0,0,0) entry is set to one instead of zero. This
33 bit should probably be moved somewhere else ...
35 Parameters
36 ----------
37 q_c: ndarray
38 Offset for the reciprocal lattice vectors (in scaled coordinates of the
39 reciprocal lattice vectors, i.e. array with index ``c``). When
40 specified, the returned array contains the values of (q+G)^2 where G
41 denotes the reciprocal lattice vectors.
43 """
45 assert gd.pbc_c.all(), 'Works only with periodic boundary conditions!'
47 q_c = np.zeros((3, 1), dtype=float) if q_c is None else q_c.reshape((3, 1))
49 # Calculate reciprocal lattice vectors
50 N_c1 = gd.N_c[:, None]
51 i_cq = np.indices(gd.n_c, dtype=float).reshape((3, -1)) # offsets....
52 i_cq += gd.beg_c[:, None]
53 i_cq += N_c1 // 2
54 i_cq %= N_c1
55 i_cq -= N_c1 // 2
57 i_cq += q_c
59 # Convert from scaled to absolute coordinates
60 B_vc = 2.0 * np.pi * gd.icell_cv.T
61 k_vq = np.dot(B_vc, i_cq)
63 k_vq *= k_vq
64 k2_Q = k_vq.sum(axis=0).reshape(gd.n_c)
66 # Avoid future divide-by-zero by setting k2_Q[G=(0,0,0)] = 1.0 if needed
67 if k2_Q[0, 0, 0] < 1e-10:
68 k2_Q[0, 0, 0] = 1.0 # Only make sense iff
69 assert gd.comm.rank == 0 # * on rank 0 (G=(0,0,0) is only there)
70 assert abs(q_c).sum() < 1e-8 # * q_c is (almost) zero
72 assert k2_Q.min() > 0.0 # Now there should be no zero left
74 # Determine N^3
75 #
76 # Why do we need to calculate and return this? The caller already
77 # has access to gd and thus knows how many points there are.
78 N3 = gd.n_c[0] * gd.n_c[1] * gd.n_c[2]
79 return k2_Q, N3
82def coordinates(gd, origin=None, tiny=1e-12):
83 """Constructs and returns matrices containing cartesian coordinates,
84 and the square of the distance from the origin.
86 The origin can be given explicitely (in Bohr units, not Anstroms).
87 Otherwise the origin is placed in the center of the box described
88 by the given grid-descriptor 'gd'.
89 """
91 if origin is None:
92 origin = 0.5 * gd.cell_cv.sum(0)
93 r0_v = np.array(origin)
95 r_vG = gd.get_grid_point_distance_vectors(r0_v)
96 r2_G = np.sum(r_vG**2, axis=0)
97 # Remove singularity at origin and replace with small number
98 r2_G = np.where(r2_G < tiny, tiny, r2_G)
100 # Return r^2 matrix
101 return r_vG, r2_G
104def pick(a_ix, i):
105 """Take integer index of a, or a linear combination of the elements of a"""
106 if isinstance(i, int):
107 return a_ix[i]
108 shape = a_ix.shape
109 a_x = np.dot(i, a_ix[:].reshape(shape[0], -1))
110 return a_x.reshape(shape[1:])
113def dagger(a, copy=True):
114 """Return Hermitian conjugate of input
116 If copy is False, the original array might be overwritten. This is faster,
117 but use with care.
118 """
119 if copy:
120 return np.conj(a.T)
121 else:
122 a = a.T
123 if a.dtype == complex:
124 a.imag *= -1
125 return a
128def project(a, b):
129 """Return the projection of b onto a."""
130 return a * (np.dot(a.conj(), b) / np.linalg.norm(a))
133def normalize(U):
134 """Normalize columns of U."""
135 for col in U.T:
136 col /= np.linalg.norm(col)
139def gram_schmidt(U):
140 """Orthonormalize columns of U according to the Gram-Schmidt procedure."""
141 for i, col in enumerate(U.T):
142 for col2 in U.T[:i]:
143 col -= col2 * np.dot(col2.conj(), col)
144 col /= np.linalg.norm(col)
147def lowdin(U, S=None):
148 """Orthonormalize columns of U according to the Lowdin procedure.
150 If the overlap matrix is know, it can be specified in S.
151 """
152 if S is None:
153 S = np.dot(dagger(U), U)
154 eig, rot = np.linalg.eigh(S)
155 rot = np.dot(rot / np.sqrt(eig), dagger(rot))
156 U[:] = np.dot(U, rot)
159def symmetrize(matrix):
160 """Symmetrize input matrix."""
161 np.add(dagger(matrix), matrix, matrix)
162 np.multiply(.5, matrix, matrix)
163 return matrix
166def tri2full(H_nn, UL='L', map=np.conj):
167 """Fill in values of hermitian or symmetric matrix.
169 Fill values in lower or upper triangle of H_nn based on the opposite
170 triangle, such that the resulting matrix is symmetric/hermitian.
172 UL='U' will copy (conjugated) values from upper triangle into the
173 lower triangle.
175 UL='L' will copy (conjugated) values from lower triangle into the
176 upper triangle.
178 The map parameter can be used to specify a different operation than
179 conjugation, which should work on 1D arrays. Example::
181 def antihermitian(src, dst):
182 np.conj(-src, dst)
184 tri2full(H_nn, map=antihermitian)
186 """
187 N, tmp = H_nn.shape
188 assert N == tmp, 'Matrix must be square'
189 if UL != 'L':
190 H_nn = H_nn.T
192 for n in range(N - 1):
193 map(H_nn[n + 1:, n], H_nn[n, n + 1:])
196def cutoff2gridspacing(E):
197 """Convert planewave energy cutoff to a real-space gridspacing."""
198 return np.pi / np.sqrt(2 * E / Hartree) * Bohr