From 705bbc589b08507c512f115a300c950f13e22d12 Mon Sep 17 00:00:00 2001
From: Guillaume Huguet <57917099+guillaumehu@users.noreply.github.com>
Date: Thu, 29 Feb 2024 22:13:47 -0500
Subject: [PATCH 1/6] add cheb approx & graph & heat dit

---
 torchcfm/cheb_approx.py        |  37 ++++++++++
 torchcfm/diffusion_distance.py | 123 +++++++++++++++++++++++++++++++++
 2 files changed, 160 insertions(+)
 create mode 100644 torchcfm/cheb_approx.py
 create mode 100644 torchcfm/diffusion_distance.py

diff --git a/torchcfm/cheb_approx.py b/torchcfm/cheb_approx.py
new file mode 100644
index 00000000..4f3ec2ec
--- /dev/null
+++ b/torchcfm/cheb_approx.py
@@ -0,0 +1,37 @@
+import typing as T
+import numpy as np
+import torch
+from scipy.special import ive
+
+
+def expm_multiply(
+    L: torch.Tensor,
+    X: torch.Tensor,
+    coeff: torch.Tensor,
+    eigval: T.Union[torch.Tensor, np.ndarray],
+):
+    """Matrix exponential with chebyshev polynomial approximation."""
+
+    def body(carry, c):
+        T0, T1, Y = carry
+        T2 = (2.0 / eigval) * torch.matmul(L, T1) - 2.0 * T1 - T0
+        Y = Y + c * T2
+        return (T1, T2, Y)
+
+    T0 = X
+    Y = 0.5 * coeff[0] * T0
+    T1 = (1.0 / eigval) * torch.matmul(L, X) - T0
+    Y = Y + coeff[1] * T1
+
+    initial_state = (T0, T1, Y)
+    for c in coeff[2:]:
+        initial_state = body(initial_state, c)
+
+    _, _, Y = initial_state
+
+    return Y
+
+
+@torch.no_grad()
+def compute_chebychev_coeff_all(eigval, t, K):
+    return 2.0 * ive(torch.arange(0, K + 1), -t * eigval)
diff --git a/torchcfm/diffusion_distance.py b/torchcfm/diffusion_distance.py
new file mode 100644
index 00000000..4e65baae
--- /dev/null
+++ b/torchcfm/diffusion_distance.py
@@ -0,0 +1,123 @@
+import torch
+from torchcfm.cheb_approx import compute_chebychev_coeff_all, expm_multiply
+
+try:
+    import scanpy as sc
+except ImportError:
+    pass
+
+EPS_LOG = 1e-6
+EPS_HEAT = 1e-4
+
+def norm_sym_laplacian(A: torch.Tensor):
+    deg = A.sum(dim=1)
+    deg_sqrt_inv = torch.diag(1.0 / torch.sqrt(deg + EPS_LOG))
+    return deg_sqrt_inv @ A @ deg_sqrt_inv
+
+
+def laplacian_from_data(data: torch.Tensor, sigma: float, alpha: int = 20):
+    affinity = torch.exp(-(torch.cdist(data, data) / (2 * sigma)).pow(alpha))
+    return norm_sym_laplacian(affinity)
+
+
+def torch_knn_from_data(
+    data: torch.Tensor, k: int, projection: bool = False, proj_dim: int = 100
+):
+    if projection:
+        _, _, V = torch.pca_lowrank(data, q=proj_dim, center=True)
+        data = data @ V
+    dist = torch.cdist(data, data)
+    _, indices = torch.topk(dist, k, largest=False)
+    affinity = torch.zeros(data.shape[0], data.shape[0])
+    affinity.scatter_(1, indices, 1)
+    return norm_sym_laplacian(affinity)
+
+
+def scanpy_knn_from_data(
+    data: torch.Tensor, k: int, projection: bool = False, proj_dim: int = 100
+):
+    adata = sc.AnnData(data.numpy())
+    if projection:
+        sc.pp.pca(adata, n_comps=proj_dim)
+    sc.pp.neighbors(
+        adata, n_neighbors=k, use_rep="X_pca" if projection else None
+    )
+    return norm_sym_laplacian(
+        torch.tensor(adata.obsp["connectivities"].toarray())
+    )
+
+
+def var_fn(x, t):
+    outer = torch.outer(torch.diag(x), torch.ones(x.shape[0]))
+    vol_approx = (outer + outer.T) * 0.5
+    return -t * torch.log(x + EPS_LOG) + t * torch.log(vol_approx + EPS_LOG)
+
+
+class BaseHeatKernel:
+    def __init__(self, t: float = 1.0, order: int = 30):
+        self.t = t
+        self.order = order
+        self.dist_fn = var_fn
+        self.graph_fn = None
+
+    def __call__(self, data: torch.Tensor):
+        if self.graph_fn is None:
+            raise NotImplementedError("graph_fn is not implemented")
+        L = self.graph_fn(data)
+        heat_kernel = self.compute_heat_from_laplacian(L)
+        heat_kernel = self.sym_clip(heat_kernel)
+        return heat_kernel
+
+    def compute_heat_from_laplacian(self, L: torch.Tensor):
+        n = L.shape[0]
+        val = torch.linalg.eigvals(L).real
+        max_eigval = val.max()
+        cheb_coeff = compute_chebychev_coeff_all(
+            0.5 * max_eigval, self.t, self.order
+        )
+        heat_kernel = expm_multiply(
+            L, torch.eye(n), cheb_coeff, 0.5 * max_eigval
+        )
+        return heat_kernel
+
+    def sym_clip(self, heat_kernel: torch.Tensor):
+        heat_kernel = (heat_kernel + heat_kernel.T) / 2
+        heat_kernel[heat_kernel < 0] = 0.0 + EPS_HEAT
+        return heat_kernel
+
+    def fit(self, data: torch.Tensor, dist_type: str = "var"):
+        assert dist_type in self.dist_fn
+        heat_kernel = self(data)
+        return self.dist_fn[dist_type](heat_kernel, self.t)
+    
+
+class HeatKernelKNN(BaseHeatKernel):
+    """Approximation of the heat kernel with a graph from a k-nearest neighbors affinity matrix.
+    Uses Chebyshev polynomial approximation.
+    """
+
+    _is_differentiable = False
+    _implemented_graph = {
+        "torch": torch_knn_from_data,
+        "scanpy": scanpy_knn_from_data,
+    }
+
+    def __init__(
+        self,
+        k: int = 10,
+        order: int = 30,
+        t: float = 1.0,
+        projection: bool = False,
+        proj_dim: int = 100,
+        graph_type: str = "torch",
+    ):
+        super().__init__(t=t, order=order)
+        assert (
+            graph_type in self._implemented_graph
+        ), f"Type must be in {self._implemented_graph}"
+        self.k = k
+        self.projection = projection
+        self.proj_dim = proj_dim
+        self.graph_fn = lambda x: self._implemented_graph[graph_type](
+            x, self.k, projection=self.projection, proj_dim=self.proj_dim
+        )

