QSpace: fix the Schur fill between repeated irreps, and make the q-space path linear in memory

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)

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)

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():

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())

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):