Repository: scverse/scanpy

Issue:
Adding Moran's I calculation to Scanpy
Could you add Moran's I calculation to Scanpy? It could be used in scIB and to also find variable genes across embedding (could be an alternative to SEMITONES that takes a while to be computed). 


Discussion:
hi @Hrovatin 
so was it useful for your task? Curious to hear. Couple of questions:
- why having a tool redundant between two packages? 
- what is SEMITONES?
thank you!
Yes, it works great for me. I can compare scores obtained on individual samples and on integrated data to show that genes that are spatially variable in samples remain such on integrated data. However, with default n_iters most pvalues were 0 (for my known marker set), but calculating more iters would take too long, so I might just use the I-score where possible. 
I would like to add Moran's I as an bio conservation integration metric to scIB - this is for me the only metric that does not require cell subtype annotation (which is cumbersome and unreliable procedure) and it performs similar to current scIB metrics. However, scIB has as dependency only scanpy, not squidpy. It seems a bit of an overkill to add package dependency to scIB for a single function. 

Semitones (https://www.biorxiv.org/content/10.1101/2020.11.17.386664v1) is a package for finding genes linearly variable across embedding and I think Moran's I would also give me similar genes (must try it out) - Moran's I might be even better for the task and quicker + less complicated. This is another reason why it would be neat to have Moran's I directly in scanpy. You may not have spatial data, so not really needing squidpy. But finding gene patterns may be useful when you have continuous effects but no trajectories - this is what my main beta cell subtype analysis is currently based on.
I think Morans I would be a good addition to scanpy, fits well with Gearys C (https://github.com/theislab/scanpy/pull/915)
I see, now I understand the usage better as a metric, thank you!
I'd suggest that this is more of a feature request for Scanpy then @Hrovatin , feel free to open one or contribute to (theislab/scanpy#915) !
Transfered to scanpy
👀 crazy didn't know it was possible!
This would be a good fit for a potential `sc.metrics` module as discussed recently (can't find the issue anymore)
@LuckyMD I think we talked about it in the PR linked above.

So, what needs to happen here? Is it just copying and pasting some code from squidpy, or is there more implementation work to do?

Resolving patch:
diff --git a/scanpy/metrics/__init__.py b/scanpy/metrics/__init__.py
--- a/scanpy/metrics/__init__.py
+++ b/scanpy/metrics/__init__.py
@@ -1,2 +1,3 @@
 from ._gearys_c import gearys_c
 from ._metrics import confusion_matrix
+from ._morans_i import morans_i
diff --git a/scanpy/metrics/_gearys_c.py b/scanpy/metrics/_gearys_c.py
--- a/scanpy/metrics/_gearys_c.py
+++ b/scanpy/metrics/_gearys_c.py
@@ -166,7 +166,7 @@ def _gearys_c_vec_W(data, indices, indptr, x, W):
 # https://github.com/numba/numba/issues/6774#issuecomment-788789663
 
 
-@numba.njit
+@numba.njit(cache=True)
 def _gearys_c_inner_sparse_x_densevec(g_data, g_indices, g_indptr, x, W):
     x_bar = x.mean()
     total = 0.0
@@ -182,7 +182,7 @@ def _gearys_c_inner_sparse_x_densevec(g_data, g_indices, g_indptr, x, W):
     return C
 
 
-@numba.njit
+@numba.njit(cache=True)
 def _gearys_c_inner_sparse_x_sparsevec(
     g_data, g_indices, g_indptr, x_data, x_indices, N, W
 ):
diff --git a/scanpy/metrics/_morans_i.py b/scanpy/metrics/_morans_i.py
new file mode 100644
--- /dev/null
+++ b/scanpy/metrics/_morans_i.py
@@ -0,0 +1,252 @@
+"""Moran's I global spatial autocorrelation."""
+from typing import Union, Optional
+from functools import singledispatch
+from anndata import AnnData
+
+import numpy as np
+from scipy import sparse
+from numba import njit, prange
+
+from scanpy.get import _get_obs_rep
+from scanpy.metrics._gearys_c import _resolve_vals
+
+
+@singledispatch
+def morans_i(
+    adata: AnnData,
+    *,
+    vals: Optional[Union[np.ndarray, sparse.spmatrix]] = None,
+    use_graph: Optional[str] = None,
+    layer: Optional[str] = None,
+    obsm: Optional[str] = None,
+    obsp: Optional[str] = None,
+    use_raw: bool = False,
+) -> Union[np.ndarray, float]:
+    r"""
+    Calculate Moran’s I Global Autocorrelation Statistic.
+
+    Moran’s I is a global autocorrelation statistic for some measure on a graph. It is commonly used in
+    spatial data analysis to assess autocorrelation on a 2D grid. It is closely related to Geary's C,
+    but not identical. More info can be found `here <https://en.wikipedia.org/wiki/Moran%27s_I>`_.
+
+    .. math::
+
+        I =
+            \frac{
+                N \sum_{i, j} w_{i, j} z_{i} z_{j}
+            }{
+                S_{0} \sum_{i} z_{i}^{2}
+            }
+
+    Params
+    ------
+    adata
+    vals
+        Values to calculate Moran's I for. If this is two dimensional, should
+        be of shape `(n_features, n_cells)`. Otherwise should be of shape
+        `(n_cells,)`. This matrix can be selected from elements of the anndata
+        object by using key word arguments: `layer`, `obsm`, `obsp`, or
+        `use_raw`.
+    use_graph
+        Key to use for graph in anndata object. If not provided, default
+        neighbors connectivities will be used instead.
+    layer
+        Key for `adata.layers` to choose `vals`.
+    obsm
+        Key for `adata.obsm` to choose `vals`.
+    obsp
+        Key for `adata.obsp` to choose `vals`.
+    use_raw
+        Whether to use `adata.raw.X` for `vals`.
+
+
+    This function can also be called on the graph and values directly. In this case
+    the signature looks like:
+
+    Params
+    ------
+    g
+        The graph
+    vals
+        The values
+
+
+    See the examples for more info.
+
+    Returns
+    -------
+    If vals is two dimensional, returns a 1 dimensional ndarray array. Returns
+    a scalar if `vals` is 1d.
+
+
+    Examples
+    --------
+
+    Calculate Morans I for each components of a dimensionality reduction:
+
+    .. code:: python
+
+        import scanpy as sc, numpy as np
+
+        pbmc = sc.datasets.pbmc68k_processed()
+        pc_c = sc.metrics.morans_i(pbmc, obsm="X_pca")
+
+
+    It's equivalent to call the function directly on the underlying arrays:
+
+    .. code:: python
+
+        alt = sc.metrics.morans_i(pbmc.obsp["connectivities"], pbmc.obsm["X_pca"].T)
+        np.testing.assert_array_equal(pc_c, alt)
+    """
+    if use_graph is None:
+        # Fix for anndata<0.7
+        if hasattr(adata, "obsp") and "connectivities" in adata.obsp:
+            g = adata.obsp["connectivities"]
+        elif "neighbors" in adata.uns:
+            g = adata.uns["neighbors"]["connectivities"]
+        else:
+            raise ValueError("Must run neighbors first.")
+    else:
+        raise NotImplementedError()
+    if vals is None:
+        vals = _get_obs_rep(adata, use_raw=use_raw, layer=layer, obsm=obsm, obsp=obsp).T
+    return morans_i(g, vals)
+
+
+###############################################################################
+# Calculation
+###############################################################################
+# This is done in a very similar way to gearys_c. See notes there for details.
+
+
+@njit(cache=True)
+def _morans_i_vec_W_sparse(
+    g_data: np.ndarray,
+    g_indices: np.ndarray,
+    g_indptr: np.ndarray,
+    x_data: np.ndarray,
+    x_indices: np.ndarray,
+    N: int,
+    W: np.float_,
+) -> float:
+    x = np.zeros(N, dtype=x_data.dtype)
+    x[x_indices] = x_data
+    return _morans_i_vec_W(g_data, g_indices, g_indptr, x, W)
+
+
+@njit(cache=True)
+def _morans_i_vec_W(
+    g_data: np.ndarray,
+    g_indices: np.ndarray,
+    g_indptr: np.ndarray,
+    x: np.ndarray,
+    W: np.float_,
+) -> float:
+    z = x - x.mean()
+    z2ss = (z * z).sum()
+    N = len(x)
+    inum = 0.0
+
+    for i in prange(N):
+        s = slice(g_indptr[i], g_indptr[i + 1])
+        i_indices = g_indices[s]
+        i_data = g_data[s]
+        inum += (i_data * z[i_indices]).sum() * z[i]
+
+    return len(x) / W * inum / z2ss
+
+
+@njit(cache=True, parallel=True)
+def _morans_i_vec(
+    g_data: np.ndarray,
+    g_indices: np.ndarray,
+    g_indptr: np.ndarray,
+    x: np.ndarray,
+) -> float:
+    W = g_data.sum()
+    return _morans_i_vec_W(g_data, g_indices, g_indptr, x, W)
+
+
+@njit(cache=True, parallel=True)
+def _morans_i_mtx(
+    g_data: np.ndarray,
+    g_indices: np.ndarray,
+    g_indptr: np.ndarray,
+    X: np.ndarray,
+) -> np.ndarray:
+    M, N = X.shape
+    assert N == len(g_indptr) - 1
+    W = g_data.sum()
+    out = np.zeros(M, dtype=np.float_)
+    for k in prange(M):
+        x = X[k, :]
+        out[k] = _morans_i_vec_W(g_data, g_indices, g_indptr, x, W)
+    return out
+
+
+@njit(
+    cache=True,
+    parallel=True,
+)
+def _morans_i_mtx_csr(
+    g_data: np.ndarray,
+    g_indices: np.ndarray,
+    g_indptr: np.ndarray,
+    X_data: np.ndarray,
+    X_indices: np.ndarray,
+    X_indptr: np.ndarray,
+    X_shape: tuple,
+) -> np.ndarray:
+    M, N = X_shape
+    W = g_data.sum()
+    out = np.zeros(M, dtype=np.float_)
+    x_data_list = np.split(X_data, X_indptr[1:-1])
+    x_indices_list = np.split(X_indices, X_indptr[1:-1])
+    for k in prange(M):
+        out[k] = _morans_i_vec_W_sparse(
+            g_data,
+            g_indices,
+            g_indptr,
+            x_data_list[k],
+            x_indices_list[k],
+            N,
+            W,
+        )
+    return out
+
+
+###############################################################################
+# Interface (taken from gearys C)
+###############################################################################
+
+
+@morans_i.register(sparse.csr_matrix)
+def _morans_i(g, vals) -> np.ndarray:
+    assert g.shape[0] == g.shape[1], "`g` should be a square adjacency matrix"
+    vals = _resolve_vals(vals)
+    g_data = g.data.astype(np.float_, copy=False)
+    if isinstance(vals, sparse.csr_matrix):
+        assert g.shape[0] == vals.shape[1]
+        return _morans_i_mtx_csr(
+            g_data,
+            g.indices,
+            g.indptr,
+            vals.data.astype(np.float_, copy=False),
+            vals.indices,
+            vals.indptr,
+            vals.shape,
+        )
+    elif isinstance(vals, np.ndarray) and vals.ndim == 1:
+        assert g.shape[0] == vals.shape[0]
+        return _morans_i_vec(g_data, g.indices, g.indptr, vals)
+    elif isinstance(vals, np.ndarray) and vals.ndim == 2:
+        assert g.shape[0] == vals.shape[1]
+        return _morans_i_mtx(
+            g_data,
+            g.indices,
+            g.indptr,
+            vals.astype(np.float_, copy=False),
+        )
+    else:
+        raise NotImplementedError()


Test patch:
diff --git a/scanpy/tests/test_metrics.py b/scanpy/tests/test_metrics.py
--- a/scanpy/tests/test_metrics.py
+++ b/scanpy/tests/test_metrics.py
@@ -34,6 +34,7 @@ def test_gearys_c_consistency():
 
     np.testing.assert_allclose(all_genes[0], first_gene)
 
+    # Test that results are similar for sparse and dense reps of same data
     np.testing.assert_allclose(
         sc.metrics.gearys_c(pbmc, layer="raw"),
         sc.metrics.gearys_c(pbmc, vals=pbmc.layers["raw"].T.toarray()),
@@ -63,6 +64,63 @@ def test_gearys_c_correctness():
         assert sc.metrics.gearys_c(adata, vals=connected) == 0.0
 
 
+def test_morans_i_consistency():
+    pbmc = sc.datasets.pbmc68k_reduced()
+    pbmc.layers["raw"] = pbmc.raw.X.copy()
+    g = pbmc.obsp["connectivities"]
+
+    assert eq(
+        sc.metrics.morans_i(g, pbmc.obs["percent_mito"]),
+        sc.metrics.morans_i(pbmc, vals=pbmc.obs["percent_mito"]),
+    )
+
+    assert eq(  # Test that series and vectors return same value
+        sc.metrics.morans_i(g, pbmc.obs["percent_mito"]),
+        sc.metrics.morans_i(g, pbmc.obs["percent_mito"].values),
+    )
+
+    np.testing.assert_array_equal(
+        sc.metrics.morans_i(pbmc, obsm="X_pca"),
+        sc.metrics.morans_i(g, pbmc.obsm["X_pca"].T),
+    )
+
+    all_genes = sc.metrics.morans_i(pbmc, layer="raw")
+    first_gene = sc.metrics.morans_i(
+        pbmc, vals=pbmc.obs_vector(pbmc.var_names[0], layer="raw")
+    )
+
+    np.testing.assert_allclose(all_genes[0], first_gene, rtol=1e-5)
+
+    # Test that results are similar for sparse and dense reps of same data
+    np.testing.assert_allclose(
+        sc.metrics.morans_i(pbmc, layer="raw"),
+        sc.metrics.morans_i(pbmc, vals=pbmc.layers["raw"].T.toarray()),
+    )
+
+
+def test_morans_i_correctness():
+    # Test case with perfectly seperated groups
+    connected = np.zeros(100)
+    connected[np.random.choice(100, size=50, replace=False)] = 1
+    graph = np.zeros((100, 100))
+    graph[np.ix_(connected.astype(bool), connected.astype(bool))] = 1
+    graph[np.ix_(~connected.astype(bool), ~connected.astype(bool))] = 1
+    graph = sparse.csr_matrix(graph)
+
+    assert sc.metrics.morans_i(graph, connected) == 1.0
+    assert eq(
+        sc.metrics.morans_i(graph, connected),
+        sc.metrics.morans_i(graph, sparse.csr_matrix(connected)),
+    )
+    # Check for anndata > 0.7
+    if hasattr(sc.AnnData, "obsp"):
+        # Checking that obsp works
+        adata = sc.AnnData(
+            sparse.csr_matrix((100, 100)), obsp={"connectivities": graph}
+        )
+        assert sc.metrics.morans_i(adata, vals=connected) == 1.0
+
+
 def test_confusion_matrix():
     mtx = sc.metrics.confusion_matrix(["a", "b"], ["c", "d"], normalize=False)
     assert mtx.loc["a", "c"] == 1