From 1ccf33d99eac875da085827c8ec3cc25319f51c1 Mon Sep 17 00:00:00 2001 From: julpe Date: Tue, 9 Jun 2026 12:16:37 +0200 Subject: [PATCH 1/5] Added calculation of kinetic and potential energy and fixed index flip in local diagrams for Eliashberg equation. --- README.md | 9 ++++----- dgamore/DGAmore.py | 14 +------------- dgamore/eliashberg_solver.py | 3 ++- dgamore/greens_function.py | 32 +++++++++++++++++++++----------- dgamore/nonlocal_sde.py | 18 ++++++++++++++++-- 5 files changed, 44 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index b5a4d6f5..514a28dd 100644 --- a/README.md +++ b/README.md @@ -133,7 +133,6 @@ mpirun/mpiexec -np $SLURM_NTASKS DGAmore.py -p "" -c " Tuple[LocalFourPoi gchi0_loc_pp_w0 = ( BubbleGenerator.create_generalized_chi0_pp_w0(g_dmft, gchi_ud_loc_pp_w0.niv, config.sys.beta) .extend_vn_to_diagonal() + .permute_orbitals("abcd->acbd") .flip_frequency_axis(-1, False) ) gamma_ud_loc_pp_w0 = config.sys.beta**2 * ( (gchi_ud_loc_pp_w0 - gchi0_loc_pp_w0).invert() + gchi0_loc_pp_w0.invert() - ) + ).permute_orbitals("abcd->acbd") gamma_ud_loc_pp_w0.save(output_dir=config.output.eliashberg_path, name="gamma_ud_loc_pp_w0") diff --git a/dgamore/greens_function.py b/dgamore/greens_function.py index 003aadbb..3099165c 100644 --- a/dgamore/greens_function.py +++ b/dgamore/greens_function.py @@ -105,16 +105,6 @@ def __init__( # config.sys.n, config.sys.occ = self._get_fill() config.sys.n, config.sys.occ, config.sys.occ_k = self.get_fill_nonlocal() - @property - def e_kin(self): - """ - Returns the kinetic energy of the system, see Eq (22) in G. Rohringer & A. Toschi - PHYSICAL REVIEW B 94, 125144 (2016). - """ - ekin = 2 / config.sys.beta * np.sum(np.mean(self._ek[..., None] * self.mat, axis=(0, 1, 2))) - assert np.abs(ekin.imag) < 1e-8, "Kinetic energy must be real." - return ekin.real - @property def ek(self) -> np.ndarray: """ @@ -198,7 +188,7 @@ def map_to_full_bz(self, k_grid: KGrid, nq: tuple = None): def transpose_orbitals(self): r""" - Transposes the orbitals of the Green's function object like :math:`G_{ab}^k -> G_{ba}^k. + Transposes the orbitals of the Green's function object like :math:`G_{ab}^k -> G_{ba}^k`. """ return self.permute_orbitals("ab->ba") @@ -239,6 +229,26 @@ def get_fill_nonlocal(self) -> tuple[float, np.ndarray, np.ndarray]: n_el = 2.0 * np.trace(occ_mean).real return n_el, occ_mean, occ_k + def get_ekin_total(self) -> float: + r""" + Returns the kinetic energy calculated from the Green's function and the k-dependent occupation. + :math:`E_{kin} = \sum_{\sigma\vec{k}ab} \varepsilon(\vec{k})_{ab} n(\vec{k})_{ba}`. + """ + return 2 * np.sum(self._ek * config.sys.occ_k.swapaxes(-1, -2)).real / config.lattice.k_grid.nk_tot + + def get_epot_total(self) -> float: + r""" + Returns the potential energy calculated from the Green's function and the self-energy. + :math:`E_{pot} = 1/(2\beta) \sum_{\sigma\vec{k}\nu ab} \Sigma(\vec{k},\nu)_{ab} G(\vec{k},\nu)_{ba}`. + """ + return ( + (self._sigma.decompress_q_dimension().mat * self.decompress_q_dimension().transpose_orbitals().mat) + .sum() + .real + / config.sys.beta + / config.lattice.k_grid.nk_tot + ) + def _get_gfull_mat(self) -> np.ndarray: """ Returns the full Green's function. diff --git a/dgamore/nonlocal_sde.py b/dgamore/nonlocal_sde.py index e8ba1a6f..0f558a0a 100644 --- a/dgamore/nonlocal_sde.py +++ b/dgamore/nonlocal_sde.py @@ -760,7 +760,7 @@ def calculate_self_energy_q( config.sys.n, config.sys.occ, config.sys.occ_k = giwk_full.get_fill_nonlocal() giwk_full.cut_niv(niv_cut) - if sigma_old is sigma_dmft: + if np.allclose(sigma_old.cut_niv(config.box.niv_core).mat, sigma_dmft.cut_niv(config.box.niv_core).mat): giwk_full.save(output_dir=config.output.output_path, name="g_latt_dmft") config.sys.n, config.sys.occ, config.sys.occ_k = comm.bcast( (config.sys.n, config.sys.occ, config.sys.occ_k), root=0 @@ -774,7 +774,13 @@ def calculate_self_energy_q( v_nonloc_full = deepcopy(v_nonloc) v_nonloc = v_nonloc.reduce_q(my_irr_q_list) + # old_mixing_history_length = config.self_consistency.mixing_history_length + # config.self_consistency.mixing_history_length = -1 + for current_iter in range(starting_iter + 1, starting_iter + config.self_consistency.max_iter + 1): + # if config.self_consistency.mixing_history_length < old_mixing_history_length: + # config.self_consistency.mixing_history_length += 1 + logger.info("----------------------------------------") logger.info(f"Starting iteration {current_iter}.") logger.info("----------------------------------------") @@ -916,9 +922,17 @@ def calculate_self_energy_q( # calculate new occupation matrix from new Green's function (outside asympt region it is the DMFT # lattice Green's function) _, config.sys.occ, config.sys.occ_k = giwk_occ.get_fill_nonlocal() # n should not change + + ekin = giwk_occ.get_ekin_total() + logger.info(f"Kinetic energy: {ekin:.4f} [t or eV].") + + epot = giwk_occ.get_epot_total() + logger.info(f"Potential energy: {epot:.4f} [t or eV].") + logger.info(f"Total energy: {(ekin + epot):.4f} [t or eV].") config.sys.occ, config.sys.occ_k = comm.bcast((config.sys.occ, config.sys.occ_k), root=0) + if config.self_consistency.max_iter > 1: - logger.info("Updated occupation matrix with new Green's function.") + logger.info("Updated occupation matrix from new Green's function.") if comm.rank == 0: sigma_new.save(name=f"sigma_dga_iteration_{current_iter}", output_dir=config.output.output_path) From e8d40ce12cf3e87580777d11ef5a04b97f6bd930 Mon Sep 17 00:00:00 2001 From: julpe Date: Tue, 9 Jun 2026 12:37:33 +0200 Subject: [PATCH 2/5] Fixed tests. --- tests/test_eliashberg_end_to_end.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_eliashberg_end_to_end.py b/tests/test_eliashberg_end_to_end.py index 79010ee8..3947aa61 100644 --- a/tests/test_eliashberg_end_to_end.py +++ b/tests/test_eliashberg_end_to_end.py @@ -129,5 +129,5 @@ def test_eliashberg_equation_with_local_part(setup, niw_core, niv_core, niv_shel lambdas_sing, lambdas_trip, gaps_sing, gaps_trip = eliashberg_solver.solve( g_dga, g_dmft, u_loc, v_nonloc, comm_mock ) - assert np.allclose(lambdas_sing, np.array([15.80373132, 14.68344248, 12.59236782, 10.8154743]), atol=1e-4) + assert np.allclose(lambdas_sing, np.array([15.80373154, 14.68344217, 12.59137005, 10.81553018]), atol=1e-4) assert np.allclose(lambdas_trip, np.array([6.70800083, 6.70799431, 6.45388305, 6.45387905]), atol=1e-4) From b7b7d5130f984445081669a5dd632a2967e4aa39 Mon Sep 17 00:00:00 2001 From: julpe Date: Tue, 9 Jun 2026 15:02:35 +0200 Subject: [PATCH 3/5] Fixed tests. --- dgamore/greens_function.py | 51 ++++++++++++++++++++++++++++------- dgamore/nonlocal_sde.py | 54 ++++++++++++++++++++++---------------- 2 files changed, 72 insertions(+), 33 deletions(-) diff --git a/dgamore/greens_function.py b/dgamore/greens_function.py index 3099165c..5de9b425 100644 --- a/dgamore/greens_function.py +++ b/dgamore/greens_function.py @@ -229,26 +229,57 @@ def get_fill_nonlocal(self) -> tuple[float, np.ndarray, np.ndarray]: n_el = 2.0 * np.trace(occ_mean).real return n_el, occ_mean, occ_k - def get_ekin_total(self) -> float: + def get_ekin(self) -> float: r""" Returns the kinetic energy calculated from the Green's function and the k-dependent occupation. :math:`E_{kin} = \sum_{\sigma\vec{k}ab} \varepsilon(\vec{k})_{ab} n(\vec{k})_{ba}`. """ return 2 * np.sum(self._ek * config.sys.occ_k.swapaxes(-1, -2)).real / config.lattice.k_grid.nk_tot - def get_epot_total(self) -> float: + def get_epot(self, niv_asympt: int = 50000) -> float: r""" - Returns the potential energy calculated from the Green's function and the self-energy. - :math:`E_{pot} = 1/(2\beta) \sum_{\sigma\vec{k}\nu ab} \Sigma(\vec{k},\nu)_{ab} G(\vec{k},\nu)_{ba}`. + Galitskii-Migdal potential energy, moment-corrected. + + E_pot = sum_k Tr[Sigma_inf rho_k] # Hartree (exact, conv. factor) + + 1/beta sum_{k,v} Tr[(Sigma - Sigma_inf) G] # in-box correlation part + + 1/beta [ sum_big - sum_box ] Tr[(Sigma_mod - Sigma_inf) G_mod] # analytic 1/v^2 tail + + with Sigma_mod - Sigma_inf = -smom1/(iv) and G_mod = [iv + mu - e_k - Sigma_inf]^-1. + The model subtraction cancels the 1/v^2 tail of the correlation sum (remainder ~1/v^4), + while sum_big supplies the part beyond the stored box. """ - return ( - (self._sigma.decompress_q_dimension().mat * self.decompress_q_dimension().transpose_orbitals().mat) - .sum() - .real - / config.sys.beta - / config.lattice.k_grid.nk_tot + smom0, smom1 = self._sigma.smom # Sigma_inf, first tail coeff; both [o1, o2] + + # 1) Hartree: physical (tail-corrected) occupation, convergence factor exact. + e_hartree = np.sum(smom0[None, None, None] * config.sys.occ_k.swapaxes(-1, -2)).real + + # 2) In-box correlation part: Tr[(Sigma - Sigma_inf) G], Sigma_inf already counted above. + dsigma = self._sigma.decompress_q_dimension().mat - smom0[..., None] + g_ba = self.decompress_q_dimension().transpose_orbitals().mat + e_corr = (dsigma * g_ba).sum().real / config.sys.beta + + # 3) Analytic 1/v^2 model tail: replace the truncated box value by the large-box one. + e_tail = self._model_epot(smom0, smom1, niv_asympt, config.sys.beta) - self._model_epot( + smom0, smom1, self.niv, config.sys.beta ) + return (e_hartree + e_corr + e_tail) / config.lattice.k_grid.nk_tot + + def _model_epot(self, smom0, smom1, niv, beta): + """ + Analytic 1/v^2 model tail + """ + h = (self._ek + smom0[None, None, None]).reshape(self.nq_tot, self.n_bands, self.n_bands) + lam, u = np.linalg.eig(h) # once per k + u_inv = np.linalg.inv(u) + smom1_rot = u_inv @ smom1 @ u # rotate tail coeff into eigenbasis + + iv = 1j * MFHelper.vn(niv, beta) + g_diag = 1.0 / (iv[None, :] + config.sys.mu - lam[:, :, None]) # [k, band, v] + # Tr[(-smom1/iv) G_mod] = -sum_i (smom1_rot)_ii * g_diag_i / iv + integrand = -np.einsum("kii,kiv->kv", smom1_rot, g_diag) / iv[None, :] + return integrand.sum().real / beta + def _get_gfull_mat(self) -> np.ndarray: """ Returns the full Green's function. diff --git a/dgamore/nonlocal_sde.py b/dgamore/nonlocal_sde.py index 0f558a0a..49196d14 100644 --- a/dgamore/nonlocal_sde.py +++ b/dgamore/nonlocal_sde.py @@ -774,13 +774,9 @@ def calculate_self_energy_q( v_nonloc_full = deepcopy(v_nonloc) v_nonloc = v_nonloc.reduce_q(my_irr_q_list) - # old_mixing_history_length = config.self_consistency.mixing_history_length - # config.self_consistency.mixing_history_length = -1 + prev_residual = np.inf for current_iter in range(starting_iter + 1, starting_iter + config.self_consistency.max_iter + 1): - # if config.self_consistency.mixing_history_length < old_mixing_history_length: - # config.self_consistency.mixing_history_length += 1 - logger.info("----------------------------------------") logger.info(f"Starting iteration {current_iter}.") logger.info("----------------------------------------") @@ -896,7 +892,7 @@ def calculate_self_energy_q( sigma_new = sigma_new.concatenate_self_energies(sigma_dmft) logger.info("Applying mixing strategy to the self-energy.") - sigma_new = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter) + sigma_new, sigma_residual = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter) old_mu = mu_history[-1] if comm.rank == 0: @@ -923,10 +919,10 @@ def calculate_self_energy_q( # lattice Green's function) _, config.sys.occ, config.sys.occ_k = giwk_occ.get_fill_nonlocal() # n should not change - ekin = giwk_occ.get_ekin_total() + ekin = giwk_occ.get_ekin() logger.info(f"Kinetic energy: {ekin:.4f} [t or eV].") - epot = giwk_occ.get_epot_total() + epot = giwk_occ.get_epot() logger.info(f"Potential energy: {epot:.4f} [t or eV].") logger.info(f"Total energy: {(ekin + epot):.4f} [t or eV].") config.sys.occ, config.sys.occ_k = comm.bcast((config.sys.occ, config.sys.occ_k), root=0) @@ -949,15 +945,15 @@ def calculate_self_energy_q( logger.info("Checking self-consistency convergence.") if comm.rank == 0 and current_iter > starting_iter + 1: - niv_start = sigma_new.niv - niv_end = niv_start + int(np.ceil(config.box.niv_core / 5)) - - sigma_converged = np.allclose( - sigma_old.compress_q_dimension()[..., niv_start:niv_end], - sigma_new.compress_q_dimension()[..., niv_start:niv_end], - atol=config.self_consistency.epsilon, + residual_change = ( + abs(sigma_residual - prev_residual) / abs(sigma_residual) if sigma_residual != 0 else np.inf + ) + sigma_converged = residual_change < config.self_consistency.epsilon + logger.info( + f"Self-energy convergence: {sigma_converged} " + f"(residual_change={residual_change:.3e}, rel_res={sigma_residual:.3e}, " + f"epsilon={config.self_consistency.epsilon:.3e})." ) - logger.info(f"Self-energy convergence: {sigma_converged}.") mu_converged = ( abs(mu_history[-1] - mu_history[-2]) < np.pi / (10 * config.sys.beta) @@ -968,6 +964,7 @@ def calculate_self_energy_q( else: converged = False converged = comm.bcast(converged) + prev_residual = sigma_residual sigma_old = sigma_new if converged: @@ -992,7 +989,7 @@ def calculate_self_energy_q( def apply_mixing_strategy( sigma_new: SelfEnergy, sigma_old: SelfEnergy, sigma_dmft: SelfEnergy, current_iter: int -) -> SelfEnergy: +) -> tuple[SelfEnergy, float]: """ Applies the mixing strategy for the self-consistency loop. The mixing strategy is defined in the config file and is either 'linear' or 'pulay'. @@ -1052,7 +1049,7 @@ def get_result(idx: int): mask = s > cutoff if not np.any(mask): logger.warning("Pulay SVD ill-conditioned - falling back to linear mixing.") - return alpha * sigma_new + (1 - alpha) * sigma_old + return alpha * sigma_new + (1 - alpha) * sigma_old, rel_res coeffs = vh[mask].T @ ((u[:, mask].T @ f_i) / s[mask]) # Pulay update: x_{n+1} = x_n + alpha*f_i - (R + alpha*F) @ c @@ -1070,7 +1067,7 @@ def get_result(idx: int): f"Pulay mixing applied (m={n_hist}, alpha={alpha:.3f}, norm_f={norm_f:.3e}, rel_res={rel_res:.3e})." ) - return sigma_new + return sigma_new, rel_res if config.self_consistency.mixing_strategy.lower() == "anderson" and current_iter > n_hist: last_sigmas = read_last_n_sigmas_from_files( n_hist, config.output.output_path, config.self_consistency.previous_sc_path @@ -1132,7 +1129,7 @@ def get_result(idx: int): except np.linalg.LinAlgError: logger.warning("Anderson SVD failed - falling back to linear mixing.") - return alpha * sigma_new + (1 - alpha) * sigma_old + return alpha * sigma_new + (1 - alpha) * sigma_old, rel_res # Undamped Anderson proposal: x_n + f_n - (dX + dF) @ c x_n = flat(last_proposals[-1]) @@ -1156,10 +1153,21 @@ def get_result(idx: int): f"Anderson acceleration applied (m={n_hist}, alpha={alpha:.3f}, norm_f={norm_f:.3e}, rel_res={rel_res:.3e})." ) - return sigma_new + return sigma_new, rel_res + + niv = sigma_new.current_shape[-1] // 2 + niv_core = config.box.niv_core + sl = slice(niv - niv_core, niv + niv_core) + f_lin = (sigma_new.mat[..., sl] - sigma_old.mat[..., sl]).flatten() + f_vec = np.concatenate([f_lin.real, f_lin.imag]) + norm_f = np.linalg.norm(f_vec) + x_lin = sigma_old.mat[..., sl].flatten() + norm_x = np.linalg.norm(np.concatenate([x_lin.real, x_lin.imag])) + rel_res = norm_f / norm_x if norm_x > 0 else np.inf sigma_new = config.self_consistency.mixing * sigma_new + (1 - config.self_consistency.mixing) * sigma_old logger.info( - f"Sigma linearly mixed with previous iteration using a mixing parameter of {config.self_consistency.mixing}." + f"Sigma linearly mixed with previous iteration using a mixing parameter of " + f"{config.self_consistency.mixing} (norm_f={norm_f:.3e}, rel_res={rel_res:.3e})." ) - return sigma_new + return sigma_new, rel_res From 26f214ccfd159eb2094e36cbac7b1d8843db4049 Mon Sep 17 00:00:00 2001 From: julpe Date: Tue, 9 Jun 2026 15:39:12 +0200 Subject: [PATCH 4/5] Fixed tests. --- tests/test_mixing.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/tests/test_mixing.py b/tests/test_mixing.py index 19bdd43e..25595b6c 100644 --- a/tests/test_mixing.py +++ b/tests/test_mixing.py @@ -14,7 +14,6 @@ from dgamore.self_energy import SelfEnergy from dgamore.nonlocal_sde import apply_mixing_strategy - BETA = 10.0 NB = 1 NK = (1, 1, 1) @@ -86,7 +85,7 @@ def run_pulay( patch_config(strategy="pulay", mixing=mixing, n_hist=n_hist, niv_core=niv_core, nk_tot=nk_tot), patch("dgamore.nonlocal_sde.read_last_n_sigmas_from_files", return_value=file_sigmas), ): - return apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=current_iter) + return apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=current_iter)[0] def run_anderson( @@ -104,7 +103,7 @@ def run_anderson( patch_config(strategy="anderson", mixing=mixing, n_hist=n_hist, niv_core=niv_core, nk_tot=nk_tot), patch("dgamore.nonlocal_sde.read_last_n_sigmas_from_files", return_value=file_sigmas), ): - return apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=current_iter) + return apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=current_iter)[0] def test_linear_mixing_basic(): @@ -114,7 +113,7 @@ def test_linear_mixing_basic(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=0.5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] np.testing.assert_allclose(result.mat, 1.0, atol=1e-5) @@ -126,7 +125,7 @@ def test_linear_mixing_alpha_zero(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=0.0): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] np.testing.assert_allclose(result.mat, 1.0, atol=1e-5) @@ -138,7 +137,7 @@ def test_linear_mixing_alpha_one(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=1.0): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] np.testing.assert_allclose(result.mat, 5.0, atol=1e-5) @@ -150,7 +149,7 @@ def test_linear_mixing_complex(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=0.5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] np.testing.assert_allclose(result.mat, 1.0 + 1.0j, atol=1e-5) @@ -162,7 +161,7 @@ def test_linear_mixing_returns_self_energy_instance(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=0.5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] assert isinstance(result, SelfEnergy) @@ -174,7 +173,7 @@ def test_pulay_falls_back_to_linear_when_iter_too_small(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="pulay", mixing=0.5, n_hist=5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=3) + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=3)[0] np.testing.assert_allclose(result.mat, 1.0, atol=1e-5) @@ -281,7 +280,7 @@ def test_anderson_falls_back_to_linear_when_iter_too_small(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="anderson", mixing=0.5, n_hist=5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=3) + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=3)[0] np.testing.assert_allclose(result.mat, 1.0, atol=1e-5) From baf8283eb2e51c380d1276c2b708e805263306fb Mon Sep 17 00:00:00 2001 From: julpe Date: Thu, 11 Jun 2026 10:57:02 +0200 Subject: [PATCH 5/5] Self-energy will now be updated differently and convergence will be tested on the relative residual of the new and old iteration. --- dgamore/n_point_base.py | 15 ++++- dgamore/nonlocal_sde.py | 66 ++++++++-------------- tests/test_mixing.py | 18 +++--- tests/test_n_point_base.py | 112 ++++++++++++++++++++++++------------- 4 files changed, 121 insertions(+), 90 deletions(-) diff --git a/dgamore/n_point_base.py b/dgamore/n_point_base.py index d4856730..f3995bdd 100644 --- a/dgamore/n_point_base.py +++ b/dgamore/n_point_base.py @@ -440,7 +440,7 @@ def find_q(self, q: tuple[int, int, int] = (0, 0, 0)): def filter_q_index(self, index: int = 0): r""" Filters the object to the given index of the momentum dimension and returns a copy. Acts like a filter. - Makes it possible to use e.g. only the first component of a non-local object. + Makes it possible to use e.g. only the n-th q-component of a non-local object. """ if not self.has_compressed_q_dimension: self.compress_q_dimension() @@ -451,6 +451,19 @@ def filter_q_index(self, index: int = 0): copy._nq = (1, 1, 1) return copy + def q_mean(self): + r""" + Averages over the momentum dimension and returns a copy of the object with nq = (1,1,1). + """ + copy = deepcopy(self) + if self.has_compressed_q_dimension: + copy.mat = np.mean(copy.mat, axis=0)[None, ...] + else: + copy.mat = np.mean(copy.mat, axis=(0, 1, 2))[None, None, None, ...] + copy.update_original_shape() + copy._nq = (1, 1, 1) + return copy + def map_to_full_bz(self, k_grid: KGrid, nq: tuple = None): """ Maps to full BZ using k_grid's inverse map and precomputed orbital rotation tensors. diff --git a/dgamore/nonlocal_sde.py b/dgamore/nonlocal_sde.py index 49196d14..0a9268e7 100644 --- a/dgamore/nonlocal_sde.py +++ b/dgamore/nonlocal_sde.py @@ -774,8 +774,6 @@ def calculate_self_energy_q( v_nonloc_full = deepcopy(v_nonloc) v_nonloc = v_nonloc.reduce_q(my_irr_q_list) - prev_residual = np.inf - for current_iter in range(starting_iter + 1, starting_iter + config.self_consistency.max_iter + 1): logger.info("----------------------------------------") logger.info(f"Starting iteration {current_iter}.") @@ -890,9 +888,20 @@ def calculate_self_energy_q( # calculated in this code and add the smooth dmft self-energy sigma_new += delta_sigma sigma_new = sigma_new.concatenate_self_energies(sigma_dmft) + delta_sigma = sigma_dmft.cut_niv(config.box.niv_core) - sigma_new.q_mean().cut_niv(config.box.niv_core) logger.info("Applying mixing strategy to the self-energy.") - sigma_new, sigma_residual = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter) + sigma_new = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter) + + # Canonical self-consistency residual. This is the relative (L2) residual used for the convergence check + sigma_new = sigma_new.compress_q_dimension() + sigma_old = sigma_old.compress_q_dimension() + + sigma_new_test = sigma_new.mat[..., sigma_new.niv : sigma_new.niv + config.box.niv_core] + sigma_old_test = sigma_old.mat[..., sigma_new.niv : sigma_old.niv + config.box.niv_core] + diff = (sigma_new_test - sigma_old_test).ravel() + norm_x = np.linalg.norm(np.concatenate([sigma_old_test.real.ravel(), sigma_old_test.imag.ravel()])) + relative_residual = np.linalg.norm(np.concatenate([diff.real, diff.imag])) / norm_x old_mu = mu_history[-1] if comm.rank == 0: @@ -945,14 +954,10 @@ def calculate_self_energy_q( logger.info("Checking self-consistency convergence.") if comm.rank == 0 and current_iter > starting_iter + 1: - residual_change = ( - abs(sigma_residual - prev_residual) / abs(sigma_residual) if sigma_residual != 0 else np.inf - ) - sigma_converged = residual_change < config.self_consistency.epsilon + sigma_converged = abs(relative_residual) < config.self_consistency.epsilon logger.info( f"Self-energy convergence: {sigma_converged} " - f"(residual_change={residual_change:.3e}, rel_res={sigma_residual:.3e}, " - f"epsilon={config.self_consistency.epsilon:.3e})." + f"(relative residual={relative_residual:.3e}, epsilon={config.self_consistency.epsilon:.3e})." ) mu_converged = ( @@ -964,7 +969,6 @@ def calculate_self_energy_q( else: converged = False converged = comm.bcast(converged) - prev_residual = sigma_residual sigma_old = sigma_new if converged: @@ -989,7 +993,7 @@ def calculate_self_energy_q( def apply_mixing_strategy( sigma_new: SelfEnergy, sigma_old: SelfEnergy, sigma_dmft: SelfEnergy, current_iter: int -) -> tuple[SelfEnergy, float]: +) -> SelfEnergy: """ Applies the mixing strategy for the self-consistency loop. The mixing strategy is defined in the config file and is either 'linear' or 'pulay'. @@ -1040,8 +1044,6 @@ def get_result(idx: int): f_i[:n_total] = iter_diff.real f_i[n_total:] = iter_diff.imag norm_f = np.linalg.norm(f_i) - norm_x = np.linalg.norm(get_proposal(-1)) - rel_res = norm_f / norm_x if norm_x > 0 else np.inf # Solve min||F @ c - f_i|| via truncated-SVD pseudoinverse (drops collinear directions) u, s, vh = np.linalg.svd(f_matrix, full_matrices=False) @@ -1049,7 +1051,7 @@ def get_result(idx: int): mask = s > cutoff if not np.any(mask): logger.warning("Pulay SVD ill-conditioned - falling back to linear mixing.") - return alpha * sigma_new + (1 - alpha) * sigma_old, rel_res + return alpha * sigma_new + (1 - alpha) * sigma_old coeffs = vh[mask].T @ ((u[:, mask].T @ f_i) / s[mask]) # Pulay update: x_{n+1} = x_n + alpha*f_i - (R + alpha*F) @ c @@ -1063,11 +1065,9 @@ def get_result(idx: int): # Update the new self energy sigma_new.mat[..., niv - niv_core : niv + niv_core] = get_proposal(-1).reshape(shape) + update.reshape(shape) - logger.info( - f"Pulay mixing applied (m={n_hist}, alpha={alpha:.3f}, norm_f={norm_f:.3e}, rel_res={rel_res:.3e})." - ) + logger.info(f"Pulay mixing applied (m={n_hist}, alpha={alpha:.3f}, norm_f={norm_f:.3e}).") - return sigma_new, rel_res + return sigma_new if config.self_consistency.mixing_strategy.lower() == "anderson" and current_iter > n_hist: last_sigmas = read_last_n_sigmas_from_files( n_hist, config.output.output_path, config.self_consistency.previous_sc_path @@ -1092,9 +1092,6 @@ def get_result(idx: int): f_curr = flat(last_results[-1]) - flat(last_proposals[-1]) f_vec = np.concatenate([f_curr.real, f_curr.imag]) norm_f = np.linalg.norm(f_vec) - x_curr = flat(last_proposals[-1]) - norm_x = np.linalg.norm(np.concatenate([x_curr.real, x_curr.imag])) - rel_res = norm_f / norm_x if norm_x > 0 else np.inf # Build dX and dF matrices (n_hist columns) # dX[:,i] = x_{n-i} - x_{n-i-1} (proposal differences) @@ -1129,7 +1126,7 @@ def get_result(idx: int): except np.linalg.LinAlgError: logger.warning("Anderson SVD failed - falling back to linear mixing.") - return alpha * sigma_new + (1 - alpha) * sigma_old, rel_res + return alpha * sigma_new + (1 - alpha) * sigma_old # Undamped Anderson proposal: x_n + f_n - (dX + dF) @ c x_n = flat(last_proposals[-1]) @@ -1149,25 +1146,10 @@ def get_result(idx: int): sigma_new.mat[..., sl] = candidate.reshape(shape) - logger.info( - f"Anderson acceleration applied (m={n_hist}, alpha={alpha:.3f}, norm_f={norm_f:.3e}, rel_res={rel_res:.3e})." - ) + logger.info(f"Anderson acceleration applied (m={n_hist}, alpha={alpha:.3f}, norm_f={norm_f:.3e}).") - return sigma_new, rel_res + return sigma_new - niv = sigma_new.current_shape[-1] // 2 - niv_core = config.box.niv_core - sl = slice(niv - niv_core, niv + niv_core) - f_lin = (sigma_new.mat[..., sl] - sigma_old.mat[..., sl]).flatten() - f_vec = np.concatenate([f_lin.real, f_lin.imag]) - norm_f = np.linalg.norm(f_vec) - x_lin = sigma_old.mat[..., sl].flatten() - norm_x = np.linalg.norm(np.concatenate([x_lin.real, x_lin.imag])) - rel_res = norm_f / norm_x if norm_x > 0 else np.inf - - sigma_new = config.self_consistency.mixing * sigma_new + (1 - config.self_consistency.mixing) * sigma_old - logger.info( - f"Sigma linearly mixed with previous iteration using a mixing parameter of " - f"{config.self_consistency.mixing} (norm_f={norm_f:.3e}, rel_res={rel_res:.3e})." - ) - return sigma_new, rel_res + sigma_new = alpha * sigma_new + (1 - alpha) * sigma_old + logger.info(f"Sigma linearly mixed (m=1, alpha={alpha}).") + return sigma_new diff --git a/tests/test_mixing.py b/tests/test_mixing.py index 25595b6c..cc3c0824 100644 --- a/tests/test_mixing.py +++ b/tests/test_mixing.py @@ -85,7 +85,7 @@ def run_pulay( patch_config(strategy="pulay", mixing=mixing, n_hist=n_hist, niv_core=niv_core, nk_tot=nk_tot), patch("dgamore.nonlocal_sde.read_last_n_sigmas_from_files", return_value=file_sigmas), ): - return apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=current_iter)[0] + return apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=current_iter) def run_anderson( @@ -103,7 +103,7 @@ def run_anderson( patch_config(strategy="anderson", mixing=mixing, n_hist=n_hist, niv_core=niv_core, nk_tot=nk_tot), patch("dgamore.nonlocal_sde.read_last_n_sigmas_from_files", return_value=file_sigmas), ): - return apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=current_iter)[0] + return apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=current_iter) def test_linear_mixing_basic(): @@ -113,7 +113,7 @@ def test_linear_mixing_basic(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=0.5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) np.testing.assert_allclose(result.mat, 1.0, atol=1e-5) @@ -125,7 +125,7 @@ def test_linear_mixing_alpha_zero(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=0.0): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) np.testing.assert_allclose(result.mat, 1.0, atol=1e-5) @@ -137,7 +137,7 @@ def test_linear_mixing_alpha_one(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=1.0): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) np.testing.assert_allclose(result.mat, 5.0, atol=1e-5) @@ -149,7 +149,7 @@ def test_linear_mixing_complex(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=0.5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) np.testing.assert_allclose(result.mat, 1.0 + 1.0j, atol=1e-5) @@ -161,7 +161,7 @@ def test_linear_mixing_returns_self_energy_instance(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="linear", mixing=0.5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1)[0] + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=1) assert isinstance(result, SelfEnergy) @@ -173,7 +173,7 @@ def test_pulay_falls_back_to_linear_when_iter_too_small(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="pulay", mixing=0.5, n_hist=5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=3)[0] + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=3) np.testing.assert_allclose(result.mat, 1.0, atol=1e-5) @@ -280,7 +280,7 @@ def test_anderson_falls_back_to_linear_when_iter_too_small(): sigma_dmft = make_sigma(0.0) with patch_config(strategy="anderson", mixing=0.5, n_hist=5): - result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=3)[0] + result = apply_mixing_strategy(sigma_new, sigma_old, sigma_dmft, current_iter=3) np.testing.assert_allclose(result.mat, 1.0, atol=1e-5) diff --git a/tests/test_n_point_base.py b/tests/test_n_point_base.py index aec4d05d..e7676bee 100644 --- a/tests/test_n_point_base.py +++ b/tests/test_n_point_base.py @@ -7,11 +7,12 @@ import os import sys import types +from unittest.mock import patch import numpy as np import pytest -from dgamore import brillouin_zone +import dgamore.symmetry_reduction as sr from dgamore import brillouin_zone as bz from dgamore.n_point_base import IHaveChannel, IHaveMat, IAmNonLocal, SpinChannel, FrequencyNotation @@ -20,7 +21,7 @@ def test_initializes_with_correct_matrix_and_shape(): mat = np.array([[1, 2], [3, 4]]) obj = IHaveMat(mat) - assert np.allclose(obj.mat, mat, rtol=1e-2) + assert np.allclose(obj.mat, mat, rtol=1e-5) assert obj.original_shape == mat.shape @@ -43,7 +44,7 @@ def test_multiplies_with_scalar_correctly(): mat = np.array([[1, 2], [3, 4]]) obj = IHaveMat(mat) result = obj * 2 - assert np.allclose(result.mat, mat * 2, rtol=1e-2) + assert np.allclose(result.mat, mat * 2, rtol=1e-5) def test_raises_error_when_multiplying_with_invalid_type(): @@ -57,21 +58,21 @@ def test_performs_right_multiplication_with_scalar_correctly(): mat = np.array([[1, 2], [3, 4]]) obj = IHaveMat(mat) result = 2 * obj - assert np.allclose(result.mat, mat * 2, rtol=1e-2) + assert np.allclose(result.mat, mat * 2, rtol=1e-5) def test_negates_matrix_correctly(): mat = np.array([[1, -2], [-3, 4]]) obj = IHaveMat(mat) result = -obj - assert np.allclose(result.mat, -mat, rtol=1e-2) + assert np.allclose(result.mat, -mat, rtol=1e-5) def test_divides_by_scalar_correctly(): mat = np.array([[2, 4], [6, 8]]) obj = IHaveMat(mat) result = obj / 2 - assert np.allclose(result.mat, mat / 2, rtol=1e-2) + assert np.allclose(result.mat, mat / 2, rtol=1e-5) def test_raises_error_when_dividing_by_invalid_type(): @@ -95,7 +96,7 @@ def test_performs_einsum_contraction_correctly(): obj1 = IHaveMat(mat1) obj2 = IHaveMat(mat2) result = obj1.times("ij,jk->ik", obj2) - assert np.allclose(result, np.dot(mat1, mat2), rtol=1e-2) + assert np.allclose(result, np.dot(mat1, mat2), rtol=1e-5) def test_performs_einsum_contraction_with_multiple_matrices(): @@ -106,7 +107,7 @@ def test_performs_einsum_contraction_with_multiple_matrices(): obj2 = IHaveMat(mat2) obj3 = IHaveMat(mat3) result = obj1.times("ij,jk,kl->il", obj2, obj3) - assert np.allclose(result, np.dot(np.dot(mat1, mat2), mat3), rtol=1e-2) + assert np.allclose(result, np.dot(np.dot(mat1, mat2), mat3), rtol=1e-5) def test_raises_error_when_contraction_argument_is_invalid(): @@ -130,7 +131,7 @@ def test_performs_einsum_contraction_with_numpy_array(): mat2 = np.array([[5, 6], [7, 8]]) obj = IHaveMat(mat1) result = obj.times("ij,jk->ik", mat2) - assert np.allclose(result, np.dot(mat1, mat2), rtol=1e-2) + assert np.allclose(result, np.dot(mat1, mat2), rtol=1e-5) def test_raises_error_when_contraction_string_is_invalid(): @@ -218,7 +219,7 @@ def test_initializes_with_correct_matrix_and_momentum_dimensions(): mat = np.zeros((4, 4, 4)) nq = (4, 4, 4) obj = IAmNonLocal(mat, nq) - assert np.allclose(obj.mat, mat, rtol=1e-2) + assert np.allclose(obj.mat, mat, rtol=1e-5) assert obj.nq == nq assert obj.has_compressed_q_dimension is False @@ -227,7 +228,7 @@ def test_initializes_with_compressed_q_dimension(): mat = np.zeros((64,)) nq = (4, 4, 4) obj = IAmNonLocal(mat, nq, has_compressed_q_dimension=True) - assert np.allclose(obj.mat, mat, rtol=1e-2) + assert np.allclose(obj.mat, mat, rtol=1e-5) assert obj.nq == nq assert obj.has_compressed_q_dimension is True @@ -236,7 +237,7 @@ def test_shifts_momentum_by_zero_correctly(): mat = np.zeros((4, 4, 4)) obj = IAmNonLocal(mat, (4, 4, 4)) shifted = obj.shift_k_by_q((0, 0, 0)) - assert np.allclose(shifted.mat, mat, rtol=1e-2) + assert np.allclose(shifted.mat, mat, rtol=1e-5) def test_shifts_momentum_by_positive_values_correctly(): @@ -244,7 +245,7 @@ def test_shifts_momentum_by_positive_values_correctly(): obj = IAmNonLocal(mat, (4, 4, 4)) shifted = obj.shift_k_by_q((1, 1, 1)) expected = np.roll(mat, shift=(-1, -1, -1), axis=(0, 1, 2)) - assert np.allclose(shifted.mat, expected, rtol=1e-2) + assert np.allclose(shifted.mat, expected, rtol=1e-5) def test_shifts_momentum_by_negative_values_correctly(): @@ -252,7 +253,7 @@ def test_shifts_momentum_by_negative_values_correctly(): obj = IAmNonLocal(mat, (4, 4, 4)) shifted = obj.shift_k_by_q((-1, -1, -1)) expected = np.roll(mat, shift=(1, 1, 1), axis=(0, 1, 2)) - assert np.allclose(shifted.mat, expected, rtol=1e-2) + assert np.allclose(shifted.mat, expected, rtol=1e-5) def test_shifts_momentum_with_compressed_q_dimension_correctly(): @@ -274,7 +275,7 @@ def test_shifts_momentum_by_pi_correctly(): obj = IAmNonLocal(mat, (4, 4, 4)) shifted = obj.shift_k_by_pi() expected = np.roll(mat, shift=(2, 2, 2), axis=(0, 1, 2)) - assert np.allclose(shifted.mat, expected, rtol=1e-2) + assert np.allclose(shifted.mat, expected, rtol=1e-5) def test_shifts_momentum_by_pi_with_compressed_q_dimension(): @@ -383,7 +384,7 @@ def test_reduces_q_dimension_to_specified_momenta_and_values(): reduced = obj.reduce_q(q_list) expected_values = mat[1, 1, 1], mat[2, 2, 2] assert reduced.mat.shape == (2,) - assert np.allclose(reduced.mat, expected_values, rtol=1e-2) + assert np.allclose(reduced.mat, expected_values, rtol=1e-5) assert reduced.has_compressed_q_dimension is True @@ -428,7 +429,7 @@ def test_maps_to_full_bz_correctly_with_valid_inverse_map(): nq = (4, 4, 4) np.array([0, 1, 2, 3, 4, 5, 6, 7]) obj = IAmNonLocal(mat, nq, has_compressed_q_dimension=True) - grid = brillouin_zone.KGrid(nk=(2, 2, 2), symmetries=[]) + grid = bz.KGrid(nk=(2, 2, 2), symmetries=[]) obj.map_to_full_bz(grid, nq=(2, 2, 2)) assert obj.mat.shape == (8,) assert obj.nq == (2, 2, 2) @@ -447,7 +448,7 @@ def test_updates_nq_correctly_when_provided(): mat = np.arange(64) nq = (4, 4, 4) obj = IAmNonLocal(mat, nq, has_compressed_q_dimension=True) - grid = brillouin_zone.KGrid(nk=(2, 2, 2), symmetries=[]) + grid = bz.KGrid(nk=(2, 2, 2), symmetries=[]) obj.map_to_full_bz(grid, nq=(2, 2, 2)) assert obj.nq == (2, 2, 2) @@ -455,7 +456,7 @@ def test_updates_nq_correctly_when_provided(): def test_retains_original_nq_when_not_provided(): mat = np.arange(64) nq = (4, 4, 4) - grid = brillouin_zone.KGrid(nk=nq, symmetries=[]) + grid = bz.KGrid(nk=nq, symmetries=[]) obj = IAmNonLocal(mat, nq, has_compressed_q_dimension=True) obj.map_to_full_bz(grid) assert obj.nq == (4, 4, 4) @@ -467,7 +468,7 @@ def test_performs_fft_correctly_on_decompressed_matrix(): obj = IAmNonLocal(mat, nq) result = obj.fft() expected = np.fft.fftn(mat, axes=(0, 1, 2)) - assert np.allclose(result.mat, expected, rtol=1e-2) + assert np.allclose(result.mat, expected, rtol=1e-5) assert result.has_compressed_q_dimension is False @@ -478,7 +479,7 @@ def test_performs_fft_correctly_on_compressed_matrix(): result = obj.fft() decompressed_mat = mat.reshape(nq) expected = np.fft.fftn(decompressed_mat, axes=(0, 1, 2)).reshape(64) - assert np.allclose(result.mat, expected, rtol=1e-2) + assert np.allclose(result.mat, expected, rtol=1e-5) assert result.has_compressed_q_dimension is True @@ -496,7 +497,7 @@ def test_performs_ifft_correctly_on_decompressed_matrix(): obj = IAmNonLocal(mat, nq) result = obj.ifft() expected = np.fft.ifftn(mat, axes=(0, 1, 2)) - assert np.allclose(result.mat, expected, rtol=1e-2) + assert np.allclose(result.mat, expected, rtol=1e-5) assert result.has_compressed_q_dimension is False @@ -507,7 +508,7 @@ def test_performs_ifft_correctly_on_compressed_matrix(): result = obj.ifft() decompressed_mat = mat.reshape(nq) expected = np.fft.ifftn(decompressed_mat, axes=(0, 1, 2)).reshape(64) - assert np.allclose(result.mat, expected, rtol=1e-2) + assert np.allclose(result.mat, expected, rtol=1e-5) assert result.has_compressed_q_dimension is True @@ -525,7 +526,7 @@ def test_flips_momentum_axis_correctly_for_decompressed_matrix(): obj = IAmNonLocal(mat, nq) flipped = obj.flip_momentum_axis() expected = np.roll(np.flip(mat, axis=(0, 1, 2)), shift=1, axis=(0, 1, 2)) - assert np.allclose(flipped.mat, expected, rtol=1e-2) + assert np.allclose(flipped.mat, expected, rtol=1e-5) assert flipped.has_compressed_q_dimension is False @@ -536,7 +537,7 @@ def test_flips_momentum_axis_correctly_for_compressed_matrix(): flipped = obj.flip_momentum_axis() decompressed_mat = mat.reshape(nq) expected = np.roll(np.flip(decompressed_mat, axis=(0, 1, 2)), shift=1, axis=(0, 1, 2)).reshape(64) - assert np.allclose(flipped.mat, expected, rtol=1e-2) + assert np.allclose(flipped.mat, expected, rtol=1e-5) assert flipped.has_compressed_q_dimension is True @@ -819,7 +820,7 @@ def test_filter_q_index_returns_correct_index(): obj = IAmNonLocal(mat, nq) result = obj.filter_q_index(5) assert result.mat.shape == (1, 2) - assert np.allclose(result.mat[0], mat.reshape(64, 2)[5], rtol=1e-2) + assert np.allclose(result.mat[0], mat.reshape(64, 2)[5], rtol=1e-5) def test_filter_q_index_default_index_is_zero(): @@ -828,7 +829,7 @@ def test_filter_q_index_default_index_is_zero(): obj = IAmNonLocal(mat, nq) result = obj.filter_q_index() assert result.mat.shape == (1, 2) - assert np.allclose(result.mat[0], mat.reshape(64, 2)[0], rtol=1e-2) + assert np.allclose(result.mat[0], mat.reshape(64, 2)[0], rtol=1e-5) def test_filter_q_index_compresses_q_dimension_if_not_already_compressed(): @@ -846,7 +847,7 @@ def test_filter_q_index_does_not_modify_original_when_already_compressed(): obj = IAmNonLocal(mat, nq, has_compressed_q_dimension=True) original_mat = obj.mat.copy() _ = obj.filter_q_index(3) - assert np.allclose(obj.mat, original_mat, rtol=1e-2) + assert np.allclose(obj.mat, original_mat, rtol=1e-5) def test_filter_q_index_sets_nq_to_one(): @@ -879,7 +880,7 @@ def test_filter_q_index_returns_deep_copy(): obj = IAmNonLocal(mat, nq) result = obj.filter_q_index(0) result.mat[0, 0] = 9999 - assert not np.allclose(obj.mat.reshape(64, 2)[0, 0], 9999, rtol=1e-2) + assert not np.allclose(obj.mat.reshape(64, 2)[0, 0], 9999, rtol=1e-5) def test_filter_q_index_last_index(): @@ -887,7 +888,7 @@ def test_filter_q_index_last_index(): nq = (4, 4, 4) obj = IAmNonLocal(mat, nq) result = obj.filter_q_index(63) - assert np.allclose(result.mat[0], mat.reshape(64, 2)[63], rtol=1e-2) + assert np.allclose(result.mat[0], mat.reshape(64, 2)[63], rtol=1e-5) def test_filter_q_index_raises_for_out_of_bounds_index(): @@ -898,12 +899,47 @@ def test_filter_q_index_raises_for_out_of_bounds_index(): obj.filter_q_index(64) -# ============================================================================= -# _map_to_full_bz: auto-mode branch -# ============================================================================= -from unittest.mock import patch +def test_q_mean_averages_full_momentum_grid_to_single_q_point(): + mat = np.arange(64 * 2).reshape((4, 4, 4, 2)) + nq = (4, 4, 4) + obj = IAmNonLocal(mat, nq) + + result = obj.q_mean() -import dgamore.symmetry_reduction as _sr + expected = np.mean(mat, axis=(0, 1, 2))[None, None, None, ...] + assert result.mat.shape == (1, 1, 1, 2) + assert np.allclose(result.mat, expected, rtol=1e-5) + assert result.nq == (1, 1, 1) + assert result.original_shape == (1, 1, 1, 2) + assert np.allclose(obj.mat, mat, rtol=1e-5) + + +def test_q_mean_averages_compressed_momentum_grid_to_single_q_point(): + mat = np.arange(64 * 2).reshape((64, 2)) + nq = (4, 4, 4) + obj = IAmNonLocal(mat, nq, has_compressed_q_dimension=True) + + result = obj.q_mean() + + expected = np.mean(mat, axis=0)[None, ...] + assert result.mat.shape == (1, 2) + assert np.allclose(result.mat, expected, rtol=1e-5) + assert result.nq == (1, 1, 1) + assert result.original_shape == (1, 2) + assert np.allclose(obj.mat, mat, rtol=1e-5) + + +def test_q_mean_preserves_values_for_single_momentum_point(): + mat = np.array([[[[3.0, -2.0j]]]], dtype=complex) + nq = (1, 1, 1) + obj = IAmNonLocal(mat, nq) + + result = obj.q_mean() + + assert result.mat.shape == (1, 1, 1, 2) + assert np.allclose(result.mat[0, 0, 0], mat[0, 0, 0], rtol=1e-5) + assert result.nq == (1, 1, 1) + assert np.allclose(obj.mat, mat, rtol=1e-5) def _build_auto_kgrid(nx=4, ny=4, nz=4, nb=1, hopping=1.0, include_antiunitary=False): @@ -1042,7 +1078,7 @@ def test_map_to_full_bz_auto_delegates_to_apply_auto_orbital_transform(): obj = _DoublePrecisionNonLocal(mat=H_ibz, nq=(4, 4, 4), has_compressed_q_dimension=True) # Patch so we can assert it gets called with the right shapes and args - with patch.object(_sr, "apply_auto_orbital_transform", wraps=_sr.apply_auto_orbital_transform) as spy: + with patch.object(sr, "apply_auto_orbital_transform", wraps=sr.apply_auto_orbital_transform) as spy: obj._map_to_full_bz(grid, num_orbital_dimensions=2) assert spy.call_count == 1 _, kwargs = spy.call_args @@ -1059,7 +1095,7 @@ def test_map_to_full_bz_auto_passes_num_orbital_dimensions_4_for_vertex(): Gamma_full = np.einsum("...ab,...cd->...abcd", H, H) Gamma_ibz = Gamma_full.reshape(-1, nb, nb, nb, nb)[grid.irrk_ind].copy() obj = _DoublePrecisionNonLocal(mat=Gamma_ibz, nq=(3, 3, 3), has_compressed_q_dimension=True) - with patch.object(_sr, "apply_auto_orbital_transform", wraps=_sr.apply_auto_orbital_transform) as spy: + with patch.object(sr, "apply_auto_orbital_transform", wraps=sr.apply_auto_orbital_transform) as spy: obj._map_to_full_bz(grid, num_orbital_dimensions=4) _, kwargs = spy.call_args assert kwargs["num_orbital_dimensions"] == 4 @@ -1072,7 +1108,7 @@ def test_map_to_full_bz_legacy_kgrid_does_not_call_orbital_transform(): nb = 1 ibz_payload = np.arange(grid.nk_irr).astype(np.complex128).reshape(grid.nk_irr, nb, nb) obj = _DoublePrecisionNonLocal(mat=ibz_payload.copy(), nq=(4, 4, 1), has_compressed_q_dimension=True) - with patch.object(_sr, "apply_auto_orbital_transform", wraps=_sr.apply_auto_orbital_transform) as spy: + with patch.object(sr, "apply_auto_orbital_transform", wraps=sr.apply_auto_orbital_transform) as spy: obj._map_to_full_bz(grid, num_orbital_dimensions=2) assert spy.call_count == 0