added more test functions and files, implementing assert frame equals from pandas

Author: anushka255

src/paste/__init__.py

@@ -1,3 +1,3 @@
-from .PASTE import pairwise_align, center_align, center_ot, center_NMF
-from .helper import match_spots_using_spatial_heuristic, filter_for_common_genes, apply_trsf, intersect
+from .PASTE import pairwise_align, center_align, center_ot, center_NMF, my_fused_gromov_wasserstein, solve_gromov_linesearch
+from .helper import match_spots_using_spatial_heuristic, filter_for_common_genes, apply_trsf, intersect,extract_data_matrix, to_dense_array, kl_divergence_backend
 from .visualization import plot_slice, stack_slices_pairwise, stack_slices_center

tests/data/output/H_center.csv

@@ -3,12 +3,22 @@
 
 import numpy as np
 import ot.backend
+from ot.lp import emd
 import pandas as pd
 import tempfile
 
 from paste import pairwise_align, center_align
-from paste.PASTE import center_ot, intersect, center_NMF
-
+from paste.PASTE import (
+    center_ot,
+    intersect,
+    center_NMF,
+    extract_data_matrix,
+    kl_divergence_backend,
+    to_dense_array,
+    my_fused_gromov_wasserstein,
+    solve_gromov_linesearch,
+)
+from pandas.testing import assert_frame_equal
 
 test_dir = Path(__file__).parent
 input_dir = test_dir / "data/input"
@@ -61,17 +71,24 @@ def test_center_alignment(slices):
         dissimilarity="kl",
         distributions=[slices[i].obsm["weights"] for i in range(len(slices))],
     )
