From 975f2edd7deecf16f206e2bcfb9fca3e1c6f2528 Mon Sep 17 00:00:00 2001 From: julpe Date: Tue, 18 Aug 2026 17:13:27 +0200 Subject: [PATCH] Sharing of more objects in memory because some rank-copied objects became really large in testing --- dgamore/DGAmore.py | 43 ++-------- dgamore/eliashberg_solver.py | 51 +++++++++--- dgamore/greens_function.py | 79 ++++++++++++++---- dgamore/max_ent.py | 29 ++++--- dgamore/memory_estimator.py | 7 +- dgamore/mpi_utils.py | 89 ++++++++++++++++---- dgamore/nonlocal_sde.py | 140 +++++++++++++++++++++++++++----- dgamore/self_energy.py | 46 ++++++++++- tests/test_autodetect_memory.py | 4 +- tests/test_eliashberg_solver.py | 72 +++++++++++++++- tests/test_greens_function.py | 64 ++++++++++++++- tests/test_max_ent.py | 16 ++++ tests/test_memory_estimator.py | 11 +++ tests/test_mpi_utils.py | 33 ++++++++ tests/test_nonlocal_sde.py | 102 +++++++++++++++++++++++ tests/test_self_energy.py | 23 ++++++ 16 files changed, 680 insertions(+), 129 deletions(-) diff --git a/dgamore/DGAmore.py b/dgamore/DGAmore.py index ba501d6..d7b5816 100644 --- a/dgamore/DGAmore.py +++ b/dgamore/DGAmore.py @@ -33,6 +33,7 @@ import dgamore.eliashberg_solver as eliashberg_solver import dgamore.local_sde as local_sde import dgamore.memory_estimator as memory_estimator +import dgamore.mpi_utils as mpi_utils import dgamore.nonlocal_sde as nonlocal_sde import dgamore.plotting as plotting from dgamore import max_ent @@ -47,42 +48,7 @@ # Fraction of a node's memory budget a run may occupy at any branch's peak. The remainder is the sole overhead # margin (OS/allocator/transients); the estimator's OVERHEAD_FACTOR is 1.0, so the two do not compound. -NODE_MEMORY_FRACTION: float = 0.97 - - -def _cgroup_memory_limit() -> int | None: - """ - Returns this process's effective cgroup memory limit in bytes, or ``None`` when no limit is set (or none is - readable). Batch schedulers such as slurm enforce a job's memory request through a cgroup, which can be far - below the node's physical memory, so the node budget must honor it. The cgroup path is taken from - ``/proc/self/cgroup``; on cgroup v2 every ancestor's ``memory.max`` is read up to the root (the limit may sit on - the job level rather than the process's own leaf) and the smallest set value wins, on v1 the controller's - ``memory.limit_in_bytes`` is read directly. Values at or above ``2**62`` mean "unlimited" and are ignored. - - :return: The smallest configured limit in bytes, or ``None`` if unlimited or undeterminable. - """ - limits = [] - try: - entries = dict( - (line.split(":", 2)[1], line.split(":", 2)[2].strip()) for line in open("/proc/self/cgroup", "r") - ) - if "" in entries: # cgroup v2: one unified hierarchy, limits possibly on an ancestor - path = os.path.normpath("/sys/fs/cgroup" + entries[""]) - while path.startswith("/sys/fs/cgroup"): - try: - value = open(os.path.join(path, "memory.max"), "r").read().strip() - if value != "max": - limits.append(int(value)) - except OSError: - pass - path = os.path.dirname(path) - elif "memory" in entries: # cgroup v1: the memory controller's own hierarchy - value = open("/sys/fs/cgroup/memory" + entries["memory"] + "/memory.limit_in_bytes", "r").read() - limits.append(int(value)) - except (OSError, ValueError, IndexError): - return None - limits = [limit for limit in limits if limit < 2**62] - return min(limits) if limits else None +NODE_MEMORY_FRACTION: float = 0.95 def main(): @@ -591,7 +557,7 @@ def autodetect_memory_settings(comm: MPI.Comm) -> None: transient is held by every rank at once, a *single-rank* transient by one rank while the others idle), minus ``(r - 1) * giwk_shareable`` for the branch's node-shared ``giwk_full`` window, and this must not exceed ``NODE_MEMORY_FRACTION`` times the node budget - ``psutil.virtual_memory().available``, capped by the cgroup - memory limit when the scheduler sets one (see :func:`_cgroup_memory_limit`). Each + memory limit when the scheduler sets one (see :func:`dgamore.mpi_utils.cgroup_memory_limit`). Each node's rank count and available memory are collected with a single ``allgather`` of ``(hostname, available_bytes)``; a branch's path is judged to "fit" only if it fits on **every** node (the flags are process-wide, so the tightest node governs, and a single-rank transient may land on any node). The @@ -608,7 +574,7 @@ def autodetect_memory_settings(comm: MPI.Comm) -> None: # Gather (hostname, available bytes) from every rank in a single collective and reduce to one entry per node: # the rank count and the (minimum, conservative) available memory on that node. node_available = psutil.virtual_memory().available - cgroup_limit = _cgroup_memory_limit() + cgroup_limit = mpi_utils.cgroup_memory_limit() if cgroup_limit is not None: node_available = min(node_available, cgroup_limit) hostname = socket.gethostname() @@ -630,6 +596,7 @@ def autodetect_memory_settings(comm: MPI.Comm) -> None: niv_core=config.box.niv_core, niv_full=config.box.niv_full, niv_cut=niv_cut, + niv_dmft=config.box.niv_dmft, niv_pp=niv_pp, n_ranks=comm.size, with_eliashberg=config.eliashberg.perform_eliashberg, diff --git a/dgamore/eliashberg_solver.py b/dgamore/eliashberg_solver.py index 191e06c..ee9b39f 100644 --- a/dgamore/eliashberg_solver.py +++ b/dgamore/eliashberg_solver.py @@ -953,9 +953,19 @@ def solve_eliashberg_lanczos_grid( block_mat = _gather_grid_vertex_block(gamma_r_pp, comm, row_slice, col_slice, in_grid) gamma_r_pp.free() - chi0_mat = mpi_utils.bcast_rows( - comm, gchi0_q0_pp.mat if comm.rank == bubble_rank else np.empty(0), root=bubble_rank - ) + # the ranks beyond the rows x cols grid never touch the bubble: ship it inside the grid only, so the idle + # ranks do not each receive a full-BZ copy (Split keeps the rank order, so the grid root index is bubble_rank) + grid_comm = comm.Split(0 if in_grid else 1, comm.rank) + if bubble_rank < rows * cols: + chi0_mat = ( + mpi_utils.bcast_rows(grid_comm, gchi0_q0_pp.mat if comm.rank == bubble_rank else None, root=bubble_rank) + if in_grid + else None + ) + else: + chi0_mat = mpi_utils.bcast_rows( + comm, gchi0_q0_pp.mat if comm.rank == bubble_rank else np.empty(0), root=bubble_rank + ) if not in_grid: return None @@ -1652,6 +1662,7 @@ def dispatch_full_vertex_calculation( niv_pp: int, mpi_dist: MpiDistributor, chunk_bytes: int | None = None, + node_comm: MPI.Comm | None = None, ) -> FourPoint: r""" Loads the local irreducible vertex for ``channel`` and builds the full ladder pp vertex through the slice-direct @@ -1680,14 +1691,21 @@ def dispatch_full_vertex_calculation( :param niv_pp: Number of positive fermionic frequencies of the pp vertex. :param mpi_dist: MPI distributor over the irreducible BZ q-points. :param chunk_bytes: Chunk byte budget of the build (``None`` uses the floor). + :param node_comm: Optional node-local communicator; when given, the multi-GB local vertex is loaded once per + node into an MPI shared-memory window instead of once per rank (the build only reads it). :return: The full ladder pp vertex :math:`F^{\mathrm{q}}_{r}` as a :class:`FourPoint`. """ - gamma_r = LocalFourPoint.load(os.path.join(config.output.output_path, f"gamma_{channel.value}_loc.npy"), channel) + gamma_r, gamma_win = nonlocal_sde._load_node_shared_local_vertex( + node_comm, os.path.join(config.output.output_path, f"gamma_{channel.value}_loc.npy"), channel + ) if config.eliashberg.save_fq: f_q_r = create_pairing_vertex_streaming_fq(u_loc, v_nonloc, gamma_r, niv_pp, mpi_dist, chunk_bytes) else: f_q_r = create_pairing_vertex_slice_q_r(u_loc, v_nonloc, gamma_r, niv_pp, mpi_dist, chunk_bytes) - gamma_r.free() + gamma_r.mat = None + if gamma_win is None: + gamma_r.free() + nonlocal_sde._free_shared_window(gamma_win, node_comm) mpi_dist.barrier() return f_q_r @@ -1824,11 +1842,16 @@ def solve( config.eliashberg.n_eig, comm.size, ) + # the node budget honors the scheduler's cgroup memory limit (e.g. slurm --mem), like the driver's fit check + node_budget = psutil.virtual_memory().available + cgroup_limit = mpi_utils.cgroup_memory_limit() + if cgroup_limit is not None: + node_budget = min(node_budget, cgroup_limit) + # one full-BZ sector residency per rank picks the in-memory solve, otherwise the grid takes over; rank 0 # decides and broadcasts, so ranks on differently loaded nodes can never pick different solvers use_grid = FORCE_GRID_SOLVER or ( - comm.size > 1 - and per_sector_bytes + giwk_dga.mat.nbytes > psutil.virtual_memory().available * NODE_MEMORY_FRACTION + comm.size > 1 and per_sector_bytes + giwk_dga.mat.nbytes > node_budget * NODE_MEMORY_FRACTION ) use_grid = comm.bcast(use_grid, root=0) if use_grid: @@ -1850,7 +1873,7 @@ def solve( ) else: sing_ranks, trip_ranks = get_ranks_for_lanczos( - comm, len(parities), psutil.virtual_memory().available, per_sector_bytes, giwk_dga.mat.nbytes + comm, len(parities), node_budget, per_sector_bytes, giwk_dga.mat.nbytes ) bubble_rank = sing_ranks[0] n_concurrent = len(set(sing_ranks) | set(trip_ranks)) @@ -1867,11 +1890,15 @@ def solve( node_comm = comm.Split_type(MPI.COMM_TYPE_SHARED) if comm.size > 1 else None chunk_bytes = memory_estimator.dynamic_chunk_budget( - psutil.virtual_memory().total, node_comm.size if node_comm is not None else 1 + mpi_utils.job_memory_total(), node_comm.size if node_comm is not None else 1 ) - f_dens_pp = dispatch_full_vertex_calculation(SpinChannel.DENS, u_loc, v_nonloc, niv_pp, mpi_dist_irrk, chunk_bytes) - f_magn_pp = dispatch_full_vertex_calculation(SpinChannel.MAGN, u_loc, v_nonloc, niv_pp, mpi_dist_irrk, chunk_bytes) + f_dens_pp = dispatch_full_vertex_calculation( + SpinChannel.DENS, u_loc, v_nonloc, niv_pp, mpi_dist_irrk, chunk_bytes, node_comm + ) + f_magn_pp = dispatch_full_vertex_calculation( + SpinChannel.MAGN, u_loc, v_nonloc, niv_pp, mpi_dist_irrk, chunk_bytes, node_comm + ) delete_files(config.output.eliashberg_path, f"gchi0_q_inv_rank_{comm.rank}.npy") @@ -1957,7 +1984,7 @@ def solve( # Fraction of a node's available host memory the sector packing may occupy (mirrors DGAmore.NODE_MEMORY_FRACTION). -NODE_MEMORY_FRACTION: float = 0.97 +NODE_MEMORY_FRACTION: float = 0.95 def get_ranks_for_lanczos( diff --git a/dgamore/greens_function.py b/dgamore/greens_function.py index f86c7ca..7875dc2 100644 --- a/dgamore/greens_function.py +++ b/dgamore/greens_function.py @@ -101,6 +101,34 @@ def root_fun( return get_total_fill(mu, ek, sigma_mat, beta, smom0) - target_filling +# Largest filling residual accepted from a "converged" Newton result; genuine roots sit orders of magnitude below +# this, while a secant that stalled in a flat filling region leaves a residual of order the filling itself. +_FILL_RESIDUAL_TOL: float = 1e-3 + + +def _find_mu_bracket( + mu0: float, args: tuple, initial_width: float = 0.05, max_width: float = 64.0 +) -> tuple[float, float] | None: + r""" + Expands a symmetric interval around ``mu0``, doubling its half-width each step, until the filling residual + :func:`root_fun` changes sign across it. Starting narrow makes the search return an interval around the root + closest to ``mu0``, which keeps a self-consistency trajectory in its current basin. + + :param mu0: Center of the search interval. + :param args: The :func:`root_fun` arguments after ``mu`` (target filling, dispersion, self-energy, beta, moment). + :param initial_width: Half-width of the first interval. + :param max_width: Half-width beyond which the search gives up. + :return: A bracketing interval ``(lo, hi)``, or ``None`` if no sign change was found. + """ + width = initial_width + while width <= max_width: + lo, hi = mu0 - width, mu0 + width + if root_fun(lo, *args) * root_fun(hi, *args) < 0: + return lo, hi + width *= 2 + return None + + def update_mu( mu0: float, target_filling: float, @@ -113,7 +141,9 @@ def update_mu( ) -> float: r""" Updates the chemical potential to match the target filling by using Newton's method to find the optimal - :math:`\mu`. On failure to converge the starting value is returned unchanged. + :math:`\mu`. A Newton result is only accepted if its filling residual is small; when Newton fails or lands + away from an actual root, the root nearest the starting value is found instead with a bracketed Brent search + (see :func:`_find_mu_bracket`). The starting value is returned unchanged only if no bracket exists. :param mu0: Initial guess for the chemical potential. :param target_filling: Desired total filling. @@ -121,18 +151,29 @@ def update_mu( :param sigma_mat: Self-energy array, shape ``[k, o1, o2, v]``. :param beta: Inverse temperature :math:`\beta`. :param smom0: Zeroth moment :math:`\Sigma_\infty` of the self-energy. - :param logger: Optional logger; if given, a failed root search is logged at debug level. - :param tol: Newton tolerance for the root search. - :return: The updated (real) chemical potential, or ``mu0`` if the root search did not converge. + :param logger: Optional logger; if given, the bracketed fallback is logged at info level and a fully failed + root search at warning level. + :param tol: Root search tolerance for the chemical potential. + :return: The updated (real) chemical potential, or ``mu0`` if no root was found. :raises ValueError: If the converged chemical potential has a non-negligible imaginary part. """ mu = mu0 + args = (target_filling, ek, sigma_mat, beta, smom0) try: - mu = opt.newton(root_fun, mu, args=(target_filling, ek, sigma_mat, beta, smom0), tol=tol) + mu = opt.newton(root_fun, mu, args=args, tol=tol) + # the secant step criterion can also "converge" inside a flat filling region far from any root, so the + # residual is verified before the value is accepted + if np.abs(root_fun(mu, *args)) > _FILL_RESIDUAL_TOL: + raise RuntimeError except RuntimeError: + bracket = _find_mu_bracket(mu0, args) + if bracket is None: + if logger is not None: + logger.warning("Root finding for chemical potential failed; keeping the previous value.") + return mu0 if logger is not None: - logger.debug("Root finding for chemical potential failed.") - return mu0 + logger.info("Newton did not find a chemical potential root; using a bracketed root search.") + mu = opt.brentq(root_fun, *bracket, args=args, xtol=tol) if np.abs(mu.imag) < 1e-8: mu = mu.real @@ -233,7 +274,9 @@ def occ_k(self) -> np.ndarray: def get_g_full(siw: SelfEnergy, mu: float, ek: np.ndarray, beta: float): r""" Builds the full momentum-dependent Green's function :math:`G^{\mathrm{k}} = [(\imath\nu + \mu) - - \varepsilon(\mathbf{k}) - \Sigma^{\mathrm{k}}]^{-1}`. + \varepsilon(\mathbf{k}) - \Sigma^{\mathrm{k}}]^{-1}`. The Dyson matrix is assembled and inverted in bounded + fermionic-frequency chunks, so beyond the result only one chunk-sized transient is alive at a time (large + boxes such as the full DMFT one would otherwise triple the peak). :param siw: The :class:`SelfEnergy` :math:`\Sigma`. :param mu: Chemical potential :math:`\mu`. @@ -243,15 +286,17 @@ def get_g_full(siw: SelfEnergy, mu: float, ek: np.ndarray, beta: float): """ eye_bands = np.eye(siw.n_bands, siw.n_bands) iv = 1j * MFHelper.vn(siw.niv, beta) - iv_bands = iv[None, None, :] * eye_bands[..., None] - mu_bands = mu * eye_bands[:, :, None] - mat = ( - iv_bands[None, None, None, ...] - + mu_bands[None, None, None, ...] - - ek[..., None] - - siw.decompress_q_dimension().mat - ) - mat = GreensFunction._invert_last_orbital_block(mat) + static_bands = (mu * eye_bands)[None, None, None, :, :, None] - ek[..., None] + sigma_mat = siw.decompress_q_dimension().mat + + # a momentum-local sigma ([1, 1, 1, ...]) broadcasts against the dispersion, so the result is always full-k + mat = np.empty((*ek.shape[:3], siw.n_bands, siw.n_bands, 2 * siw.niv), dtype=sigma_mat.dtype) + step = max(1, _MODEL_EPOT_CHUNK_ELEMENTS // (int(np.prod(ek.shape[:3])) * siw.n_bands**2)) + for start in range(0, 2 * siw.niv, step): + stop = min(2 * siw.niv, start + step) + dyson = (iv[start:stop][None, None, :] * eye_bands[..., None])[None, None, None, ...] + static_bands + dyson -= sigma_mat[..., start:stop] + mat[..., start:stop] = GreensFunction._invert_last_orbital_block(dyson) return GreensFunction(mat, siw, ek, siw.full_niv_range, False, False, nk=ek.shape[:3], beta=beta, mu=mu) @staticmethod diff --git a/dgamore/max_ent.py b/dgamore/max_ent.py index f69c0b5..6fd0b75 100644 --- a/dgamore/max_ent.py +++ b/dgamore/max_ent.py @@ -79,21 +79,24 @@ def perform_maxent_giwk(giwk: GreensFunction, name: str, comm: MPI.Comm): logger = config.logger logger.info(f"Starting analytic continuation of the {name} Green's function using the maximum entropy method.") - giwk_maxent = giwk.cut_niv(config.box.niv_core).to_half_niv_range() - - # Rotate G into the band (H(k)-eigen) basis first, so the band-diagonal below is the band-resolved spectral function - # (invariant under lattice symmetries, so the naive irrk_inv unfold is correct); done on the full BZ. - hk = config.lattice.hamiltonian.get_ek(config.lattice.k_grid) - giwk_maxent = giwk_maxent.decompress_q_dimension() - orbital_to_band_basis(hk, giwk_maxent.mat) - irrq_list = config.lattice.k_grid.get_irrq_list() - mpi_dist = MpiDistributor(ntasks=len(irrq_list), comm=comm, name="Maxent_G", output_path=config.output.output_path) - giwk_maxent = giwk_maxent.reduce_q(irrq_list) + # The full-BZ preparation runs on rank 0 only: every rank used to build the identical core-cut, band-rotated, + # irr-reduced Green's function just to keep its own scatter slice - multi-GB transients times the rank count. + g_irr_mat = None + if comm.rank == 0: + giwk_maxent = giwk.cut_niv(config.box.niv_core).to_half_niv_range() + + # Rotate G into the band (H(k)-eigen) basis first, so the band-diagonal below is the band-resolved spectral + # function (invariant under lattice symmetries, so the naive irrk_inv unfold is correct); on the full BZ. + hk = config.lattice.hamiltonian.get_ek(config.lattice.k_grid) + giwk_maxent = giwk_maxent.decompress_q_dimension() + orbital_to_band_basis(hk, giwk_maxent.mat) + g_irr_mat = giwk_maxent.reduce_q(irrq_list).mat + logger.info("Scattering Green's function in the IBZ to all ranks.") - giwk_maxent.mat = mpi_dist.scatter(giwk_maxent.mat) # each rank now has a slice of the irr BZ + g_irr_slice = mpi_dist.scatter(g_irr_mat) # each rank now has a slice of the irr BZ wn = np.pi / config.sys.beta * (2 * np.arange(config.box.niv_core) + 1) w = ( @@ -109,7 +112,7 @@ def perform_maxent_giwk(giwk: GreensFunction, name: str, comm: MPI.Comm): for band in range(config.sys.n_bands): logger.info(f"Processing analytic continuation of band {band+1}.") - for k in range(giwk_maxent.mat.shape[0]): + for k in range(g_irr_slice.shape[0]): # Capture the vendored solver's stdout so its print() diagnostics go through the logger instead of # leaking to the output; re-logged (prefixed) below whether the continuation succeeds or fails. captured_output = io.StringIO() @@ -121,7 +124,7 @@ def perform_maxent_giwk(giwk: GreensFunction, name: str, comm: MPI.Comm): # The alpha-fit curve_fit inside the solver harmlessly fails to estimate its covariance; mute it. warnings.simplefilter("ignore", OptimizeWarning) probl_maxent = AnalyticContinuationProblem( - im_axis=wn, re_axis=w, im_data=giwk_maxent[k, band, band], beta=config.sys.beta + im_axis=wn, re_axis=w, im_data=g_irr_slice[k, band, band], beta=config.sys.beta ) result = probl_maxent.solve(model=model, stdev=stdev)[0] spectral_function[k, band] = result.A_opt.astype(np.float32) diff --git a/dgamore/memory_estimator.py b/dgamore/memory_estimator.py index ddabb46..c395503 100644 --- a/dgamore/memory_estimator.py +++ b/dgamore/memory_estimator.py @@ -228,6 +228,7 @@ def estimate_peaks( niv_core: int, niv_full: int, niv_cut: int, + niv_dmft: int, niv_pp: int, n_ranks: int, with_eliashberg: bool, @@ -263,6 +264,8 @@ def estimate_peaks( :param niv_full: Number of positive fermionic full-region frequencies. :param niv_cut: Number of positive fermionic frequencies the full-grid ``giwk_full`` is built at (``min(niw_core + niv_full + 10, niv_dmft)`` in :func:`dgamore.nonlocal_sde.calculate_self_energy_q`). + :param niv_dmft: Number of positive fermionic frequencies of the DMFT input box (the rank-0 occupation and + energy step of every iteration concatenates the self-energy back to it). :param niv_pp: Number of positive fermionic frequencies of the pp (Eliashberg) box. :param n_ranks: Number of MPI ranks the q-points are distributed over. :param with_eliashberg: Whether the Eliashberg step runs (adds the ``"fq"`` and ``"lanczos"`` branches). @@ -343,7 +346,9 @@ def estimate_peaks( # chunk-capped exchange transients (floor modeled - the dynamic budget stays below an eighth of the fair share). sde_chunk = min(SLICE_CHUNK_BYTES, DTYPE_BYTES * _bubble_block(qt, nb, wp, vc)) sde_distributed = scale * _bubble_block(qi, nb, wp, vc) + overhead * SDE_CHUNK_FACTOR * sde_chunk - sde_single = scale * 2 * _giwk_rspace(nk_tot, nb, vc) + # rank-0 single: the sigma finalize buffers, or the occupation/energy step's DMFT-box sigma + giwk pair + # (its concatenation and Dyson-build transients are broadcast-assigned and v-chunked, so only the pair counts) + sde_single = scale * max(2 * _giwk_rspace(nk_tot, nb, vc), 2 * _giwk_rspace(nk_tot, nb, 2 * niv_dmft)) peaks["sde"] = BranchPeak( baseline=baseline_sde, giwk_shareable=2 * giwk_sde, diff --git a/dgamore/mpi_utils.py b/dgamore/mpi_utils.py index 84a831e..9397763 100644 --- a/dgamore/mpi_utils.py +++ b/dgamore/mpi_utils.py @@ -32,6 +32,7 @@ import h5py import mpi4py.MPI as MPI +import psutil import numpy as np import scipy.fft as fft @@ -83,6 +84,55 @@ def build_node_shared_array(node_comm, compute_fn, dtype=DTYPE): return shared, win +def cgroup_memory_limit(proc_file: str = "/proc/self/cgroup", cgroup_root: str = "/sys/fs/cgroup") -> int | None: + """ + Returns this process's effective cgroup memory limit in bytes, or ``None`` when no limit is set (or none is + readable). Batch schedulers such as slurm enforce a job's memory request through a cgroup, which can be far + below the node's physical memory, so every memory budget must honor it. The cgroup path is taken from + ``proc_file``; on cgroup v2 every ancestor's ``memory.max`` is read up to the root (the limit may sit on the + job level rather than the process's own leaf) and the smallest set value wins, on v1 the memory controller's + ``memory.limit_in_bytes`` is read directly. Values at or above ``2**62`` mean "unlimited" and are ignored. + + :param proc_file: Path of the process's cgroup membership file. + :param cgroup_root: Mount point of the cgroup filesystem. + :return: The smallest configured limit in bytes, or ``None`` if unlimited or undeterminable. + """ + limits = [] + try: + entries = dict((line.split(":", 2)[1], line.split(":", 2)[2].strip()) for line in open(proc_file, "r")) + if "" in entries: # cgroup v2: one unified hierarchy, limits possibly on an ancestor + path = os.path.normpath(cgroup_root + entries[""]) + while path.startswith(cgroup_root): + try: + value = open(os.path.join(path, "memory.max"), "r").read().strip() + if value != "max": + limits.append(int(value)) + except OSError: + pass + path = os.path.dirname(path) + elif "memory" in entries: # cgroup v1: the memory controller's own hierarchy + value = open(cgroup_root + "/memory" + entries["memory"] + "/memory.limit_in_bytes", "r").read() + limits.append(int(value)) + except (OSError, ValueError, IndexError): + return None + limits = [limit for limit in limits if limit < 2**62] + return min(limits) if limits else None + + +def job_memory_total() -> int: + """ + Returns the memory a job may plan with on this node: the hardware total, capped by the cgroup limit when the + scheduler sets one (see :func:`cgroup_memory_limit`). Chunk budgets derive from this instead of the free + memory, so the chunking - and with it the floating-point reduction order - stays reproducible across reruns + (the cgroup limit is part of the job specification, unlike the machine's momentary load). + + :return: The plannable memory of this node in bytes. + """ + total = psutil.virtual_memory().total + limit = cgroup_memory_limit() + return total if limit is None else min(total, limit) + + def count_nodes(comm, node_comm) -> int: r""" Returns the number of distinct nodes ``comm`` spans, i.e. how many copies of a per-node shared array (see @@ -1108,7 +1158,11 @@ def get_pencil_indices(rank: int, size: int, nq: tuple[int, int, int], layout: s def _redistribute_p2p(mat, nq, comm, source_layout, target_layout): """ Peer-to-peer redistributes the rows of ``mat`` (indexed by flattened q) from one pencil/flat layout to another, - exchanging only the rows each rank pair shares (in below-2 GB byte chunks). + exchanging only the rows each rank pair shares (in below-2 GB byte chunks). Every pair's transfer is posted + non-blocking at once and completed by a single ``Waitall``: a one-round-per-pair schedule would serialize + ``size`` blocking rounds per redistribution, which dominates the FFT step at high rank counts. The staging stays + bounded because each source row belongs to exactly one target (and vice versa), so all send copies together hold + at most one local slab and all receive stagings at most one target slab. :param mat: The local array slice, with the q-index on axis 0. :param nq: Number of momenta per spatial direction ``(nx, ny, nz)``. @@ -1127,15 +1181,18 @@ def _redistribute_p2p(mat, nq, comm, source_layout, target_layout): src_map = {g_idx: l_idx for l_idx, g_idx in enumerate(src_indices)} tgt_map = {g_idx: l_idx for l_idx, g_idx in enumerate(tgt_indices)} - for shift in range(size): - if shift == 0: - # Self-overlap: rows this rank both owns (source layout) and needs (target layout). Copy locally instead - # of round-tripping the data through MPI to itself. - common = np.intersect1d(src_indices, tgt_indices, assume_unique=True) - if len(common) > 0: - res_mat[[tgt_map[g] for g in common]] = mat[[src_map[g] for g in common]] - continue + # Self-overlap: rows this rank both owns (source layout) and needs (target layout). Copy locally instead + # of round-tripping the data through MPI to itself. + common = np.intersect1d(src_indices, tgt_indices, assume_unique=True) + if len(common) > 0: + res_mat[[tgt_map[g] for g in common]] = mat[[src_map[g] for g in common]] + + # every rank pair's transfer is posted at once and completed by the single Waitall below (see the docstring) + reqs = [] + send_bufs = [] # keep alive until Waitall + recv_stagings = [] + for shift in range(1, size): target_rank = (rank + shift) % size source_rank = (rank - shift) % size @@ -1145,29 +1202,25 @@ def _redistribute_p2p(mat, nq, comm, source_layout, target_layout): remote_src_indices = get_pencil_indices(source_rank, size, nq, source_layout) to_recv_g = np.intersect1d(tgt_indices, remote_src_indices, assume_unique=True) - reqs = [] - send_buf = None # keep alive until Waitall - recv_staging = None - if len(to_send_g) > 0: send_l = [src_map[g] for g in to_send_g] send_buf = np.ascontiguousarray(mat[send_l]) + send_bufs.append(send_buf) send_view = send_buf.view(np.byte).reshape(-1) for i in range(0, send_view.nbytes, MAX_MPI_BYTES): reqs.append(comm.Isend(send_view[i : i + MAX_MPI_BYTES], dest=target_rank, tag=shift)) if len(to_recv_g) > 0: recv_staging = np.empty((len(to_recv_g),) + mat.shape[1:], dtype=mat.dtype) + recv_stagings.append((recv_staging, [tgt_map[g] for g in to_recv_g])) recv_view = recv_staging.view(np.byte).reshape(-1) for i in range(0, recv_view.nbytes, MAX_MPI_BYTES): reqs.append(comm.Irecv(recv_view[i : i + MAX_MPI_BYTES], source=source_rank, tag=shift)) - MPI.Request.Waitall(reqs) + MPI.Request.Waitall(reqs) - # Now copy from staging into res_mat at the right rows - if len(to_recv_g) > 0: - recv_l = [tgt_map[g] for g in to_recv_g] - res_mat[recv_l] = recv_staging + for recv_staging, recv_l in recv_stagings: + res_mat[recv_l] = recv_staging return res_mat diff --git a/dgamore/nonlocal_sde.py b/dgamore/nonlocal_sde.py index b3ed4c2..6c62a78 100644 --- a/dgamore/nonlocal_sde.py +++ b/dgamore/nonlocal_sde.py @@ -1351,6 +1351,92 @@ def _free_shared_window(win, node_comm) -> None: win.Free() +def _share_sigma_per_node(sigma: SelfEnergy, node_comm, roots_comm) -> "MPI.Win | None": + r""" + Replaces the mixed self-energy's array on every rank by a view of one per-node MPI shared-memory window: the + node roots receive rank 0's array through a chunked broadcast over ``roots_comm``, then expose it to their + node's other ranks (see :func:`dgamore.mpi_utils.build_node_shared_array`). The self-consistency loop thereby + holds the full-BZ self-energy **once per node instead of once per rank**. Every rank keeps its own + :class:`SelfEnergy` object; only the buffer is shared and must be treated as read-only (every consumer copies + via ``cut_niv``/``copy``/``concatenate`` before writing). + + :param sigma: The :class:`SelfEnergy` holding rank 0's mixed array (every rank's ``mat`` is replaced in place + by the shared view). + :param node_comm: The node-local (shared-memory) communicator. + :param roots_comm: Communicator over exactly the node-root ranks (global rank 0 first). + :return: The MPI shared-memory window (``None`` for a single-rank node); free it via + :func:`_free_shared_window` once no rank reads the buffer anymore. + """ + if node_comm.rank == 0: + sigma.mat = mpi_utils.bcast_rows(roots_comm, sigma.mat, root=0) + mat, win = mpi_utils.build_node_shared_array(node_comm, lambda: sigma.mat) + sigma.mat = mat + return win + + +def _update_occ_and_energies_distributed( + sigma_new: SelfEnergy, sigma_dmft_full: SelfEnergy, mpi_dist_fullbz: MpiDistributor, mu: float +) -> tuple[float, np.ndarray, np.ndarray, float, float]: + r""" + Computes the occupation and the kinetic and potential energies of the mixed self-energy on the DMFT frequency + box, distributed over the full-BZ momenta: every rank concatenates and Dyson-inverts only its own momentum + slice, evaluates the occupation and energy sums there, and the results are recombined (the former evaluation + built the whole DMFT-box Green's function and its asymptotic tail sums on rank 0 while every other rank idled). + The self-energy moments are fitted from the momentum-averaged concatenated self-energy, allreduced first so + they match the full-box fit on every rank; the k-resolved occupation is allgathered and the k-summed scalars + are recombined with each rank's momentum count as weight. + + :param sigma_new: The mixed :class:`SelfEnergy` (full BZ, compressed momenta, identical on every rank). + :param sigma_dmft_full: The DMFT :class:`SelfEnergy` supplying the shell frequencies (momentum-local). + :param mpi_dist_fullbz: MPI distributor over the full BZ q-points. + :param mu: Chemical potential :math:`\mu`. + :return: The tuple ``(n, occ, occ_k, ekin, epot)``: total filling, k-averaged occupation ``[o1, o2]``, + k-resolved occupation ``[kx, ky, kz, o1, o2]``, and the kinetic and potential energies per site. + """ + nk_tot = config.lattice.k_grid.nk_tot + n_bands = config.sys.n_bands + n_my = mpi_dist_fullbz.my_size + + sigma_slice = SelfEnergy( + sigma_new.compress_q_dimension().mat[mpi_dist_fullbz.my_slice], + (n_my, 1, 1), + has_compressed_q_dimension=True, + calc_smom=False, + beta=config.sys.beta, + ) + sigma_occ = sigma_slice.concatenate_self_energies(sigma_dmft_full) + + # sigma_new is replicated, so every rank fits the full-box moments locally from the k-mean fit window; the fit + # is bit-identical to the momentum-resolved concatenation's fit and needs no reduction + sigma_occ._smom0, sigma_occ._smom1 = sigma_new.fit_smom_concatenated(sigma_dmft_full) + + ek = config.lattice.hamiltonian.get_ek() + ek_slice = ek.reshape(nk_tot, n_bands, n_bands)[mpi_dist_fullbz.my_slice].reshape(n_my, 1, 1, n_bands, n_bands) + giwk_occ = GreensFunction.get_g_full(sigma_occ, mu, ek_slice, config.sys.beta) + _, _, occ_k_slice = giwk_occ.get_fill_nonlocal() + ekin, epot = giwk_occ.get_ekin(), giwk_occ.get_epot() + giwk_occ.free() + sigma_occ.free() + + # assembled through the chunked Allreduce of zero-padded slices: each momentum is contributed by exactly one + # rank, so the sum is bit-exact and no point-to-point matching is involved. The dtype MUST be pinned: the + # fill's dtype is value-dependent (real-matrix eig returns float when a slice's eigenvalues are all real), and + # ranks entering a collective with different element types abort with truncated messages. + occ_k = np.zeros((nk_tot, n_bands, n_bands), dtype=np.complex128) + occ_k[mpi_dist_fullbz.my_slice] = occ_k_slice.reshape(n_my, n_bands, n_bands) + if mpi_dist_fullbz.mpi_size > 1: + occ_k = mpi_dist_fullbz.allreduce(occ_k) + occ_k = occ_k.reshape(*config.lattice.k_grid.nk, n_bands, n_bands) + scalars = np.array([ekin, epot]) * n_my + if mpi_dist_fullbz.mpi_size > 1: + scalars = mpi_dist_fullbz.comm.allreduce(scalars) + ekin, epot = scalars / nk_tot + + occ = np.mean(occ_k, axis=(0, 1, 2)) + occ.real[np.abs(occ) < 1e-12] = 0.0 + return 2.0 * np.trace(occ).real, occ, occ_k, float(ekin), float(epot) + + def calculate_sigma_proposal( sigma_in: SelfEnergy, mu: float, @@ -1455,10 +1541,8 @@ def calculate_sigma_proposal( if config.eliashberg.perform_eliashberg: gchi0_q_core_inv.save(name=f"gchi0_q_inv_rank_{comm.rank}", output_dir=config.output.eliashberg_path) - import psutil - node_ranks = shared_node_comm.size if shared_node_comm is not None else 1 - chunk_bytes = memory_estimator.dynamic_chunk_budget(psutil.virtual_memory().total, node_ranks) + chunk_bytes = memory_estimator.dynamic_chunk_budget(mpi_utils.job_memory_total(), node_ranks) gamma_dens, gamma_dens_win = _load_node_shared_local_vertex( shared_node_comm, os.path.join(config.output.output_path, "gamma_dens_loc.npy"), SpinChannel.DENS @@ -1652,6 +1736,12 @@ def calculate_self_energy_q( ntasks=config.lattice.k_grid.nk_tot, comm=comm, name="FBZ", output_path=config.output.output_path ) + # the loop's full-BZ self-energy is held once per NODE (mixed on rank 0, broadcast to the node roots, then + # window-shared read-only); single-rank runs keep the plain per-rank broadcast + sc_node_comm = comm.Split_type(MPI.COMM_TYPE_SHARED) if comm.size > 1 else None + sc_roots_comm = comm.Split(0 if sc_node_comm.rank == 0 else 1) if sc_node_comm is not None else None + sigma_win = None + sigma_old, starting_iter = get_starting_sigma(sigma_dmft) if starting_iter > 0: logger.info( @@ -1729,16 +1819,24 @@ def calculate_self_energy_q( # delta_sigma = sigma_dmft.cut_niv(config.box.niv_core) - sigma_new.q_mean().cut_niv(config.box.niv_core) sigma_old = sigma_old.cut_niv(config.box.niv_core) + # the cut copied everything still needed from the previous iteration's shared sigma buffer + _free_shared_window(sigma_win, sc_node_comm) + sigma_win = None logger.info("Applying mixing strategy to the self-energy.") sigma_old = sigma_old.concatenate_self_energies(sigma_dmft) history_cap = _mixing_history_cap(current_iter, release_iter, anneal_reset_iter) - # mixing runs on rank 0 only (all ranks computed identical results before) and the mixed sigma is broadcast + # mixing runs on rank 0 only (all ranks computed identical results before); the mixed sigma then reaches the + # other ranks once per node through a shared window instead of once per rank if comm.rank == 0: sigma_new = apply_mixing_strategy( sigma_new, sigma_old, sigma_dmft, current_iter, history_cap, sigma_history ) - sigma_new = mpi_dist_fullbz.bcast_npoint(sigma_new) + if sc_node_comm is None: + sigma_new = mpi_dist_fullbz.bcast_npoint(sigma_new) + else: + # compressing first pins one momentum layout on every rank, so the shared buffer matches all metadata + sigma_win = _share_sigma_per_node(sigma_new.compress_q_dimension(), sc_node_comm, sc_roots_comm) if sigma_history is not None: # the in-memory analog of what read_last_n_sigmas_from_files reproduced from the file just saved below sigma_history.append(sigma_new.decompress_q_dimension().cut_niv(config.box.niv_core).mat) @@ -1766,22 +1864,14 @@ def calculate_self_energy_q( mu_history.append(config.sys.mu) logger.info(f"Updated mu from {old_mu} to {config.sys.mu}.") - if comm.rank == 0: - sigma_occ = sigma_new.copy().concatenate_self_energies(sigma_dmft_full) - giwk_occ = GreensFunction.get_g_full( - sigma_occ, config.sys.mu, config.lattice.hamiltonian.get_ek(), config.sys.beta - ) - # 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) + # new occupation matrix and energies from the new Green's function (outside the asympt region it is the + # DMFT lattice Green's function); k-distributed, so no rank builds the whole DMFT-box Green's function + _, config.sys.occ, config.sys.occ_k, ekin, epot = _update_occ_and_energies_distributed( + sigma_new, sigma_dmft_full, mpi_dist_fullbz, config.sys.mu + ) # n should not change + logger.info(f"Kinetic energy: {ekin:.4f} [t or eV].") + logger.info(f"Potential energy: {epot:.4f} [t or eV].") + logger.info(f"Total energy: {(ekin + epot):.4f} [t or eV].") if config.self_consistency.max_iter > 1: logger.info("Updated occupation matrix from new Green's function.") @@ -1872,6 +1962,14 @@ def calculate_self_energy_q( else: logger.info("Self-consistency not reached.") + # the caller-owned result must outlive the loop's shared window; hand back a private copy and release the window + if sigma_win is not None: + sigma_old.mat = sigma_old.mat.copy() + _free_shared_window(sigma_win, sc_node_comm) + if sc_node_comm is not None: + sc_roots_comm.Free() + sc_node_comm.Free() + mpi_dist_irrk.delete_file() mpi_dist_fullbz.delete_file() diff --git a/dgamore/self_energy.py b/dgamore/self_energy.py index 1f4788d..67c3f13 100644 --- a/dgamore/self_energy.py +++ b/dgamore/self_energy.py @@ -245,14 +245,52 @@ def concatenate_self_energies(self, other: "SelfEnergy") -> "SelfEnergy": self.compress_q_dimension() other = other.compress_q_dimension() - other_mat = np.tile(other.mat, (self.nq_tot, 1, 1, 1)) if other.nq_tot == 1 else other.mat - result_mat = np.concatenate( - (other_mat[..., :niv_diff], self.mat, other_mat[..., niv_diff + 2 * self.niv :]), axis=-1 - ) + # the shell donor's momentum axis broadcasts into the result, so a local tail is never tiled over the BZ + result_mat = np.empty(self.current_shape[:-1] + (2 * other.niv,), dtype=self.mat.dtype) + result_mat[..., :niv_diff] = other.mat[..., :niv_diff] + result_mat[..., niv_diff : niv_diff + 2 * self.niv] = self.mat + result_mat[..., niv_diff + 2 * self.niv :] = other.mat[..., niv_diff + 2 * self.niv :] return SelfEnergy( result_mat, self.nq, self.full_niv_range, self.has_compressed_q_dimension, False, beta=self._beta ) + def fit_smom_concatenated(self, other: "SelfEnergy") -> tuple[np.ndarray, np.ndarray]: + """ + Fits the high-frequency moments that :meth:`fit_smom` would report for + ``self.concatenate_self_energies(other)``, without building the momentum-resolved concatenation: only the + fit window (the top fifth of ``other``'s positive frequencies) is assembled from the same column sources, + so the moments are bit-identical to the full concatenation's fit. + + :param other: The self-energy supplying the shell frequencies; must have at least as many frequencies as ``self``. + :return: The tuple ``(mom0, mom1)`` of moments, each of shape ``[o1, o2]``. + :raises ValueError: If ``other`` has fewer frequencies than ``self``. + """ + if self.niv > other.niv: + raise ValueError("Can not concatenate with a self-energy that has less frequencies.") + niv_res = other.niv + niv_diff = niv_res - self.niv + n_freq_fit = max(int(0.2 * niv_res), 4) + + self.compress_q_dimension() + other = other.compress_q_dimension() + + # first absolute fit column and the core/shell boundary inside the window (the window never reaches the + # negative shell, since it spans at most a fifth of the positive half) + lo = 2 * niv_res - n_freq_fit + split = min(max(niv_res + self.niv, lo), 2 * niv_res) + + mat_fit = np.empty(self.current_shape[:-1] + (n_freq_fit,), dtype=self.mat.dtype) + mat_fit[..., : split - lo] = self.mat[..., lo - niv_diff : split - niv_diff] + mat_fit[..., split - lo :] = other.mat[..., split:] + + fitdata = np.mean(mat_fit.reshape(*self.nq, self.n_bands, self.n_bands, n_freq_fit), axis=(0, 1, 2)) + iv = 1j * MFHelper.vn(niv_res, self._beta, return_only_positive=True) + iwfit = iv[niv_res - n_freq_fit :][None, None, :] + + mom0 = np.mean(fitdata.real, axis=-1) + mom1 = np.mean(fitdata.imag * iwfit.imag, axis=-1) + return mom0, mom1 + def fit_polynomial(self, n_fit: int = 4, degree: int = 3, niv_core: int = 0) -> "SelfEnergy": """ Replaces the self-energy by a per-(momentum, orbital) polynomial fit of the positive-frequency data, diff --git a/tests/test_autodetect_memory.py b/tests/test_autodetect_memory.py index b424d1e..4be5c02 100644 --- a/tests/test_autodetect_memory.py +++ b/tests/test_autodetect_memory.py @@ -14,7 +14,9 @@ # the q-grid / box parameters the fake_system fixture installs, so tests can reproduce the driver's estimate # niv_cut == min(niw_core + niv_full + 10, niv_dmft) == min(32, 50) == 32 with the fixture's niv_dmft below. -FIXTURE_PARAMS = dict(n_bands=1, nk_tot=256, nk_irr=40, niw_core=10, niv_core=10, niv_full=12, niv_cut=32, niv_pp=5) +FIXTURE_PARAMS = dict( + n_bands=1, nk_tot=256, nk_irr=40, niw_core=10, niv_core=10, niv_full=12, niv_cut=32, niv_dmft=50, niv_pp=5 +) @pytest.fixture diff --git a/tests/test_eliashberg_solver.py b/tests/test_eliashberg_solver.py index c927047..48bd55c 100644 --- a/tests/test_eliashberg_solver.py +++ b/tests/test_eliashberg_solver.py @@ -2021,7 +2021,8 @@ def test_streaming_fq_file_has_gather_layout_and_matches_pp_band(setup): def test_dispatch_selects_slice_or_streaming_construction_by_save_fq(setup, monkeypatch, save_fq): """dispatch_full_vertex_calculation routes to the slice constructor, or to the streaming one under save_fq.""" config.eliashberg.save_fq = save_fq - monkeypatch.setattr(LocalFourPoint, "load", staticmethod(lambda *a, **k: MagicMock(spec=LocalFourPoint))) + loader = MagicMock(return_value=(MagicMock(spec=LocalFourPoint), None)) + monkeypatch.setattr(es.nonlocal_sde, "_load_node_shared_local_vertex", loader) slice_mock = MagicMock(return_value="slice") streaming_mock = MagicMock(return_value="streamed") monkeypatch.setattr(es, "create_pairing_vertex_slice_q_r", slice_mock) @@ -2036,6 +2037,26 @@ def test_dispatch_selects_slice_or_streaming_construction_by_save_fq(setup, monk assert slice_mock.call_count == (0 if save_fq else 1) +def test_dispatch_loads_the_vertex_node_shared_and_frees_its_window(setup, monkeypatch): + """dispatch_full_vertex_calculation loads the local vertex through the node-shared loader and frees the window.""" + config.eliashberg.save_fq = False + win = MagicMock() + loader = MagicMock(return_value=(MagicMock(spec=LocalFourPoint), win)) + free_mock = MagicMock() + monkeypatch.setattr(es.nonlocal_sde, "_load_node_shared_local_vertex", loader) + monkeypatch.setattr(es.nonlocal_sde, "_free_shared_window", free_mock) + monkeypatch.setattr(es, "create_pairing_vertex_slice_q_r", MagicMock(return_value="slice")) + node_comm = MagicMock() + + result = es.dispatch_full_vertex_calculation( + SpinChannel.DENS, MagicMock(), MagicMock(), 2, _make_single_rank_distributor(), node_comm=node_comm + ) + + assert result == "slice" + assert loader.call_args.args[0] is node_comm + free_mock.assert_called_once_with(win, node_comm) + + def test_solver_grid_shape_degenerates_and_caps(): """The solver grid uses row-only splitting up to the frequency count, then widens by divisor columns.""" assert es.solver_grid_shape(1, 80) == (1, 1) @@ -2093,6 +2114,55 @@ def test_grid_solver_on_one_rank_matches_in_memory_solver(): assert np.allclose(np.abs(gap_new.mat), np.abs(gap_ref.mat), atol=1e-4) +def test_grid_solver_ships_the_bubble_inside_the_grid_only(monkeypatch): + """With idle ranks beyond the solver grid, the pp bubble broadcast runs over the grid communicator only.""" + from copy import deepcopy + + monkeypatch.setattr(es, "MPI", conftest.FAKE_MPI) + + def fake_eigsh(op, k, tol, v0, which, maxiter): + # thread-unsafe ARPACK is densified via n identical matvecs per rank, keeping the lockstep collectives aligned + n = op.shape[0] + dense = np.column_stack([op.matvec(np.eye(n, dtype=np.complex64)[:, i]) for i in range(n)]) + lam, vec = np.linalg.eig(dense) + order = np.argsort(lam.real)[::-1][:k] + return lam.real[order], vec[:, order] + + monkeypatch.setattr("dgamore.eliashberg_solver.sp.sparse.linalg.eigsh", fake_eigsh) + seen = [] + real_bcast = mu.bcast_rows + monkeypatch.setattr( + es.mpi_utils, + "bcast_rows", + lambda comm, arr, root, **k: seen.append(comm.Get_size()) or real_bcast(comm, arr, root, **k), + ) + nq, o, niv_pp, size = (4, 2, 1), 1, 1, 3 + _grid_test_config(nq, niv_pp) + gamma, chi0 = _random_pairing_vertex_and_bubble(nq, o, niv_pp, 5) + nq_tot = int(np.prod(nq)) + bounds = np.linspace(0, nq_tot, size + 1).astype(int) + + def worker(comm, rank): + _grid_test_config(nq, niv_pp) + local = FourPoint( + gamma.mat[bounds[rank] : bounds[rank + 1]].copy(), + SpinChannel.SING, + nq, + 0, + 2, + False, + True, + True, + FrequencyNotation.PP, + ) + return es.solve_eliashberg_lanczos_grid(local, deepcopy(chi0) if rank == 0 else None, comm, 0) + + conftest.run_parallel(size, worker) + # the vertex-slice gather is 3 sources x 3 participants; the bubble must add grid-sized (2) calls, never full (3) + assert sum(1 for s in seen if s == size) == size * size + assert sum(1 for s in seen if s == 2) == 2 + + @pytest.mark.parametrize("size, niv_pp", [(2, 2), (4, 1), (3, 2), (3, 1)]) def test_grid_solver_multi_rank_matches_in_memory_solver(monkeypatch, size, niv_pp): """Row-split, 2x2, uneven-row and idle-rank grids reproduce the in-memory eigenvalues on the fake MPI.""" diff --git a/tests/test_greens_function.py b/tests/test_greens_function.py index 1afb970..9c1111e 100644 --- a/tests/test_greens_function.py +++ b/tests/test_greens_function.py @@ -244,13 +244,71 @@ def test_update_mu_without_logger_is_silent_on_failure(): assert out == mu -def test_update_mu_with_logger_logs_on_failure(): - """update_mu logs a debug message and returns the input mu when root-finding fails.""" +def test_update_mu_with_logger_warns_on_failure(): + """update_mu logs a warning and returns the input mu when no root exists for the target filling.""" nk, ek, sig, beta, mu = _toy_inputs() logger = MagicMock() out = update_mu(mu, 1e9, ek, sig.mat, beta, sig.smom[0], logger=logger) assert out == mu - logger.debug.assert_called_once() + logger.warning.assert_called_once() + + +def test_update_mu_falls_back_to_bracketed_search_when_newton_fails(monkeypatch): + """update_mu recovers the true root through the bracketed Brent search when the Newton solver fails.""" + from dgamore.greens_function import root_fun + + nk, ek, sig, beta, mu = _toy_inputs() + target = root_fun(mu + 1.3, 0.0, ek, sig.mat, beta, sig.smom[0]) + with monkeypatch.context() as mp: + mp.setattr("dgamore.greens_function.opt.newton", MagicMock(side_effect=RuntimeError)) + out = update_mu(mu, target, ek, sig.mat, beta, sig.smom[0]) + assert np.allclose(out, mu + 1.3, atol=1e-5) + + +def test_update_mu_rejects_a_false_newton_root_via_the_residual_check(monkeypatch): + """update_mu discards a Newton result whose filling residual is large and finds the actual root instead.""" + from dgamore.greens_function import root_fun + + nk, ek, sig, beta, mu = _toy_inputs() + target = root_fun(mu + 1.3, 0.0, ek, sig.mat, beta, sig.smom[0]) + with monkeypatch.context() as mp: + mp.setattr("dgamore.greens_function.opt.newton", MagicMock(return_value=mu - 60.0)) + out = update_mu(mu, target, ek, sig.mat, beta, sig.smom[0]) + assert np.allclose(out, mu + 1.3, atol=1e-5) + + +def test_update_mu_logs_bracketed_fallback_at_info_level(monkeypatch): + """update_mu logs the bracketed fallback at info level, not warning, when Newton fails but a root exists.""" + from dgamore.greens_function import root_fun + + nk, ek, sig, beta, mu = _toy_inputs() + target = root_fun(mu + 1.3, 0.0, ek, sig.mat, beta, sig.smom[0]) + logger = MagicMock() + with monkeypatch.context() as mp: + mp.setattr("dgamore.greens_function.opt.newton", MagicMock(side_effect=RuntimeError)) + update_mu(mu, target, ek, sig.mat, beta, sig.smom[0], logger=logger) + logger.info.assert_called_once() + logger.warning.assert_not_called() + + +def test_find_mu_bracket_encloses_the_nearest_root(monkeypatch): + """_find_mu_bracket returns an interval around the root closest to the start, excluding the farther one.""" + import dgamore.greens_function as gf_module + from dgamore.greens_function import _find_mu_bracket + + mu = 0.3 + with monkeypatch.context() as mp: + mp.setattr(gf_module, "root_fun", lambda x, *a: (x - (mu + 0.1)) * (x - (mu + 3.0))) + lo, hi = _find_mu_bracket(mu, ()) + assert lo < mu + 0.1 < hi and hi < mu + 3.0 + + +def test_find_mu_bracket_returns_none_without_sign_change(): + """_find_mu_bracket gives up with None when the filling residual never changes sign.""" + from dgamore.greens_function import _find_mu_bracket + + nk, ek, sig, beta, mu = _toy_inputs() + assert _find_mu_bracket(mu, (1e9, ek, sig.mat, beta, sig.smom[0])) is None def test_update_mu_forwards_newton_tolerance(monkeypatch): diff --git a/tests/test_max_ent.py b/tests/test_max_ent.py index 09ccb67..965e58f 100644 --- a/tests/test_max_ent.py +++ b/tests/test_max_ent.py @@ -180,6 +180,22 @@ def patch_maxent_mpi(monkeypatch): monkeypatch.setattr(max_ent, "AnalyticContinuationProblem", _fake_problem) +def test_perform_maxent_giwk_prepares_the_continued_g_on_rank0_only(tmp_path, patch_maxent_mpi, monkeypatch): + """perform_maxent_giwk builds the full-BZ band-basis input on rank 0 only; other ranks receive their slice.""" + nk, n_bands, niv_core, w_count = (4, 4, 1), 2, 3, 7 + _setup_maxent_config(tmp_path, nk, n_bands, niv_core, w_count, seed=7) + mat = _build_giwk_mat(nk, n_bands, niv=4, seed=11) + calls = [] + real_rotation = max_ent.orbital_to_band_basis + monkeypatch.setattr(max_ent, "orbital_to_band_basis", lambda *a: calls.append(1) or real_rotation(*a)) + + def fn(comm, rank): + return max_ent.perform_maxent_giwk(GreensFunction(mat.copy(), nk=config.lattice.nk), "TEST", comm) + + run_parallel(2, fn) + assert len(calls) == 1 + + @pytest.mark.parametrize("size", [1, 2]) def test_perform_maxent_giwk_continues_band_diagonal_and_unfolds(tmp_path, patch_maxent_mpi, size): """perform_maxent_giwk continues the band-diagonal G and unfolds it, unlike orbital-diagonal continuation.""" diff --git a/tests/test_memory_estimator.py b/tests/test_memory_estimator.py index d57eae1..097584b 100644 --- a/tests/test_memory_estimator.py +++ b/tests/test_memory_estimator.py @@ -31,6 +31,7 @@ niv_core=30, niv_full=40, niv_cut=80, # min(niw_core + niv_full + 10, niv_dmft) == 80 here + niv_dmft=120, niv_pp=15, n_ranks=4, with_eliashberg=False, @@ -44,6 +45,7 @@ niv_core=5, niv_full=6, niv_cut=15, + niv_dmft=20, niv_pp=2, n_ranks=4, with_eliashberg=False, @@ -249,6 +251,15 @@ def test_sde_transient_is_the_irr_kernel_plus_the_bounded_exchange_chunks(): assert estimate_peaks(**TINY)["sde"].off_distributed == pytest.approx(expected) +def test_sde_single_covers_the_rank0_occupation_step(): + """The sde single-rank slot grows to the DMFT-box sigma/giwk pair once that exceeds the finalize buffers.""" + small = estimate_peaks(**{**TINY, "niv_dmft": TINY["niv_core"]})["sde"].off_single + big = estimate_peaks(**{**TINY, "niv_dmft": 100 * TINY["niv_core"]})["sde"].off_single + nb = TINY["n_bands"] + assert big == pytest.approx(SCALE * 2 * TINY["nk_tot"] * nb**2 * 2 * 100 * TINY["niv_core"]) + assert small < big + + def test_sde_chunk_term_is_capped_by_the_byte_budget(monkeypatch): """Once the per-rank full-BZ kernel exceeds the byte budget, the sde transient stops growing with the grid.""" monkeypatch.setattr(memory_estimator, "SLICE_CHUNK_BYTES", 64) diff --git a/tests/test_mpi_utils.py b/tests/test_mpi_utils.py index 8952c11..0424046 100644 --- a/tests/test_mpi_utils.py +++ b/tests/test_mpi_utils.py @@ -266,6 +266,39 @@ def test_get_pencil_indices_partition(layout, size): assert np.array_equal(np.sort(allidx), np.arange(n_tot)) +def test_cgroup_memory_limit_walks_v2_ancestors_to_the_smallest_set_limit(tmp_path): + """The v2 reader skips a "max" leaf and returns the smallest configured ancestor memory.max.""" + (tmp_path / "proc_cgroup").write_text("0::/a/b\n") + (tmp_path / "cg" / "a" / "b").mkdir(parents=True) + (tmp_path / "cg" / "a" / "b" / "memory.max").write_text("max\n") + (tmp_path / "cg" / "a" / "memory.max").write_text("1234567\n") + assert mu.cgroup_memory_limit(str(tmp_path / "proc_cgroup"), str(tmp_path / "cg")) == 1234567 + + +def test_cgroup_memory_limit_reads_v1_controller_and_ignores_unlimited(tmp_path): + """The v1 fallback reads memory.limit_in_bytes and treats huge sentinel values as unlimited.""" + (tmp_path / "proc_cgroup").write_text("9:memory:/slurm/job1\n") + (tmp_path / "cg" / "memory" / "slurm" / "job1").mkdir(parents=True) + (tmp_path / "cg" / "memory" / "slurm" / "job1" / "memory.limit_in_bytes").write_text("2222\n") + assert mu.cgroup_memory_limit(str(tmp_path / "proc_cgroup"), str(tmp_path / "cg")) == 2222 + (tmp_path / "cg" / "memory" / "slurm" / "job1" / "memory.limit_in_bytes").write_text(str(2**63 - 4096)) + assert mu.cgroup_memory_limit(str(tmp_path / "proc_cgroup"), str(tmp_path / "cg")) is None + + +def test_cgroup_memory_limit_returns_none_without_cgroup_information(tmp_path): + """A missing cgroup file (or one without limits) yields None, so the caller falls back to host memory.""" + assert mu.cgroup_memory_limit(str(tmp_path / "absent"), str(tmp_path / "cg")) is None + + +def test_job_memory_total_caps_the_hardware_total_by_the_cgroup_limit(monkeypatch): + """job_memory_total returns the hardware total capped by the cgroup limit, or the plain total without one.""" + monkeypatch.setattr(mu.psutil, "virtual_memory", lambda: SimpleNamespace(total=1000)) + monkeypatch.setattr(mu, "cgroup_memory_limit", lambda: 200) + assert mu.job_memory_total() == 200 + monkeypatch.setattr(mu, "cgroup_memory_limit", lambda: None) + assert mu.job_memory_total() == 1000 + + def test_get_pencil_indices_is_cached(): """Repeated get_pencil_indices calls with the same arguments are served from the cache with equal content.""" first = mu.get_pencil_indices(1, 3, (4, 4, 2), "y_pencil") diff --git a/tests/test_nonlocal_sde.py b/tests/test_nonlocal_sde.py index f150120..38e5047 100644 --- a/tests/test_nonlocal_sde.py +++ b/tests/test_nonlocal_sde.py @@ -823,6 +823,11 @@ def _setup_self_energy_loop(monkeypatch, tmp_path, proposal_step, max_iter=10, e ) monkeypatch.setattr(nonlocal_sde, "GreensFunction", SimpleNamespace(get_g_full=lambda *a, **k: gf_stub)) monkeypatch.setattr(nonlocal_sde, "update_mu", lambda *a, **k: 0.5) + monkeypatch.setattr( + nonlocal_sde, + "_update_occ_and_energies_distributed", + lambda *a: (1.0, np.zeros((1, 1)), np.zeros((1, 1, 1, 1, 1)), 0.0, 0.0), + ) calls = [] @@ -1026,6 +1031,103 @@ def test_create_auxiliary_chi_r_q_sum_is_chunk_size_invariant(monkeypatch): assert np.allclose(chunked.mat, whole.mat, atol=1e-6) +def test_update_occ_and_energies_distributed_matches_the_full_box_evaluation(monkeypatch): + """The k-distributed occupation/energy evaluation matches the single-rank full-box reference.""" + monkeypatch.setattr(mpi_utils, "MPI", FAKE_MPI) + nk, o, niv, niv_dmft, beta, mu = (4, 2, 1), 2, 3, 8, 9.0, 0.4 + nk_tot = int(np.prod(nk)) + rng = np.random.default_rng(9) + config.sys.beta, config.sys.n_bands, config.sys.mu = beta, o, mu + config.lattice.nk = nk + config.lattice.k_grid = SimpleNamespace(nk_tot=nk_tot, nk=nk) + ek = rng.standard_normal((*nk, o, o)) + ek = ek + ek.swapaxes(-1, -2) + config.lattice.hamiltonian = MagicMock(get_ek=MagicMock(return_value=ek)) + sig_mat = (rng.standard_normal((nk_tot, o, o, 2 * niv)) * 0.1 + 0.3j).astype(np.complex64) + dmft_mat = (rng.standard_normal((1, 1, 1, o, o, 2 * niv_dmft)) * 0.1 + 0.2j).astype(np.complex64) + sigma_new = SelfEnergy(sig_mat.copy(), nk, has_compressed_q_dimension=True, beta=beta) + sigma_dmft_full = SelfEnergy(dmft_mat.copy(), (1, 1, 1), beta=beta) + + sigma_ref = sigma_new.copy().concatenate_self_energies(sigma_dmft_full) + giwk_ref = GreensFunction.get_g_full(sigma_ref, mu, ek, beta) + _, occ_ref, occ_k_ref = giwk_ref.get_fill_nonlocal() + ekin_ref, epot_ref = giwk_ref.get_ekin(), giwk_ref.get_epot() + + def fn(comm, rank): + d_full = mpi_utils.MpiDistributor(ntasks=nk_tot, comm=comm) + return nonlocal_sde._update_occ_and_energies_distributed(sigma_new, sigma_dmft_full, d_full, mu) + + _, res = run_parallel(2, fn) + # occupations reproduce the full-box reference bit-for-bit; only the energy scalars regroup their k-sums + for _, occ, occ_k, ekin, epot in res: + assert np.array_equal(occ, occ_ref) and np.array_equal(occ_k, occ_k_ref) + assert np.allclose([ekin, epot], [ekin_ref, epot_ref], atol=1e-5) + + +def test_update_occ_and_energies_distributed_pins_the_occupation_dtype_across_ranks(monkeypatch): + """Ranks whose fill comes out float (all-real eigenvalues) and ranks with complex fill reduce to one complex128 occ_k.""" + monkeypatch.setattr(mpi_utils, "MPI", FAKE_MPI) + nk, o, niv, niv_dmft, beta, mu = (4, 2, 1), 2, 3, 8, 9.0, 0.4 + nk_tot = int(np.prod(nk)) + rng = np.random.default_rng(11) + config.sys.beta, config.sys.n_bands, config.sys.mu = beta, o, mu + config.lattice.nk = nk + config.lattice.k_grid = SimpleNamespace(nk_tot=nk_tot, nk=nk) + ek = rng.standard_normal((*nk, o, o)) + config.lattice.hamiltonian = MagicMock(get_ek=MagicMock(return_value=ek + ek.swapaxes(-1, -2))) + sigma_new = SelfEnergy( + (rng.standard_normal((nk_tot, o, o, 2 * niv)) * 0.1 + 0.3j).astype(np.complex64), + nk, + has_compressed_q_dimension=True, + beta=beta, + ) + sigma_dmft_full = SelfEnergy( + (rng.standard_normal((1, 1, 1, o, o, 2 * niv_dmft)) * 0.1 + 0.2j).astype(np.complex64), (1, 1, 1), beta=beta + ) + + def fake_get_g_full(sigma_occ, mu_in, ek_slice, beta_in): + # the first slice (rows starting at ek row 0) plays the all-real-eigenvalue rank and returns a float fill + n_my = ek_slice.shape[0] + is_first = np.allclose(ek_slice[0, 0, 0], config.lattice.hamiltonian.get_ek().reshape(nk_tot, o, o)[0]) + dtype, value = (np.float64, 0.25) if is_first else (np.complex128, 0.75 + 0.5j) + g = MagicMock() + g.get_fill_nonlocal.return_value = (0.0, None, np.full((n_my, 1, 1, o, o), value, dtype=dtype)) + g.get_ekin.return_value = 1.0 + g.get_epot.return_value = 2.0 + return g + + def fn(comm, rank): + d_full = mpi_utils.MpiDistributor(ntasks=nk_tot, comm=comm) + return nonlocal_sde._update_occ_and_energies_distributed(sigma_new, sigma_dmft_full, d_full, mu) + + with monkeypatch.context() as mp: + mp.setattr(nonlocal_sde.GreensFunction, "get_g_full", fake_get_g_full) + _, res = run_parallel(2, fn) + for _, occ, occ_k, ekin, epot in res: + assert occ_k.dtype == np.complex128 + assert np.allclose(occ_k.reshape(nk_tot, o, o)[0], 0.25) and np.allclose( + occ_k.reshape(nk_tot, o, o)[-1], 0.75 + 0.5j + ) + + +def test_share_sigma_per_node_gives_each_node_one_shared_buffer_with_rank0_values(monkeypatch): + """_share_sigma_per_node leaves every rank viewing its node's single shared buffer holding rank 0's array.""" + monkeypatch.setattr(mpi_utils, "MPI", FAKE_MPI) + mixed = (np.arange(2 * 2 * 6).reshape(2, 2, 6) + 1j).astype(np.complex64) + + def fn(comm, rank): + sigma = SimpleNamespace(mat=mixed.copy() if rank == 0 else np.zeros_like(mixed)) + node_comm = comm.Split_type(0) + roots_comm = comm.Split(0 if node_comm.rank == 0 else 1) + win = nonlocal_sde._share_sigma_per_node(sigma, node_comm, roots_comm) + return sigma.mat.copy(), sigma.mat.__array_interface__["data"][0], win is not None + + _, res = run_parallel(4, fn, hostnames=["n0", "n0", "n1", "n1"]) + pointers = {r[1] for r in res} + assert all(np.array_equal(r[0], mixed) for r in res) + assert len(pointers) == 2 and all(r[2] for r in res) + + @pytest.mark.parametrize("negative_w", [False, True]) def test_fft_sde_pass_is_invariant_under_the_w_chunk_size(negative_w, monkeypatch): """A one-byte chunk budget reproduces the all-w-at-once result of the chunked FFT self-energy pass.""" diff --git a/tests/test_self_energy.py b/tests/test_self_energy.py index 1564eb0..6ca8b12 100644 --- a/tests/test_self_energy.py +++ b/tests/test_self_energy.py @@ -314,6 +314,29 @@ def test_raises_error_when_concatenating_with_smaller_niv(): self_energy1.concatenate_self_energies(self_energy2) +@pytest.mark.parametrize("niv_core, niv_shell", [(3, 12), (8, 10), (10, 10)]) +def test_fit_smom_concatenated_is_bit_identical_to_the_full_concatenation_fit(niv_core, niv_shell): + """fit_smom_concatenated equals the momentum-resolved concatenation's own moment fit bit-for-bit.""" + nk_odd = (3, 2, 1) + rng = np.random.default_rng(7) + core_mat = (rng.standard_normal((int(np.prod(nk_odd)), 2, 2, 2 * niv_core)) * 0.1 + 0.3j).astype(np.complex64) + shell_mat = (rng.standard_normal((1, 1, 1, 2, 2, 2 * niv_shell)) * 0.1 + 0.2j).astype(np.complex64) + core = _se(core_mat, nk=nk_odd, has_compressed_q_dimension=True) + shell = _se(shell_mat, nk=(1, 1, 1)) + ref0, ref1 = core.copy().concatenate_self_energies(shell).smom + mom0, mom1 = core.fit_smom_concatenated(shell) + assert np.array_equal(mom0, ref0) and np.array_equal(mom1, ref1) + + +def test_fit_smom_concatenated_raises_for_smaller_shell(): + """fit_smom_concatenated raises when the shell donor has fewer frequencies than the core.""" + self_energy1 = _se(mat_decompressed, nk=nk, has_compressed_q_dimension=False) + smaller_mat = np.random.rand(*nk, 2, 2, 2 * (niv - 1)) + self_energy2 = _se(smaller_mat, nk=nk, has_compressed_q_dimension=False) + with pytest.raises(ValueError, match="Can not concatenate with a self-energy that has less frequencies."): + self_energy1.fit_smom_concatenated(self_energy2) + + def test_concatenates_self_energies_correctly_with_larger_niv(): """concatenate_self_energies wraps the core in the larger self-energy's tails.""" self_energy1 = _se(mat_decompressed, nk=nk, has_compressed_q_dimension=False)