test density functions and multiscale

tests/test_util_density.py

@@ -0,0 +1,73 @@
+from anndata._core.anndata import AnnData
+from pandas.core.frame import DataFrame
+import pytest
+import pandas as pd
+import scanpy as sc
+import numpy as np
+
+from palantir.utils import (
+    run_density,
+    run_low_density_variability,
+    run_density_evaluation,
+)
+
+
+@pytest.fixture
+def mock_data():
+    n_cells = 50
+    n_genes = 500
+    return pd.DataFrame(
+        np.random.rand(n_cells, n_genes),
+        columns=[f"gene_{i}" for i in range(n_genes)],
+        index=[f"cell_{i}" for i in range(n_cells)],
+    )
+
+
+@pytest.fixture
+def mock_anndata(mock_data: DataFrame):
+    ad = sc.AnnData(X=mock_data)
+    ad.obsm["DM_EigenVectors"] = mock_data.iloc[:, :10].copy()
+    ad.obsm["branch_masks"] = pd.DataFrame(
+        columns=["branch_0", "branch_1"],
+        index=mock_data.index,
+        data=np.random.choice([True, False], size=(mock_data.shape[0], 2)),
+    )
+    ad.obs["other_density"] = np.random.rand(mock_data.shape[0])
+    ad.layers["local_variability"] = np.random.rand(*mock_data.shape)
+    return ad
+
+
+@pytest.fixture
+def mock_anndata_custom(mock_data: DataFrame):
+    ad = sc.AnnData(X=mock_data)
+    ad.obsm["DM_EigenVectors_custom"] = mock_data.iloc[:, :10].copy()
+    return ad
+
+
+def test_run_density(mock_anndata: AnnData):
+    run_density(mock_anndata)
+    assert "mellon_log_density" in mock_anndata.obs.keys()
+    assert "mellon_log_density_clipped" in mock_anndata.obs.keys()
+
+
+def test_run_density_custom_keys(mock_anndata_custom: AnnData):
+    run_density(
+        mock_anndata_custom, repr_key="DM_EigenVectors_custom", density_key="custom_key"
+    )
+    assert "custom_key" in mock_anndata_custom.obs.keys()
+    assert "custom_key_clipped" in mock_anndata_custom.obs.keys()
+
+
+def test_run_low_density_variability(mock_anndata: AnnData):
+    run_low_density_variability(mock_anndata, density_key="other_density")
+    for branch in mock_anndata.obsm["branch_masks"].columns:
+        assert f"low_density_gene_variability_{branch}" in mock_anndata.var.keys()
+
+
+def test_run_density_evaluation(mock_anndata: AnnData, mock_anndata_custom: AnnData):
+    run_density(mock_anndata)
+    run_density_evaluation(
+        mock_anndata, mock_anndata_custom, repr_key="DM_EigenVectors_custom"
+    )
+    assert "cross_log_density" in mock_anndata_custom.obs.keys()
+    assert "cross_log_density_clipped" in mock_anndata_custom.obs.keys()

tests/utils_diffusion_maps_from_kernel.py

@@ -5,18 +5,18 @@
 from scipy.sparse.linalg import eigs
 from pytest import approx
 
-from palantir.utils import diffusion_maps_from_kernel
+from palantir.utils import diffusion_maps_from_kernel, determine_multiscale_space
 
 
-def create_mock_kernel(size):
-    # Creating a mock symmetric positive definite kernel matrix
+@pytest.fixture
+def mock_kernel():
+    size = 50
     A = np.random.rand(size, size)
     return csr_matrix((A + A.T) / 2)
 
 
-def test_diffusion_maps_basic():
-    kernel = create_mock_kernel(50)
-    result = diffusion_maps_from_kernel(kernel)
+def test_diffusion_maps_basic(mock_kernel):
+    result = diffusion_maps_from_kernel(mock_kernel)
 
     assert isinstance(result, dict)
     assert "T" in result and "EigenVectors" in result and "EigenValues" in result
@@ -26,30 +26,35 @@ def test_diffusion_maps_basic():
     assert result["EigenValues"].shape == (10,)
 
 
-def test_diffusion_maps_n_components():
-    kernel = create_mock_kernel(50)
-    result = diffusion_maps_from_kernel(kernel, n_components=5)
+def test_diffusion_maps_n_components(mock_kernel):
+    result = diffusion_maps_from_kernel(mock_kernel, n_components=5)
 
     assert result["EigenVectors"].shape == (50, 5)
     assert result["EigenValues"].shape == (5,)
 
 
-def test_diffusion_maps_seed():
-    kernel = create_mock_kernel(50)
-    result1 = diffusion_maps_from_kernel(kernel, seed=0)
-    result2 = diffusion_maps_from_kernel(kernel, seed=0)
+def test_diffusion_maps_seed(mock_kernel):
+    result1 = diffusion_maps_from_kernel(mock_kernel, seed=0)
+    result2 = diffusion_maps_from_kernel(mock_kernel, seed=0)
 
     # Seed usage should yield the same result
     assert np.allclose(result1["EigenValues"], result2["EigenValues"])
 
 
-def test_diffusion_maps_eigen():
-    kernel = create_mock_kernel(50)
-    result = diffusion_maps_from_kernel(kernel)
+def test_diffusion_maps_eigen(mock_kernel):
+    result = diffusion_maps_from_kernel(mock_kernel)
 
     T = result["T"].toarray()
     e_values, e_vectors = eigs(T, 10, tol=1e-4, maxiter=1000)
 
     assert np.allclose(
         result["EigenValues"], np.real(sorted(e_values, reverse=True)[:10]), atol=1e-4
     )
+
+
+def test_determine_multiscale_space(mock_kernel):
+    result = diffusion_maps_from_kernel(mock_kernel)
+    determine_multiscale_space(result)
+    assert "EigenValues" in result.keys()
+    assert "EigenValues" in result.keys()
+    assert "T" in result.keys()