Resolution of identity

.gitignore

@@ -1,4 +1,6 @@
 __pycache__/
+.pytest_cache/
 .DS_STORE
 .vscode
-*.egg-info
+*.egg-info
+build/

adcgen/__init__.py

@@ -20,6 +20,7 @@
 from .sympy_objects import (AntiSymmetricTensor, SymmetricTensor, Amplitude,
                             NonSymmetricTensor, KroneckerDelta, SymbolicTensor)
 from .tensor_names import tensor_names
+from .resolution_of_identity import apply_resolution_of_identity
 from . import sort_expr as sort
 
 
@@ -35,6 +36,7 @@
            "Intermediates", "reduce_expr", "factor_intermediates",
            "sort",
            "transform_to_spatial_orbitals",
+           "apply_resolution_of_identity",
            "apply_cvs_approximation",
            "generate_code", "optimize_contractions",
            "unoptimized_contraction", "Contraction",

adcgen/expression/expr_container.py

@@ -271,6 +271,29 @@ def rename_tensor(self, current: str, new: str) -> "ExprContainer":
         self._inner = renamed
         return self
 
+    def expand_coulomb_ri(self, factorisation: str = 'sym'
+                          ) -> 'ExprContainer':
+        """
+        Expands the Coulomb operators (pq | rs) into RI format
+
+        Parameters
+        ----------
+        factorisation : str, optional
+            The type of factorisation ('sym' or 'asym'), by default 'sym'
+
+        Returns
+        -------
+        ExprContainer
+            The factorised expression.
+        """
+        res = S.Zero
+        for term in self.terms:
+            res += term.expand_coulomb_ri(factorisation=factorisation,
+                                          wrap_result=False)
+        assert isinstance(res, Expr)
+        self._inner = res
+        return self
+
     def expand_antisym_eri(self) -> "ExprContainer":
         """
         Expands the antisymmetric ERI using chemists notation

adcgen/expression/object_container.py

@@ -445,6 +445,11 @@ def allowed_spin_blocks(self) -> tuple[str, ...] | None:
                 ]))
             elif name == tensor_names.coulomb:  # ERI in chemist notation
                 return ("aaaa", "aabb", "bbaa", "bbbb")
+            elif name in (tensor_names.ri_sym, tensor_names.ri_asym_eri,
+                          tensor_names.ri_asym_factor):
+                return ("aaa", "abb")
+            elif name == tensor_names.fock:
+                return ("aa", "bb")
         elif isinstance(obj, KroneckerDelta):  # delta
             # spins have to be equal
             return ("aa", "bb")
@@ -499,6 +504,8 @@ def format_indices(indices) -> str:
                 ),
                 # coulomb integral chemist notation
                 tensor_names.coulomb: lambda up, lo: f"({up}\\vert {lo})",
+                # 2e3c integral in asymmetric RI
+                tensor_names.ri_asym_eri: lambda up, lo: f"({up}\\vert {lo})",
                 # orbital energy
                 tensor_names.orb_energy: lambda _, lo: f"\\varepsilon_{{{lo}}}"
             }
@@ -724,6 +731,84 @@ def rename_tensor(self, current: str, new: str,
             obj = ExprContainer(obj, **self.assumptions)
         return obj
 
+    def expand_coulomb_ri(self, factorisation: str = 'sym',
+                          wrap_result: bool = True) -> "ExprContainer | Expr":
+        """
+        Expands the Coulomb operators (pq | rs) into RI format
+
+        Parameters
+        ----------
+        factorisation : str, optional
+            The type of factorisation ('sym' or 'asym'), by default 'sym'
+        wrap_result : bool, optional
+            Whether to wrap the result in an ExprContainer, by default True
+
+        Returns
+        -------
+        ExprContainer | Expr
+            The factorised expression.
+        """
+        from .expr_container import ExprContainer
+
+        if factorisation not in ('sym', 'asym'):
+            raise NotImplementedError("Only symmetric ('sym') and asymmetric "
+                                      "('asym') factorisation of the Coulomb "
+                                      "integral is implemented")
+
+        res = self.inner
+        base, exponent = self.base_and_exponent
+        if isinstance(base, SymmetricTensor) and \
+                base.name == tensor_names.coulomb:
+            if base.bra_ket_sym != 1:
+                raise NotImplementedError("Can only apply RI approximation to "
+                                          "coulomb integrals with "
+                                          "bra-ket symmetry.")
+            # we dont expand negative exponents, because the result
+            # (ab)^-n will evaluate to a^-n b^-n, which is
+            # only correct if the product ab has no contracted
+            # indices
+            if not exponent.is_Integer or exponent < S.Zero:
+                raise NotImplementedError("Can only apply RI approximation to "
+                                          "coulomb integrals "
+                                          "with positive integer exponents. "
+                                          f"{self} has an invalid exponent.")
+            # setup the assumptions for the aux index:
+            # assign alpha spin if represented in spatial orbitals
+            idx = self.idx
+            has_spin = bool(idx[0].spin)
+            if any(bool(s) != has_spin for s in idx):
+                raise NotImplementedError(f"The coulomb integral {self} has "
+                                          "to be represented either in spatial"
+                                          " or spin orbitals. A mixture is not"
+                                          " valid.")
+            assumptions = {"aux": True}
+            if has_spin:
+                assumptions["alpha"] = True
+            # actually do the expansion
+            p, q, r, s = idx
+            res = S.One
+            for _ in range(int(exponent)):  # exponent has to be positive
+                aux_idx = Index("P", **assumptions)
+                if factorisation == "sym":
+                    res *= SymmetricTensor(
+                        tensor_names.ri_sym, (aux_idx,), (p, q), 0
+                    )
+                    res *= SymmetricTensor(
+                        tensor_names.ri_sym, (aux_idx,), (r, s), 0
+                    )
+                else:
+                    assert factorisation == "asym"
+                    res *= SymmetricTensor(
+                        tensor_names.ri_asym_factor, (aux_idx,), (p, q), 0
+                    )
+                    res *= SymmetricTensor(
+                        tensor_names.ri_asym_eri, (aux_idx,), (r, s), 0
+                    )
+        if wrap_result:
+            kwargs = self.assumptions
+            res = ExprContainer(res, **kwargs)
+        return res
+
     def expand_antisym_eri(self, wrap_result: bool = True
                            ) -> "ExprContainer | Expr":
         """
