revert t2 indexing convention in documentation (#466)

Author: kevinsung

docs/explanations/lucj.ipynb

@@ -113,9 +113,9 @@
     "$$\n",
     "t_{ijab} = i \\sum_{k=0}^{L - 1} \\sum_{pq}\n",
     "    Z^{(k)}_{pq}\n",
-    "    U^{(k)}_{(\\eta + a)p} U^{(k)*}_{ip} U^{(k)}_{(\\eta + b)q} U^{(k)*}_{jq}\n",
+    "    U^{(k)}_{ap} U^{(k)*}_{ip} U^{(k)}_{bq} U^{(k)*}_{jq}.\n",
     "$$\n",
-    "Here, each $Z^{(k)}$ is a real symmetric matrix representing a diagonal Coulomb operator, each $U^{(k)}$ is a unitary matrix representing an orbital rotation, and the shape of the t2 amplitudes tensor is $(\\eta, \\eta, N - \\eta, N - \\eta)$, where $N$ is the number of orbitals and $\\eta$ is the number of orbitals that are occupied. If the number of terms in the sum is truncated, or if some diagonal Coulomb interactions are dropped (by zeroing out the corresponding entries of the $Z^{(k)}$ matrices), then the equality no longer holds. However, given any two tensors $\\bar{U}^{(k)}_{ij}$ and $\\bar{Z}^{(k)}_{ij}$, we can plug them into the expression to obtain some $t_2$-amplitudes $\\bar{t}_{ijab}$. The compressed double factorization attempts to find tensors that minimize the least-squares objective function\n",
+    "Here, each $Z^{(k)}$ is a real symmetric matrix representing a diagonal Coulomb operator, each $U^{(k)}$ is a unitary matrix representing an orbital rotation, $i$ and $j$ run over occupied orbitals, and $a$ and $b$ run over virtual orbitals. If the number of terms in the sum is truncated, or if some diagonal Coulomb interactions are dropped (by zeroing out the corresponding entries of the $Z^{(k)}$ matrices), then the equality no longer holds. However, given any two tensors $\\bar{U}^{(k)}_{ij}$ and $\\bar{Z}^{(k)}_{ij}$, we can plug them into the expression to obtain some $t_2$-amplitudes $\\bar{t}_{ijab}$. The compressed double factorization attempts to find tensors that minimize the least-squares objective function\n",
     "$$\n",
     "\\frac12 \\sum_{ijab} \\lvert \\bar{t}_{ijab} - t_{ijab} \\rvert ^2.\n",
     "$$\n",

python/ffsim/linalg/double_factorized_decomposition.py

@@ -505,21 +505,21 @@ def double_factorized_t2(
     tuple[np.ndarray, np.ndarray]
     | tuple[np.ndarray, np.ndarray, scipy.optimize.OptimizeResult]
 ):
-    r"""Double-factorized decomposition of t2 amplitudes.
+    r"""Double-factorized decomposition of :math:`t_2` amplitudes.
 
-    The double-factorized decomposition of a t2 amplitudes tensor :math:`t_{ijab}` is
+    The double-factorized decomposition of a :math:`t_2` amplitudes tensor
+    :math:`t_{ijab}` is
 
     .. math::
 
         t_{ijab} = i \sum_{k=0}^{L - 1} \sum_{pq}
             Z^{(k)}_{pq}
-            U^{(k)}_{(\eta + a)p} U^{(k)*}_{ip} U^{(k)}_{(\eta + b)q} U^{(k)*}_{jq}
+            U^{(k)}_{ap} U^{(k)*}_{ip} U^{(k)}_{bq} U^{(k)*}_{jq},
 
-    Here each :math:`Z^{(k)}` is a real symmetric matrix, referred to as a
-    "diagonal Coulomb matrix," each :math:`U^{(k)}` is a unitary matrix,
-    referred to as an "orbital rotation," and the shape of the t2 amplitudes tensor is
-    :math:`(\eta, \eta, N - \eta, N - \eta)`, where :math:`N` is the number of orbitals
-    and :math:`\eta` is the number of orbitals that are occupied.
+    where each :math:`Z^{(k)}` is a real symmetric matrix representing a diagonal
+    Coulomb operator, each :math:`U^{(k)}` is a unitary matrix representing an orbital
+    rotation, :math:`i` and :math:`j` run over occupied orbitals, and :math:`a` and
+    :math:`b` run over virtual orbitals.
 
     The number of terms :math:`L` in the decomposition depends on the allowed
     error threshold. A larger error threshold may yield a smaller number of terms.
@@ -529,7 +529,7 @@ def double_factorized_t2(
     error threshold.
 
     The default behavior of this routine is to perform a straightforward
-    "exact" factorization of the t2 amplitudes tensor based on a nested
+    "exact" factorization of the :math:`t_2` amplitudes tensor based on a nested
     eigenvalue decomposition, and then truncate the terms based on the values of `tol`
     and `max_terms`.
     If `optimize` is set to ``True``, then the entries of the resulting tensors
@@ -552,16 +552,38 @@ def double_factorized_t2(
     automatically: `multi_stage_start` defaults to the total number of terms returned by
     the exact factorization, and `multi_stage_step` defaults to 1.
 
-    Note: Currently, only real-valued t2 amplitudes are supported.
+    The original :math:`t_2` amplitudes tensor can be reconstructed, up to the error
+    tolerance, using the following function:
+
+    .. code::
+
+        def reconstruct_t2(
+            diag_coulomb_mats: np.ndarray, orbital_rotations: np.ndarray, nocc: int
+        ) -> np.ndarray:
+            return (
+                1j
+                * np.einsum(
+                    "kpq,kap,kip,kbq,kjq->ijab",
+                    diag_coulomb_mats,
+                    orbital_rotations,
+                    orbital_rotations.conj(),
+                    orbital_rotations,
+                    orbital_rotations.conj(),
+                )[:nocc, :nocc, nocc:, nocc:]
+            )
+
+    Note: Currently, only real-valued :math:`t_2` amplitudes are supported.
 
     Args:
-        t2_amplitudes: The t2 amplitudes tensor.
+        t2_amplitudes: The doubles amplitudes tensor of shape
+            ``(nocc, nocc, nvrt, nvrt)``, where ``nocc`` is the number of occupied
+            orbitals and ``nvrt`` is the number of virtual orbitals.
         tol: Tolerance for error in the decomposition.
             The error is defined as the maximum absolute difference between
             an element of the original tensor and the corresponding element of
             the reconstructed tensor.
         max_terms: An optional limit on the number of terms to keep in the decomposition
-            of the t2 amplitudes tensor. This argument overrides `tol`.
+            of the :math:`t_2` amplitudes tensor. This argument overrides `tol`.
         optimize: Whether to optimize the tensors returned by the decomposition to
             to minimize the error in the factorization.
         method: The optimization method. See the documentation of
@@ -743,10 +765,10 @@ def fun(x: np.ndarray):
 def double_factorized_t2_alpha_beta(
     t2_amplitudes: np.ndarray, *, tol: float = 1e-8, max_terms: int | None = None
 ) -> tuple[np.ndarray, np.ndarray]:
-    r"""Double-factorized decomposition of alpha-beta t2 amplitudes.
+    r"""Double-factorized decomposition of alpha-beta :math:`t_2` amplitudes.
 
-    Decompose alpha-beta t2 amplitudes into diagonal Coulomb matrices with orbital
-    rotations. This function returns two arrays:
+    Decompose alpha-beta :math:`t_2` amplitudes into diagonal Coulomb matrices with
+    orbital rotations. This function returns two arrays:
 
     - `diagonal_coulomb_mats`, with shape `(n_terms, 3, norb, norb)`.
     - `orbital_rotations`, with shape `(n_terms, 2, norb, norb)`.
@@ -755,8 +777,8 @@ def double_factorized_t2_alpha_beta(
     tolerance might yield a smaller value for `n_terms`. You can also set an optional
     upper bound on `n_terms` using the `max_terms` argument.
 
-    The original t2 amplitudes tensor can be reconstructed, up to the error tolerance,
-    using the following function:
+    The original :math:`t_2` amplitudes tensor can be reconstructed, up to the error
+    tolerance, using the following function:
 
     .. code::
 
@@ -783,7 +805,7 @@ def reconstruct_t2_alpha_beta(
                 )
             return (
                 2j
-                * contract(
+                * np.einsum(
                     "kpq,kap,kip,kbq,kjq->ijab",
                     expanded_diag_coulomb_mats,
                     expanded_orbital_rotations,
@@ -793,16 +815,20 @@ def reconstruct_t2_alpha_beta(
                 )[:nocc_a, norb : norb + nocc_b, nocc_a:norb, norb + nocc_b :]
             )
 
-    Note: Currently, only real-valued t2 amplitudes are supported.
+    Note: Currently, only real-valued :math:`t_2` amplitudes are supported.
 
     Args:
-        t2_amplitudes: The t2 amplitudes tensor.
+        t2_amplitudes: The doubles amplitudes tensor of shape
+            ``(nocc_a, nocc_b, nvrt_a, nvrt_b)``, where ``nocc_a`` is the number of
+            occupied spin-alpha orbitals, ``nvrt_a`` is the number of virtual spin-alpha
+            orbitals, ``nocc_b`` is the number of occupied spin-beta orbitals, and
+            ``nvrt_b`` is the number of virtual spin-beta orbitals.
         tol: Tolerance for error in the decomposition.
             The error is defined as the maximum absolute difference between
             an element of the original tensor and the corresponding element of
             the reconstructed tensor.
         max_terms: An optional limit on the number of terms to keep in the decomposition
-            of the t2 amplitudes tensor. This argument overrides `tol`.
+            of the :math:`t_2` amplitudes tensor. This argument overrides `tol`.
 
 
     Returns: