diff --git a/cellconstructor/Phonons.py b/cellconstructor/Phonons.py index b50cef6..3f15f50 100644 --- a/cellconstructor/Phonons.py +++ b/cellconstructor/Phonons.py @@ -3134,12 +3134,19 @@ def AdjustQStar(self, use_spglib = False): for iq, q in enumerate(self.q_tot): if np.max(np.abs(q)) < __EPSILON__: i_gamma = iq + break + + if i_gamma < 0: + raise ValueError("Error, no Gamma point found among the q points; cannot reorder the q_tot list.") if i_gamma != 0: + # Swap the Gamma block, not the last q visited by the loop above: + # iq keeps its final value from the loop, so using it here moved the + # wrong dynamical matrix whenever Gamma was neither first nor last. mydyn = self.dynmats[0].copy() - self.dynmats[0] = self.dynmats[iq].copy() - self.dynmats[iq] = mydyn - self.q_tot[iq] = self.q_tot[0].copy() + self.dynmats[0] = self.dynmats[i_gamma].copy() + self.dynmats[i_gamma] = mydyn + self.q_tot[i_gamma] = self.q_tot[0].copy() self.q_tot[0][:] = 0 @@ -3696,10 +3703,25 @@ def DiagonalizeSupercell_slow(self, verbose = False, lo_to_split = None, return_ self.dynmats[iq] = re_part # Check if this is gamma (to apply the LO-TO splitting) - if Methods.get_min_dist_into_cell(bg, q, np.zeros(3)) < 1e-16 and lo_to_split is not None: - if self.effective_charges is None: - warnings.warn("WARNING: Requested LO-TO splitting without effective charges. LO-TO ignored.") - + apply_lo_to = Methods.get_min_dist_into_cell(bg, q, np.zeros(3)) < 1e-16 and lo_to_split is not None + if apply_lo_to and isinstance(lo_to_split, str) and lo_to_split.lower() != "random": + # Validate before the effective-charges check, otherwise a typo in + # lo_to_split is silently accepted whenever the charges are missing + # (the branch that used to raise is skipped in that case). + raise ValueError("Error, lo_to_split argument '%s' not recognized" % lo_to_split) + if apply_lo_to and self.effective_charges is None: + warnings.warn("WARNING: Requested LO-TO splitting without effective charges. LO-TO ignored.") + # Honour the warning instead of running the branch anyway. Without + # effective charges ForceTensor.Tensor2.Interpolate returns the plain + # short-range Fourier sum -- the whole non-analytic part is guarded by + # "if self.effective_charges is not None" -- so this branch rebuilds the + # very same dynamical matrix that DyagDinQ already has, after paying for + # a full supercell Tensor2: a (3*nat_sc, 3*nat_sc) complex array, hundreds + # of MB on a large supercell. Falling through to DyagDinQ gives the same + # frequencies and keeps the Gamma polarization gauge consistent with all + # the other q points. + apply_lo_to = False + if apply_lo_to: # TIMER: LO-TO splitting computation t_lo_to_start = time.time() # Initialize the Force Constant @@ -3887,7 +3909,7 @@ def DiagonalizeSupercell_slow(self, verbose = False, lo_to_split = None, return_ return w_array, e_pols_sc - def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmodes = False, timer=None): + def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmodes = False, timer=None, q_only = False): r""" DIAGONALIZE THE DYNAMICAL MATRIX IN THE SUPERCELL (FAST VERSION) ================================================================ @@ -3908,16 +3930,36 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode If LO-TO is specified but no effective charges are present, then a warning is print and it is ignored. - return_qmodes : bool If true, frequencies and polarizations in q space are returned. + - q_only : bool + If True the (3N, 3N) supercell polarization matrix is never + allocated. The mode selection is performed exactly as in the + standard path, but instead of storing the supercell + eigenvector of each accepted mode only its origin + (iq, band) is recorded, in two O(3N) integer arrays. The + returned w_q / pols_q are bitwise identical to the ones + obtained with return_qmodes = True. + Note that this path also skips the pre-computation of the + phase factors (OPTIMIZATION 2 below), recomputing the phase + one q at a time instead; the two agree to BLAS rounding + (gemv vs gemm) and w_q / pols_q do not depend on it. That + table is (3N, nq), i.e. O(N^2/nat), and would dominate the memory + budget of a path whose whole point is to stay O(N). Results ------- - w_mu : ndarray( size = (n_modes), dtype = np.double) Frequencies in the supercell - e_mu : ndarray( size = (3*Nat_sc, n_modes), dtype = np.double, order = "F") - Polarization vectors in the supercell + Polarization vectors in the supercell (not returned if q_only is True) + - mode_iq : ndarray( size = (n_modes), dtype = np.intp) + For each sorted supercell mode, the index of the q point it comes + from (only if q_only is True) + - mode_band : ndarray( size = (n_modes), dtype = np.intp) + For each sorted supercell mode, the band index within that q point, + so that w_q[mode_band[k], mode_iq[k]] == w_mu[k] (only if q_only is True) - w_q : ndarray( size = (3*Nat, nq), dtype = np.double, order = "F") - Frequencies in the q space (only if return_qmodes is True) + Frequencies in the q space (only if return_qmodes or q_only is True) - e_q : ndarray( size = (3*Nat, 3*Nat, nq), dtype = np.complex128, order = "F") - Polarization vectors in the q space (only if return_qmodes is True) + Polarization vectors in the q space (only if return_qmodes or q_only is True) """ supercell_size = len(self.q_tot) @@ -3927,7 +3969,17 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode nat_sc = nat*supercell_size w_array = np.zeros( nmodes, dtype = np.double) - e_pols_sc = np.zeros( (nmodes, nmodes), dtype = np.double, order = "F") + # q_only path: never allocate the (3N, 3N) supercell polarization matrix. + # Only the O(3N) index arrays mapping each accepted supercell mode back to + # the (iq, band) of the per-q diagonalization that produced it. + if q_only: + e_pols_sc = None + mode_iq = np.full(nmodes, -1, dtype = np.intp) + mode_band = np.full(nmodes, -1, dtype = np.intp) + else: + e_pols_sc = np.zeros( (nmodes, nmodes), dtype = np.double, order = "F") + mode_iq = None + mode_band = None nq = len(self.q_tot) w_q = np.zeros((3*nat, nq), dtype = np.double, order = "F") @@ -3993,8 +4045,15 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode # OPTIMIZATION 2: Pre-compute phase factors for all q-points at once # TIMER: Phase factor computation t_start = time.time() - # This avoids computing R_vec.dot(q) repeatedly in the inner loop - phase_factors = np.exp(1j * 2 * np.pi * R_vec.dot(q_array.T)) # (nmodes, nq) + # This avoids computing R_vec.dot(q) repeatedly in the inner loop. + # The table is (nmodes, nq) complex128 = O(N^2/nat): on the q_only path + # it would be the single largest allocation of a routine that is meant + # to stay O(N), so there the phase is recomputed per q point instead + # (one (nmodes,) vector at a time, same numbers). + if q_only: + phase_factors = None + else: + phase_factors = np.exp(1j * 2 * np.pi * R_vec.dot(q_array.T)) # (nmodes, nq) t_end = time.time() if timer is not None: timer.add_timer("Compute phase factors", t_end - t_start) @@ -4008,7 +4067,7 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode # TIMER: Process skipped q-point t_skip_start = time.time() # Check if we must return anyway the polarization in q space - if return_qmodes: + if return_qmodes or q_only: if timer is not None: wq, eq = timer.execute_timed_function(self.DyagDinQ, iq) else: @@ -4038,10 +4097,25 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode self.dynmats[iq] = re_part # Check if this is gamma (to apply the LO-TO splitting) - if Methods.get_min_dist_into_cell(bg, q, np.zeros(3)) < 1e-16 and lo_to_split is not None: - if self.effective_charges is None: - warnings.warn("WARNING: Requested LO-TO splitting without effective charges. LO-TO ignored.") - + apply_lo_to = Methods.get_min_dist_into_cell(bg, q, np.zeros(3)) < 1e-16 and lo_to_split is not None + if apply_lo_to and isinstance(lo_to_split, str) and lo_to_split.lower() != "random": + # Validate before the effective-charges check, otherwise a typo in + # lo_to_split is silently accepted whenever the charges are missing + # (the branch that used to raise is skipped in that case). + raise ValueError("Error, lo_to_split argument '%s' not recognized" % lo_to_split) + if apply_lo_to and self.effective_charges is None: + warnings.warn("WARNING: Requested LO-TO splitting without effective charges. LO-TO ignored.") + # Honour the warning instead of running the branch anyway. Without + # effective charges ForceTensor.Tensor2.Interpolate returns the plain + # short-range Fourier sum -- the whole non-analytic part is guarded by + # "if self.effective_charges is not None" -- so this branch rebuilds the + # very same dynamical matrix that DyagDinQ already has, after paying for + # a full supercell Tensor2: a (3*nat_sc, 3*nat_sc) complex array, hundreds + # of MB on a large supercell. Falling through to DyagDinQ gives the same + # frequencies and keeps the Gamma polarization gauge consistent with all + # the other q points. + apply_lo_to = False + if apply_lo_to: # TIMER: LO-TO splitting computation t_lo_to_start = time.time() # Initialize the Force Constant @@ -4086,7 +4160,15 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode e_contracted = tilde_e_qnu_all[itau_modes, :] # (nmodes, nmodes_unit) # Get phase factors for this q-point - phase_q = phase_factors[:, iq] # (nmodes,) + if q_only: + # One column at a time instead of the pre-computed table. This + # is a gemv where the table is a gemm, so the two agree to BLAS + # rounding (~1e-15) rather than bitwise; the mode selection has + # a margin of nine orders of magnitude on that, and w_q / pols_q + # do not depend on the phase at all. + phase_q = np.exp(1j * 2 * np.pi * R_vec.dot(q_array[iq])) # (nmodes,) + else: + phase_q = phase_factors[:, iq] # (nmodes,) # Broadcast and multiply: e_sc = e_contracted * phase / sqrt(N_q) c_e_sc = e_contracted * phase_q[:, np.newaxis] / np.sqrt(supercell_size) @@ -4137,11 +4219,19 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode # Add the vectors if add_1: w_array[i_mu] = w_qnu - e_pols_sc[:, i_mu] = evec_1 / np.sqrt(norm1) + if q_only: + mode_iq[i_mu] = iq + mode_band[i_mu] = i_qnu + else: + e_pols_sc[:, i_mu] = evec_1 / np.sqrt(norm1) i_mu += 1 if add_2: w_array[i_mu] = w_qnu - e_pols_sc[:, i_mu] = evec_2 / np.sqrt(norm2) + if q_only: + mode_iq[i_mu] = iq + mode_band[i_mu] = i_qnu + else: + e_pols_sc[:, i_mu] = evec_2 / np.sqrt(norm2) i_mu += 1 t2 = time.time() @@ -4162,7 +4252,58 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode # Sort the frequencies sort_mask = np.argsort(w_array) w_array = w_array[sort_mask] - e_pols_sc = e_pols_sc[:, sort_mask] + + if q_only: + # No (3N, 3N) has ever been allocated on this path: only the O(3N) + # index arrays are reordered, so that + # w_q[mode_band[k], mode_iq[k]] == w_array[k] + # for every sorted supercell mode k. The normalization assert below + # is skipped because it checks e_pols_sc, which does not exist here. + # That assert is also the only thing that made a malformed q_tot + # (e.g. an exactly duplicated q point) fail loudly on the standard + # path; without a check here q_only would return silently with + # unassigned modes, i.e. mode_iq = -1 and spurious zero frequencies. + assert i_mu == nmodes, "Error, only {} / {} supercell modes were assigned".format(i_mu, nmodes) + mode_iq = mode_iq[sort_mask] + mode_band = mode_band[sort_mask] + t_sort_end = time.time() + if timer is not None: + timer.add_timer("Sort and validate", t_sort_end - t_sort_start) + return w_array, mode_iq, mode_band, w_q, pols_q + + # Reorder the polarization columns according to sort_mask, in place. + # + # The result is the same as e_pols_sc[:, sort_mask], i.e. column k of + # the sorted matrix is column sort_mask[k] of the original one. Writing + # it that way, however, allocates the whole sorted copy while the + # original is still alive, so the peak doubles on the largest array of + # the routine (2 x (3N)^2 doubles). Doing it in place costs one spare + # column instead. + # + # A permutation decomposes into disjoint cycles, and a cycle can be + # applied in place by saving one element and shifting the others onto + # it. Starting from a column _i not yet moved, save it in _tmp_col, + # then walk the cycle _j -> sort_mask[_j]: each step overwrites the + # column _j with the one it must receive. When the walk comes back to + # _i its column is no longer in the matrix, so the saved copy is + # written there and the cycle is closed. The _done mask marks the + # columns already placed, so every cycle is walked exactly once and + # each column is written exactly once. + _tmp_col = np.empty(nmodes, dtype = e_pols_sc.dtype) + _done = np.zeros(nmodes, dtype = bool) + for _i in range(nmodes): + if _done[_i]: + continue + _j = _i + _tmp_col[:] = e_pols_sc[:, _i] # save the head of the cycle + while True: + _done[_j] = True + _src = sort_mask[_j] # column _j must receive _src + if _src == _i: # back to the head: close it + e_pols_sc[:, _j] = _tmp_col + break + e_pols_sc[:, _j] = e_pols_sc[:, _src] + _j = _src # Get the check for the polarization vector normalization assert np.max(np.abs(np.einsum("ab, ab->b", e_pols_sc, e_pols_sc) - 1)) < __EPSILON__