Phonons: q-only supercell diagonalization, in-place sort, LO-TO and AdjustQStar fixes

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__