diff --git a/README.md b/README.md index b5a4d6f..514a28d 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 003aadb..5de9b42 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,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(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(self, niv_asympt: int = 50000) -> float: + r""" + 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. + """ + 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/n_point_base.py b/dgamore/n_point_base.py index d485673..f3995bd 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 e8ba1a6..0a9268e 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 @@ -888,10 +888,21 @@ 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 = 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: mu_finding_failed = False @@ -916,9 +927,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() + logger.info(f"Kinetic energy: {ekin:.4f} [t or eV].") + + 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) + 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) @@ -935,15 +954,11 @@ 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, + sigma_converged = abs(relative_residual) < config.self_consistency.epsilon + logger.info( + f"Self-energy convergence: {sigma_converged} " + f"(relative residual={relative_residual:.3e}, 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) @@ -1029,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) @@ -1052,9 +1065,7 @@ 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 if config.self_consistency.mixing_strategy.lower() == "anderson" and current_iter > n_hist: @@ -1081,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) @@ -1138,14 +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 - 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}." - ) + 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_eliashberg_end_to_end.py b/tests/test_eliashberg_end_to_end.py index 79010ee..3947aa6 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) diff --git a/tests/test_mixing.py b/tests/test_mixing.py index 19bdd43..cc3c082 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) diff --git a/tests/test_n_point_base.py b/tests/test_n_point_base.py index aec4d05..e7676be 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