From 705b2afd4265b3dbb11e8912f77a975f54f15995 Mon Sep 17 00:00:00 2001
From: Guillaume Huguet <57917099+guillaumehu@users.noreply.github.com>
Date: Thu, 29 Feb 2024 22:15:46 -0500
Subject: [PATCH 2/6] add tests

---
 tests/test_heat.py | 52 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 tests/test_heat.py

diff --git a/tests/test_heat.py b/tests/test_heat.py
new file mode 100644
index 00000000..b73b7bd1
--- /dev/null
+++ b/tests/test_heat.py
@@ -0,0 +1,52 @@
+import pytest
+import torch
+from torchcfm.diffusion_distance import laplacian_from_data, HeatKernelKNN, torch_knn_from_data
+
+def gt_heat_kernel_knn(
+    data,
+    t,
+    k,
+):
+    L = torch_knn_from_data(data, k=k, projection=False, proj_dim=10)
+    # eigendecomposition
+    eigvals, eigvecs = torch.linalg.eigh(L)
+    # compute the heat kernel
+    heat_kernel = eigvecs @ torch.diag(torch.exp(-t * eigvals)) @ eigvecs.T
+    heat_kernel = (heat_kernel + heat_kernel.T) / 2
+    heat_kernel[heat_kernel < 0] = 0.0
+    return heat_kernel
+
+
+def test_laplacian():
+    data = torch.randn(100, 5)
+    sigma = 1.0
+    L = laplacian_from_data(data, sigma)
+    assert torch.allclose(L, L.T)
+    # compute the largest eigenvalue
+    eigvals = torch.linalg.eigvals(L).real
+    max_eigval = eigvals.max()
+    min_eigval = eigvals.min()
+    assert max_eigval <= 2.0
+
+
+@pytest.mark.parametrize("t", [0.1, 1.0,])
+@pytest.mark.parametrize("order", [10, 30, 50])
+@pytest.mark.parametrize("k", [10, 20])
+def test_heat_kernel_knn(t, order, k):
+    tol = 2e-1 if t > 1.0 else 1e-1
+    data = torch.randn(100, 5)
+    heat_op = HeatKernelKNN(k=k, t=t, order=order, graph_type="scanpy")
+    heat_kernel = heat_op(data)
+    
+    # test if symmetric
+    assert torch.allclose(heat_kernel, heat_kernel.T)
+
+    # test if positive
+    assert torch.all(heat_kernel >= 0)
+
+    # test if the heat kernel is close to the ground truth
+    gt_heat_kernel = gt_heat_kernel_knn(data, t=t, k=k)
+    assert torch.allclose(heat_kernel, gt_heat_kernel, atol=tol, rtol=tol)
+
+if __name__ == "__main__":
+    pytest.main([__file__])
\ No newline at end of file

