refactor preprocessing pipelines for spatial_domain tasks

dance/data/base.py

@@ -81,6 +81,9 @@ def __init__(self, data: Union[AnnData, MuData], train_size: Optional[int] = Non
  if "dance_config" not in self._data.uns:
  self._data.uns["dance_config"] = dict()
 
+ def __repr__(self) -> str:
+ return f"{self.__class__.__name__} object that wraps (.data):\n{self.data}"
+
  def _setup_splits(self, train_size: Optional[Union[int, str]], val_size: int, test_size: int):
  if train_size is None:
  return

dance/datasets/base.py

@@ -129,10 +129,15 @@ def _maybe_load_cache(self, transform, cache, redo_cache) -> Union[Data, str]:
 
  cache_file_path = osp.join(cache_dir, f"{md5_hash}.pkl")
  if osp.isfile(cache_file_path) and cache:
- logger.info(f"Loading cached data at {cache_file_path}\n{'Cache data info':-^100}\n"
- f"Dataset: {self!r}\nTransformation: {transform!r}\n{'End of cache data info':-^100}")
  with open(cache_file_path, "rb") as f:
  data = pickle.load(f)
+ terminal_width = os.get_terminal_size().columns
+ logger.info(f"Loading cached data at {cache_file_path}\n"
+ f"{'Cache data info':=^{terminal_width}}\n"
+ f"{'Dataset object info':-^{terminal_width}}\n{self!r}\n"
+ f"{'Transformation info':-^{terminal_width}}\n{transform!r}\n"
+ f"{'Loaded data info':-^{terminal_width}}\n{data!r}\n"
+ f"{'End of cache data info':=^{terminal_width}}")
  return data
  else:
  return cache_file_path

dance/datasets/spatial.py

@@ -8,25 +8,13 @@
 import scanpy as sc
 
 from dance import logger
+from dance.data import Data
+from dance.datasets.base import BaseDataset
+from dance.registers import register_dataset
 from dance.utils.download import download_file, download_unzip, unzip_file
 
 IGNORED_FILES = ["readme.txt"]
 
-dataset = {
- "151510": "https://www.dropbox.com/sh/41h9brsk6my546x/AADa18mkJge-KQRTndRelTpMa?dl=1",
- "151507": "https://www.dropbox.com/sh/m3554vfrdzbwv2c/AACGsFNVKx8rjBgvF7Pcm2L7a?dl=1",
- "151508": "https://www.dropbox.com/sh/tm47u3fre8692zt/AAAJJf8-za_Lpw614ft096qqa?dl=1",
- "151509": "https://www.dropbox.com/sh/hihr7906vyirjet/AACslV5mKIkF2CF5QqE1LE6ya?dl=1",
- "151669": "https://www.dropbox.com/sh/ulw2nnnmgtbswvc/AAC0fT549EwtxKZWWoB89gb4a?dl=1",
- "151670": "https://www.dropbox.com/sh/8fw44zyyjgh0ddc/AAA1asGAmyDiMmvhRmL7pN1Na?dl=1",
- "151671": "https://www.dropbox.com/sh/9g5qzd5ykx2mpk3/AAD3xjx1i2h0RhYBc-Vft6CEa?dl=1",
- "151672": "https://www.dropbox.com/sh/l6519tr280krd4p/AAAWefCSp2iKhVmLgytlyxTta?dl=1",
- "151673": "https://www.dropbox.com/sh/qc64ps6gd64dm0c/AAC_5_mP4AczKj8lORLLKcIba?dl=1",
- "151674": "https://www.dropbox.com/sh/q7io99psd2xuqgw/AABske8dgX_kc1oaDSxuiqjpa?dl=1",
- "151675": "https://www.dropbox.com/sh/uahka2h5klnrzvj/AABe7K0_ewqOcqKUxHebE6qLa?dl=1",
- "151676": "https://www.dropbox.com/sh/jos5jjurezy5zp1/AAB2uaVm3-Us1a4mDkS1Q-iAa?dl=1",
-}
-
 cellDeconvo_dataset = {
  "CARD_synthetic": "https://www.dropbox.com/sh/v0vpv0jsnfexj7f/AADpizLGOrF7M8EesDihgbBla?dl=1",
  "GSE174746": "https://www.dropbox.com/sh/spfv06yfttetrab/AAAgORS6ocyoZEyxiRYKTymCa?dl=1",
@@ -38,81 +26,112 @@
 }
 
 
-class SpotDataset:
+@register_dataset("spatiallibd")
+class SpatialLIBDDataset(BaseDataset):
+
+ _DISPLAY_ATTRS = ("data_id", )
+ url_dict = {
+ "151510": "https://www.dropbox.com/sh/41h9brsk6my546x/AADa18mkJge-KQRTndRelTpMa?dl=1",
+ "151507": "https://www.dropbox.com/sh/m3554vfrdzbwv2c/AACGsFNVKx8rjBgvF7Pcm2L7a?dl=1",
+ "151508": "https://www.dropbox.com/sh/tm47u3fre8692zt/AAAJJf8-za_Lpw614ft096qqa?dl=1",
+ "151509": "https://www.dropbox.com/sh/hihr7906vyirjet/AACslV5mKIkF2CF5QqE1LE6ya?dl=1",
+ "151669": "https://www.dropbox.com/sh/ulw2nnnmgtbswvc/AAC0fT549EwtxKZWWoB89gb4a?dl=1",
+ "151670": "https://www.dropbox.com/sh/8fw44zyyjgh0ddc/AAA1asGAmyDiMmvhRmL7pN1Na?dl=1",
+ "151671": "https://www.dropbox.com/sh/9g5qzd5ykx2mpk3/AAD3xjx1i2h0RhYBc-Vft6CEa?dl=1",
+ "151672": "https://www.dropbox.com/sh/l6519tr280krd4p/AAAWefCSp2iKhVmLgytlyxTta?dl=1",
+ "151673": "https://www.dropbox.com/sh/qc64ps6gd64dm0c/AAC_5_mP4AczKj8lORLLKcIba?dl=1",
+ "151674": "https://www.dropbox.com/sh/q7io99psd2xuqgw/AABske8dgX_kc1oaDSxuiqjpa?dl=1",
+ "151675": "https://www.dropbox.com/sh/uahka2h5klnrzvj/AABe7K0_ewqOcqKUxHebE6qLa?dl=1",
+ "151676": "https://www.dropbox.com/sh/jos5jjurezy5zp1/AAB2uaVm3-Us1a4mDkS1Q-iAa?dl=1",
+ }
+
+ def __init__(self, root=".", full_download=False, data_id="151673", data_dir="data/spot"):
+ super().__init__(root, full_download)
 
- def __init__(self, data_id="151673", data_dir="data/spot", build_graph_fn="default"):
  self.data_id = data_id
  self.data_dir = data_dir + "/{}".format(data_id)
- self.data_url = dataset[data_id]
- self._load_data()
- self.adj = None
-
- def get_all_data(self):
- # provide an interface to get all data at one time
- print("All data includes {} datasets: {}".format(len(dataset), ",".join(dataset.keys())))
- res = {}
- for each_dataset in dataset.keys():
- res[each_dataset] = SpotDataset(each_dataset)
- return res
-
- def download_data(self):
- # judge whether a file exists or not
- isdownload = download_file(self.data_url, self.data_dir + "/{}.zip".format(self.data_id))
- if isdownload:
- unzip_file(self.data_dir + "/{}.zip".format(self.data_id), self.data_dir + "/")
- return self
+
+ def download_all(self):
+ logger.info(f"All data includes {len(self.url_dict)} datasets: {list(self.url_dict)}")
+ _data_id = self.data_id
+ for data_id in self.url_dict:
+ self.data_id = data_id
+ self.download()
+ self.data_id = _data_id
+
+ def is_complete_all(self):
+ _data_id = self.data_id
+ for data_id in self.url_dict:
+ self.data_id = data_id
+ if not self.is_complete():
+ self.data_id = _data_id
+ return False
+ self.data_id = _data_id
+ return True
+
+ def download(self):
+ out_path = osp.join(self.data_dir, f"{self.data_id}.zip")
+ if download_file(self.url_dict[self.data_id], out_path):
+ unzip_file(out_path, self.data_dir)
 
  def is_complete(self):
- # data.h5ad
- # histology.tif
- # positions.txt
- # judge whether data is complete or not
  check = [
- self.data_dir + "/{}_raw_feature_bc_matrix.h5".format(self.data_id),
- self.data_dir + "/{}_full_image.tif".format(self.data_id), self.data_dir + "/tissue_positions_list.txt"
+ osp.join(self.data_dir, f"{self.data_id}_raw_feature_bc_matrix.h5"), # expression
+ osp.join(self.data_dir, f"{self.data_id}_full_image.tif"), # histology
+ osp.join(self.data_dir, "tissue_positions_list.txt"), # positions
  ]
 
  for i in check:
  if not os.path.exists(i):
- print("lack {}".format(i))
+ logger.info(f"lack {i}")
  return False
+
  return True
 
- def _load_data(self):
- if self.is_complete():
- pass
- else:
- self.download_data()
- self.data = sc.read_10x_h5(self.data_dir + "/{}_raw_feature_bc_matrix.h5".format(self.data_id))
- self.img = cv2.imread(self.data_dir + "/{}_full_image.tif".format(self.data_id))
- label = pd.read_csv(self.data_dir + "/cluster_labels.csv")
- classes = {layer_class: idx for idx, layer_class in enumerate(set(label["ground_truth"].tolist()))}
- self.spatial = pd.read_csv(self.data_dir + "/tissue_positions_list.txt", sep=",", header=None, na_filter=False,
- index_col=0)
- self.data.obs["x1"] = self.spatial[1]
- self.data.obs["x2"] = self.spatial[2]
- self.data.obs["x3"] = self.spatial[3]
- self.data.obs["x4"] = self.spatial[4]
- self.data.obs["x5"] = self.spatial[5]
- self.data.obs["x"] = self.data.obs["x2"]
- self.data.obs["y"] = self.data.obs["x3"]
- self.data.obs["x_pixel"] = self.data.obs["x4"]
- self.data.obs["y_pixel"] = self.data.obs["x5"]
-
- self.data = self.data[self.data.obs["x1"] == 1]
- self.data.var_names = [i.upper() for i in list(self.data.var_names)]
- self.data.var["genename"] = self.data.var.index.astype("str")
- self.data.obs["label"] = list(map(lambda x: classes[x], label["ground_truth"].tolist()))
- self.data.obs["ground_truth"] = label["ground_truth"].tolist()
- return self
-
- def load_data(self):
- adata = self.data
- spatial = adata.obs[["x", "y"]]
- spatial_pixel = adata.obs[["x_pixel", "y_pixel"]]
- image = self.img
- label = adata.obs[["label"]]
- return image, adata, spatial, spatial_pixel, label
+ def _load_raw_data(self):
+ image_path = osp.join(self.data_dir, f"{self.data_id}_full_image.tif")
+ data_path = osp.join(self.data_dir, f"{self.data_id}_raw_feature_bc_matrix.h5")
+ spatial_path = osp.join(self.data_dir, "tissue_positions_list.txt")
+ meta_path = osp.join(self.data_dir, "cluster_labels.csv")
+
+ logger.info(f"Loading image data from {image_path}")
+ img = cv2.imread(image_path)
+
+ logger.info(f"Loading expression data from {data_path}")
+ adata = sc.read_10x_h5(data_path)
+
+ logger.info(f"Loading spatial info from {spatial_path}")
+ spatial = pd.read_csv(spatial_path, header=None, index_col=0).loc[adata.obs_names]
+
+ logger.info(f"Loading label info from {meta_path}")
+ meta_df = pd.read_csv(meta_path)
+
+ # Restrict to captured spots
+ indicator = spatial[1].values == 1
+ adata = adata[indicator]
+ spatial = spatial.iloc[indicator]
+
+ # Prepare spatial info tables
+ xy = spatial[[2, 3]].rename(columns={2: "x", 3: "y"})
+ xy_pixel = spatial[[4, 5]].rename(columns={4: "x_pixel", 5: "y_pixel"})
+
+ # Prepare meta data and create a column with indexed label info
+ label_classes = {j: i for i, j in enumerate(meta_df["ground_truth"].unique())}
+ meta_df["label"] = list(map(label_classes.get, meta_df["ground_truth"]))
+
+ return img, adata, xy, xy_pixel, meta_df
+
+ def _raw_to_dance(self, raw_data):
+ img, adata, xy, xy_pixel, meta_df = raw_data
+ adata.var_names_make_unique()
+
+ adata.obs = meta_df.set_index(adata.obs_names)
+ adata.obsm["spatial"] = xy.set_index(adata.obs_names)
+ adata.obsm["spatial_pixel"] = xy_pixel.set_index(adata.obs_names)
+ adata.uns["image"] = img
+
+ data = Data(adata, train_size="all")
+ return data
 
 
 class CellTypeDeconvoDatasetLite:

dance/modules/spatial/spatial_domain/louvain.py

@@ -14,6 +14,11 @@
 
 import networkx as nx
 import numpy as np
+import scanpy as sc
+
+from dance.transforms import AnnDataTransform, CellPCA, Compose, FilterGenesMatch, SetConfig
+from dance.transforms.graph import NeighborGraph
+from dance.typing import LogLevel
 
 PASS_MAX = -1
 MIN = 0.0000001
@@ -333,6 +338,23 @@ class Louvain:
  def __init__(self, resolution: float = 1):
  self.resolution = resolution
 
+ @staticmethod
+ def preprocessing_pipeline(dim: int = 50, n_neighbors: int = 17, log_level: LogLevel = "INFO"):
+ return Compose(
+ FilterGenesMatch(prefixes=["ERCC", "MT-"]),
+ AnnDataTransform(sc.pp.normalize_total, target_sum=1e4),
+ AnnDataTransform(sc.pp.log1p),
+ CellPCA(n_components=dim),
+ NeighborGraph(n_neighbors=n_neighbors),
+ SetConfig({
+ "feature_channel": "NeighborGraph",
+ "feature_channel_type": "obsp",
+ "label_channel": "label",
+ "label_channel_type": "obs"
+ }),
+ log_level=log_level,
+ )
+
  def fit(self, adj, partition=None, weight="weight", randomize=None, random_state=None):
  """Fit function for model training.
 

dance/modules/spatial/spatial_domain/spagcn.py

@@ -23,6 +23,9 @@
 from torch.nn.parameter import Parameter
 
 from dance import utils
+from dance.transforms import AnnDataTransform, CellPCA, Compose, FilterGenesMatch, SetConfig
+from dance.transforms.graph import SpaGCNGraph, SpaGCNGraph2D
+from dance.typing import LogLevel
 
 
 def refine(sample_id, pred, dis, shape="hexagon"):
@@ -447,9 +450,27 @@ class SpaGCN:
 
  def __init__(self, l=None):
  super().__init__()
- self.l = l or None
+ self.l = l
  self.res = None
 
+ @staticmethod
+ def preprocessing_pipeline(alpha: float = 1, beta: int = 49, dim: int = 50, log_level: LogLevel = "INFO"):
+ return Compose(
+ FilterGenesMatch(prefixes=["ERCC", "MT-"]),
+ AnnDataTransform(sc.pp.normalize_total, target_sum=1e4),
+ AnnDataTransform(sc.pp.log1p),
+ SpaGCNGraph(alpha=alpha, beta=beta),
+ SpaGCNGraph2D(),
+ CellPCA(n_components=dim),
+ SetConfig({
+ "feature_channel": ["CellPCA", "SpaGCNGraph", "SpaGCNGraph2D"],
+ "feature_channel_type": ["obsm", "obsp", "obsp"],
+ "label_channel": "label",
+ "label_channel_type": "obs"
+ }),
+ log_level=log_level,
+ )
+
  def search_l(self, p, adj, start=0.01, end=1000, tol=0.01, max_run=100):
  """Search best l.
 
@@ -625,5 +646,4 @@ def score(self, y_true):
  """
  from sklearn.metrics.cluster import adjusted_rand_score
  score = adjusted_rand_score(y_true, self.y_pred)
- print("ARI {}".format(adjusted_rand_score(y_true, self.y_pred)))
  return score

dance/modules/spatial/spatial_domain/stagate.py

@@ -25,6 +25,10 @@
 from torch_sparse import SparseTensor, set_diag
 from tqdm import tqdm
 
+from dance.transforms import AnnDataTransform, Compose, SetConfig
+from dance.transforms.graph import StagateGraph
+from dance.typing import LogLevel
+
 
 def transfer_pytorch_data(adata, adj):
  edgeList = adj
@@ -175,6 +179,23 @@ def __init__(self, hidden_dims):
  self.conv3 = GATConv(out_dim, num_hidden, heads=1, concat=False, dropout=0, add_self_loops=False, bias=False)
  self.conv4 = GATConv(num_hidden, in_dim, heads=1, concat=False, dropout=0, add_self_loops=False, bias=False)
 
+ @staticmethod
+ def preprocessing_pipeline(hvg_flavor: str = "seurat_v3", n_top_hvgs: int = 3000, model_name: str = "radius",
+ radius: float = 150, n_neighbors: int = 5, log_level: LogLevel = "INFO"):
+ return Compose(
+ AnnDataTransform(sc.pp.highly_variable_genes, flavor=hvg_flavor, n_top_genes=n_top_hvgs, subset=True),
+ AnnDataTransform(sc.pp.normalize_total, target_sum=1e4),
+ AnnDataTransform(sc.pp.log1p),
+ StagateGraph(model_name, radius=radius, n_neighbors=n_neighbors),
+ SetConfig({
+ "feature_channel": "StagateGraph",
+ "feature_channel_type": "obsp",
+ "label_channel": "label",
+ "label_channel_type": "obs"
+ }),
+ log_level=log_level,
+ )
+
  def forward(self, features, edge_index):
  """Forward function for training.