@@ -782,6 +867,12 @@ def expand_intermediates(self, target: Sequence[Index],
             False: The intermediate is only expanded once, e.g., n'th
               order MP t-amplitudes are expressed by means of (n-1)'th order
               MP t-amplitudes and ERI.
+        braket_sym_tensors: Sequence[str], optional
+            Add bra-ket-symmetry to the given tensors of the expanded
+            expression (after expansion of the intermediates).
+        braket_antisym_tensors: Sequence[str], optional
+            Add bra-ket-antisymmetry to the given tensors of the expanded
+            expression (after expansion of the intermediates).
         """
         from ..intermediates import Intermediates
         from .expr_container import ExprContainer
@@ -800,18 +891,31 @@ def expand_intermediates(self, target: Sequence[Index],
         itmd = Intermediates().available.get(longname, None)
         expanded = self.inner
         if itmd is not None:
+            # for negative exponents we would have to ensure that the
+            # intermediate is a "long" intermediate that consists of
+            # multiple terms. Or if it consists of a single term
+            # that it does not have any contracted indices
+            # However, this can only be checked by calling ".expand()"
+            # on the contributions in the for loop below, which seems bad.
+            # A short intermediates will be expanded as
+            # X^-2 = (ab * cd)^-1 -> a^-1 b^-1 c^-1 d^-1
+            # where the last step is not correct if the intermediate
+            # has contracted indices.
+            exponent = self.exponent
+            if not exponent.is_Integer or exponent < S.Zero:
+                raise NotImplementedError(
+                    "Can only expand intermediates with positive "
+                    f"integer exponents. {self} has an invalid exponent."
+                )
             # Use a for loop to obtain different contracted itmd indices
             # for each x in: x * x * ...
             expanded = S.One
-            exponent = self.exponent
             assert exponent.is_Integer
             for _ in range(abs(int(exponent))):
                 expanded *= itmd.expand_itmd(
                     indices=self.idx, wrap_result=False,
                     fully_expand=fully_expand
                 )
-            if exponent < S.Zero:
-                expanded = Pow(expanded, -1)
         # apply assumptions to the expanded object
         if braket_sym_tensors or braket_antisym_tensors:
             expanded = ExprContainer(expanded).add_bra_ket_sym(

adcgen/expression/polynom_container.py

@@ -2,7 +2,7 @@
 from functools import cached_property
 from typing import Any, TYPE_CHECKING
 
-from sympy import Add, Expr, Pow, S
+from sympy import Add, Expr, Mul, Pow, S
 
 from ..indices import Index, sort_idx_canonical
 from .container import Container
@@ -188,6 +188,46 @@ def rename_tensor(self, current: str, new: str,
             renamed = ExprContainer(inner=renamed, **self.assumptions)
         return renamed
 
+    def expand_coulomb_ri(self, factorisation: str = 'sym',
+                          wrap_result: bool = True) -> "Expr | ExprContainer":
+        """
+        Expands the Coulomb operators (pq | rs) into RI format
+
+        Parameters
+        ----------
+        factorisation : str, optional
+            The type of factorisation ('sym' or 'asym'), by default 'sym'
+        wrap_result : bool, optional
+            Whether to wrap the result in an ExprContainer, by default True
+        """
+        from .expr_container import ExprContainer
+
+        exponent = self.exponent
+        if not exponent.is_Integer:
+            raise ValueError("Can only apply RI approximation to Polynomials "
+                             "with integer exponents. "
+                             f"{self} has an invalid exponent.")
+        # use a for loop so the contracted aux indices for each x in
+        # x * x * ... = x^n are different.
+        expanded = S.One
+        for _ in range(abs(int(exponent))):
+            contrib = S.Zero
+            for term in self.terms:
+                contrib += term.expand_coulomb_ri(
+                    factorisation=factorisation, wrap_result=False
+                )
+            assert isinstance(contrib, Expr)
+            if exponent < S.Zero:
+                # a mul object would be simplified as
+                # (ab)^-1 -> a^-1 b^-1
+                # which is only correct if a and b have no contracted indices.
+                assert not isinstance(contrib, Mul)
+                contrib = Pow(contrib, -1)
+            expanded *= contrib
+        if wrap_result:
+            expanded = ExprContainer(inner=expanded, **self.assumptions)
+        return expanded
+
     def expand_antisym_eri(self, wrap_result: bool = True):
         """
         Expands the antisymmetric ERI using chemists notation
@@ -217,16 +257,30 @@ def expand_intermediates(self, target: Sequence[Index],
         """Expands all known intermediates in the polynom."""
         from .expr_container import ExprContainer
 
-        expanded = S.Zero
-        for term in self.terms:
-            expanded += term.expand_intermediates(
-                target, wrap_result=False, fully_expand=fully_expand,
-                braket_sym_tensors=braket_sym_tensors,
-                braket_antisym_tensors=braket_antisym_tensors
-            )
-        assert isinstance(expanded, Expr)
-        expanded = Pow(expanded, self.exponent)
-
+        exponent = self.exponent
+        if not exponent.is_Integer:
+            raise NotImplementedError("Can only expand intermediates for "
+                                      "polynoms with integer exponents."
+                                      f"{self} has an invalid exponent.")
+        # use a for loop so the contracted itmd indices for each x in
+        # x * x * ... = x^n are different.
+        expanded = S.One
+        for _ in range(abs(int(exponent))):
+            contrib = S.Zero
+            for term in self.terms:
+                contrib += term.expand_intermediates(
+                    target, wrap_result=False, fully_expand=fully_expand,
+                    braket_sym_tensors=braket_sym_tensors,
+                    braket_antisym_tensors=braket_antisym_tensors
+                )
+            assert isinstance(contrib, Expr)
+            if exponent < S.Zero:
+                # a mul object would be simplified as
+                # (ab)^-1 -> a^-1 b^-1
+                # which is only correct if a and b have no contracted indices.
+                assert not isinstance(contrib, Mul)
+                contrib = Pow(contrib, -1)
+            expanded *= contrib
         if wrap_result:
             assumptions = self.assumptions
             assumptions["target_idx"] = target

adcgen/expression/term_container.py

@@ -583,6 +583,35 @@ def rename_tensor(self, current: str, new: str, wrap_result: bool = True
             renamed = ExprContainer(renamed, **self.assumptions)
         return renamed
 
+    def expand_coulomb_ri(self, factorisation: str = 'sym',
+                          wrap_result: bool = True) -> "Expr | ExprContainer":
+        """
+        Expands the Coulomb operators (pq | rs) into RI format
+
+        Parameters
+        ----------
+        factorisation : str, optional
+            The type of factorisation ('sym' or 'asym'), by default 'sym'
+        wrap_result : bool, optional
+            Whether to wrap the result in an ExprContainer, by default True
+
+        Returns
+        -------
+        ExprContainer | Expr
+            The factorised expression.
+        """
+        from .expr_container import ExprContainer
+
+        factorised = S.One
+        for obj in self.objects:
+            factorised *= obj.expand_coulomb_ri(factorisation=factorisation,
+                                                wrap_result=False)
+
+        if wrap_result:
+            assumptions = self.assumptions
+            factorised = ExprContainer(factorised, **assumptions)
+        return factorised
+
     def expand_antisym_eri(self, wrap_result: bool = True
                            ) -> "ExprContainer | Expr":
         """

adcgen/func.py

@@ -159,7 +159,10 @@ def import_tensor(tensor: str) -> Expr:
             # ADC-Amplitude or t-amplitudes
             if is_adc_amplitude(name) or is_t_amplitude(name):
                 base: Expr = Amplitude(name, upper, lower)
-            elif name == tensor_names.coulomb:  # eri in chemist notation
+            elif name in (tensor_names.coulomb, tensor_names.ri_sym,
+                          tensor_names.ri_asym_eri,
+                          tensor_names.ri_asym_factor):
+                # eri in chemist notation or RI tensor
                 base: Expr = SymmetricTensor(name, upper, lower)
             else:
                 base: Expr = AntiSymmetricTensor(name, upper, lower)

adcgen/generate_code/config.json

@@ -2,6 +2,7 @@
   "sizes": {
     "core": 5,
     "occ": 20,
-    "virt": 200
+    "virt": 200,
+    "aux": 250
   }
 }

adcgen/generate_code/contraction.py

@@ -23,6 +23,7 @@ class Sizes:
     occ: int = 0
     virt: int = 0
     general: int = 0
+    aux: int = 0
 
     @staticmethod
     def from_dict(input: dict[str, int]) -> "Sizes":
@@ -32,7 +33,7 @@ def from_dict(input: dict[str, int]) -> "Sizes":
         if not provided.
         """
         if "general" not in input:
-            input["general"] = sum(input.values())
+            input["general"] = sum(v for k, v in input.items() if k != "aux")
         return Sizes(**input)
 
     @staticmethod
@@ -232,6 +233,7 @@ class ScalingComponent:
     virt: int
     occ: int
     core: int
+    aux: int
 
     def evaluate_costs(self, sizes: Sizes) -> int:
         """