计算某个区域中细胞的复杂性

在某系空间组图谱文章中，我们经常会看到一种图片，是去计算某个区域中细胞的复杂性，我们以今年在 Cell 上发表的一篇文章为例，来看看如何计算某个区域中细胞的复杂性。文章名字为Single-cell spatial transcriptome reveals cell-type organization in the macaque cortex，文章链接为：https://linkinghub.elsevier.com/retrieve/pii/S0092867423006797 ，图二中有这么一张图片：

I图的意思是不同的 Layer 中的每个细胞的复杂度分布，每个细胞的复杂的计算为，以每个细胞为圆心，200个像素为半径，画一个圆圈，计算这个圆圈中和这个细胞的类型不同的细胞类型个数。根据这个理解，我们便是可以写程序了。

# 首先是包的引入
import anndata as ad
import numpy as np
import pandas as pd
import scanpy as sc
from matplotlib.ticker import FuncFormatter
import seaborn as sns
import numpy as np
from collections import Counter
import matplotlib.pyplot as plt

这里我写了三个函数，三个函数的计算速度是一个比一个快，但是易理解程度是一个比一个难

# 第一个函数是非常常见的我们所理解的遍历所有的细胞
def calculate_neighborhood_complexity(cell_coordinates, cluster_labels, radius):
    neighborhood_complexities = []

    for central_cell, central_cluster_label in zip(cell_coordinates, cluster_labels):
        complexity = 0

        for cell, cluster_label in zip(cell_coordinates, cluster_labels):
            distance = np.linalg.norm(cell - central_cell)
            if 0 < distance <= radius and cluster_label not in cluster_labels[:complexity]:
                complexity += 1

        neighborhood_complexities.append(complexity)

    complexity_counts = Counter(neighborhood_complexities)
    total_cells = len(neighborhood_complexities)
    complexity_probabilities = {
        complexity: count / total_cells for complexity, count in complexity_counts.items()
    }
    return np.array(neighborhood_complexities), complexity_probabilities

# 第二个函数则是利用了一些 numpy 的函数，速度比第一个函数快了很多
def calculate_neighborhood_complexity(cell_coordinates, cluster_labels, radius):
    unique_labels, label_indices = np.unique(cluster_labels, return_inverse=True)
    complexity_counts = np.zeros(len(cell_coordinates), dtype=int)
    
    for i, central_cell in enumerate(cell_coordinates):
        distances = np.linalg.norm(cell_coordinates - central_cell, axis=1)
        mask = np.logical_and(distances > 0, distances <= radius)
        previous_complexities = set()
        
        for cell_idx, cluster_idx in enumerate(label_indices[mask]):
            if cluster_idx not in previous_complexities:
                complexity_counts[i] += 1
                previous_complexities.add(cluster_idx)
    
    total_cells = len(cell_coordinates)
    complexity_probabilities = np.bincount(complexity_counts) / total_cells
    
    return complexity_counts, complexity_probabilities

# 第三个函数则是利用了 KD 树，在建树的时候就把边进行的维护，速度比第二个函数快了很多
from scipy.spatial import cKDTree

def calculate_neighborhood_complexity(cell_coordinates, cluster_labels, radius):
    unique_labels, label_indices = np.unique(cluster_labels, return_inverse=True)
    complexity_counts = np.zeros(len(cell_coordinates), dtype=int)

    # Build a KD-tree from the cell coordinates
    kdtree = cKDTree(cell_coordinates)

    for i, central_cell in enumerate(cell_coordinates):
        # Query the KD-tree to find the cell indices within the radius
        neighbor_indices = kdtree.query_ball_point(central_cell, radius)
        previous_complexities = set()

        for cell_idx, cluster_idx in zip(neighbor_indices, label_indices[neighbor_indices]):
            if cluster_idx not in previous_complexities:
                complexity_counts[i] += 1
                previous_complexities.add(cluster_idx)

    total_cells = len(cell_coordinates)
    complexity_probabilities = np.bincount(complexity_counts) / total_cells

    return complexity_counts, complexity_probabilities

那么应该如何是用，以及如何绘制图片呢？

# adata 为空间组的数据
# obs 中的 region 为脑区的信息
# obs 中的 celltype 为细胞类型的信息
# obsm 中的 spatial 为细胞的坐标信息
# layer 为各个层的信息
# 实测速度 第一个函数对于 20k 细胞，大概 15 min 左右， 第二个函数大概 5s 左右，第三个函数大概 1.5s 左右
cell_coordinates_list = [
    adata[adata.obs['region']=='layer1'].obsm['spatial'],
    adata[adata.obs['region']=='layer2'].obsm['spatial'],
    adata[adata.obs['region']=='layer3'].obsm['spatial'],
    adata[adata.obs['region']=='layer4'].obsm['spatial'],
]
cluster_labels_list = [
    adata[adata.obs['region']=='layer1'].obs['celltype'].tolist(),
    adata[adata.obs['region']=='layer2'].obs['celltype'].tolist(),
    adata[adata.obs['region']=='layer3'].obs['celltype'].tolist(),
    adata[adata.obs['region']=='layer4'].obs['celltype'].tolist(),
]

name = ['layer1', 'layer2', 'layer3', 'layer4']

radius_list = [200, 200, 200, 200]

fig, ax = plt.subplots(figsize=(6, 6))

for i, (cell_coordinates, cluster_labels, radius) in enumerate(zip(cell_coordinates_list, cluster_labels_list, radius_list)):
    complexities, complexity_probabilities = calculate_neighborhood_complexity(cell_coordinates, cluster_labels, radius)

    sns.kdeplot(complexities, cumulative=True, fill=False, ax=ax, label=name[i])

ax.set_xlabel('Neighborhood Complexity')
ax.set_ylabel('Probability (%)')
ax.set_title('Brain Neighborhood Complexity Distribution')
ax.legend(loc='lower right')

# Format y-axis labels as percentages
ax.yaxis.set_major_formatter(FuncFormatter(lambda x, _: f'{x*100:.0f}'))
plt.tight_layout()
plt.savefig('Brain_Neighborhood_Complexity_Distribution.pdf')
plt.show()

画出来的结果便是 Cell 的图片的样子了