From 47552a0562b2d4c83c86fcaa3f2c1f47e9993b5b Mon Sep 17 00:00:00 2001 From: SorBalda Date: Sat, 1 Aug 2026 20:53:40 +0200 Subject: [PATCH 1/2] Phonons: q-only supercell diagonalization, in-place sort, LO-TO and AdjustQStar fixes Four independent changes to Phonons.py, all in the fast DiagonalizeSupercell except the last. 1. New kwarg q_only=False. When True the (3N, 3N) supercell polarization matrix is never allocated: mode selection is unchanged, but for each accepted mode only its origin (iq, band) is recorded, in two O(3N) integer arrays, so that w_q[mode_band[k], mode_iq[k]] == w_mu[k]. Callers that only need the q-space eigenmodes (the q-space TDSCHA path) then run in memory linear in the supercell size instead of quadratic. This path also skips the pre-computed phase table, which is (3N, nq) and would be the largest allocation left; the phase is recomputed one q at a time, agreeing with the table to BLAS rounding, and w_q/pols_q do not depend on it. Verified on a 27-q gold dyn: w_mu, w_q and pols_q are bitwise identical to return_qmodes=True, and the (iq, band) map is exact. 2. Sort the polarization columns in place. e_pols_sc[:, sort_mask] is fancy indexing, so the sorted copy is built while the original is still alive, doubling the peak of the routine's largest allocation. Replaced by a cycle-following permutation with a single length-3N temporary column; verified bitwise equal to the fancy-indexed result over 200 random permutations and on real dyns with heavily degenerate spectra. 3. Honour the LO-TO warning. With lo_to_split set but no effective charges the code warned "LO-TO ignored" and then ran the branch anyway. Since ForceTensor.Tensor2.Interpolate guards the whole non-analytic term behind "effective_charges is not None", that branch rebuilt the dynamical matrix DyagDinQ already has, after allocating a full supercell Tensor2 through SetupFromPhonons -> GetRealSpaceFC (measured: 132 MB at a 8x8x8 supercell, growing as N^2). Falling through to DyagDinQ gives the same frequencies and keeps the Gamma polarization gauge consistent with the other q points. The lo_to_split string validation is moved above the check so a typo still raises when the charges are missing. Applied to both copies of the routine. 4. AdjustQStar swapped the block of the last q visited by the search loop instead of the Gamma one: the loop had no break, so iq kept its final value, and the reordering was corrupted whenever Gamma was neither first nor last. It also proceeded silently when no Gamma point was present, in which case the swap zeroed q_tot[0] and fabricated a Gamma that did not exist; it now raises. --- cellconstructor/Phonons.py | 175 ++++++++++++++++++++++++++++++++----- 1 file changed, 152 insertions(+), 23 deletions(-) diff --git a/cellconstructor/Phonons.py b/cellconstructor/Phonons.py index b50cef65..725d1a9d 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,46 @@ 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 + + # Sort the supercell polarization columns IN PLACE. Fancy indexing + # (e_pols_sc[:, sort_mask]) builds a second (3N, 3N) array while the first + # is still alive, doubling the peak of the largest allocation of the + # routine. Cycle-following permutation reproducing exactly + # e_pols_sc[:, k] <- e_pols_sc_presort[:, sort_mask[k]] + # with a single length-3N temporary column and an O(3N) boolean mask. + _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] + while True: + _done[_j] = True + _src = sort_mask[_j] + if _src == _i: + 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__ From 4a79f03015a64e0afb7dfd4db79a977c1266f223 Mon Sep 17 00:00:00 2001 From: SorBalda Date: Sun, 2 Aug 2026 12:09:52 +0200 Subject: [PATCH 2/2] Explain the in-place column permutation The comment described what was wrong with the fancy-indexed version and left the cycle-following loop itself unexplained. Describe the algorithm instead: why a permutation can be applied in place one cycle at a time, what the saved column and the _done mask are for, and what each step of the walk does. --- cellconstructor/Phonons.py | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/cellconstructor/Phonons.py b/cellconstructor/Phonons.py index 725d1a9d..3f15f500 100644 --- a/cellconstructor/Phonons.py +++ b/cellconstructor/Phonons.py @@ -4271,23 +4271,35 @@ def DiagonalizeSupercell(self, verbose = False, lo_to_split = None, return_qmode timer.add_timer("Sort and validate", t_sort_end - t_sort_start) return w_array, mode_iq, mode_band, w_q, pols_q - # Sort the supercell polarization columns IN PLACE. Fancy indexing - # (e_pols_sc[:, sort_mask]) builds a second (3N, 3N) array while the first - # is still alive, doubling the peak of the largest allocation of the - # routine. Cycle-following permutation reproducing exactly - # e_pols_sc[:, k] <- e_pols_sc_presort[:, sort_mask[k]] - # with a single length-3N temporary column and an O(3N) boolean mask. + # 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] + _tmp_col[:] = e_pols_sc[:, _i] # save the head of the cycle while True: _done[_j] = True - _src = sort_mask[_j] - if _src == _i: + _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]