グラフニューラルネットワークを使って、細胞コミュニティをクラスタリングするというこの論文。教師無しクラスタリングと教師有りクラスタリングが可能。クラスター数は自分で設定しなければならない。
教師有りクラスタリングの場合、各サンプルがどんな背景か(例えば健常者か病気か)を渡してやると、グラフクラスタリングはサンプルの分類タスクに沿っても訓練されていくため、群特異的なコミュニティが見つけやすくなる。
原著(Nature Methods, 2024): https://www.nature.com/articles/s41592-023-02124-2
原著より抜粋
TCN: tissue cellular neighborhood
チュートリアルも用意されており、デモデータ(細胞種とその座標)もGitHub内にある。
GitHub: https://github.com/huBioinfo/CytoCommunity
チュートリアル:https://cytocommunity.readthedocs.io/en/latest/
チュートリアルのGitHub: https://github.com/yafeixu-xidian/CytoCommunity_Tutorials
① soft TCN assignment
各ノードが指定クラスター数に切れ味よく振り分けられるようなグラフ畳み込みモデルを訓練する。そして、モデルにデータを通して出力を得る。この出力は、各ノードがどのクラスターかの確率値である。
これを論文ではsoftなクラスター割り当てと称している。
② hard TCN assignment
①のsoftな割り当てクラスター結果からhardな割り当てを行って最終的なクラスター番号を確定する。
「モデルの訓練➔予測」の一連の処理を1度ではなく、何度も行うことでアンサンブル学習を行う。(同じ訓練データで同じGNNモデルの訓練を初めからやり直す。)
アンサンブル学習の結果、細胞数*クラスター数*繰り返し数の情報が得られる。チュートリアルでは20回繰り返していた。
ここから最終的なクラスター番号を確定する。
これにはdiceRパッケージを使う。複数のsoftなクラスタリング結果を比較し、majority votingで最終クラスターを確定する。
細胞コミュニティ
細胞と細胞の座標位置からk最近傍法により、接続先を決定する。
チュートリアルではk=70ぐらいを用いていた。さすがに1つの細胞に接している細胞は70個も無いので、1つの細胞の隣の隣の隣ぐらいという風にある程度の範囲の細胞の情報を使うことを想定しているようだ。
k最近傍法では有向グラフとなる。これを無向グラフ版に変換してさらに接続を増やしている。
損失の設定
目的は、コミュニティ(細胞集合)の構成に応じてクラスター分類することである。このクラスタリングの過程には正解データが無い。
そこでモデルは次の項目を考慮してパラメーターを更新する。
- 細胞がどのクラスターなのかの割り当て確率の切れ味が鋭い。(どれか1つのクラスターの確率値が高い。)
- 接続している細胞は同じクラスターになる
- クラスターごとのノード数は偏らないようにする
接続しているからと言ってどんどん同じクラスターに割り当てることが無いように、最後の正則項は重要である。
0. 事前準備
本記事では以下を使用する。
import numpy as np
import pandas as pd
import tqdm.auto as tqdm
from sklearn.neighbors import kneighbors_graph
import torch
import torch_geometric
from torch_geometric.data import Data
from tqdm.auto import tqdm
from torch_geometric.data import InMemoryDataset
import torch_geometric.transforms as T
from torch_geometric.loader import DenseDataLoader
import torch.nn.functional as F
from torch.nn import Linear
from torch_geometric.nn import DenseGraphConv
from torch_geometric.nn.dense import dense_mincut_pool
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
デモデータ
次のGitHubページをクローンしておくとよい。Tutrialフォルダに教師無しクラスタリング用、教師有りクラスタリング用のデモデータが入っている。