From d4e31f2855fdf85215945a58cef60b0fc542001e Mon Sep 17 00:00:00 2001
From: Guillaume Huguet <57917099+guillaumehu@users.noreply.github.com>
Date: Thu, 29 Feb 2024 22:23:42 -0500
Subject: [PATCH 3/6] rm `laplacian_from_data`

---
 tests/test_heat.py             | 14 +-------------
 torchcfm/diffusion_distance.py |  5 -----
 2 files changed, 1 insertion(+), 18 deletions(-)

diff --git a/tests/test_heat.py b/tests/test_heat.py
index b73b7bd1..24bc3f34 100644
--- a/tests/test_heat.py
+++ b/tests/test_heat.py
@@ -1,6 +1,6 @@
 import pytest
 import torch
-from torchcfm.diffusion_distance import laplacian_from_data, HeatKernelKNN, torch_knn_from_data
+from torchcfm.diffusion_distance import HeatKernelKNN, torch_knn_from_data
 
 def gt_heat_kernel_knn(
     data,
@@ -17,18 +17,6 @@ def gt_heat_kernel_knn(
     return heat_kernel
 
 
-def test_laplacian():
-    data = torch.randn(100, 5)
-    sigma = 1.0
-    L = laplacian_from_data(data, sigma)
-    assert torch.allclose(L, L.T)
-    # compute the largest eigenvalue
-    eigvals = torch.linalg.eigvals(L).real
-    max_eigval = eigvals.max()
-    min_eigval = eigvals.min()
-    assert max_eigval <= 2.0
-
-
 @pytest.mark.parametrize("t", [0.1, 1.0,])
 @pytest.mark.parametrize("order", [10, 30, 50])
 @pytest.mark.parametrize("k", [10, 20])
diff --git a/torchcfm/diffusion_distance.py b/torchcfm/diffusion_distance.py
index 4e65baae..e9807ca2 100644
--- a/torchcfm/diffusion_distance.py
+++ b/torchcfm/diffusion_distance.py
@@ -15,11 +15,6 @@ def norm_sym_laplacian(A: torch.Tensor):
     return deg_sqrt_inv @ A @ deg_sqrt_inv
 
 
-def laplacian_from_data(data: torch.Tensor, sigma: float, alpha: int = 20):
-    affinity = torch.exp(-(torch.cdist(data, data) / (2 * sigma)).pow(alpha))
-    return norm_sym_laplacian(affinity)
-
-
 def torch_knn_from_data(
     data: torch.Tensor, k: int, projection: bool = False, proj_dim: int = 100
 ):

From 7fa245d40723bd072b4e48355f0f8514eefebb4b Mon Sep 17 00:00:00 2001
From: Guillaume Huguet <57917099+guillaumehu@users.noreply.github.com>
Date: Thu, 29 Feb 2024 22:40:03 -0500
Subject: [PATCH 4/6] change default

---
 torchcfm/diffusion_distance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchcfm/diffusion_distance.py b/torchcfm/diffusion_distance.py
index e9807ca2..ea237f67 100644
--- a/torchcfm/diffusion_distance.py
+++ b/torchcfm/diffusion_distance.py
@@ -104,7 +104,7 @@ def __init__(
         t: float = 1.0,
         projection: bool = False,
         proj_dim: int = 100,
-        graph_type: str = "torch",
+        graph_type: str = "scanpy",
     ):
         super().__init__(t=t, order=order)
         assert (

From 75c63379b8ffa333866e96eea79aacbc43beebd6 Mon Sep 17 00:00:00 2001
From: Guillaume Huguet <57917099+guillaumehu@users.noreply.github.com>
Date: Fri, 1 Mar 2024 15:07:15 -0500
Subject: [PATCH 5/6] add devices

---
 tests/test_heat.py             | 8 +++++++-
 torchcfm/cheb_approx.py        | 3 ++-
 torchcfm/diffusion_distance.py | 2 +-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tests/test_heat.py b/tests/test_heat.py
index 24bc3f34..ab59427d 100644
--- a/tests/test_heat.py
+++ b/tests/test_heat.py
@@ -2,6 +2,10 @@
 import torch
 from torchcfm.diffusion_distance import HeatKernelKNN, torch_knn_from_data
 
+DEVICES = ["cpu"]
+if torch.cuda.is_available():
+    DEVICES.append("cuda")
+
 def gt_heat_kernel_knn(
     data,
     t,
@@ -20,9 +24,11 @@ def gt_heat_kernel_knn(
 @pytest.mark.parametrize("t", [0.1, 1.0,])
 @pytest.mark.parametrize("order", [10, 30, 50])
 @pytest.mark.parametrize("k", [10, 20])
-def test_heat_kernel_knn(t, order, k):
+@pytest.mark.parametrize("device", DEVICES)
+def test_heat_kernel_knn(t, order, k, device):
     tol = 2e-1 if t > 1.0 else 1e-1
     data = torch.randn(100, 5)
+    data = data.to(device)
     heat_op = HeatKernelKNN(k=k, t=t, order=order, graph_type="scanpy")
     heat_kernel = heat_op(data)
     
diff --git a/torchcfm/cheb_approx.py b/torchcfm/cheb_approx.py
index 4f3ec2ec..2f5ba3d7 100644
--- a/torchcfm/cheb_approx.py
+++ b/torchcfm/cheb_approx.py
@@ -34,4 +34,5 @@ def body(carry, c):
 
 @torch.no_grad()
 def compute_chebychev_coeff_all(eigval, t, K):
-    return 2.0 * ive(torch.arange(0, K + 1), -t * eigval)
+    eigval = eigval.detach().cpu()
+    return 2.0 * ive(torch.arange(0, K + 1, device=eigval.device), -t * eigval)
diff --git a/torchcfm/diffusion_distance.py b/torchcfm/diffusion_distance.py
index ea237f67..f91197b0 100644
--- a/torchcfm/diffusion_distance.py
+++ b/torchcfm/diffusion_distance.py
@@ -38,7 +38,7 @@ def scanpy_knn_from_data(
         adata, n_neighbors=k, use_rep="X_pca" if projection else None
     )
     return norm_sym_laplacian(
-        torch.tensor(adata.obsp["connectivities"].toarray())
+        torch.tensor(adata.obsp["connectivities"].toarray(), device=data.device)
     )
 
 

From d14e4e5626b9c257c15f89c2ce2c61b675f8ced2 Mon Sep 17 00:00:00 2001
From: Guillaume Huguet <57917099+guillaumehu@users.noreply.github.com>
Date: Thu, 25 Jul 2024 17:04:09 -0400
Subject: [PATCH 6/6] fix laplacian

---
 torchcfm/diffusion_distance.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/torchcfm/diffusion_distance.py b/torchcfm/diffusion_distance.py
index f91197b0..127b8642 100644
--- a/torchcfm/diffusion_distance.py
+++ b/torchcfm/diffusion_distance.py
@@ -12,7 +12,8 @@
 def norm_sym_laplacian(A: torch.Tensor):
     deg = A.sum(dim=1)
     deg_sqrt_inv = torch.diag(1.0 / torch.sqrt(deg + EPS_LOG))
-    return deg_sqrt_inv @ A @ deg_sqrt_inv
+    id = torch.eye(A.shape[0], device=A.device, dtype=A.dtype)
+    return id - deg_sqrt_inv @ A @ deg_sqrt_inv
 
 
 def torch_knn_from_data(