From 8c241d3a7ff00451ffff79076d0a66fb8d65a593 Mon Sep 17 00:00:00 2001
From: Andrew Fitzgibbon <awf@graphcore.ai>
Date: Wed, 11 Oct 2023 21:24:56 +0100
Subject: [PATCH] Comment, and add binom_factor_direct to motivate the other
 cases

---
 notebooks/binom_factor_table.ipynb | 28 ++++++++++------------------
 pyscf_ipu/experimental/special.py  | 22 +++++++++++++++++++---
 2 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/notebooks/binom_factor_table.ipynb b/notebooks/binom_factor_table.ipynb
index b096c0f..0f4287d 100644
--- a/notebooks/binom_factor_table.ipynb
+++ b/notebooks/binom_factor_table.ipynb
@@ -61,16 +61,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
@@ -157,25 +150,24 @@
     "\n",
     "\n",
     "def binom_factor(i: int, j: int, a: float, b: float, s: int):\n",
-    "    out = 0\n",
-    "    for t in range(max(s - i, 0), j + 1):\n",
-    "        assert ((s - i) <= t) & (t <= j)\n",
-    "        val = binom(i, s - t) * binom(j, t) * a ** (i - (s - t)) * b ** (j - t)\n",
-    "        out += val\n",
-    "    return out\n",
+    "    return sum(\n",
+    "        binom(i, s - t) * binom(j, t) * a ** (i - (s - t)) * b ** (j - t)\n",
+    "        for t in range(max(s - i, 0), j + 1)\n",
+    "    )\n",
+    "\n",
     "\n",
     "def binom_factor_sym(i: int, j: int, a: Symbol, b: Symbol, s: int):\n",
     "    return Poly(binom_factor(i, j, a, b, s), a, b)\n",
     "\n",
     "\n",
     "a, b = symbols(\"a b\", real=True)\n",
-    "for i in range(0,7,3):\n",
-    "    for j in range(0,7,2):\n",
-    "        for s in range(0,7,2):\n",
+    "for i in range(0, 7, 3):\n",
+    "    for j in range(0, 7, 2):\n",
+    "        for s in range(0, 7, 2):\n",
     "            bf = binom_factor_sym(i, j, a, b, s)\n",
     "            print((i, j, s), bf)\n",
     "\n",
-    "display(bf, bf.coeff_monomial(a**3*b**3))"
+    "display(bf, bf.coeff_monomial(a**3 * b**3))"
    ]
   },
   {
diff --git a/pyscf_ipu/experimental/special.py b/pyscf_ipu/experimental/special.py
index c277eb7..32dd6eb 100644
--- a/pyscf_ipu/experimental/special.py
+++ b/pyscf_ipu/experimental/special.py
@@ -58,6 +58,8 @@ def factorial2_lookup(n: IntN, nmax: int = 2 * LMAX) -> IntN:
 
 factorial2 = factorial2_lookup
 
+# Various binom implementations
+
 
 def binom_beta(x: IntN, y: IntN) -> IntN:
     approx = 1.0 / ((x + 1) * jnp.exp(betaln(x - y + 1, y + 1)))
@@ -78,6 +80,8 @@ def binom_lookup(x: IntN, y: IntN, nmax: int = LMAX) -> IntN:
 
 binom = binom_lookup
 
+# Various gammanu implementations
+
 
 def gammanu_gamma(nu: IntN, t: FloatN, epsilon: float = 1e-10) -> FloatN:
     """
@@ -117,15 +121,26 @@ def gammanu_series(nu: IntN, t: FloatN, num_terms: int = 128) -> FloatN:
 gammanu = gammanu_series
 
 
-def binom_factor_segment_sum(
-    i: int, j: int, a: float, b: float, lmax: int = LMAX
-) -> FloatN:
+# Several binom_factor implementations
+
+
+def binom_factor_direct(i: int, j: int, a: float, b: float, s: int):
     """
     Eq. 15 from Augspurger JD, Dykstra CE. General quantum mechanical operators. An
     open-ended approach for one-electron integrals with Gaussian bases. Journal of
     computational chemistry. 1990 Jan;11(1):105-11.
     <https://doi.org/10.1002/jcc.540110113>
     """
+    return sum(
+        binom_beta(i, s - t) * binom_beta(j, t) * a ** (i - (s - t)) * b ** (j - t)
+        for t in range(max(s - i, 0), j + 1)
+    )
+
+
+def binom_factor_segment_sum(
+    i: int, j: int, a: float, b: float, lmax: int = LMAX
+) -> FloatN:
+    # Vectorized version of above
     s, t = jnp.tril_indices(lmax + 1)
     out = binom(i, s - t) * binom(j, t) * a ** (i - (s - t)) * b ** (j - t)
     mask = ((s - i) <= t) & (t <= j)
@@ -141,6 +156,7 @@ def binom_factor__via_segment_sum(i: int, j: int, a: float, b: float, s: int):
 
 
 def binom_factor__via_lookup(i: int, j: int, a: float, b: float, s: int) -> FloatN:
+    # Lookup-table version of above -- see binom_factor_table.ipynb for the derivation
     monomials = jnp.array(binom_factor_table.get_monomials(a, b))
     coeffs = binom_factor_table_W[i, j, s]
     return coeffs @ monomials