-    pd.DataFrame(center_slice.uns["paste_W"], index=center_slice.obs.index).to_csv(
-        temp_dir / "W_center.csv"
+    assert_frame_equal(
+        pd.DataFrame(
+            center_slice.uns["paste_W"],
+            index=center_slice.obs.index,
+            columns=[str(i) for i in range(15)],
+        ),
+        pd.read_csv(output_dir / "W_center.csv", index_col=0),
+        check_names=False,
+        rtol=1e-05,
+        atol=1e-08,
     )
-    pd.DataFrame(center_slice.uns["paste_H"], columns=center_slice.var.index).to_csv(
-        temp_dir / "H_center.csv"
+    assert_frame_equal(
+        pd.DataFrame(center_slice.uns["paste_H"], columns=center_slice.var.index),
+        pd.read_csv(output_dir / "H_center.csv"),
+        rtol=1e-05,
+        atol=1e-08,
     )
 
-    # TODO: The following computations seem to be architecture dependent (need to look into as for how)
-    # assert_checksum_equals(temp_dir, "W_center.csv")
-    # assert_checksum_equals(temp_dir, "H_center.csv")
-
     for i, pi in enumerate(pairwise_info):
         pd.DataFrame(
             pi, index=center_slice.obs.index, columns=slices[i].obs.index
@@ -118,7 +135,6 @@ def test_center_ot(slices):
 
 
 def test_center_NMF(intersecting_slices):
-    temp_dir = Path(tempfile.mkdtemp())
     n_slices = len(intersecting_slices)
 
     pairwise_info = [
@@ -136,8 +152,106 @@ def test_center_NMF(intersecting_slices):
         random_seed=0,
     )
 
-    pd.DataFrame(_W).to_csv(temp_dir / "W_center_NMF.csv")
-    pd.DataFrame(_H).to_csv(temp_dir / "H_center_NMF.csv")
-    # TODO: The following computations seem to be architecture dependent (need to look into as for how)
-    # assert_checksum_equals(temp_dir, "W_center_NMF.csv")
-    # assert_checksum_equals(temp_dir, "H_center_NMF.csv")
+    assert_frame_equal(
+        pd.DataFrame(
+            _W,
+            index=intersecting_slices[0].obs.index,
+            columns=[str(i) for i in range(15)],
+        ),
+        pd.read_csv(output_dir / "W_center_NMF.csv", index_col=0),
+        rtol=1e-05,
+        atol=1e-08,
+    )
+    assert_frame_equal(
+        pd.DataFrame(_H, columns=intersecting_slices[0].var.index),
+        pd.read_csv(output_dir / "H_center_NMF.csv"),
+        rtol=1e-05,
+        atol=1e-08,
+    )
+
+
+def test_fused_gromov_wasserstein(slices):
+    temp_dir = Path(tempfile.mkdtemp())
+
+    common_genes = intersect(slices[0].var.index, slices[1].var.index)
+    sliceA = slices[0][:, common_genes]
+    sliceB = slices[1][:, common_genes]
+
+    nx = ot.backend.NumpyBackend()
+    slice1_dist = ot.dist(
+        nx.from_numpy(sliceA.obsm["spatial"]),
+        nx.from_numpy(sliceA.obsm["spatial"]),
+        metric="euclidean",
+    )
+    slice2_dist = ot.dist(
+        nx.from_numpy(sliceB.obsm["spatial"]),
+        nx.from_numpy(sliceB.obsm["spatial"]),
+        metric="euclidean",
+    )
+    slice1_distr = nx.ones((sliceA.shape[0],)) / sliceA.shape[0]
+    slice2_distr = nx.ones((sliceB.shape[0],)) / sliceB.shape[0]
+
+    slice1_X = nx.from_numpy(to_dense_array(extract_data_matrix(sliceA, None)))
+    slice2_X = nx.from_numpy(to_dense_array(extract_data_matrix(sliceB, None)))
+
+    M = nx.from_numpy(kl_divergence_backend(slice1_X + 0.01, slice2_X + 0.01))
+
+    pairwise_info, log = my_fused_gromov_wasserstein(
+        M,
+        slice1_dist,
+        slice2_dist,
+        slice1_distr,
+        slice2_distr,
+        G_init=None,
+        loss_fun="square_loss",
+        alpha=0.1,
+        log=True,
+        numItermax=200,
+    )
+    pd.DataFrame(pairwise_info).to_csv(temp_dir / "fused_gromov_wasserstein.csv")
+    # TODO: Need to figure out where the randomness is coming from
+    # assert_checksum_equals(temp_dir, "fused_gromov_wasserstein.csv")
+
+
+def test_gromov_linesearch(slices):
+    common_genes = intersect(slices[1].var.index, slices[2].var.index)
+    sliceA = slices[1][:, common_genes]
+    sliceB = slices[2][:, common_genes]
+
+    nx = ot.backend.NumpyBackend()
+    slice1_dist = ot.dist(
+        nx.from_numpy(sliceA.obsm["spatial"]),
+        nx.from_numpy(sliceA.obsm["spatial"]),
+        metric="euclidean",
+    )
+    slice2_dist = ot.dist(
+        nx.from_numpy(sliceB.obsm["spatial"]),
+        nx.from_numpy(sliceB.obsm["spatial"]),
+        metric="euclidean",
+    )
+    slice1_distr = nx.ones((sliceA.shape[0],)) / sliceA.shape[0]
+    slice2_distr = nx.ones((sliceB.shape[0],)) / sliceB.shape[0]
+
+    slice1_X = nx.from_numpy(to_dense_array(extract_data_matrix(sliceA, None)))
+    slice2_X = nx.from_numpy(to_dense_array(extract_data_matrix(sliceB, None)))
+
+    M = nx.from_numpy(kl_divergence_backend(slice1_X + 0.01, slice2_X + 0.01))
+    slice1_distr, slice2_distr = ot.utils.list_to_array(slice1_distr, slice2_distr)
+
+    constC, hC1, hC2 = ot.gromov.init_matrix(
+        slice1_dist, slice2_dist, slice1_distr, slice2_distr, loss_fun="square_loss"
+    )
+
+    G = slice1_distr[:, None] * slice2_distr[None, :]
+    Mi = M + 0.1 + ot.gromov.gwggrad(constC, hC1, hC2, G)
+    Mi = Mi + nx.min(Mi)
+
+    Gc = emd(slice1_distr, slice2_distr, Mi)
+    deltaG = Gc - G
+    costG = nx.sum(M * G) + 0.1 * ot.gromov.gwloss(constC, hC1, hC2, G)
+    alpha, fc, cost_G = solve_gromov_linesearch(
+        G, deltaG, costG, slice1_dist, slice2_dist, M=0.0, reg=1.0, nx=nx
+    )
+    assert alpha == 1.0
+    assert fc == 1
+    assert round(cost_G,6) == -11.419226