Merge pull request #99 from EliHei2/cosmx_run

CosMx Data Support and Transcript Processing Improvements

Author: EliHei2

README.md

@@ -9,7 +9,7 @@
 # How segger Works
 
 ![Segger Model](docs/images/Segger_model_08_2024.png)
-Some illustrations (cells and data) are borrowed from [Biorender](https://www.biorender.com/) and [BIDCell's paper](https://www.nature.com/articles/s41467-023-44560-w). 
+Some illustrations (cells and data) are borrowed from [Biorender](https://www.biorender.com/) and [BIDCell's paper](https://www.nature.com/articles/s41467-023-44560-w).
 
 ---
 

scripts/create_data_cosmx.py

@@ -0,0 +1,99 @@
+from segger.data.parquet.sample import STSampleParquet, STInMemoryDataset
+from path import Path
+from segger.data.utils import calculate_gene_celltype_abundance_embedding
+import scanpy as sc
+import pandas as pd
+import math
+import numpy as np
+from segger.data.parquet._utils import get_polygons_from_xy
+
+"""
+This script preprocesses Xenium spatial transcriptomics data for SEGGER cell segmentation model.
+
+Key steps:
+1. Data Loading:
+   - Loads scRNA-seq reference data to create gene-celltype embeddings
+   - Imports Xenium transcripts and nucleus boundaries
+   
+2. Parameter Optimization:
+   - Calculates optimal neighborhood parameters based on tissue characteristics
+   - dist_tx: Sets transcript neighbor search radius to 1/4 of typical nucleus size
+   - k_tx: Determines number of transcripts to sample based on local density
+   
+3. Dataset Creation:
+   - Filters transcripts to those overlapping nuclei
+   - Creates graph connections between nearby transcripts
+   - Splits data into training/validation sets
+   - Saves in PyG format for SEGGER training
+
+Usage:
+- Input: Raw Xenium data (transcripts.parquet, nucleus_boundaries.parquet)
+- Output: Processed dataset with graph structure and embeddings
+"""
+
+# Define data paths
+# XENIUM_DATA_DIR = Path('/omics/odcf/analysis/OE0606_projects_temp/xenium_projects/20241209_Xenium5k_CNSL_BrM/20241209_Xenium5k_CNSL_BrM/output-XETG00078__0041719__Region_1__20241203__142052')
+# SEGGER_DATA_DIR = Path('data_tidy/pyg_datasets/CNSL_5k')
+# # SCRNASEQ_FILE = Path('/omics/groups/OE0606/internal/tangy/tasks/schier/data/atals_filtered.h5ad')
+# CELLTYPE_COLUMN = 'celltype_minor'
+
+
+XENIUM_DATA_DIR = Path("data_raw/cosmx/human_pancreas/processed/")
+SEGGER_DATA_DIR = Path("data_tidy/pyg_datasets/cosmx_pancreas_50")
+# SCRNASEQ_FILE = Path('/omics/groups/OE0606/internal/mimmo/Xenium/notebooks/data/scData/bh/bh_mng_scdata_20250306.h5ad')
+# CELLTYPE_COLUMN = 'annot_v1'
+
+# Calculate gene-celltype embeddings from reference data
+# gene_celltype_abundance_embedding = calculate_gene_celltype_abundance_embedding(
+#     sc.read(SCRNASEQ_FILE),
+#     CELLTYPE_COLUMN
+# )
+
+# Initialize spatial transcriptomics sample object
+sample = STSampleParquet(
+    base_dir=XENIUM_DATA_DIR,
+    n_workers=4,
+    sample_type="cosmx",
+    buffer_ratio=1,
+    # weights=gene_celltype_abundance_embedding
+)
+
+# Load and filter data
+transcripts = pd.read_parquet(XENIUM_DATA_DIR / "transcripts.parquet")
+boundaries = pd.read_parquet(XENIUM_DATA_DIR / "nucleus_boundaries.parquet")
+
+# Calculate optimal neighborhood parameters
+transcript_counts = transcripts.groupby("cell").size()
+nucleus_polygons = get_polygons_from_xy(boundaries, "x_global_px", "y_global_px", "cell")
+
+transcript_densities = nucleus_polygons[transcript_counts.index].area / transcript_counts
+nucleus_diameter = nucleus_polygons.minimum_bounding_radius().median() * 2
+
+# Set neighborhood parameters
+dist_tx = nucleus_diameter / 4  # Search radius = 1/4 nucleus diameter
+k_tx = math.ceil(
+    np.quantile(dist_tx**2 * np.pi * transcript_densities, 0.9)
+)  # Sample size based on 90th percentile density
+
+print(f"Calculated parameters: k_tx={k_tx}, dist_tx={dist_tx:.2f}")
+
+# Save processed dataset for SEGGER
+# Parameters:
+# - k_bd/dist_bd: Control nucleus boundary point connections
+# - k_tx/dist_tx: Control transcript neighborhood connections
+# - tile_width/height: Size of spatial tiles for processing
+# - neg_sampling_ratio: Ratio of negative to positive samples
+# - val_prob: Fraction of data for validation
+sample.save(
+    data_dir=SEGGER_DATA_DIR,
+    k_bd=3,  # Number of boundary points to connect
+    dist_bd=15,  # Maximum distance for boundary connections
+    k_tx=20,  # Use calculated optimal transcript neighbors
+    dist_tx=70,  # Use calculated optimal search radius
+    tile_width=500,  # Tile size for processing
+    tile_height=500,
+    neg_sampling_ratio=5.0,  # 5:1 negative:positive samples
+    frac=1.0,  # Use all data
+    val_prob=0.3,  # 30% validation set
+    test_prob=0,  # No test set
+)

scripts/create_data_fast_sample.py

@@ -7,97 +7,93 @@
 import numpy as np
 from segger.data.parquet._utils import get_polygons_from_xy
 
-xenium_data_dir = Path('data_raw/breast_cancer/Xenium_FFPE_Human_Breast_Cancer_Rep1/outs/')
-segger_data_dir = Path('data_tidy/pyg_datasets/bc_rep1_emb_200_final')
-
-
-scrnaseq_file = Path('/omics/groups/OE0606/internal/tangy/tasks/schier/data/atals_filtered.h5ad')
-celltype_column = 'celltype_minor'
-gene_celltype_abundance_embedding = calculate_gene_celltype_abundance_embedding(
-    sc.read(scrnaseq_file),
-    celltype_column
+"""
+This script preprocesses Xenium spatial transcriptomics data for SEGGER cell segmentation model.
+
+Key steps:
+1. Data Loading:
+   - Loads scRNA-seq reference data to create gene-celltype embeddings
+   - Imports Xenium transcripts and nucleus boundaries
+   
+2. Parameter Optimization:
+   - Calculates optimal neighborhood parameters based on tissue characteristics
+   - dist_tx: Sets transcript neighbor search radius to 1/4 of typical nucleus size
+   - k_tx: Determines number of transcripts to sample based on local density
+   
+3. Dataset Creation:
+   - Filters transcripts to those overlapping nuclei
+   - Creates graph connections between nearby transcripts
+   - Splits data into training/validation sets
+   - Saves in PyG format for SEGGER training
+
+Usage:
+- Input: Raw Xenium data (transcripts.parquet, nucleus_boundaries.parquet)
+- Output: Processed dataset with graph structure and embeddings
+"""
+
+# Define data paths
+# XENIUM_DATA_DIR = Path('/omics/odcf/analysis/OE0606_projects_temp/xenium_projects/20241209_Xenium5k_CNSL_BrM/20241209_Xenium5k_CNSL_BrM/output-XETG00078__0041719__Region_1__20241203__142052')
+# SEGGER_DATA_DIR = Path('data_tidy/pyg_datasets/CNSL_5k')
+# # SCRNASEQ_FILE = Path('/omics/groups/OE0606/internal/tangy/tasks/schier/data/atals_filtered.h5ad')
+# CELLTYPE_COLUMN = 'celltype_minor'
+
+
+XENIUM_DATA_DIR = Path(
+    "/omics/odcf/analysis/OE0606_projects_temp/oncolgy_data_exchange/analysis_domenico/project_24/output-XETG00423__0053177__mng_04_TMA__20250306__170821"
 )
+SEGGER_DATA_DIR = Path("data_tidy/pyg_datasets/MNG_0053177")
+SCRNASEQ_FILE = Path("/omics/groups/OE0606/internal/mimmo/Xenium/notebooks/data/scData/bh/bh_mng_scdata_20250306.h5ad")
+CELLTYPE_COLUMN = "annot_v1"
 
+# Calculate gene-celltype embeddings from reference data
+# gene_celltype_abundance_embedding = calculate_gene_celltype_abundance_embedding(
+#     sc.read(SCRNASEQ_FILE),
+#     CELLTYPE_COLUMN
+# )
+
+# Initialize spatial transcriptomics sample object
 sample = STSampleParquet(
-    base_dir=xenium_data_dir,
+    base_dir=XENIUM_DATA_DIR,
     n_workers=4,
     sample_type="xenium",
-    weights=gene_celltype_abundance_embedding,  # uncomment if gene-celltype embeddings are available
-)
-
-transcripts = pd.read_parquet(xenium_data_dir / "transcripts.parquet", filters=[[("overlaps_nucleus", "=", 1)]])
-boundaries = pd.read_parquet(xenium_data_dir / "nucleus_boundaries.parquet")
-
-sizes = transcripts.groupby("cell_id").size()
-polygons = get_polygons_from_xy(boundaries, "vertex_x", "vertex_y", "cell_id")
-densities = polygons[sizes.index].area / sizes
-bd_width = polygons.minimum_bounding_radius().median() * 2
-
-# 1/4 median boundary diameter
-dist_tx = bd_width / 4
-# 90th percentile density of bounding circle with radius=dist_tx
-k_tx = math.ceil(np.quantile(dist_tx**2 * np.pi * densities, 0.9))
-
-print(k_tx)
-print(dist_tx)
-
-
-sample.save(
-      data_dir=segger_data_dir,
-      k_bd=3,
-      dist_bd=15,
-      k_tx=3,
-      dist_tx=5,
-      tile_width=200,
-      tile_height=200,
-      neg_sampling_ratio=5.0,
-      frac=1.0,
-      val_prob=0.3,
-      test_prob=0,
-)
-
-
-xenium_data_dir = Path('data_tidy/bc_5k')
-segger_data_dir = Path('data_tidy/pyg_datasets/bc_5k_emb_new')
-
-
-
-sample = STSampleParquet(
-    base_dir=xenium_data_dir,
-    n_workers=8,
-    sample_type='xenium',
-    weights=gene_celltype_abundance_embedding, # uncomment if gene-celltype embeddings are available
+    # weights=gene_celltype_abundance_embedding
 )
 
-
-transcripts = pd.read_parquet(xenium_data_dir / "transcripts.parquet", filters=[[("overlaps_nucleus", "=", 1)]])
-boundaries = pd.read_parquet(xenium_data_dir / "nucleus_boundaries.parquet")
-
-sizes = transcripts.groupby("cell_id").size()
-polygons = get_polygons_from_xy(boundaries, "vertex_x", "vertex_y", "cell_id")
-densities = polygons[sizes.index].area / sizes
-bd_width = polygons.minimum_bounding_radius().median() * 2
-
-# 1/4 median boundary diameter
-dist_tx = bd_width / 4
-# 90th percentile density of bounding circle with radius=dist_tx
-k_tx = math.ceil(np.quantile(dist_tx**2 * np.pi * densities, 0.9))
-
-print(k_tx)
-print(dist_tx)
-
-
+# Load and filter data
+transcripts = pd.read_parquet(XENIUM_DATA_DIR / "transcripts.parquet", filters=[[("overlaps_nucleus", "=", 1)]])
+boundaries = pd.read_parquet(XENIUM_DATA_DIR / "nucleus_boundaries.parquet")
+
+# Calculate optimal neighborhood parameters
+transcript_counts = transcripts.groupby("cell_id").size()
+nucleus_polygons = get_polygons_from_xy(boundaries, "vertex_x", "vertex_y", "cell_id")
+transcript_densities = nucleus_polygons[transcript_counts.index].area / transcript_counts
+nucleus_diameter = nucleus_polygons.minimum_bounding_radius().median() * 2
+
+# Set neighborhood parameters
+dist_tx = nucleus_diameter / 4  # Search radius = 1/4 nucleus diameter
+k_tx = math.ceil(
+    np.quantile(dist_tx**2 * np.pi * transcript_densities, 0.9)
+)  # Sample size based on 90th percentile density
+
+print(f"Calculated parameters: k_tx={k_tx}, dist_tx={dist_tx:.2f}")
+
+# Save processed dataset for SEGGER
+# Parameters:
+# - k_bd/dist_bd: Control nucleus boundary point connections
+# - k_tx/dist_tx: Control transcript neighborhood connections
+# - tile_width/height: Size of spatial tiles for processing
+# - neg_sampling_ratio: Ratio of negative to positive samples
+# - val_prob: Fraction of data for validation
 sample.save(
-      data_dir=segger_data_dir,
-      k_bd=3,
-      dist_bd=15.0,
-      k_tx=15,
-      dist_tx=3,
-      tile_size=50_000, 
-      neg_sampling_ratio=5.0,
-      frac=0.1,
-      val_prob=0.1,
-      test_prob=0.1,
+    data_dir=SEGGER_DATA_DIR,
+    k_bd=3,  # Number of boundary points to connect
+    dist_bd=15,  # Maximum distance for boundary connections
+    k_tx=k_tx,  # Use calculated optimal transcript neighbors
+    dist_tx=dist_tx,  # Use calculated optimal search radius
+    tile_width=100,  # Tile size for processing
+    tile_height=100,
+    neg_sampling_ratio=5.0,  # 5:1 negative:positive samples
+    frac=1.0,  # Use all data
+    val_prob=0.3,  # 30% validation set
+    test_prob=0,  # No test set
 )
-
-

scripts/predict_model_sample.py

@@ -22,8 +22,8 @@
 seg_tag = "bc_fast_data_emb_major"
 model_version = 1
 
-segger_data_dir = Path('data_tidy/pyg_datasets') / seg_tag
-models_dir = Path("./models") / seg_tag 
+segger_data_dir = Path("data_tidy/pyg_datasets") / seg_tag
+models_dir = Path("./models") / seg_tag
 benchmarks_dir = Path("/dkfz/cluster/gpu/data/OE0606/elihei/segger_experiments/data_tidy/benchmarks/xe_rep1_bc")
 transcripts_file = "data_raw/xenium/Xenium_FFPE_Human_Breast_Cancer_Rep1/transcripts.parquet"
 # Initialize the Lightning data module

scripts/train_model_sample.py

@@ -15,16 +15,16 @@
 import os
 
 
-segger_data_dir = segger_data_dir = Path('data_tidy/pyg_datasets/bc_rep1_emb_final_200')
-models_dir = Path("./models/bc_rep1_emb_final_200")
+segger_data_dir = segger_data_dir = Path("data_tidy/pyg_datasets/cosmx_pancreas")
+models_dir = Path("./models/cosmx_pancreas")
 
 # Base directory to store Pytorch Lightning models
 # models_dir = Path('models')
 
 # Initialize the Lightning data module
 dm = SeggerDataModule(
     data_dir=segger_data_dir,
-    batch_size=1,
+    batch_size=2,
     num_workers=2,
 )
 
@@ -35,37 +35,34 @@
 
 # If you use custom gene embeddings, use the following two lines instead:
 is_token_based = False
-num_tx_tokens = dm.train[0].x_dict["tx"].shape[1] # Set the number of tokens to the number of genes
+num_tx_tokens = len(dm.train[0].x_dict["tx"])  # Set the number of tokens to the number of genes
 
 
 num_bd_features = dm.train[0].x_dict["bd"].shape[1]
 
 # Initialize the Lightning model
 ls = LitSegger(
-    is_token_based = is_token_based,
-    num_node_features = {"tx": num_tx_tokens, "bd": num_bd_features},
-    init_emb=8,    
+    is_token_based=is_token_based,
+    num_node_features={"tx": num_tx_tokens, "bd": num_bd_features},
+    init_emb=8,
     hidden_channels=64,
     out_channels=16,
     heads=4,
     num_mid_layers=3,
-    aggr='sum',
-    learning_rate=1e-3
+    aggr="sum",
+    learning_rate=1e-3,
 )
 
 # Initialize the Lightning trainer
 trainer = Trainer(
-    accelerator='cuda',
-    strategy='auto',
-    precision='16-mixed',
-    devices=2, # set higher number if more gpus are available
-    max_epochs=400,
+    accelerator="cpu",
+    strategy="auto",
+    precision="16-mixed",
+    devices=4,  # set higher number if more gpus are available
+    max_epochs=100,
     default_root_dir=models_dir,
     logger=CSVLogger(models_dir),
 )
 
 
-trainer.fit(
-    model=ls,
-    datamodule=dm
-)
+trainer.fit(model=ls, datamodule=dm)