# SPDX-FileCopyrightText: 2025-2026 Julian Peil <julian.peil@tuwien.ac.at>
# SPDX-License-Identifier: MIT
#
# DGAmore — Multi-Orbital Ladder Dynamical Vertex Approximation (LDGA) &
# Eliashberg Equation Solver for Strongly Correlated Electron Systems
r"""
Momentum-dependent four-point objects. :class:`FourPoint` extends :class:`LocalFourPoint` with one momentum axis
(see :class:`IAmNonLocal`) to represent quantities such as the ladder susceptibility
:math:`\chi_{abcd}^{q\omega\nu\nu'}` and vertex :math:`F^{q}` over the (irreducible) BZ. The arithmetic and
compound-index machinery mirrors :class:`LocalFourPoint` but accounts for the extra momentum dimension; these
operations are the performance and memory bottleneck of the non-local ladder DGA step, so several variants are
provided that trade speed for footprint. Notation mirrors the thesis (Chapters 3 & 4).
"""
import gc
from copy import deepcopy
import numpy as np
from dgamore.brillouin_zone import KGrid
from dgamore.interaction import Interaction, LocalInteraction
from dgamore.local_four_point import LocalFourPoint
from dgamore.n_point_base import IAmNonLocal, SpinChannel, FrequencyNotation, DTYPE
[docs]
class FourPoint(IAmNonLocal, LocalFourPoint):
"""
This class is used to represent a non-local four-point object in a given channel with one momentum dimension,
four orbital dimensions and variable bosonic and fermionic frequency dimensions. Calculations using these objects
are the bottleneck of the DGA algorithm and need to be optimized for performance and memory usage.
"""
def __init__(
self,
mat: np.ndarray,
channel: SpinChannel = SpinChannel.NONE,
nq: tuple[int, int, int] = (1, 1, 1),
num_wn_dimensions: int = 1,
num_vn_dimensions: int = 2,
full_niw_range: bool = True,
full_niv_range: bool = True,
has_compressed_q_dimension: bool = False,
frequency_notation: FrequencyNotation = FrequencyNotation.PH,
):
r"""
Initializes the momentum-dependent four-point object from an array and its channel/momentum/frequency metadata.
:param mat: Underlying array with one momentum axis (compressed or not), four orbital axes, then frequency axes.
:param channel: Spin channel of the object (see :class:`SpinChannel`).
:param nq: Number of momenta per spatial direction.
:param num_wn_dimensions: Number of bosonic frequency axes (0–1).
:param num_vn_dimensions: Number of fermionic frequency axes (0–2).
:param full_niw_range: Whether the object spans the full (signed) bosonic range or only :math:`\omega \geq 0`.
:param full_niv_range: Whether the object spans the full (signed) fermionic range or only :math:`\nu \geq 0`.
:param has_compressed_q_dimension: Whether the momentum is stored as a single compressed axis ``[q, ...]``
(True) or as ``[qx, qy, qz, ...]`` (False).
:param frequency_notation: Frequency convention (see :class:`FrequencyNotation`).
"""
LocalFourPoint.__init__(
self,
mat,
channel,
num_wn_dimensions,
num_vn_dimensions,
full_niw_range,
full_niv_range,
frequency_notation,
)
IAmNonLocal.__init__(self, mat, nq, has_compressed_q_dimension)
def __add__(self, other) -> "FourPoint":
"""
Operator form of :meth:`add` (``A + B``). See :meth:`add`.
"""
return self.add(other)
def __radd__(self, other) -> "FourPoint":
"""
Reflected operator form of :meth:`add` (``B + A``). See :meth:`add`.
"""
return self.add(other)
def __sub__(self, other) -> "FourPoint":
"""
Operator form of :meth:`sub` (``A - B``). See :meth:`sub`.
"""
return self.sub(other)
def __rsub__(self, other) -> "FourPoint":
"""
Reflected operator form of :meth:`sub` (``B - A``), returning ``-(A - B)``. See :meth:`sub`.
"""
return -self.sub(other)
def __mul__(self, other) -> "FourPoint":
"""
Operator form of :meth:`mul` (``A * B``). See :meth:`mul` for the (element-wise, not matrix) semantics.
"""
return self.mul(other)
def __rmul__(self, other) -> "FourPoint":
"""
Reflected operator form of :meth:`mul` (``B * A``). See :meth:`mul`.
"""
return self.mul(other)
def __matmul__(self, other) -> "FourPoint":
"""
Operator form of :meth:`matmul` with ``self`` on the left (``A @ B``). See :meth:`matmul`.
"""
return self.matmul(other, left_hand_side=True)
def __rmatmul__(self, other) -> "FourPoint":
"""
Operator form of :meth:`matmul` with ``self`` on the right (``B @ A``). See :meth:`matmul`.
"""
return self.matmul(other, left_hand_side=False)
def __invert__(self):
"""
Operator form of :meth:`invert` (``~A``). See :meth:`invert`.
"""
return self.invert()
def __pow__(self, power, modulo=None) -> "FourPoint":
"""
Operator form of :meth:`pow` (``A ** n``), supplying the matching identity. See :meth:`pow`.
"""
return self.pow(power, FourPoint.identity_like(self))
[docs]
def sum_over_vn(self, beta: float, axis: tuple = (-1,), copy: bool = True) -> "FourPoint":
r"""
Sums over the given fermionic frequency axes and applies the Matsubara prefactor :math:`1/\beta^{n}` (with
:math:`n` the number of summed axes). The in-place branch frees the old array before allocating the result to
cap peak memory.
:param beta: Inverse temperature :math:`\beta`.
:param axis: Fermionic axes to sum over (negative indices into the frequency tail).
:param copy: If True, operate on and return a deep copy; if False, mutate and return ``self`` in place.
:return: A :class:`FourPoint` with the summed axes removed and ``num_vn_dimensions`` reduced accordingly.
:raises ValueError: If more axes are requested than the object has fermionic frequency dimensions.
"""
if len(axis) > self.num_vn_dimensions:
raise ValueError(f"Cannot sum over more fermionic axes than available in {self.current_shape}.")
if not copy:
summed = np.sum(self.mat, axis=axis)
summed *= 1.0 / beta ** len(axis) # in-place scale avoids a second full-size temporary
self.mat = None
gc.collect()
self.mat = summed
self._num_vn_dimensions -= len(axis)
self.update_original_shape()
return self
copy_mat = np.sum(self.mat, axis=axis)
copy_mat *= 1.0 / beta ** len(axis) # in-place scale avoids a second full-size temporary
self.update_original_shape()
return FourPoint(
copy_mat,
self.channel,
self.nq,
self.num_wn_dimensions,
self.num_vn_dimensions - len(axis),
self.full_niw_range,
self.full_niv_range,
self.has_compressed_q_dimension,
self.frequency_notation,
)
[docs]
def sum_over_orbitals(self, orbital_contraction: str = "abcd->ad") -> "FourPoint":
"""
Sums over orbital indices according to an einsum-style contraction on the four orbital axes (the leading
momentum axis and trailing frequency axes are preserved automatically). Mutates ``self`` in place.
:param orbital_contraction: Contraction of the form ``"abcd->..."`` whose target side is a subset of ``abcd``.
:return: ``self``, with the summed orbital axes removed.
:raises ValueError: If the left side does not have four orbitals or the right side has more indices than the left.
"""
split = orbital_contraction.split("->")
if len(split[0]) != 4 or len(split[1]) > len(split[0]):
raise ValueError("Invalid orbital contraction.")
permutation = (
f"i{split[0]}...->i{split[1]}..."
if self.has_compressed_q_dimension
else f"ijk{split[0]}...->ijk{split[1]}..."
)
self.mat = np.einsum(permutation, self.mat)
diff = len(split[0]) - len(split[1])
self.update_original_shape()
self._num_orbital_dimensions -= diff
return self
[docs]
def to_compound_indices(self) -> "FourPoint":
r"""
Converts the indices of the FourPoint object :math:`F^{wvv'}_{lmm'l'}` to compound indices :math:`F^{w}_{c_1, c_2}`
by transposing the object to [q, w, o1, o2, v, o4, o3, v'] (if the object has any fermionic frequency dimension,
otherwise the compound indices are built from orbital dimensions only) and grouping {o1, o2, v} and {o4, o3, v'}
to the new compound index. Always returns the object with a compressed momentum dimension and in the same niw
range as the original object.
:return: ``self`` with shape ``[q, w, c1, c2]`` (compound indices, compressed momentum).
:raises NotImplementedError: If the frequency notation is neither ph nor pp.
"""
if self.frequency_notation == FrequencyNotation.PH:
return self._to_compound_indices_ph()
elif self.frequency_notation == FrequencyNotation.PP:
return self._to_compound_indices_pp()
else:
raise NotImplementedError(
f"Frequency notation {self.frequency_notation} not supported for transformation to "
f"compound indices."
)
def _to_compound_indices_ph(self) -> "FourPoint":
"""
Converts the indices of the FourPoint object in the ph notation to compound indices (see
:meth:`to_compound_indices`).
:return: ``self`` in compound-index layout with compressed momentum.
:raises ValueError: If the object has no bosonic frequency dimension but not exactly two fermionic ones.
"""
if len(self.current_shape) == 4: # [q, w, x1, x2]
return self
if not self.has_compressed_q_dimension:
self.compress_q_dimension()
self.update_original_shape()
if (
self.num_wn_dimensions == 0 # [q, o1, o2, o3, o4, v, vp]
): # special case for objects without any bosonic frequency dimension (such as the pairing vertex)
if self.num_vn_dimensions != 2:
raise ValueError(
"Object must have 2 fermionic frequency dimensions if it does not have any w dimension."
)
self.mat = self.mat.reshape(self.nq_tot, self.n_bands**2 * 2 * self.niv, self.n_bands**2 * 2 * self.niv)
return self
w_dim = self.current_shape[5]
if self.num_vn_dimensions == 0: # [q, o1, o2, o3, o4, w]
self.mat = self.mat.transpose(0, 5, 1, 2, 4, 3).reshape(
self.nq_tot, w_dim, self.n_bands**2, self.n_bands**2
) # reshaping to [q,w,o1,o2,o4,o3] and then collecting {o1,o2} and {o4,o3} into two indices
return self
if self.num_vn_dimensions == 1: # [q, o1, o2, o3, o4, w, v]
self.extend_vn_to_diagonal()
# [q, o1, o2, o3, o4, w, v, vp]
self.mat = self.mat.transpose(0, 5, 1, 2, 6, 4, 3, 7).reshape(
self.nq_tot, w_dim, self.n_bands**2 * 2 * self.niv, self.n_bands**2 * 2 * self.niv
) # reshaping to [q,w,o1,o2,v,o4,o3,vp] and then collecting {o1,o2,v} and {o4,o3,vp} into two indices
return self
def _to_compound_indices_pp(self) -> "FourPoint":
"""
Converts the indices of the FourPoint object in the pp notation to compound indices. The difference to the ph
case is that the orbital indices are permuted first, since the ordering of pp quantities is not "1234" but
rather "1324".
:return: ``self`` in compound-index layout.
"""
if len(self.current_shape) == 3: # [q, w, x1, x2]
return self
return self.permute_orbitals("abcd->acbd", copy=False)._to_compound_indices_ph()
[docs]
def to_full_indices(self, shape: tuple = None) -> "FourPoint":
"""
Converts an object stored with compound indices to an object that has unraveled momentum,
orbital and frequency axes. Always returns the object with a compressed momentum dimension. This is the inverse
transformation of :meth:`to_compound_indices`. Will make use of the ``original_shape`` the object was
created or last modified with. If the ``original_shape`` is not set or is hard to obtain, the ``shape``
argument can be used to specify the original shape of the object.
:param shape: Optional override for the stored ``original_shape`` used to unravel the compound axes.
:return: ``self`` with unraveled orbital and frequency axes (compressed momentum).
:raises NotImplementedError: If the frequency notation is neither ph nor pp.
"""
if self.frequency_notation == FrequencyNotation.PH:
return self._to_full_indices_ph(shape)
elif self.frequency_notation == FrequencyNotation.PP:
return self._to_full_indices_pp(shape)
else:
raise NotImplementedError(
f"Frequency notation {self.frequency_notation} not supported for transformation to full indices."
)
def _to_full_indices_ph(self, shape: tuple = None) -> "FourPoint":
"""
Converts the indices of the FourPoint object in the ph notation back to full indices (see
:meth:`to_full_indices`).
:param shape: Optional override for the stored ``original_shape``.
:return: ``self`` with unraveled orbital and frequency axes.
:raises ValueError: If the current shape is not a compound-index layout, or there is no bosonic frequency axis.
"""
if (
len(self.current_shape) == 1 + self.num_orbital_dimensions + self.num_wn_dimensions + self.num_vn_dimensions
and self.has_compressed_q_dimension
):
return self
elif (
len(self.current_shape) == 3 + self.num_orbital_dimensions + self.num_wn_dimensions + self.num_vn_dimensions
and not self.has_compressed_q_dimension
):
return self
if (len(self.current_shape) != 4 and self.has_compressed_q_dimension) or (
len(self.current_shape) != 6 and not self.has_compressed_q_dimension
): # (q,w,x1,x2) or (qx,qy,qz,w,x1,x2)
raise ValueError(f"Converting to full indices with shape {self.current_shape} not supported.")
if self.num_wn_dimensions != 1:
raise ValueError("Number of bosonic frequency dimensions must be 1.")
self.original_shape = shape if shape is not None else self.original_shape
w_dim = self.original_shape[5] if self.has_compressed_q_dimension else self.original_shape[7]
if self.num_vn_dimensions == 0: # original was [q,o1,o2,o4,o3,w]
self.mat = self.mat.reshape(
(self.nq_tot,) + (w_dim,) + (self.n_bands,) * self.num_orbital_dimensions
).transpose(0, 2, 3, 5, 4, 1)
self._has_compressed_q_dimension = True
return self
compound_index_shape = (self.n_bands, self.n_bands, 2 * self.niv)
# original was [q,o1,o2,o4,o3,w,v,v']
self.mat = self.mat.reshape((self.nq_tot,) + (w_dim,) + compound_index_shape * 2).transpose(
0, 2, 3, 6, 5, 1, 4, 7
)
if self.num_vn_dimensions == 1: # original was [q,o1,o2,o4,o3,w,v]
# ``.copy()`` since ``diagonal`` returns a read-only view that also keeps the larger parent alive.
self.mat = self.mat.diagonal(axis1=-2, axis2=-1).copy()
return self
def _to_full_indices_pp(self, shape: tuple = None) -> "FourPoint":
"""
Converts the indices of the FourPoint object in the pp notation back to full indices. The difference to the ph
case is that the orbital indices are permuted back, since the ordering of pp quantities is not "1234" but
rather "1324".
:param shape: Optional override for the stored ``original_shape``.
:return: ``self`` with unraveled orbital and frequency axes.
"""
return self._to_full_indices_ph(shape).permute_orbitals("abcd->acbd", copy=False)
[docs]
def permute_orbitals(self, permutation: str = "abcd->abcd", copy: bool = True) -> "FourPoint":
"""
Permutes the four orbital axes according to an einsum-style string (the momentum and frequency axes are kept
fixed). Summing over orbitals is not allowed (both sides must list all four orbitals).
:param permutation: A permutation of the form ``"abcd->..."`` using exactly the four orbital labels.
:param copy: If True, operate on and return a deep copy; if False, mutate and return ``self`` in place.
:return: The orbital-permuted :class:`FourPoint` (``self`` unchanged if the permutation is the identity).
:raises ValueError: If the permutation is malformed or does not list all four orbitals on both sides.
"""
split = permutation.split("->")
if len(split) != 2 or len(split[0]) != 4 or len(split[1]) != 4:
raise ValueError("Invalid permutation.")
if split[0] == split[1]:
return self
if copy:
return deepcopy(self).permute_orbitals(permutation, copy=False)
permutation = (
f"i{split[0]}...->i{split[1]}..."
if self.has_compressed_q_dimension
else f"ijk{split[0]}...->ijk{split[1]}..."
)
self.mat = np.einsum(permutation, self.mat, optimize=True)
return self
[docs]
def map_to_full_bz(self, grid: KGrid, nq: tuple = None):
"""
Unfolds the object from the irreducible BZ to the full BZ using the grid's symmetry index map (see
:meth:`IAmNonLocal._map_to_full_bz`), with four orbital dimensions.
:param grid: The :class:`KGrid` providing the irreducible-to-full BZ index mapping.
:param nq: Optional number of momenta per direction for the unfolded grid; defaults to the object's ``nq``.
:return: ``self`` defined on the full BZ.
"""
return self._map_to_full_bz(grid, 4, nq)
[docs]
def add(self, other) -> "FourPoint":
"""
Adds ``other`` to this object (operator ``+``); see :meth:`_add` for the accepted operands and the niw-range
handling.
:param other: A :class:`FourPoint`, :class:`LocalFourPoint`, :class:`Interaction`, :class:`LocalInteraction`,
numpy array, or number.
:return: A new :class:`FourPoint` holding the sum.
"""
return self._add(other)
def _add(self, other, subtract: bool = False) -> "FourPoint":
"""
Helper method that allows for addition of FourPoint objects and other FourPoint, LocalFourPoint, Interaction or
LocalInteraction objects. Additions with numpy arrays, floats, ints or complex numbers are also supported.
Depending on the number of frequency and momentum dimensions, the vertices have to be added slightly different.
If the objects have different niw ranges, they will be converted to the half niw range before the addition.
Objects will always be returned in the half niw range to save memory.
:param other: A :class:`FourPoint`, :class:`LocalFourPoint`, :class:`Interaction`, :class:`LocalInteraction`,
numpy array, or number. Local operands are broadcast over the momentum axis.
:param subtract: If True, subtract ``other`` instead of adding it (used by :meth:`sub` to avoid a negated copy).
:return: A new :class:`FourPoint` (in the half niw range for the vertex-vertex case).
:raises ValueError: If ``other`` has an unsupported type.
"""
if not isinstance(
other, (FourPoint, LocalFourPoint, Interaction, LocalInteraction, np.ndarray, float, int, complex)
):
raise ValueError(f"Operations '+/-' for {type(self)} and {type(other)} not supported.")
op = np.subtract if subtract else np.add
if isinstance(other, (np.ndarray, float, int, complex)):
return FourPoint(
op(self.mat, other),
self.channel,
self.nq,
self.num_wn_dimensions,
self.num_vn_dimensions,
self.full_niw_range,
self.full_niv_range,
self.has_compressed_q_dimension,
self.frequency_notation,
)
channel = self.channel if self.channel != SpinChannel.NONE else other.channel
if isinstance(other, (Interaction, LocalInteraction)):
self.compress_q_dimension()
other_mat = other.mat[None, ...] if not isinstance(other, Interaction) else other.compress_q_dimension().mat
other_mat = other_mat.reshape(other.mat.shape + (1,) * (self.num_wn_dimensions + self.num_vn_dimensions))
return FourPoint(
op(self.mat, other_mat),
self.channel,
self.nq,
self.num_wn_dimensions,
self.num_vn_dimensions,
self.full_niw_range,
self.full_niv_range,
self.has_compressed_q_dimension,
self.frequency_notation,
)
self_full_niw_range = self.full_niw_range
other_full_niw_range = other.full_niw_range
self.to_half_niw_range()
other = other.to_half_niw_range()
if not isinstance(other, FourPoint):
# if other is LocalFourPoint
other, self_extended, other_extended = self._align_frequency_dimensions_for_operation(other)
result = FourPoint(
op(
self.mat,
(other.mat[None, ...] if self.has_compressed_q_dimension else other.mat[None, None, None, ...]),
),
channel,
self.nq,
self.num_wn_dimensions,
max(self.num_vn_dimensions, other.num_vn_dimensions),
False,
self.full_niv_range,
self.has_compressed_q_dimension,
self.frequency_notation,
)
if self_full_niw_range:
self.to_full_niw_range()
if other_full_niw_range:
other = other.to_full_niw_range()
other = self._revert_frequency_dimensions_after_operation(other, other_extended, self_extended)
return result
other = self._align_q_dimensions_for_operations(other)
other, self_extended, other_extended = self._align_frequency_dimensions_for_operation(other)
result = FourPoint(
op(self.mat, other.mat),
channel,
self.nq,
self.num_wn_dimensions,
self.num_vn_dimensions,
False,
self.full_niv_range,
self.has_compressed_q_dimension,
self.frequency_notation,
)
if self_full_niw_range:
self.to_full_niw_range()
if other_full_niw_range:
other = other.to_full_niw_range()
other = self._revert_frequency_dimensions_after_operation(other, other_extended, self_extended)
return result
[docs]
def sub(self, other) -> "FourPoint":
"""
Helper method that allows for subtraction of FourPoint objects and other FourPoint, LocalFourPoint, Interaction or
LocalInteraction objects. Subtractions with numpy arrays, floats, ints or complex numbers are also supported.
Depending on the number of frequency and momentum dimensions, the vertices have to be subtracted slightly different.
If the objects have different niw ranges, they will be converted to the half niw range before the subtraction.
Objects will always be returned in the half niw range to save memory.
:param other: A :class:`FourPoint`, :class:`LocalFourPoint`, :class:`Interaction`, :class:`LocalInteraction`,
numpy array, or number.
:return: The difference, implemented as ``self._add(other, subtract=True)`` (see :meth:`_add`).
:raises ValueError: Propagated from :meth:`_add` for unsupported operands.
"""
return self._add(other, subtract=True)
[docs]
def mul(self, other) -> "FourPoint":
r"""
Allows for the multiplication with a number, a numpy array or another FourPoint object. This is different from
the :meth:`matmul` method, which is used for matrix multiplication.
In the case the other object is a FourPoint object, we require that both objects have only one fermionic
frequency dimension, such that :math:`A_{abcd}^{qv} * B_{dcef}^{qv'} = C_{abef}^{qvv'}`. This is needed to
construct the full vertex, see Eq. (3.139) in my thesis. Returns the object in the half niw range.
:param other: A number, numpy array, or :class:`FourPoint`.
:return: A new :class:`FourPoint` (in the half niw range for the four-point case).
:raises ValueError: If ``other`` has an unsupported type, or either four-point operand does not have exactly
one fermionic frequency dimension.
"""
if not isinstance(other, (int, float, complex, np.ndarray, FourPoint)):
raise ValueError("Multiplication only supported with numbers, numpy arrays or FourPoint objects.")
if not isinstance(other, FourPoint):
copy = deepcopy(self)
copy.mat *= other
return copy
if self.num_vn_dimensions != 1 or other.num_vn_dimensions != 1:
raise ValueError("Both objects must have only one fermionic frequency dimension.")
is_self_full_niw_range = self.full_niw_range
is_other_full_niw_range = other.full_niw_range
self.to_half_niw_range()
other = other.to_half_niw_range()
result_mat = self.times("qabcdwv,qdcefwp->qabefwvp", other)
if is_self_full_niw_range:
self.to_full_niw_range()
if is_other_full_niw_range:
other = other.to_full_niw_range()
return FourPoint(result_mat, self.channel, self.nq, 1, 2, False, True, True, self.frequency_notation)
[docs]
def matmul(self, other, left_hand_side: bool = True) -> "FourPoint":
"""
Helper method that allows for a matrix multiplication between FourPoint and FourPoint, LocalFourPoint,
Interaction and LocalInteraction objects. Depending on the number of frequency and momentum dimensions,
the objects have to be multiplied differently. The use of einsum is very crucial for memory efficiency here,
as a regular matrix multiplication in compound index space would create large intermediate arrays if one of both
partaking objects has less than two fermionic frequency dimensions. Result objects will always be returned in
half of their niw range to save memory.
:param other: A :class:`FourPoint`, :class:`LocalFourPoint`, :class:`Interaction`, or :class:`LocalInteraction`.
Local operands are broadcast over the momentum axis.
:param left_hand_side: If True, compute ``self @ other``; if False, compute ``other @ self``.
:return: A new :class:`FourPoint` in the half bosonic frequency range, carrying the non-NONE channel.
:raises ValueError: If ``other`` has an unsupported type.
"""
if not isinstance(other, (FourPoint, LocalFourPoint, Interaction, LocalInteraction)):
raise ValueError(f"Multiplication {type(self)} @ {type(other)} not supported.")
if isinstance(other, (LocalInteraction, Interaction)):
is_local = not isinstance(other, Interaction)
q_prefix = "" if is_local else "q"
self.compress_q_dimension()
left_orbs = "abij"
right_orbs = "jief"
final_orbs = "abef"
suffix = {0: "w", 1: "wv", 2: "wvp"}.get(self.num_vn_dimensions, "")
einsum_str = (
f"q{left_orbs}{suffix},{q_prefix}{right_orbs}->q{final_orbs}{suffix}"
if left_hand_side
else f"{q_prefix}{left_orbs},q{right_orbs}{suffix}->q{final_orbs}{suffix}"
)
return FourPoint(
(
np.einsum(einsum_str, self.mat, other.mat, optimize=True)
if left_hand_side
else np.einsum(einsum_str, other.mat, self.mat, optimize=True)
),
self.channel,
self.nq,
self.num_wn_dimensions,
self.num_vn_dimensions,
self.full_niw_range,
self.full_niv_range,
self.has_compressed_q_dimension,
self.frequency_notation,
)
is_local = not isinstance(other, FourPoint)
channel = self.channel if self.channel != SpinChannel.NONE else other.channel
if self.num_vn_dimensions in (0, 1) or other.num_vn_dimensions in (0, 1):
# special cases if both objects do not have two fermionic frequency dimensions each. Straightforward
# contraction is saving memory as we do not have to add fermionic frequency dimensions to artificially
# create compound indices
q_prefix = "" if is_local else "q"
self.compress_q_dimension()
if not is_local:
other = other.compress_q_dimension()
self.to_half_niw_range()
other.to_half_niw_range()
suffix_other, suffix_result, suffix_self = self._get_frequency_suffixes_for_matmul(other, left_hand_side)
einsum_str = (
f"qabcd{suffix_self},{q_prefix}dcef{suffix_other}->qabef{suffix_result}"
if left_hand_side
else f"{q_prefix}abcd{suffix_other},qdcef{suffix_self}->qabef{suffix_result}"
)
return FourPoint(
(
np.einsum(einsum_str, self.mat, other.mat, optimize=True)
if left_hand_side
else np.einsum(einsum_str, other.mat, self.mat, optimize=True)
),
channel,
self.nq,
self.num_wn_dimensions,
max(self.num_vn_dimensions, other.num_vn_dimensions),
self.full_niw_range,
self.full_niv_range,
self.has_compressed_q_dimension,
self.frequency_notation,
)
is_self_full_niw_range = self.full_niw_range
is_other_full_niw_range = other.full_niw_range
self.to_half_niw_range().to_compound_indices()
other = other.to_half_niw_range().to_compound_indices()
# for __matmul__ self needs to be the LHS object, for __rmatmul__ self needs to be the RHS object
new_mat = (
np.matmul(self.mat, other.mat[None, ...] if is_local else other.mat)
if left_hand_side
else np.matmul(other.mat[None, ...] if is_local else other.mat, self.mat)
)
self.to_full_indices()
if is_self_full_niw_range:
self.to_full_niw_range()
other = other.to_full_indices()
if is_other_full_niw_range:
other = other.to_full_niw_range()
return FourPoint(
new_mat,
channel,
self.nq,
self.num_wn_dimensions,
2,
False,
self.full_niv_range,
self.has_compressed_q_dimension,
self.frequency_notation,
).to_full_indices(self.original_shape)
[docs]
def invert(self, copy: bool = True):
r"""
Inverts the object in compound-index (matrix) space, per momentum. The single-fermionic-frequency case is
handled by a dedicated block-diagonal reshape; otherwise each momentum slice is inverted in a loop to keep
intermediate arrays small. The result is always returned in the half bosonic frequency range.
:param copy: If True, operate on and return a deep copy; if False, mutate and return ``self`` in place.
:return: The inverted :class:`FourPoint` in the half niw range.
"""
if copy:
return deepcopy(self).invert(copy=False)
self.to_half_niw_range()
if self.num_vn_dimensions == 1:
w_dim = self.original_shape[5] if self.has_compressed_q_dimension else self.original_shape[7]
self.compress_q_dimension()
self.mat = self.mat.transpose(0, 5, 6, 1, 2, 4, 3).reshape(
(self.current_shape[0], w_dim, 2 * self.niv, self.n_bands**2, self.n_bands**2)
) # transpose to [q,w,v,o1,o2,o4,o3] and collecting [q,w,v,x1,x2]
# invert in [x1,x2] per q (still batched over [w,v]); writing each q-slice in place caps the peak at the
# input size instead of allocating the full batched-inverse output (~one extra full array).
for i in range(self.current_shape[0]):
self.mat[i] = np.linalg.inv(self.mat[i])
# reshape to [q,w,v,o1,o2,o4,o3] and transpose to [q,o1,o2,o3,o4,w,v]
self.mat = self.mat.reshape(
(self.current_shape[0], w_dim, 2 * self.niv, self.n_bands, self.n_bands, self.n_bands, self.n_bands)
).transpose(0, 3, 4, 6, 5, 1, 2)
return self
self.to_compound_indices()
for i in range(self.current_shape[0]):
self.mat[i] = np.linalg.inv(self.mat[i])
return self.to_full_indices()
[docs]
def invert_and_sum_over_last_vn(self, beta: float):
r"""
Inverts the object in compound-index space per momentum and bosonic frequency, then sums over the last
fermionic frequency axis (with the :math:`1/\beta` prefactor). This computes the auxiliary susceptibility
used in the ladder construction in a single fused pass. Mutates ``self`` in place.
:param beta: Inverse temperature :math:`\beta`.
:return: ``self`` with the last fermionic axis summed out (``num_vn_dimensions`` reduced to 1).
:raises NotImplementedError: If the object does not have exactly two fermionic frequency dimensions.
"""
if self.num_vn_dimensions != 2:
raise NotImplementedError("Method only implemented for objects with two fermionic frequency dimensions.")
compound_index_shape = (self.n_bands, self.n_bands, 2 * self.niv)
size = np.prod(compound_index_shape)
self.to_half_niw_range().compress_q_dimension()
w_dim = self.original_shape[5] if self.has_compressed_q_dimension else self.original_shape[7]
new_arr = np.empty(self.original_shape[:-1], dtype=self.mat.dtype)
for i in range(self.current_shape[0]):
compound_arr = self.mat[i].transpose(4, 0, 1, 5, 3, 2, 6).reshape(w_dim, size, size)
new_arr[i] = (
np.linalg.inv(compound_arr).reshape((w_dim,) + compound_index_shape * 2).transpose(1, 2, 5, 4, 0, 3, 6)
).sum(axis=-1)
new_arr /= beta # in-place scale: new_arr is freshly allocated and unaliased, so no full-size temporary
self.mat = new_arr
self._num_vn_dimensions = 1
self.update_original_shape()
return self
[docs]
def invert_and_sum_over_last_vn_v2(self, beta: float):
r"""
Helper method that explicitly handles the calculation of the sum over the auxiliary susceptibility while
being highly memory-efficient. Does not invert the matrix directly but uses a linear solver to avoid the
creation of large intermediate arrays. This is especially important for objects with a large number of
orbital degrees of freedom, where the matrix in compound index space can become very large. Is up to
numerical precision the same as :meth:`invert_and_sum_over_last_vn`.
:param beta: Inverse temperature :math:`\beta`.
:return: ``self`` with the last fermionic axis summed out (``num_vn_dimensions`` reduced to 1).
"""
o = self.n_bands
vn = 2 * self.niv
compound_size = o * o * vn
self.to_half_niw_range().compress_q_dimension()
w_dim = self.original_shape[5] if self.has_compressed_q_dimension else self.original_shape[7]
new_arr = np.empty(self.original_shape[:-1], dtype=self.mat.dtype)
idx = np.arange(compound_size)
# decode flat index -> (o4,o3)
idx_o4 = idx // (o * vn)
idx_o3 = (idx // vn) % o
rhs = np.zeros((compound_size, o * o), dtype=self.mat.dtype)
rhs[idx, idx_o4 * o + idx_o3] = 1.0
rhs = np.ascontiguousarray(rhs)
rhs_batched = np.broadcast_to(rhs, (w_dim, compound_size, o * o))
for i in range(self.current_shape[0]):
compound_arr = np.ascontiguousarray(self.mat[i].transpose(4, 0, 1, 5, 3, 2, 6)).reshape(
w_dim, compound_size, compound_size
)
new_arr[i] = (
np.linalg.solve(compound_arr, rhs_batched).reshape((w_dim, o, o, vn, o, o)).transpose(1, 2, 5, 4, 0, 3)
)
new_arr /= beta # in-place scale: new_arr is freshly allocated and unaliased, so no full-size temporary
self.mat = new_arr
self._num_vn_dimensions = 1
self.update_original_shape()
return self
[docs]
@staticmethod
def load(
filename: str,
channel: SpinChannel = SpinChannel.NONE,
nq: tuple[int, int, int] = (1, 1, 1),
num_wn_dimensions: int = 1,
num_vn_dimensions: int = 2,
full_niw_range: bool = False,
full_niv_range: bool = True,
has_compressed_q_dimension: bool = True,
frequency_notation: FrequencyNotation = FrequencyNotation.PH,
) -> "FourPoint":
r"""
Loads a :class:`FourPoint` from a ``.npy`` file.
:param filename: Path to the ``.npy`` file (loaded with ``allow_pickle=False``).
:param channel: Spin channel of the object (see :class:`SpinChannel`).
:param nq: Number of momenta per spatial direction.
:param num_wn_dimensions: Number of bosonic frequency axes (0–1).
:param num_vn_dimensions: Number of fermionic frequency axes (0–2).
:param full_niw_range: Whether the stored array spans the full (signed) bosonic range or only :math:`\omega \geq 0`.
:param full_niv_range: Whether the stored array spans the full (signed) fermionic range or only :math:`\nu \geq 0`.
:param has_compressed_q_dimension: Whether the momentum is stored as a single compressed axis (True) or as
``[qx, qy, qz, ...]`` (False).
:param frequency_notation: Frequency convention (see :class:`FrequencyNotation`).
:return: The loaded :class:`FourPoint`.
"""
return FourPoint(
np.load(filename, allow_pickle=False),
channel,
nq,
num_wn_dimensions,
num_vn_dimensions,
full_niw_range,
full_niv_range,
has_compressed_q_dimension,
frequency_notation,
)
[docs]
@staticmethod
def identity(
n_bands: int,
niw: int,
niv: int,
nq_tot: int = 1,
nq: tuple[int, int, int] = (1, 1, 1),
num_vn_dimensions: int = 2,
frequency_notation: FrequencyNotation = FrequencyNotation.PH,
) -> "FourPoint":
r"""
Creates a :class:`FourPoint` that is the identity in compound-index (matrix) space at each momentum, returned
in the half bosonic frequency range.
:param n_bands: Number of orbitals/bands per orbital axis.
:param niw: Number of positive bosonic frequencies.
:param niv: Number of positive fermionic frequencies.
:param nq_tot: Total number of momenta (product over directions).
:param nq: Number of momenta per spatial direction.
:param num_vn_dimensions: Number of fermionic frequency axes (1 or 2).
:param frequency_notation: Frequency convention (see :class:`FrequencyNotation`).
:return: The identity :class:`FourPoint` (compressed momentum, half niw range).
:raises ValueError: If ``num_vn_dimensions`` is not 1 or 2.
"""
if num_vn_dimensions not in (1, 2):
raise ValueError("Invalid number of fermionic frequency dimensions.")
full_shape = (nq_tot,) + (n_bands,) * 4 + (2 * niw + 1,) + (2 * niv,) * num_vn_dimensions
if num_vn_dimensions == 1:
mat = (
np.tile(np.eye(n_bands**2, dtype=DTYPE)[None, ..., None, None], (nq_tot, 1, 1, 2 * niw + 1, 2 * niv))
.reshape(full_shape)
.swapaxes(3, 4)
)
# swapping the last two orbital axes to get [q,o1,o2,o4,o3,w,v] since the matrix is unity in
# [o1,o2,o4,o3] (remember last two orbital indices are swapped in the compound index notation)
return FourPoint(
mat,
nq=nq,
num_vn_dimensions=num_vn_dimensions,
has_compressed_q_dimension=True,
frequency_notation=frequency_notation,
).to_half_niw_range()
compound_index_size = 2 * niv * n_bands**2
mat = np.tile(np.eye(compound_index_size, dtype=DTYPE)[None, None, ...], (nq_tot, 2 * niw + 1, 1, 1))
return (
FourPoint(
mat,
nq=nq,
num_vn_dimensions=num_vn_dimensions,
has_compressed_q_dimension=True,
frequency_notation=frequency_notation,
)
.to_full_indices(full_shape)
.to_half_niw_range()
)
[docs]
@staticmethod
def identity_like(other: "FourPoint") -> "FourPoint":
"""
Creates a compound-index identity matching the bands, frequency box, momenta, fermionic-axis count and
frequency notation of ``other`` (see :meth:`identity`).
:param other: The :class:`FourPoint` whose shape/attributes the identity should match.
:return: The matching identity :class:`FourPoint`.
"""
return FourPoint.identity(
other.n_bands,
other.niw if other.full_niw_range else 2 * other.niw,
other.niv,
other.nq_tot,
other.nq,
other.num_vn_dimensions,
other.frequency_notation,
)