From fd258b9381f5fb5c0fdb519598ce61f652e71f42 Mon Sep 17 00:00:00 2001
From: Rob Zinkov <rob@zinkov.com>
Date: Mon, 2 Jun 2025 21:42:52 +0200
Subject: [PATCH] Remove pivoted cholesky function

---
 pymc_extras/utils/pivoted_cholesky.py | 69 ---------------------------
 tests/test_pivoted_cholesky.py        | 24 ----------
 2 files changed, 93 deletions(-)
 delete mode 100644 pymc_extras/utils/pivoted_cholesky.py
 delete mode 100644 tests/test_pivoted_cholesky.py

diff --git a/pymc_extras/utils/pivoted_cholesky.py b/pymc_extras/utils/pivoted_cholesky.py
deleted file mode 100644
index 756d328db..000000000
--- a/pymc_extras/utils/pivoted_cholesky.py
+++ /dev/null
@@ -1,69 +0,0 @@
-try:
-    import torch
-
-    from gpytorch.utils.permutation import apply_permutation
-except ImportError as e:
-    raise ImportError("PyTorch and GPyTorch not found.") from e
-
-import numpy as np
-
-
-def pp(x):
-    return np.array2string(x, precision=4, floatmode="fixed")
-
-
-def pivoted_cholesky(mat: np.matrix, error_tol=1e-6, max_iter=np.inf):
-    """
-    mat: numpy matrix of N x N
-
-    This is to replicate what is done in GPyTorch verbatim.
-    """
-    n = mat.shape[-1]
-    max_iter = min(int(max_iter), n)
-
-    d = np.array(np.diag(mat))
-    orig_error = np.max(d)
-    error = np.linalg.norm(d, 1) / orig_error
-    pi = np.arange(n)
-
-    L = np.zeros((max_iter, n))
-
-    m = 0
-    while m < max_iter and error > error_tol:
-        permuted_d = d[pi]
-        max_diag_idx = np.argmax(permuted_d[m:])
-        max_diag_idx = max_diag_idx + m
-        max_diag_val = permuted_d[max_diag_idx]
-        i = max_diag_idx
-
-        # swap pi_m and pi_i
-        pi[m], pi[i] = pi[i], pi[m]
-        pim = pi[m]
-
-        L[m, pim] = np.sqrt(max_diag_val)
-
-        if m + 1 < n:
-            row = apply_permutation(
-                torch.from_numpy(mat), torch.tensor(pim), right_permutation=None
-            )  # left permutation just swaps row
-            row = row.numpy().flatten()
-            pi_i = pi[m + 1 :]
-            L_m_new = row[pi_i]  # length = 9
-
-            if m > 0:
-                L_prev = L[:m, pi_i]
-                update = L[:m, pim]
-                prod = update @ L_prev
-                L_m_new = L_m_new - prod  # np.sum(prod, axis=-1)
-
-            L_m = L[m, :]
-            L_m_new = L_m_new / L_m[pim]
-            L_m[pi_i] = L_m_new
-
-            matrix_diag_current = d[pi_i]
-            d[pi_i] = matrix_diag_current - L_m_new**2
-
-            L[m, :] = L_m
-            error = np.linalg.norm(d[pi_i], 1) / orig_error
-        m = m + 1
-    return L, pi
diff --git a/tests/test_pivoted_cholesky.py b/tests/test_pivoted_cholesky.py
deleted file mode 100644
index a30050666..000000000
--- a/tests/test_pivoted_cholesky.py
+++ /dev/null
@@ -1,24 +0,0 @@
-# try:
-#     import gpytorch
-#     import torch
-# except ImportError as e:
-#     # print(
-#     #     f"Please install Pytorch and GPyTorch to use this pivoted Cholesky implementation. Error {e}"
-#     # )
-#     pass
-# import numpy as np
-#
-# import pymc_extras as pmx
-#
-#
-# def test_match_gpytorch_linearcg_output():
-#     N = 10
-#     rank = 5
-#     np.random.seed(1234)  # nans with seed 1234
-#     K = np.random.randn(N, N)
-#     K = K @ K.T + N * np.eye(N)
-#     K_torch = torch.from_numpy(K)
-#
-#     L_gpt = gpytorch.pivoted_cholesky(K_torch, rank=rank, error_tol=1e-3)
-#     L_np, _ = pmx.utils.pivoted_cholesky(K, max_iter=rank, error_tol=1e-3)
-#     assert np.allclose(L_gpt, L_np.T)