diff --git a/Modules/DynamicalLanczos.py b/Modules/DynamicalLanczos.py index 51c8a322..4a4b8f64 100644 --- a/Modules/DynamicalLanczos.py +++ b/Modules/DynamicalLanczos.py @@ -246,16 +246,24 @@ def __init__(self, ensemble = None, mode = None, unwrap_symmetries = False, sele # ========== END OF VARIABLE DEFINITION (EACH NEW DEFINITION FROM NOW ON RESULTS IN AN ERROR) ======= - self.dyn = ensemble.current_dyn.Copy() + # Split init: a linear part (geometry/ensemble scalars/masses/structures, + # no (3N,3N) allocation) shared with the q-space subclasses, and a + # real-space part (diagonalization, pols, X/Y, psi, linops) that the + # q-space subclasses override to skip. + self._init_linear(ensemble) + self._init_realspace(ensemble, unwrap_symmetries, select_modes, lo_to_split) + + def _init_linear(self, ensemble): + """Linear-cost part of the initialization: geometry, ensemble scalars, + masses and structures. Shared by the real-space Lanczos and the q-space + subclasses; never allocates any (3N,3N)-order array.""" + self.dyn = ensemble.current_dyn.Copy() self.uci_structure = ensemble.current_dyn.structure.copy() self.super_structure = self.dyn.structure.generate_supercell(self.dyn.GetSupercell())#superdyn.structure self.T = ensemble.current_T - ws, pols = self.dyn.DiagonalizeSupercell(lo_to_split = lo_to_split) - self.nat = self.super_structure.N_atoms - n_cell = np.prod(self.dyn.GetSupercell()) self.qe_sym = CC.symmetries.QE_Symmetry(self.dyn.structure) self.qe_sym.SetupQPoint() @@ -264,6 +272,30 @@ def __init__(self, ensemble = None, mode = None, unwrap_symmetries = False, sele m = self.super_structure.get_masses_array() self.m = np.tile(m, (3,1)).T.ravel() + # Ignore v3 or v4. You can set them for testing + self.ignore_v3 = False + self.ignore_v4 = False + + # The number of configurations and the ensemble weights + self.N = ensemble.N + self.rho = ensemble.rho.copy() + self.N_eff = np.sum(self.rho) + + def _init_realspace(self, ensemble, unwrap_symmetries, select_modes, lo_to_split): + """Real-space preprocessing that allocates the (3N,3N)-order arrays: + supercell diagonalization, polarization basis, mass-rescaled + displacements/forces, X/Y projections, the psi working vector and the + L/M linear operators. Only the direct real-space Lanczos runs this; the + q-space subclasses override it to a no-op.""" + order = "C" + + ws, pols = self.dyn.DiagonalizeSupercell(lo_to_split = lo_to_split) + + n_cell = np.prod(self.dyn.GetSupercell()) + + # Get the (un-tiled) masses for the translation projector + m = self.super_structure.get_masses_array() + # Remove the translations if lo_to_split is not None and self.dyn.effective_charges is not None: trans_mask = np.zeros(len(ws), dtype=bool) diff --git a/Modules/QSpaceHessian.py b/Modules/QSpaceHessian.py index 412abf30..e2d72a16 100644 --- a/Modules/QSpaceHessian.py +++ b/Modules/QSpaceHessian.py @@ -49,6 +49,199 @@ __RyToK__ = 157887.32400374097 +def _adaptive_schur_fill(G_q, solve_schedule, rep_x, solve_column, nb, tol, + use_mode_symmetry, verbose=False, iq=None): + """Fill G_q from the solved representative columns, Schur-consistently. + + Schur's lemma (L commutes with the little group of q) fixes the + diagonal block of G on a d-dim irrep copy to c*I in ANY orthonormal + basis of that copy, but the cross block between two copies of the + SAME irrep is c*U_AB with an unknown unitary intertwiner (eigh + returns arbitrary bases in each degenerate subspace) -- NOT c*I. + The scalar shortcut is therefore valid only where the coupling + vanishes (distinct irreps). Two kinds of failure are detected, both + measured on columns that are solved anyway, so detection is free: + + * a block that is reducible on its own (a repeated irrep at the same + frequency, or an accidental degeneracy between different irreps): + the representative column of a clean single copy has zero support + on the rest of its own block, so nonzero leakage exposes it; + * a coupling between two blocks sharing an irrep. + + Both use the threshold min(50*tol, 1e-5)*scale (the cap keeps the + detection meaningful for loose solver tolerances). Every block of a + coupled group, and every self-reducible block, is then solved column + by column exactly: k blocks of dimension d cost k*(d-1) extra solves. + A false positive only costs solves. + + Known limits, all inherent to reading only the representative column: + + * a reducible block whose basis already happens to be symmetry + adapted is invisible -- the off-diagonal leakage is then exactly + zero while the two Schur constants still differ. Detecting it + needs the other columns, which is what the shortcut avoids; + * the coupling between a singleton mode and a degenerate block that + is accidentally degenerate is zeroed without being measured: Schur + forbids coupling between different irreps, not between different + frequencies; + * min(50*tol, 1e-5) is a floor on sensitivity: a coupling weaker + than that, relative to the column norm, is not detected. + + Parameters + ---------- + G_q : ndarray(nb, nb), complex -- filled in place + solve_schedule : list of (band_i, block) + rep_x : dict band_i -> solved column (length >= nb) + solve_column : callable(band) -> (x, n_iters, elapsed) + nb : int -- number of bands (R-sector size) + tol : float -- iterative solver relative tolerance + use_mode_symmetry : bool + verbose, iq : diagnostics only + + Returns + ------- + full_solve : set of representative bands whose blocks were solved + column by column (empty when no repeated irrep was detected). + """ + deg_blocks = [b for _, b in solve_schedule if len(b) >= 2] + full_solve = set() + group_rows = {} + if use_mode_symmetry and deg_blocks: + parent = {b[0]: b[0] for b in deg_blocks} + + def _find(a): + while parent[a] != a: + parent[a] = parent[parent[a]] + a = parent[a] + return a + + # Pass 1: blocks that are not a single irrep copy. A degenerate block + # can already be reducible on its own -- a repeated irrep at the same + # frequency, or an accidental degeneracy between different irreps -- + # and then c*I is wrong for it with no partner block to reveal it. + # Schur forces G on one irrep copy to be c*I in ANY basis, so the + # representative column of a clean copy has zero support on the rest + # of its own block: nonzero leakage means the block is reducible. + # This is measured on columns that are solved anyway, so it is free. + for A in deg_blocks: + others = np.array([m for m in A if m != A[0]]) + if not others.size: + continue + xA = rep_x[A[0]] + scale = max(np.linalg.norm(xA[:nb]), 1e-300) + if np.max(np.abs(xA[others])) > min(50.0 * tol, 1e-5) * scale: + full_solve.add(A[0]) + + # Pass 2: couplings between blocks. Kept separate from pass 1 because + # the dimension shortcut below reads full_solve, which must already be + # complete: deciding both in one loop makes the result depend on the + # order the blocks happen to be listed in. + for i in range(len(deg_blocks)): + for j in range(i + 1, len(deg_blocks)): + A, B = deg_blocks[i], deg_blocks[j] + # "different dimension -> different irrep -> no coupling" only + # holds when BOTH blocks are irreducible; a reducible one can + # share an irrep with a block of any size, so it is measured. + if len(A) != len(B) and \ + A[0] not in full_solve and B[0] not in full_solve: + continue + xA, xB = rep_x[A[0]], rep_x[B[0]] + # Recomputed per pair: carrying a running maximum across the + # loop would make the threshold monotonically non-decreasing, + # so a single soft-mode block (||x|| ~ 1/w^2, orders of + # magnitude larger) would raise it for every later pair and + # silently hide their couplings. + scale = max(np.linalg.norm(xA[:nb]), + np.linalg.norm(xB[:nb]), 1e-300) + coup = max(np.max(np.abs(xA[np.array(B)])), + np.max(np.abs(xB[np.array(A)]))) + if coup > min(50.0 * tol, 1e-5) * scale: + ra, rb = _find(A[0]), _find(B[0]) + if ra != rb: + parent[rb] = ra + + groups = {} + for b in deg_blocks: + groups.setdefault(_find(b[0]), []).append(b) + for blist in groups.values(): + rows = sorted(m for b in blist for m in b) + for b in blist: + # Every block of a coupled group is solved in full; a block + # that is reducible on its own is a group of one. + if len(blist) > 1 or b[0] in full_solve: + full_solve.add(b[0]) + group_rows[b[0]] = rows + + # Extra solves for the blocks with repeated irreps + extra_x = {} + if full_solve: + if verbose: + print(" Repeated irreps detected at iq={}: exact solve " + "for all columns of blocks {}".format( + iq, sorted(full_solve))) + for band_i, block in solve_schedule: + if band_i not in full_solve: + continue + extra_x[band_i] = rep_x[band_i] + for col in block: + if col == band_i: + continue + extra_x[col], it, dt = solve_column(col) + if verbose: + print(" Band {} (block {}): {} iters, " + "{:.2f}s".format(col, block, it, dt)) + + # Fill G_q using Schur's lemma for degenerate blocks. + for band_i, block in solve_schedule: + x = rep_x[band_i] + if len(block) == 1: + # Non-degenerate: full column from R-sector + G_q[:, band_i] = x[:nb] + # Zero out entries for modes in degenerate blocks + # (different irreps -> zero by Schur's lemma). + # This prevents solver noise from breaking degeneracy + # after Hermitian symmetrization. + for _, other_block in solve_schedule: + if len(other_block) >= 2: + for m in other_block: + G_q[m, band_i] = 0.0 + elif band_i in full_solve: + # Repeated irrep: every column of the coupled group was + # solved exactly; rows outside the group are zero by Schur. + rows = group_rows[band_i] + for col in block: + xc = extra_x[col] + for m in rows: + G_q[m, col] = xc[m] + else: + d = len(block) + # Extract Schur-consistent scalars only (not the full + # noisy solver column). This ensures all columns within + # the degenerate block are filled identically, preserving + # perfect block structure and preventing degeneracy + # breaking after symmetrization. + c_diag = x[band_i] # within-block diagonal constant + + # Fill ALL columns in this block (including rep) uniformly + for j in range(d): + col = block[j] + # Within-block diagonal + G_q[col, col] = c_diag + # Cross-coupling with other same-dimension blocks + # (verified uncoupled above, so this is only the + # residual solver noise on a Schur-zero entry) + for _, other_block in solve_schedule: + if other_block[0] == band_i: + continue + if len(other_block) == d and \ + other_block[0] not in full_solve: + G_q[other_block[j], col] = x[other_block[0]] + # Entries with different-dimension blocks and singlets + # are zero by Schur (different irreps), left as 0. + + return full_solve + + class QSpaceHessian: """Compute the free energy Hessian in q-space via iterative linear solves. @@ -728,8 +921,12 @@ def compute_hessian_at_q(self, iq, tol=1e-6, max_iters=500, When use_mode_symmetry=True and degenerate modes are present, exploits Schur's lemma: L_static commutes with the little group - of q, so G_q restricted to a d-dimensional irrep block is c*I_d. - Only one solve per degenerate block is needed instead of d solves. + of q, so G_q restricted to a d-dimensional irrep block is c*I_d, + and the cross block between two copies of the SAME irrep is + c*U with an unknown unitary intertwiner U. Only one solve per + degenerate block is needed for the uncoupled blocks; blocks with + a detected non-vanishing mutual coupling (repeated irreps) are + solved column by column instead (see _adaptive_schur_fill). Parameters ---------- @@ -747,8 +944,12 @@ def compute_hessian_at_q(self, iq, tol=1e-6, max_iters=500, can be very large for big supercells. Default is False. use_mode_symmetry : bool If True, exploit mode degeneracy to reduce the number of GMRES - solves. Within each degenerate block, only one solve is performed - and G_q is filled using Schur's lemma (G_block = c * I). + solves. Within each degenerate block, only one solve is + performed and G_q is filled using Schur's lemma (diagonal + block c * I). Groups of same-dimension blocks with detected + non-zero coupling (repeated irreps, where Schur only fixes + the cross block up to a unitary) fall back to exact + column-by-column solves automatically. Returns ------- @@ -796,6 +997,18 @@ def apply_M_tilde(x_tilde): non_acoustic = [nu for nu in range(nb) if self.qlanc.valid_modes_q[nu, iq]] + # With w < 0 (unstable modes) the Bose occupations become n < -1 and + # Lambda/Y_w produce finite but physically meaningless numbers with + # no other diagnostic: refuse loudly instead of returning garbage. + if any(w_qp[nu] < 0 for nu in non_acoustic): + raise ValueError( + "Negative (unstable) frequencies among the non-acoustic " + "modes at iq={} (min w = {:.6e} Ry): the free energy " + "Hessian Bose factors are meaningless for w < 0. Apply " + "ForcePositiveDefinite() to the dynamical matrix or check " + "the SSCHA convergence.".format( + iq, min(w_qp[nu] for nu in non_acoustic))) + # Build solve schedule: list of (band_to_solve, block_members) if use_mode_symmetry: blocks = self._find_degenerate_blocks(iq) @@ -819,7 +1032,9 @@ def apply_M_tilde(x_tilde): total_iters = 0 L_dense = None # Built lazily if iterative solvers fail - for band_i, block in solve_schedule: + def _solve_column(band_i): + """Solve L_static x = e_{band_i}; returns (x, n_iters, elapsed).""" + nonlocal L_dense, total_iters rhs = np.zeros(psi_size, dtype=np.complex128) rhs[band_i] = 1.0 rhs_tilde = rhs * sqrt_mask @@ -878,52 +1093,25 @@ def _count(xk): total_iters += n_iters[0] # Un-transform - x = x_tilde * inv_sqrt_mask - - # Fill G_q using Schur's lemma for degenerate blocks. - # G commutes with the little group, so between two d-dim - # copies of the same irrep, G = c_cross * I_d. - if len(block) == 1: - # Non-degenerate: full column from R-sector - G_q[:, band_i] = x[:nb] - # Zero out entries for modes in degenerate blocks - # (different irreps → zero by Schur's lemma). - # This prevents GMRES noise from breaking degeneracy - # after Hermitian symmetrization. - for _, other_block in solve_schedule: - if len(other_block) >= 2: - for m in other_block: - G_q[m, band_i] = 0.0 - else: - d = len(block) - # Extract Schur-consistent scalars only (not the full - # noisy GMRES column). This ensures all columns within - # the degenerate block are filled identically, preserving - # perfect block structure and preventing degeneracy - # breaking after symmetrization. - c_diag = x[band_i] # within-block diagonal constant - - # Fill ALL columns in this block (including rep) uniformly - for j in range(d): - col = block[j] - # Within-block diagonal - G_q[col, col] = c_diag - # Cross-coupling with other same-dimension blocks - for _, other_block in solve_schedule: - if other_block[0] == band_i: - continue - if len(other_block) == d: - # Same irrep type: shifted diagonal - G_q[other_block[j], col] = x[other_block[0]] - # Entries with different-dimension blocks and singlets - # are zero by Schur (different irreps), left as 0. + return x_tilde * inv_sqrt_mask, n_iters[0], t2 - t1 + # Phase 1: solve one representative column per block + rep_x = {} + for band_i, block in solve_schedule: + x, it, dt = _solve_column(band_i) + rep_x[band_i] = x if self.verbose: block_str = "{}".format(block) if len(block) > 1 else "" print(" Band {}{}: {} iters, {:.2f}s".format( band_i, " (block {})".format(block_str) if block_str else "", - n_iters[0], t2 - t1)) + it, dt)) + + # Phases 2-4: detect repeated irreps and fill G_q Schur-consistently + # (module-level so the fill logic is unit-testable in isolation). + _adaptive_schur_fill( + G_q, solve_schedule, rep_x, _solve_column, nb, tol, + use_mode_symmetry, verbose=self.verbose, iq=iq) # 7. Symmetrize G_q (should be Hermitian) G_q = (G_q + G_q.conj().T) / 2 @@ -1132,7 +1320,7 @@ def load_distributed_hessian(data_dir, population_id, dyn, T, lo_to_split=None, use_symmetries=True, n_configs=None, final_dyn=None, final_T=None, verbose=True, ignore_v3=False, ignore_v4=False, - **kwargs): + fourier_weights=True, **kwargs): """Load QSpaceHessian with distributed configurations across MPI ranks. Loads the ensemble on master rank only, then distributes configuration data @@ -1165,6 +1353,14 @@ def load_distributed_hessian(data_dir, population_id, dyn, T, lo_to_split=None, If True, exclude cubic (D3) anharmonic contributions. ignore_v4 : bool If True, exclude quartic (D4) anharmonic contributions. + fourier_weights : bool + If True (default), the ensemble is built in the opt-in light mode + (Ensemble(..., qspace_light=True): no (3N,3N) supercell polarization + vectors, not even at construction) and the final_dyn weight update + uses the memory-clean q-space update_weights_fourier, when the Julia + Fourier backend is available; otherwise falls back to a standard + ensemble + the real-space update_weights. Set False to force the + legacy behavior exactly. **kwargs Additional arguments passed to QSpaceLanczos. @@ -1198,6 +1394,7 @@ def load_distributed_hessian(data_dir, population_id, dyn, T, lo_to_split=None, n_configs=n_configs, final_dyn=final_dyn, final_T=final_T, + fourier_weights=fourier_weights, **kwargs ) diff --git a/Modules/QSpaceLanczos.py b/Modules/QSpaceLanczos.py index 1584b057..c6275977 100644 --- a/Modules/QSpaceLanczos.py +++ b/Modules/QSpaceLanczos.py @@ -22,6 +22,7 @@ import sys, os import time +import inspect import warnings import numpy as np @@ -59,6 +60,47 @@ except ImportError: __SPGLIB__ = False +# Capability probes for the two companion packages. The q-space path needs +# features that only exist in newer CellConstructor / python-sscha releases; +# the probes let this module degrade with an explicit warning instead of dying +# on a TypeError or an AttributeError deep inside __init__ when installed next +# to an older one. They are cheap and evaluated once at import. +# +# These will be removed in favour of a version pin in requirements.txt and +# pyproject.toml once the companion features are released: today no published +# version of either package exposes them, so a pin cannot be written yet. +# See CellConstructor PR #126 and python-sscha PR #428. + + +def _cc_has_q_only(): + try: + return "q_only" in inspect.signature( + CC.Phonons.Phonons.DiagonalizeSupercell).parameters + except (AttributeError, TypeError, ValueError): + return False + + +def _ensemble_supports_light(): + """True if python-sscha exposes the linear q-space ensemble API.""" + try: + return "qspace_light" in inspect.signature( + sscha.Ensemble.Ensemble.__init__).parameters + except Exception: + return False + + +def _ensemble_has_qspace_cache_api(ensemble_cls): + """True if the ensemble class can rebuild its q-space caches. + + Accepts the deprecated private name so that a python-sscha predating the + public rename still works. + """ + return (hasattr(ensemble_cls, "refresh_qspace_caches_from_real_space") + or hasattr(ensemble_cls, "_refresh_qspace_caches_from_real_space")) + + +_CC_HAS_Q_ONLY = _cc_has_q_only() + # Constants __EPSILON__ = 1e-12 @@ -131,6 +173,7 @@ def __init__(self, ensemble, lo_to_split=None, **kwargs): # -- Add the q-space attributes -- qspace_attrs = [ 'q_points', 'n_q', 'n_bands', 'w_q', 'pols_q', + 'mode_iq', 'mode_band', 'valid_modes_q', 'X_q', 'Y_q', 'iq_pert', 'q_pair_map', 'unique_pairs', '_psi_size', '_block_offsets_a', '_block_offsets_b', '_block_sizes', @@ -142,12 +185,40 @@ def __init__(self, ensemble, lo_to_split=None, **kwargs): # If ensemble is None, perform a bare initialization like the parent if ensemble is None: + # Declared above but never assigned on this path: the non-master + # ranks of the distributed loader go through here, so reading them + # would raise AttributeError instead of returning "not available". + self.mode_iq = None + self.mode_band = None return # == 1. Get q-space eigenmodes == - ws_sc, pols_sc, w_q, pols_q = self.dyn.DiagonalizeSupercell( - return_qmodes=True, lo_to_split=lo_to_split) - + # q_only=True never allocates the (3N,3N) supercell polarization matrix. + # It returns (w_array, mode_iq, mode_band, w_q, pols_q); w_q/pols_q are + # bitwise identical to the legacy return_qmodes=True outputs. mode_iq / + # mode_band map the sorted supercell mode index k -> (iq, band) and are + # kept for Phase F3. + # Older CellConstructor releases have no q_only: fall back to + # return_qmodes, which computes the same w_q/pols_q but does allocate + # the dense (3N,3N). Probing the signature keeps this module usable + # against an unpatched CellConstructor instead of raising TypeError. + if _CC_HAS_Q_ONLY: + w_array, mode_iq, mode_band, w_q, pols_q = self.dyn.DiagonalizeSupercell( + q_only=True, lo_to_split=lo_to_split) + else: + warnings.warn( + "This CellConstructor has no DiagonalizeSupercell(q_only=True): " + "falling back to return_qmodes=True, which allocates the dense " + "(3N,3N) supercell polarization matrix. The q-space path will " + "give the same numbers but will not be linear in memory.") + w_array, e_pols_sc, w_q, pols_q = self.dyn.DiagonalizeSupercell( + return_qmodes=True, lo_to_split=lo_to_split) + del e_pols_sc + mode_iq = None + mode_band = None + + self.mode_iq = mode_iq # (3N,) intp: supercell mode -> q index + self.mode_band = mode_band # (3N,) intp: supercell mode -> band index self.q_points = np.array(self.dyn.q_tot) # (n_q, 3) self.n_q = len(self.q_points) self.n_bands = 3 * self.uci_structure.N_atoms # uniform band count @@ -193,6 +264,39 @@ def __init__(self, ensemble, lo_to_split=None, **kwargs): self._N_eff_global = self.N_eff self._N_local = self.N + # -- Pin the real-space array slots to None -- + # + # As of Phase F2 these arrays are NEVER allocated on the q-space path: + # `_init_realspace` is overridden below to a no-op, so the base class + # skips the supercell diagonalization, `pols` (3*N_sc x n_modes), the + # `psi` working vector (n_modes + n_modes*(n_modes+1)/2 ~ (3N)^2), and + # X/Y/u_tilde/f_tilde entirely. They keep their empty scalar-block + # defaults; we pin them to None so the downstream None-guards behave: + # * self.psi must be None so run_FT / QSpaceKPM.run_KPM raise their + # explicit ValueError until reset_q() sizes psi at the q-space size. + # * load_distributed_tdscha frees X/Y via `if qlanc.X is not None`. + # Everything the q-space algorithm needs lives in X_q, Y_q, pols_q, w_q + # and rho. None of QSpaceLanczos / QSpaceKPM / QSpaceHessian ever reads + # self.pols / X / Y / u_tilde / f_tilde (verified by grep, Phase F2). + self.pols = None + self.psi = None + self.X = None + self.Y = None + self.u_tilde = None + self.f_tilde = None + + def _init_realspace(self, ensemble, unwrap_symmetries, select_modes, lo_to_split): + """Override: skip the base-class real-space preprocessing entirely. + + The direct real-space Lanczos allocates the (3N,3N) polarization matrix, + the X/Y projections and the O((3N)^2) psi vector here. The q-space path + needs none of it -- it does its own single q_only diagonalization and + Bloch-transforms the ensemble into X_q/Y_q -- so this is a no-op. This is + the seam that keeps any (3N,3N)-order array from ever being allocated on + the QSpaceLanczos / QSpaceHessian / QSpaceKPM construction path. + """ + return + def _bloch_transform_ensemble(self): """Bloch transform the ensemble displacements and forces into q-space mode basis. @@ -202,14 +306,39 @@ def _bloch_transform_ensemble(self): The forces are the anharmonic residual: f - f_SSCHA - , matching the preprocessing done in DynamicalLanczos.__init__. """ - # Ensure the ensemble has computed q-space quantities - # Check if fourier_gradient is active or force it - if self.ensemble.u_disps_qspace is None: - # Force fourier gradient initialization in the ensemble - if not self.ensemble.fourier_gradient: - print("Ensemble checking: computing Fourier transform of displacements and forces...") - self.ensemble.fourier_gradient = True - self.ensemble.init() + # Ensure the q-space caches describe the same current dynamical matrix + # as the real-space ensemble data. update_weights() refreshes the + # real-space SSCHA forces but historically leaves sscha_forces_qspace + # stale; update_weights_fourier() does the converse. Refresh from + # real-space only after a real-space update, by Fourier transforming + # the real-space arrays themselves: unlike Ensemble.init(), this keeps + # the displacements referenced to current_dyn's centroids (init() + # would reset them to dyn_0's) and does not clobber rho, + # sscha_energies or the u_disps_original baseline. On the + # Fourier/light path the q-space cache is already authoritative, so + # the refresh is skipped and memory stays O(Nq). + qspace_cache_current = bool(getattr( + self.ensemble, "_last_weight_update_fourier", False)) + if self.ensemble.u_disps_qspace is None or not qspace_cache_current: + refresh = getattr( + self.ensemble, "refresh_qspace_caches_from_real_space", + getattr(self.ensemble, + "_refresh_qspace_caches_from_real_space", None)) + if refresh is None: + raise AttributeError( + "This ensemble has no refresh_qspace_caches_from_real_space(): " + "the q-space path needs a python-sscha that provides the " + "q-space cache API (qspace_light / update_weights_fourier). " + "Install the matching python-sscha, or build the Lanczos " + "from a real-space ensemble instead.") + if not self.ensemble.fourier_gradient: + if Parallel.am_i_the_master(): + print("Ensemble checking: computing Fourier transform of " + "displacements and forces...") + self.ensemble.fourier_gradient = True + # The ensemble raises its own coherence marker: this package must + # not reach into Ensemble.__dict__ to flip private state. + refresh() # Unit conversion factors # Target: Bohr (u) and Ry/Bohr (f) @@ -219,7 +348,12 @@ def _bloch_transform_ensemble(self): u_conv = CC.Units.A_TO_BOHR f_conv = 1.0 / CC.Units.A_TO_BOHR elif self.ensemble.units == "hartree": - f_conv = 2.0 # Ha -> Ry + raise NotImplementedError( + "units='hartree' is not supported on the q-space path: the " + "Fourier caches (u_disps_qspace, forces_qspace, " + "sscha_forces_qspace) are built with Ry/Angstrom conversion " + "factors and Ensemble.convert_units does not update them. " + "Convert the ensemble back to default units first.") # Mass scaling factors (sqrt(m) for u, 1/sqrt(m) for f) # We use self.dyn.structure corresponding to unit cell @@ -227,9 +361,12 @@ def _bloch_transform_ensemble(self): sqrt_m = np.sqrt(m_uc) sqrt_m_3 = np.repeat(sqrt_m, 3) # (3*nat_uc,) - # Compute the average anharmonic force (matching parent DynamicalLanczos) - # get_average_forces returns rho-weighted in unit cell, Ry/Angstrom - f_mean_uc = self.ensemble.get_average_forces(get_error=False) # (nat_uc, 3) + # Compute the average from the SAME q-space residual used below. + # get_average_forces() reads the real-space SSCHA cache, which can be + # stale after update_weights_fourier(); get_fourier_forces() reads + # forces_qspace - sscha_forces_qspace and therefore stays coherent. + f_mean_uc = self.ensemble.get_fourier_forces( + get_error=False).reshape((-1, 3)) # Ry/Angstrom # Symmetrize the average force qe_sym = CC.symmetries.QE_Symmetry(self.dyn.structure) qe_sym.SetupQPoint() @@ -1511,7 +1648,7 @@ def _get_atom_perm(structure, R_cart, t_cart, M, Minv, tol=0.1): Returns irt such that R @ tau[kappa] + t ≡ tau[irt[kappa]] mod lattice. """ nat = structure.N_atoms - irt = np.zeros(nat, dtype=int) + irt = np.full(nat, -1, dtype=int) for kappa in range(nat): tau = structure.coords[kappa] mapped = R_cart @ tau + t_cart @@ -1522,6 +1659,13 @@ def _get_atom_perm(structure, R_cart, t_cart, M, Minv, tol=0.1): if np.linalg.norm(M @ diff_frac) < tol: irt[kappa] = kp break + if np.any(irt < 0) or len(np.unique(irt)) != nat: + raise ValueError( + "A symmetry operation does not map the atoms onto themselves " + "within tol={} A (irt={}): the structure is distorted away " + "from the detected symmetry group, or two atoms fall within " + "the matching tolerance. Refusing to build a wrong " + "symmetrization matrix.".format(tol, irt)) return irt def _build_qspace_symmetries(self, rot_frac_all, trans_frac_all, @@ -1543,13 +1687,11 @@ def _build_qspace_symmetries(self, rot_frac_all, trans_frac_all, n_total = self.n_q * self.n_bands nb = self.n_bands - n_syms = len(pg_indices) - self.n_syms_qspace = n_syms - # Build all sparse matrices in Python, then pass to Julia all_rows = [] all_cols = [] all_vals = [] + n_skipped_syms = 0 for i_sym_idx in pg_indices: R_frac = rot_frac_all[i_sym_idx].astype(float) @@ -1559,6 +1701,19 @@ def _build_qspace_symmetries(self, rot_frac_all, trans_frac_all, R_cart = M @ R_frac @ Minv t_cart = M @ t_frac + # Map all q points first: a point-group operation of the unit + # cell need not preserve an anisotropic q grid (e.g. C4 on a + # 2x2x4 supercell). The preserving operations form a subgroup, + # so averaging over them alone is still a valid projector: + # skip the others instead of crashing. + try: + iq_prime_map = [ + find_q_index(R_cart @ self.q_points[jq], self.q_points, bg) + for jq in range(self.n_q)] + except ValueError: + n_skipped_syms += 1 + continue + # Get atom permutation irt = self._get_atom_perm( self.uci_structure, R_cart, t_cart, M, Minv) @@ -1567,10 +1722,9 @@ def _build_qspace_symmetries(self, rot_frac_all, trans_frac_all, for iq in range(self.n_q): q = self.q_points[iq] - Rq = R_cart @ q - # Find iq' matching Rq - iq_prime = find_q_index(Rq, self.q_points, bg) + # iq' matching R q (precomputed above) + iq_prime = iq_prime_map[iq] q_prime = self.q_points[iq_prime] # Build P_uc with Bloch phase factor @@ -1600,6 +1754,19 @@ def _build_qspace_symmetries(self, rot_frac_all, trans_frac_all, all_cols.append(np.array(cols, dtype=np.int32)) all_vals.append(np.array(vals, dtype=np.complex128)) + n_syms = len(all_rows) + self.n_syms_qspace = n_syms + if n_skipped_syms > 0: + warnings.warn( + "{} point-group operations do not preserve the q grid and " + "were excluded from the q-space symmetrization (subgroup of " + "{} operations kept).".format(n_skipped_syms, n_syms)) + if n_syms == 0: + raise ValueError( + "No point-group operation preserves the q grid: cannot " + "build the q-space symmetrization. Use no_sym=True or fix " + "the q grid.") + # Pass to Julia for caching (convert to 1-indexed) for i in range(n_syms): all_rows[i] += 1 @@ -1643,7 +1810,8 @@ def init(self, use_symmetries=True): def load_distributed_tdscha(data_dir, population_id, dyn, T, lo_to_split=None, use_symmetries=True, n_configs=None, - final_dyn=None, final_T=None, **kwargs): + final_dyn=None, final_T=None, fourier_weights=True, + **kwargs): """Load QSpaceLanczos with distributed configurations across MPI ranks. Loads the ensemble on master rank only, then distributes configuration data @@ -1674,6 +1842,17 @@ def load_distributed_tdscha(data_dir, population_id, dyn, T, lo_to_split=None, final_T : float, optional Temperature for weight updates. Defaults to T if not specified. Use this if the final temperature differs from the ensemble temperature. + fourier_weights : bool + If True (default), build the ensemble in the opt-in light mode + (Ensemble(..., qspace_light=True): the (3N,3N) supercell polarization + vectors are never materialized) and route the final_dyn weight update + through the memory-clean q-space update_weights_fourier, when the + Julia Fourier backend is available; otherwise fall back to a standard + ensemble + real-space update_weights. Set False to force the legacy + behavior exactly (standard ensemble, real-space update, exact bitwise + legacy rho, dense Upsilon transient). Direct (non-loader) usage of the + light pipeline: Ensemble(dyn, T, qspace_light=True) + + update_weights_fourier + QSpaceHessian/QSpaceLanczos. **kwargs Additional arguments passed to QSpaceLanczos. @@ -1697,18 +1876,85 @@ def load_distributed_tdscha(data_dir, population_id, dyn, T, lo_to_split=None, rank = comm.Get_rank() n_procs = Parallel.GetNProc() + # Decide before touching the disk. Whatever fourier_weights says, the + # QSpaceLanczos built at the end of this function needs the q-space cache + # API on the ensemble: without it _bloch_transform_ensemble raises, and it + # raises *after* load_bin and a real-space update_weights have already + # spent minutes and a dense Upsilon transient. Failing here keeps the error + # cheap and says what to install. The deprecated private name is accepted + # so that a python-sscha checkout predating the public rename still works. + if not _ensemble_has_qspace_cache_api(sscha.Ensemble.Ensemble): + raise AttributeError( + "This python-sscha does not expose the q-space cache API " + "(refresh_qspace_caches_from_real_space): the q-space Lanczos " + "cannot be built from any ensemble it produces. Install the " + "q-space python-sscha.") + if Parallel.am_i_the_master(): # ========== MASTER (RANK 0) ========== - ensemble = sscha.Ensemble.Ensemble(dyn, T) + # fourier_weights=True (default): the ensemble is built in the OPT-IN + # light mode (Ensemble(..., qspace_light=True)): the (3N,3N) supercell + # polarization vectors are never materialized (not even at + # construction) and the final_dyn weight update goes through the + # memory-clean q-space update_weights_fourier -- the whole load window + # stays free of quadratic transients. The light rho is bitwise + # identical to the standard fourier path and differs from the + # real-space update only at floating-point noise (~1e-12 relative). + # It requires the Julia Fourier backend; if that is unavailable we + # warn and fall back to a standard ensemble + the real-space + # update_weights (exact bitwise legacy rho, dense (3N,3N) transients). + # Pass fourier_weights=False to force the legacy behavior exactly. + # Two independent capabilities are needed, and __JULIA_EXT__ alone + # covers neither reliably: upstream defines it as "juliacall is + # importable", not "the runtime works", and it says nothing about + # whether this python-sscha even accepts qspace_light. Probing the + # constructor signature is what actually decides, and it is what keeps + # this loader working against an unpatched python-sscha instead of + # raising AttributeError from the frozen-attribute hook. + has_light = _ensemble_supports_light() + # qspace_light also needs the CellConstructor side: Ensemble.__setattr__ + # calls DiagonalizeSupercell(q_only=True) as soon as the flag is set, so + # building a light ensemble against an unpatched CellConstructor raises + # TypeError right there -- before __init__ of this class, hence before + # its q_only fallback can do anything about it. + use_fourier = fourier_weights and has_light and _CC_HAS_Q_ONLY and \ + bool(getattr(sscha.Ensemble, "__JULIA_EXT__", False)) + if fourier_weights and not use_fourier: + if not has_light: + warnings.warn( + "fourier_weights=True requested but this python-sscha has " + "no Ensemble(qspace_light=...): falling back to a standard " + "ensemble and the real-space update_weights (dense Upsilon " + "transient). Install the q-space python-sscha for the " + "linear-memory path.") + elif not _CC_HAS_Q_ONLY: + warnings.warn( + "fourier_weights=True requested but this CellConstructor has " + "no DiagonalizeSupercell(q_only=True), so the light ensemble " + "cannot be built: falling back to a standard ensemble and the " + "real-space update_weights (dense Upsilon transient). Install " + "the q-space CellConstructor for the linear-memory path.") + else: + warnings.warn( + "fourier_weights=True requested but the Julia Fourier backend " + "is unavailable; falling back to a standard ensemble and the " + "real-space update_weights (dense Upsilon transient).") + if use_fourier: + ensemble = sscha.Ensemble.Ensemble(dyn, T, qspace_light=True) + else: + ensemble = sscha.Ensemble.Ensemble(dyn, T) if n_configs is not None: ensemble.load_bin(data_dir, population_id, n_configs=n_configs) else: ensemble.load_bin(data_dir, population_id) - # Update weights if final_dyn is provided + # Update weights if final_dyn is provided. if final_dyn is not None: T_for_update = final_T if final_T is not None else T - ensemble.update_weights(final_dyn, T_for_update) + if use_fourier: + ensemble.update_weights_fourier(final_dyn, T_for_update) + else: + ensemble.update_weights(final_dyn, T_for_update) qlanc = QSpaceLanczos(ensemble, lo_to_split=lo_to_split, **kwargs) qlanc.init(use_symmetries=use_symmetries) @@ -1775,7 +2021,13 @@ def load_distributed_tdscha(data_dir, population_id, dyn, T, lo_to_split=None, qlanc._N_eff_global = N_eff_global qlanc._N_local = N_local qlanc.N = N_local - qlanc.N_eff = int(np.sum(qlanc.rho)) + # NOTE: must stay a *float*. Julia normalises its result by + # n_syms * sum(rho_local) (float); _call_julia_qspace_distributed + # multiplies back by self.N_eff to undo exactly that division. + # The cancellation is only exact if N_eff == float(sum(rho_local)), + # so truncating to int silently corrupts the result whenever rho is + # non-integer (i.e. after ensemble.update_weights). + qlanc.N_eff = float(np.sum(qlanc.rho)) # Free unused arrays if hasattr(qlanc, 'X') and qlanc.X is not None: @@ -1830,7 +2082,13 @@ def load_distributed_tdscha(data_dir, population_id, dyn, T, lo_to_split=None, qlanc._N_eff_global = metadata['N_eff_global'] qlanc._N_local = N_local qlanc.N = N_local - qlanc.N_eff = int(np.sum(qlanc.rho)) + # NOTE: must stay a *float*. Julia normalises its result by + # n_syms * sum(rho_local) (float); _call_julia_qspace_distributed + # multiplies back by self.N_eff to undo exactly that division. + # The cancellation is only exact if N_eff == float(sum(rho_local)), + # so truncating to int silently corrupts the result whenever rho is + # non-integer (i.e. after ensemble.update_weights). + qlanc.N_eff = float(np.sum(qlanc.rho)) # Build Julia symmetry cache qlanc.prepare_symmetrization(no_sym=not use_symmetries) diff --git a/Modules/tdscha_qspace.jl b/Modules/tdscha_qspace.jl index 868205f3..9380d338 100644 --- a/Modules/tdscha_qspace.jl +++ b/Modules/tdscha_qspace.jl @@ -79,6 +79,11 @@ function get_d2v_from_R_pert_qspace( # Buffers x_buf = zeros(ComplexF64, n_total) y_buf = zeros(ComplexF64, n_total) + # NOTE: x_rot/y_rot are overwritten by mul! on every (config, sym) + # iteration and never outlive it. Do NOT parallelize that loop + # without making these buffers thread-local. + x_rot = zeros(ComplexF64, n_total) + y_rot = zeros(ComplexF64, n_total) for bigindex in start_index:end_index i_config = div(bigindex - 1, n_syms) + 1 @@ -94,8 +99,8 @@ function get_d2v_from_R_pert_qspace( end # Apply symmetry - x_rot = symmetries[j_sym] * x_buf - y_rot = symmetries[j_sym] * y_buf + mul!(x_rot, symmetries[j_sym], x_buf) + mul!(y_rot, symmetries[j_sym], y_buf) # Views at q_pert x_pert = view(x_rot, (iq_pert-1)*n_bands+1:iq_pert*n_bands) @@ -193,6 +198,11 @@ function get_d2v_from_Y_pert_qspace( # Buffers x_buf = zeros(ComplexF64, n_total) y_buf = zeros(ComplexF64, n_total) + # NOTE: x_rot/y_rot are overwritten by mul! on every (config, sym) + # iteration and never outlive it. Do NOT parallelize that loop + # without making these buffers thread-local. + x_rot = zeros(ComplexF64, n_total) + y_rot = zeros(ComplexF64, n_total) buffer_u = zeros(ComplexF64, n_q, n_bands) for bigindex in start_index:end_index @@ -209,8 +219,8 @@ function get_d2v_from_Y_pert_qspace( end # Apply symmetry - x_rot = symmetries[j_sym] * x_buf - y_rot = symmetries[j_sym] * y_buf + mul!(x_rot, symmetries[j_sym], x_buf) + mul!(y_rot, symmetries[j_sym], y_buf) # Step 1: Compute buffer_u and total_wD4 # buffer_u[iq1, nu1] = sum_nu2 alpha1[p][nu1, nu2] * x_rot[iq2, nu2] @@ -352,6 +362,11 @@ function get_f_from_Y_pert_qspace( # Buffers x_buf = zeros(ComplexF64, n_total) y_buf = zeros(ComplexF64, n_total) + # NOTE: x_rot/y_rot are overwritten by mul! on every (config, sym) + # iteration and never outlive it. Do NOT parallelize that loop + # without making these buffers thread-local. + x_rot = zeros(ComplexF64, n_total) + y_rot = zeros(ComplexF64, n_total) buffer_u = zeros(ComplexF64, n_q, n_bands) for bigindex in start_index:end_index @@ -368,8 +383,8 @@ function get_f_from_Y_pert_qspace( end # Apply symmetry - x_rot = symmetries[j_sym] * x_buf - y_rot = symmetries[j_sym] * y_buf + mul!(x_rot, symmetries[j_sym], x_buf) + mul!(y_rot, symmetries[j_sym], y_buf) # Compute buffer_u and total_sum (same as d2v function) total_sum = zero(ComplexF64) @@ -492,6 +507,11 @@ function get_perturb_averages_qspace_fused( # Buffers (reused each iteration) x_buf = zeros(ComplexF64, n_total) y_buf = zeros(ComplexF64, n_total) + # NOTE: x_rot/y_rot are overwritten by mul! on every (config, sym) + # iteration and never outlive it. Do NOT parallelize that loop + # without making these buffers thread-local. + x_rot = zeros(ComplexF64, n_total) + y_rot = zeros(ComplexF64, n_total) buffer_u = zeros(ComplexF64, n_q, n_bands) for bigindex in start_index:end_index @@ -507,8 +527,8 @@ function get_perturb_averages_qspace_fused( end end - x_rot = symmetries[j_sym] * x_buf - y_rot = symmetries[j_sym] * y_buf + mul!(x_rot, symmetries[j_sym], x_buf) + mul!(y_rot, symmetries[j_sym], y_buf) # === Step 2: D3 weights from R1 perturbation === x_pert = view(x_rot, (iq_pert-1)*n_bands+1:iq_pert*n_bands) diff --git a/tests/test_qspace/test_distributed.py b/tests/test_qspace/test_distributed.py index 2c8c6ef2..942379d4 100644 --- a/tests/test_qspace/test_distributed.py +++ b/tests/test_qspace/test_distributed.py @@ -216,7 +216,10 @@ def test_distributed_hessian(): hess = QH.QSpaceHessian.from_qspace_lanczos(qlanc, verbose=False, use_symmetries=False) pprint("Computing Hessian...") - hess.compute_full_hessian() + # This case is built with the spatial symmetries off, so the mode symmetry + # must be off as well: leaving it on would apply Schur's lemma to a + # symmetry structure the rest of the object does not use. + hess.compute_full_hessian(use_mode_symmetry=False) # Check that we have results for Gamma assert 0 in hess.H_q_dict, "Missing Gamma point in Hessian results" @@ -227,8 +230,17 @@ def test_distributed_hessian(): pprint(f"Hessian eigenvalues at Gamma: {evals}") - # All eigenvalues should be non-negative (for stable system) - assert np.all(evals >= -1e-10), "Negative eigenvalues in Hessian" + # The eigenvalues are not required to be non-negative here: this is a + # 10-configuration ensemble with the symmetries off, so the free energy + # Hessian is genuinely noisy and its lowest eigenvalues come out slightly + # negative (order 1e-6 Ry/bohr^2). What must hold is that nothing blows + # up: the values are finite and small compared to the physical scale. + # (Before the Schur cross-block fix this assert passed only because the + # scalar shortcut filled a fabricated degeneracy, replacing the true + # slightly-negative eigenvalues with zeros.) + assert np.all(np.isfinite(evals)), "Non-finite eigenvalues in Hessian" + assert np.all(evals >= -1e-3), \ + "Hessian eigenvalues far below zero: {}".format(evals) def test_distributed_kpm(): diff --git a/tests/test_qspace/test_neff_cast.py b/tests/test_qspace/test_neff_cast.py new file mode 100644 index 00000000..811844e1 --- /dev/null +++ b/tests/test_qspace/test_neff_cast.py @@ -0,0 +1,164 @@ +"""Regression test for the N_eff integer-truncation bug in +``tdscha.QSpaceLanczos.load_distributed_tdscha``. + +The bug +------- +``load_distributed_tdscha`` used to set, on every rank:: + + qlanc.N_eff = int(np.sum(qlanc.rho)) + +The Julia kernel ``get_perturb_averages_qspace`` returns a result already +divided by ``n_syms * sum(rho_local)`` -- a *float*, see ``tdscha_qspace.jl`` +(``N_eff = sum(rho)``, ``norm_factor = n_syms * N_eff``). +``_call_julia_qspace_distributed`` then multiplies by ``self.N_eff`` to undo +exactly that division, MPI-Allreduces and divides by ``N_eff_global``. The +cancellation is exact only if ``self.N_eff == float(sum(rho_local))``. +Truncating to ``int`` breaks it as soon as ``rho`` is non-integer, i.e. after +``ensemble.update_weights`` -- the normal production situation. The result is +a *silently* wrong (mis-weighted) anharmonic average, hence a wrong Hessian / +Lanczos spectrum. + +The test +-------- +Run under ``mpirun -np 2``. Load the in-repo test ensemble +(``tests/test_julia/data``, 10 configs), call ``update_weights`` with a +slightly different dynamical matrix so that ``rho != 1``, then compute the +TDSCHA Lanczos coefficients twice: + + * path A -- ordinary ``QSpaceLanczos`` (full ensemble on every rank, Julia + work split by ``GoParallel``); ``N_eff`` comes from the base class as + ``np.sum(rho)`` and is therefore correct; + * path B -- ``load_distributed_tdscha`` (config slices per rank), which is + the code path containing the bug. + +The two must agree. On the unfixed code the local weight sums are truncated +(here 4.9461 -> 4 and 4.0589 -> 4), which mis-weights each rank's contribution +by up to ~20%. + +Why the Lanczos and not the free-energy Hessian: the GOLD **size 2** ensemble +is purely harmonic (its forces are exactly odd in u, ``f(u)+f(-u) = 2.2e-15``), +so its anharmonic operator contributes nothing and a Hessian-level +serial-vs-distributed comparison on it is vacuous -- it agrees bitwise even on +the buggy code. GOLD size >= 3 *is* genuinely anharmonic and is the minimum +case for correctness validation. The ``run_FT`` Lanczos on the small in-repo +ensemble populates the two-phonon sector, is far cheaper than a size-3 +Hessian, and its averages go through exactly the same ``N_eff`` +normalisation. + +Usage +----- + mpirun -np 2 python test_neff_cast.py # the actual test + python test_neff_cast.py # re-execs itself under mpirun -np 2 + +It is also collected by pytest as ``test_neff_not_truncated``. +""" + +import os +import subprocess +import sys + +import numpy as np + +import cellconstructor as CC +import cellconstructor.Phonons +import cellconstructor.Settings as Parallel +import sscha.Ensemble +import tdscha.QSpaceLanczos as QL +from tdscha.QSpaceLanczos import load_distributed_tdscha + +# In-repo test ensemble (the one used by tests/test_qspace/test_distributed.py) +ENS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), + "..", "test_julia", "data") +POPULATION = 1 +NQIRR = 3 +T = 250.0 + +N_STEPS = 6 +IQ = 0 +SCALE = 1.05 # perturbation of the dynamical matrix used by update_weights + + +def _load_dyns(): + dyn0 = CC.Phonons.Phonons(os.path.join(ENS_DIR, "dyn_gen_pop1_"), NQIRR) + # A slightly different dyn, so that update_weights yields non-integer rho. + dyn_f = dyn0.Copy() + for i in range(len(dyn_f.dynmats)): + dyn_f.dynmats[i] = dyn_f.dynmats[i] * SCALE + return dyn0, dyn_f + + +def run_mpi(): + """Body of the test; must be executed under mpirun with >= 2 ranks.""" + n_procs = Parallel.GetNProc() + assert n_procs >= 2, "this test must be run with mpirun -np 2 (got %d)" % n_procs + + dyn0, dyn_f = _load_dyns() + + # ---------------- path A: reference (GoParallel, full ensemble) -------- + ens = sscha.Ensemble.Ensemble(dyn0, T) + ens.load_bin(ENS_DIR, POPULATION) + ens.update_weights(dyn_f, T) + + rho = np.asarray(ens.rho, dtype=np.float64) + assert np.abs(rho - np.round(rho)).max() > 1e-6, \ + "rho is (near) integer: the test would not exercise the bug" + + ref = QL.QSpaceLanczos(ens, lo_to_split=None) + ref.ignore_v3 = False + ref.ignore_v4 = False + ref.init(use_symmetries=True) + assert not ref._distributed + band = int(np.argmax(ref.w_q[:, IQ])) + ref.prepare_mode_q(IQ, band) + ref.run_FT(N_STEPS, verbose=False) + a_ref = np.array(ref.a_coeffs, dtype=np.float64) + b_ref = np.array(ref.b_coeffs, dtype=np.float64) + + # ---------------- path B: distributed (the buggy code path) ------------ + dist = load_distributed_tdscha(ENS_DIR, POPULATION, dyn0, T, lo_to_split=None, + use_symmetries=True, + final_dyn=dyn_f, final_T=T) + dist.ignore_v3 = False + dist.ignore_v4 = False + assert dist._distributed + dist.prepare_mode_q(IQ, band) + dist.run_FT(N_STEPS, verbose=False) + a_dis = np.array(dist.a_coeffs, dtype=np.float64) + b_dis = np.array(dist.b_coeffs, dtype=np.float64) + + ok = True + if Parallel.am_i_the_master(): + print("N_eff local = %r (sum rho_local = %.10f)" + % (dist.N_eff, float(np.sum(dist.rho)))) + print("a_ref = %s" % np.array2string(a_ref, precision=10)) + print("a_dis = %s" % np.array2string(a_dis, precision=10)) + for name, x, y in (("a_coeffs", a_ref, a_dis), ("b_coeffs", b_ref, b_dis)): + scale = max(np.abs(x).max(), 1e-30) + rel = np.abs(x - y).max() / scale + print("max |ref - dist| / max|ref| [%s] = %.6e" % (name, rel)) + if not np.allclose(x, y, rtol=1e-8, atol=1e-12 * scale): + ok = False + if not ok: + print("FAIL: distributed Lanczos differs from the reference. " + "N_eff must be float(sum(rho)), not int(sum(rho)).") + else: + print("OK: distributed Lanczos matches the GoParallel reference") + + return ok + + +def _run_under_mpirun(): + env = dict(os.environ) + env["OMP_NUM_THREADS"] = "1" + return subprocess.call(["mpirun", "-np", "2", sys.executable, + os.path.abspath(__file__), "--inner"], env=env) + + +def test_neff_not_truncated(): + assert _run_under_mpirun() == 0 + + +if __name__ == "__main__": + if "--inner" in sys.argv: + sys.exit(0 if run_mpi() else 1) + sys.exit(_run_under_mpirun()) diff --git a/tests/test_qspace/test_qspace_hessian.py b/tests/test_qspace/test_qspace_hessian.py index 7d2ad03b..7cbf4b17 100644 --- a/tests/test_qspace/test_qspace_hessian.py +++ b/tests/test_qspace/test_qspace_hessian.py @@ -230,12 +230,16 @@ def test_hessian_L_operator_timing(): "Hessian L-operator took {:.1f}s per call — too slow".format(t_hessian)) -def test_qspace_hessian_mode_symmetry(): +def test_qspace_hessian_mode_symmetry(capsys): """Verify that mode symmetry optimization gives same Hessian eigenvalues. For each irreducible q-point, computes the Hessian with use_mode_symmetry=False (full solves) and use_mode_symmetry=True (degenerate block reduction), then compares eigenvalues. + + SnTe has a strongly coupled pair of repeated irreps at iq=5 + (cross coupling ~0.7 relative), so this test also checks that the + adaptive repeated-irrep detection actually triggers there. """ try: import tdscha.QSpaceHessian as QH @@ -270,6 +274,14 @@ def test_qspace_hessian_mode_symmetry(): "Mode symmetry optimization changed eigenvalues at iq={}: " "max diff = {:.2e}".format(iq_irr, max_diff)) + # The repeated-irrep branch must have been exercised (SnTe iq=5): + # without this assert the test cannot distinguish the adaptive fill + # from the old scalar shortcut, which was exact here only because + # eigh happened to return aligned bases (U = I). + captured = capsys.readouterr().out + assert "Repeated irreps detected" in captured, ( + "the repeated-irrep detection was expected to trigger on SnTe") + print("=== Mode symmetry optimization test PASSED ===") def test_qspace_hessian_checkpoint(tmp_path): diff --git a/tests/test_qspace/test_schur_fill.py b/tests/test_qspace/test_schur_fill.py new file mode 100644 index 00000000..a91418eb --- /dev/null +++ b/tests/test_qspace/test_schur_fill.py @@ -0,0 +1,238 @@ +"""Unit tests for the adaptive Schur fill of QSpaceHessian. + +By Schur's lemma the cross block of G between two copies of the SAME +irrep is c*U with a unitary intertwiner U that is nontrivial whenever +eigh picks arbitrary bases in the two degenerate subspaces: the scalar +c*I shortcut is exact only for distinct irreps (zero coupling). These +tests exercise _adaptive_schur_fill directly with an exact solver, so +they need no Julia extension and no ensemble data. +""" +import numpy as np +import pytest + +try: + from tdscha.QSpaceHessian import _adaptive_schur_fill +except ModuleNotFoundError: + # Only a genuinely missing dependency is a legitimate skip. A bare + # ImportError would also swallow "cannot import name _adaptive_schur_fill", + # i.e. exactly the regression these tests exist to catch: the suite would + # then report green while never running a single assertion. + pytest.skip("tdscha.QSpaceHessian not importable", allow_module_level=True) + + +def _exact_solver(G_true): + """solve_column callable backed by L = G_true^-1 (exact columns).""" + L = np.linalg.inv(G_true) + + def solve_column(i): + e = np.zeros(G_true.shape[0], dtype=np.complex128) + e[i] = 1.0 + return np.linalg.solve(L, e), 0, 0.0 + + return solve_column + + +def _two_triplets(c12): + """G_true with two 3-dim blocks (a*I, b*I) and cross c12*U, U random.""" + rng = np.random.default_rng(42) + d, nb = 3, 6 + U = np.linalg.qr(rng.normal(size=(d, d)) + + 1j * rng.normal(size=(d, d)))[0] + G = np.zeros((nb, nb), dtype=np.complex128) + G[:d, :d] = 2.3 * np.eye(d) + G[d:, d:] = 4.1 * np.eye(d) + G[:d, d:] = c12 * U + G[d:, :d] = np.conj(c12) * U.conj().T + return G, U + + +def test_repeated_irrep_intertwiner_exact(): + """Two copies of the same irrep with a nontrivial intertwiner: the + coupling must be detected and the fill must reproduce G exactly + (the old c*I shortcut was wrong by ~7% on the eigenvalues here).""" + c12 = 0.7 * np.exp(0.6j) + G_true, U = _two_triplets(c12) + # Sanity: the intertwiner really is nontrivial, so c*I would be wrong + assert np.max(np.abs(c12 * U - c12 * np.eye(3))) > 0.1 + + solve = _exact_solver(G_true) + schedule = [(0, [0, 1, 2]), (3, [3, 4, 5])] + rep_x = {b: solve(b)[0] for b, _ in schedule} + + G = np.zeros_like(G_true) + full = _adaptive_schur_fill(G, schedule, rep_x, solve, 6, 1e-6, True) + assert full == {0, 3}, "repeated-irrep coupling not detected" + + G = (G + G.conj().T) / 2 + assert np.max(np.abs(G - G_true)) < 1e-12 + + ev = np.linalg.eigvalsh(np.linalg.inv(G)) + ev_true = np.linalg.eigvalsh(np.linalg.inv(G_true)) + assert np.max(np.abs(ev - ev_true) / np.abs(ev_true)) < 1e-12 + + +def test_distinct_irreps_keep_scalar_path(): + """Zero coupling (distinct irreps): no full solve, and the scalar + Schur fill is exact.""" + G_true, _ = _two_triplets(0.0) + solve = _exact_solver(G_true) + schedule = [(0, [0, 1, 2]), (3, [3, 4, 5])] + rep_x = {b: solve(b)[0] for b, _ in schedule} + + G = np.zeros_like(G_true) + full = _adaptive_schur_fill(G, schedule, rep_x, solve, 6, 1e-6, True) + assert full == set(), "spurious coupling detected on distinct irreps" + + G = (G + G.conj().T) / 2 + assert np.max(np.abs(G - G_true)) < 1e-12 + + +def test_singletons_and_no_mode_symmetry(): + """All-singleton schedules (use_mode_symmetry False or no degeneracy) + must reproduce the full columns verbatim.""" + rng = np.random.default_rng(7) + A = rng.normal(size=(4, 4)) + 1j * rng.normal(size=(4, 4)) + G_true = A @ A.conj().T + 5 * np.eye(4) + solve = _exact_solver(G_true) + schedule = [(i, [i]) for i in range(4)] + rep_x = {b: solve(b)[0] for b, _ in schedule} + + for ums in (True, False): + G = np.zeros_like(G_true) + full = _adaptive_schur_fill(G, schedule, rep_x, solve, 4, 1e-6, ums) + assert full == set() + assert np.max(np.abs(G - G_true)) < 1e-12 + + +def _block_diag_irrep(consts, dim, mixer=None): + """G with `len(consts)` copies of one dim-dimensional irrep. + + Each copy m contributes consts[m] * I_dim on its diagonal block; `mixer`, + if given, is the (ncopies, ncopies) hermitian matrix of cross-block + constants, so that the block (m, n) is mixer[m, n] * I_dim. This is the + exact structure Schur's lemma allows in a symmetry-adapted basis. + """ + n = len(consts) * dim + G = np.zeros((n, n), dtype=np.complex128) + for m, cm in enumerate(consts): + G[m * dim:(m + 1) * dim, m * dim:(m + 1) * dim] = cm * np.eye(dim) + if mixer is not None: + for m in range(len(consts)): + for n_ in range(len(consts)): + if m != n_: + G[m * dim:(m + 1) * dim, n_ * dim:(n_ + 1) * dim] = \ + mixer[m, n_] * np.eye(dim) + return G + + +def _rotate(G, rng, blocks): + """Rotate each degenerate block by a random unitary. + + This is what eigh does in practice: inside a degenerate subspace the basis + is arbitrary, which is exactly why the cross block is c*U and not c*I. + """ + n = G.shape[0] + U = np.eye(n, dtype=np.complex128) + for b in blocks: + k = len(b) + M = rng.normal(size=(k, k)) + 1j * rng.normal(size=(k, k)) + Q, _ = np.linalg.qr(M) + U[np.ix_(b, b)] = Q + return U.conj().T @ G @ U + + +def test_single_reducible_block_is_detected(): + """A lone degenerate block can already be reducible. + + Two copies of the same irrep degenerate at the SAME frequency land in one + block, so there is no partner block to reveal the coupling. The scalar + shortcut is wrong for it, and the leakage of the representative column + onto the rest of its own block is what exposes it. + """ + rng = np.random.default_rng(11) + block = list(range(4)) + G_true = _block_diag_irrep([2.0, 3.5], dim=2, + mixer=np.array([[0.0, 0.9], [0.9, 0.0]])) + G_true = _rotate(G_true, rng, [block]) + solve = _exact_solver(G_true) + schedule = [(0, block)] + rep_x = {0: solve(0)[0]} + + G = np.zeros_like(G_true) + full = _adaptive_schur_fill(G, schedule, rep_x, solve, 4, 1e-6, True) + assert full == {0}, "reducible single block not detected" + G = (G + G.conj().T) / 2 + assert np.max(np.abs(G - G_true)) < 1e-10 + + +def test_soft_mode_spectator_does_not_raise_the_threshold(): + """A large-norm spectator block must not hide a later coupling. + + A block with a large column norm (a soft mode: the columns of G go as + 1/w^2) must not raise the detection threshold for the pairs examined + after it. Carrying a running maximum of `scale` across the pair loop + makes the threshold monotonically non-decreasing, so such a spectator + sitting between two coupled blocks masks their coupling entirely. The + spectator is therefore at index 1, with the coupled copies at 0 and 2, + and its constant is large so that its column norm dominates. + """ + rng = np.random.default_rng(3) + n = 6 + G_true = np.zeros((n, n), dtype=np.complex128) + # two copies of a 2-dim irrep, coupled, at indices 0-1 and 4-5 + G_true[0:2, 0:2] = 1.0 * np.eye(2) + G_true[4:6, 4:6] = 1.2 * np.eye(2) + G_true[0:2, 4:6] = 3e-4 * np.eye(2) + G_true[4:6, 0:2] = 3e-4 * np.eye(2) + # spectator in between, with a column norm ~1e3 times the coupled pair + # and no coupling of its own + G_true[2:4, 2:4] = 1e3 * np.eye(2) + blocks = [[0, 1], [2, 3], [4, 5]] + G_true = _rotate(G_true, rng, blocks) + G_true = (G_true + G_true.conj().T) / 2 + solve = _exact_solver(G_true) + schedule = [(b[0], b) for b in blocks] + rep_x = {b[0]: solve(b[0])[0] for b in blocks} + + G = np.zeros_like(G_true) + full = _adaptive_schur_fill(G, schedule, rep_x, solve, n, 1e-6, True) + assert full == {0, 4}, "coupling masked by the soft-mode spectator" + G = (G + G.conj().T) / 2 + assert np.max(np.abs(G - G_true)) < 1e-10 + + +def test_detection_is_order_independent(): + """The result must not depend on the order the blocks are listed in. + + The dimension shortcut reads the set of self-reducible blocks, so that set + has to be complete before any pair is examined; deciding both in one pass + makes the outcome depend on the iteration order. + """ + rng = np.random.default_rng(5) + # a self-reducible 4-dim block coupled to a 2-dim block: the pair has + # different dimensions, so it is only examined because one is reducible + G_true = np.zeros((6, 6), dtype=np.complex128) + G_true[:4, :4] = _block_diag_irrep([2.0, 2.6], dim=2, + mixer=np.array([[0.0, 0.8], + [0.8, 0.0]])) + G_true[4:, 4:] = 1.5 * np.eye(2) + G_true[:2, 4:] = 0.4 * np.eye(2) + G_true[4:, :2] = 0.4 * np.eye(2) + blocks_a = [[0, 1, 2, 3], [4, 5]] + G_true = _rotate(G_true, rng, blocks_a) + G_true = (G_true + G_true.conj().T) / 2 + solve = _exact_solver(G_true) + + results = [] + for blocks in (blocks_a, list(reversed(blocks_a))): + schedule = [(b[0], b) for b in blocks] + rep_x = {b[0]: solve(b[0])[0] for b in blocks} + G = np.zeros_like(G_true) + full = _adaptive_schur_fill(G, schedule, rep_x, solve, 6, 1e-6, True) + G = (G + G.conj().T) / 2 + results.append((full, np.max(np.abs(G - G_true)))) + + assert results[0][0] == results[1][0], "detection depends on block order" + assert results[0][0] == {0, 4}, "expected both coupled blocks to be solved" + for _, err in results: + assert err < 1e-10