diff --git a/scripts/run_benchmark/run_test_segger_nebius.sh b/scripts/run_benchmark/run_test_segger_nebius.sh index 22d1916ba..53e9488cb 100644 --- a/scripts/run_benchmark/run_test_segger_nebius.sh +++ b/scripts/run_benchmark/run_test_segger_nebius.sh @@ -1,10 +1,11 @@ #!/bin/bash -# Test run for the segger transcript-assignment method. +# Test run for the segger and fastreseg transcript-assignment methods. # Runs on the small S3 test resources, using the DEFAULT method at every stage -# plus segger added to the transcript-assignment stage. Segger requires a GPU, -# so this runs on Nebius (the gpu labels in labels_nebius.config pin it to the -# GPU node group). +# plus segger and fastreseg added to the transcript-assignment stage. Segger +# requires a GPU, so this runs on Nebius (the gpu labels in labels_nebius.config +# pin it to the GPU node group). fastreseg is CPU-based (midcpu/midmem labels) +# and schedules on a regular node within the same run. # get the root of the directory REPO_ROOT=$(git rev-parse --show-toplevel) @@ -32,6 +33,7 @@ segmentation_methods: transcript_assignment_methods: - basic_transcript_assignment - segger + - fastreseg count_aggregation_methods: - basic_count_aggregation qc_filtering_methods: @@ -68,4 +70,4 @@ tw launch https://github.com/openproblems-bio/task_ist_preprocessing.git \ --params-file /tmp/params_segger.yaml \ --entry-name auto \ --config src/base/labels_nebius.config \ - --labels task_ist_preprocessing,test,segger + --labels task_ist_preprocessing,test,segger,fastreseg diff --git a/src/methods_transcript_assignment/fastreseg/config.vsh.yaml b/src/methods_transcript_assignment/fastreseg/config.vsh.yaml new file mode 100644 index 000000000..512828392 --- /dev/null +++ b/src/methods_transcript_assignment/fastreseg/config.vsh.yaml @@ -0,0 +1,78 @@ +__merge__: /src/api/comp_method_transcript_assignment.yaml + +name: fastreseg +label: "fastReseg transcript assignment" +summary: "Spatial segmentation correction using fastReseg method" +description: | + fastReseg is an R package designed to enhance the precision of cell segmentation in spatial transcriptomics by + leveraging transcriptomic data to correct and refine initial image-based segmentation results. +links: + documentation: "https://github.com/openproblems-bio/task_ist_preprocessing" + repository: "https://github.com/openproblems-bio/task_ist_preprocessing" +references: + doi: "10.1038/s41598-025-08733-5" + + +#arguments: +# - name: --transcripts_key +# type: string +# default: "transcripts" +# description: "Key for transcripts in the points layer" +# - name: --coordinate_system +# type: string +# default: "global" +# description: "Coordinate system for the transcripts" + + +resources: + - type: bash_script + path: orchestrator.sh + - type: python_script + path: input.py + - type: r_script + path: script.R + - type: python_script + path: output.py + +engines: + - type: docker + image: openproblems/base_python:1 + setup: + - type: apt + packages: + - libv8-dev + - libudunits2-dev + - libabsl-dev + - gdal-bin + - libgdal-dev + - r-base + - type: r + packages: + - devtools + - terra + - Matrix + - dbscan + - igraph + - matrixStats + - codetools + - data.table + # Add other CRAN packages here + github: + # Add GitHub packages here if needed + - Nanostring-Biostats/FastReseg + #- type: python + # github: + # - openproblems-bio/core#subdirectory=packages/python/openproblems + __merge__: + - /src/base/setup_txsim_partial.yaml + - /src/base/setup_spatialdata_partial.yaml + - type: native + +runners: + - type: executable + - type: nextflow + directives: + label: [ midtime, midcpu, midmem ] + + +##macOS workaround: export DOCKER_DEFAULT_PLATFORM=linux/amd64 \ No newline at end of file diff --git a/src/methods_transcript_assignment/fastreseg/input.py b/src/methods_transcript_assignment/fastreseg/input.py new file mode 100644 index 000000000..1da6f5a48 --- /dev/null +++ b/src/methods_transcript_assignment/fastreseg/input.py @@ -0,0 +1,115 @@ +import spatialdata as sd +import sys +import numpy as np +import pandas as pd +import xarray as xr +import dask +import txsim as tx +import anndata as ad +import argparse + +def parse_arguments(): + parser = argparse.ArgumentParser( + description="Process input files and generate output files for FastReseg input" + ) + parser.add_argument('input_path_ist', help='Path to the input file') + parser.add_argument('input_segmentation_path', help='Path to the input segmentation file') + parser.add_argument('input_sc_reference_path', help='Path to the input single-cell reference file') + parser.add_argument('output_path_counts', help='Path for the output TSV file') + parser.add_argument('output_path_transcripts', help='Path for the output TIF file') + parser.add_argument('output_path_cell_type', help='Output cell type specification') + + return parser.parse_args() + +### parsing arguments +args = parse_arguments() +print("args:") +print(args) +input_path = args.input_path_ist +input_segmentation_path = args.input_segmentation_path +input_sc_reference_path = args.input_sc_reference_path +print("path") +print(input_sc_reference_path) +output_path_counts = args.output_path_counts +output_path_transcripts = args.output_path_transcripts +output_path_cell_type = args.output_path_cell_type + +## potential other parameters (TODO - make configurable) +um_per_pixel = 0.5 +sc_celltype_key = 'cell_type' + +### reading the data in +sdata = sd.read_zarr(input_path) + +### reading in basic segmentation +sdata_segm = sd.read_zarr(input_segmentation_path) +segmentation_coord_systems = sd.transformations.get_transformation(sdata_segm["segmentation"], get_all=True).keys() + +# In case of a translation transformation of the segmentation (e.g. crop of the data), we need to adjust the transcript coordinates +trans = sd.transformations.get_transformation(sdata_segm["segmentation"], get_all=True)['global'].inverse() + +transcripts = sd.transform(sdata['transcripts'], to_coordinate_system='global') +transcripts = sd.transform(transcripts, trans, 'global') + +print('Assigning transcripts to cell ids', flush=True) +y_coords = transcripts.y.compute().to_numpy(dtype=np.int64) +x_coords = transcripts.x.compute().to_numpy(dtype=np.int64) +if isinstance(sdata_segm["segmentation"], xr.DataTree): + label_image = sdata_segm["segmentation"]["scale0"].image.to_numpy() +else: + label_image = sdata_segm["segmentation"].to_numpy() +cell_id_dask_series = dask.dataframe.from_dask_array( + dask.array.from_array( + label_image[y_coords, x_coords], chunks=tuple(sdata['transcripts'].map_partitions(len).compute()) + ), + index=sdata['transcripts'].index +) +sdata['transcripts']["cell_id"] = cell_id_dask_series + +### extracting transcript ids +print('Transforming transcripts coordinates', flush=True) +transcripts = sd.transform(sdata['transcripts'], to_coordinate_system='global') + +transcripts_df = transcripts.compute() +transcripts_df.rename(columns = {'feature_name': 'target', +'transcript_id': 'UMI_transID', 'cell_id': 'UMI_cellID'}, inplace = True) + +transcripts_df = transcripts_df.loc[:, ['target', 'x', 'y', 'z', 'UMI_transID', 'UMI_cellID']] +transcripts_df.to_csv(output_path_transcripts) + + +#### aggregating counts per transcript, based on +df = sdata['transcripts'].compute() +df.feature_name = df.feature_name.astype(str) + +adata_sp = tx.preprocessing.generate_adata(df, cell_id_col='cell_id', gene_col='feature_name') #TODO: x and y refers to a specific coordinate system. Decide which space we want to use here. (probably should be handled in the previous assignment step) +adata_sp.layers['counts'] = adata_sp.layers['raw_counts'] +del adata_sp.layers['raw_counts'] +adata_sp.var["gene_name"] = adata_sp.var_names +print(adata_sp.var_names[1:10]) + +# currently the function also saves the transcripts in the adata object, but this is not necessary here +del adata_sp.uns['spots'] +del adata_sp.uns['pct_noise'] + + +count_df = pd.DataFrame(adata_sp.X.toarray(), + index=adata_sp.obs_names, + columns=adata_sp.var_names) +count_df.to_csv(output_path_counts) + +#### run cell annotation with ssam +adata_sc = ad.read_h5ad(input_sc_reference_path) +adata_sc.X = adata_sc.layers["normalized"] +print(adata_sc.var_names[1:10]) + +shared_genes = [g for g in adata_sc.var_names if g in adata_sp.var_names] +adata_sp = adata_sp[:,shared_genes] + +print('Annotating cell types', flush=True) +adata_sp = tx.preprocessing.run_ssam( + adata_sp, transcripts.compute(), adata_sc, um_p_px=um_per_pixel, + cell_id_col='cell_id', gene_col='feature_name', sc_ct_key=sc_celltype_key +) +cell_type_df = adata_sp.obs["ct_ssam"].astype(str) +cell_type_df.to_csv(output_path_cell_type, header=True) \ No newline at end of file diff --git a/src/methods_transcript_assignment/fastreseg/orchestrator.sh b/src/methods_transcript_assignment/fastreseg/orchestrator.sh new file mode 100644 index 000000000..9529d5b51 --- /dev/null +++ b/src/methods_transcript_assignment/fastreseg/orchestrator.sh @@ -0,0 +1,69 @@ +#!/bin/bash + +## VIASH START +# The following code has been auto-generated by Viash. +par_input_ist='resources_test/task_ist_preprocessing/mouse_brain_combined/raw_ist.zarr' +par_input_segmentation='resources_test/task_ist_preprocessing/mouse_brain_combined/segmentation.zarr' +par_input_scrnaseq='resources_test/task_ist_preprocessing/mouse_brain_combined/scrnaseq_reference.h5ad' +par_sc_cell_type_key='cell_type' +par_output='resources_test/task_ist_preprocessing/mouse_brain_combined/transcript_assignments.zarr' +par_transcripts_key='transcripts' +par_coordinate_system='global' +meta_name='fastreseg' +meta_functionality_name='fastreseg' +meta_resources_dir='/private/tmp/viash_inject_fastreseg18170326436127140412' +meta_executable='/private/tmp/viash_inject_fastreseg18170326436127140412/fastreseg' +meta_config='/private/tmp/viash_inject_fastreseg18170326436127140412/.config.vsh.yaml' +meta_temp_dir='/var/folders/fq/ymt0vml175s4yvqxzbmlmpz80000gn/T/' +meta_cpus='123' +meta_memory_b='123' +meta_memory_kb='123' +meta_memory_mb='123' +meta_memory_gb='123' +meta_memory_tb='123' +meta_memory_pb='123' +meta_memory_kib='123' +meta_memory_mib='123' +meta_memory_gib='123' +meta_memory_tib='123' +meta_memory_pib='123' + +## VIASH END + +par_intermediate_dir=$(mktemp -d -p "$(pwd)" tmp-processing-XXXXXXXX) + +echo "running FastReseg orchestrator" + +# Create intermediate directory +mkdir -p "$par_intermediate_dir" +echo $(date +%T) + +# Step 1: Run Python script to reformat input in the first Python environment +python "$meta_resources_dir/input.py" \ + "$par_input_ist" \ + "$par_input_segmentation" \ + "$par_input_scrnaseq" \ + "$par_intermediate_dir/counts.tsv" \ + "$par_intermediate_dir/transcripts.tsv" \ + "$par_intermediate_dir/cell_types.tsv" + +head $par_intermediate_dir/cell_types.tsv + +# Step 2: RunFastReseg + +##running the R script +Rscript "$meta_resources_dir/script.R" "$par_intermediate_dir/counts.tsv" \ + "$par_intermediate_dir/transcripts.tsv" \ + "$par_intermediate_dir/cell_types.tsv" \ + "$par_intermediate_dir/cell_ids.csv" \ + "$par_intermediate_dir/gene_names.csv" \ + "$par_intermediate_dir/transcripts_out.csv" + +## python output +python "$meta_resources_dir/output.py" \ + "$par_intermediate_dir/cell_ids.csv" \ + "$par_intermediate_dir/gene_names.csv" \ + "$par_intermediate_dir/transcripts_out.csv" \ + "$par_output" + +echo $(date +%T) \ No newline at end of file diff --git a/src/methods_transcript_assignment/fastreseg/output.py b/src/methods_transcript_assignment/fastreseg/output.py new file mode 100644 index 000000000..7849d351e --- /dev/null +++ b/src/methods_transcript_assignment/fastreseg/output.py @@ -0,0 +1,58 @@ +import spatialdata as sd +import sys +import numpy as np +import pandas as pd +import xarray as xr +from scipy.sparse import coo_matrix +from pathlib import Path +import os +import dask +import anndata as ad + +def convert_to_lower_dtype(arr): + max_val = arr.max() + if max_val <= np.iinfo(np.uint8).max: + new_dtype = np.uint8 + elif max_val <= np.iinfo(np.uint16).max: + new_dtype = np.uint16 + elif max_val <= np.iinfo(np.uint32).max: + new_dtype = np.uint32 + else: + new_dtype = np.uint64 + + return arr.astype(new_dtype) + +path_to_cell_ids_out = sys.argv[1] +path_to_gene_names_out = sys.argv[2] +path_to_transcripts_out = sys.argv[3] + +output_zarr_path = sys.argv[4] + +cell_df = pd.read_csv(path_to_cell_ids_out, index_col=0) +var_names = pd.read_csv(path_to_gene_names_out, index_col=0) +df = pd.read_csv(path_to_transcripts_out, index_col=0) +print(df.head()) + +##converting to dask transcript output for spatialdata format +df = df.loc[:,["x", "y", "z", "target", "updated_cellID", "updated_celltype", "UMI_transID"]] +df.rename(columns={"updated_cellID": "cell_id", "target": "feature_name", "UMI_transID": "transcript_id"}, inplace=True) +transcripts_dask = dask.dataframe.from_pandas(df, npartitions = 1) + + +sdata_transcripts_only = sd.SpatialData( + points={ + "transcripts": sd.models.PointsModel.parse(transcripts_dask) + }, + tables={ + "table": ad.AnnData( + # The transcript-assignment output format requires a `cell_id` column in + # the table's obs (see src/api/file_transcript_assignments.yaml); expose + # fastReseg's updated segmentation under the standard column names. + obs=cell_df.loc[:,['updated_cellID', 'updated_celltype']].rename( + columns={'updated_cellID': 'cell_id', 'updated_celltype': 'cell_type'} + ), + var=pd.DataFrame(index=var_names['x']) + ) + } +) +sdata_transcripts_only.write(output_zarr_path) \ No newline at end of file diff --git a/src/methods_transcript_assignment/fastreseg/script.R b/src/methods_transcript_assignment/fastreseg/script.R new file mode 100644 index 000000000..f1e67810e --- /dev/null +++ b/src/methods_transcript_assignment/fastreseg/script.R @@ -0,0 +1,91 @@ +#install.packages("devtools") +#devtools::install_github("Nanostring-Biostats/FastReseg", build_vignettes = FALSE, ref = "main") + +library(FastReseg) +library(data.table) +library(dplyr) + +args <- commandArgs(trailingOnly = TRUE) + +# Check if arguments are provided +if (length(args) == 0) { + stop("No arguments provided") +} + +# Access individual arguments +path_to_counts <- args[1] +path_to_transcripts <- args[2] +path_to_cell_annot <- args[3] + +path_to_cell_ids_out <- args[4] +path_to_gene_names_out <- args[5] +path_to_transcripts_out <- args[6] + + +### reading in data +count_df <- as.matrix(read.csv(path_to_counts, row.names = 1)) +head(count_df) + +cell_df <- read.csv(path_to_cell_annot, row.names = 1) +cell_annot <- cell_df[row.names(count_df),] +cell_annot <- setNames(cell_annot, row.names(count_df)) + +transcriptDF <- read.csv(path_to_transcripts) +head(transcriptDF) + + +##run pipeline +refineAll_res_one_FOC <- fastReseg_full_pipeline( + counts = count_df, + clust = cell_annot, + + # Similar to `runPreprocess()`, one can use `clust = NULL` if providing `refProfiles` + + transcript_df = transcriptDF, # + transDF_fileInfo = NULL, + pixel_size = 0.18, + zstep_size = 0.8, + transID_coln = NULL, + transGene_coln = "target", + cellID_coln = "UMI_cellID", + spatLocs_colns = c("x","y","z"), + extracellular_cellID = c(-1), + + # Similar to `runPreprocess()`, one can set various cutoffs to NULL for automatic calculation from input data + + # distance cutoff for neighborhood searching at molecular and cellular levels, respectively + molecular_distance_cutoff = 2.7, + cellular_distance_cutoff = NULL, + + # cutoffs for transcript scores and number for cells under each cell type + score_baseline = NULL, + lowerCutoff_transNum = NULL, + higherCutoff_transNum= NULL, + imputeFlag_missingCTs = TRUE, + + # Settings for error detection and correction, refer to `runSegRefinement()` for more details + flagCell_lrtest_cutoff = 5, # cutoff to flag for cells with strong spatial dependcy in transcript score profiles + svmClass_score_cutoff = -2, # cutoff of transcript score to separate between high and low score classes + groupTranscripts_method = "dbscan", + spatialMergeCheck_method = "leidenCut", + cutoff_spatialMerge = 0.5, # spatial constraint cutoff for a valid merge event + + path_to_output = "res2_multiFiles", + save_intermediates = TRUE, # flag to return and write intermediate results to disk + return_perCellData = TRUE, # flag to return per cell level outputs from updated segmentation + combine_extra = FALSE # flag to include trimmed and extracellular transcripts in the exported `updated_transDF.csv` files +) + +if(is.null(refineAll_res_one_FOC$updated_transDF_list[[1]])){ + print("no transcripts assigned after refinement") + cell_df$UMI_cellID <- as.integer(row.names(cell_df)) + transcriptDF <- left_join(cell_df, transcriptDF) + names(transcriptDF)[names(transcriptDF) == "UMI_cellID"] <- "updated_cellID" + names(transcriptDF)[names(transcriptDF) == "ct_ssam"] <- "updated_celltype" + write.csv(transcriptDF, file = path_to_transcripts_out) +} else { + write.csv(refineAll_res_one_FOC$updated_transDF_list[[1]], file = path_to_transcripts_out) +} +### export outputs +write.csv(row.names(refineAll_res_one_FOC$updated_perCellExprs), file = path_to_gene_names_out) +write.csv(refineAll_res_one_FOC$updated_perCellDT, file = path_to_cell_ids_out) diff --git a/src/methods_transcript_assignment/segger/config.vsh.yaml b/src/methods_transcript_assignment/segger/config.vsh.yaml index 377150533..c77e62aaf 100644 --- a/src/methods_transcript_assignment/segger/config.vsh.yaml +++ b/src/methods_transcript_assignment/segger/config.vsh.yaml @@ -35,6 +35,17 @@ arguments: choices: [nucleus, cell, uniform] description: "Which polygon set drives the prediction graph." default: cell + - name: --node_representation_dim + type: integer + required: false + description: >- + Node embedding size for segger's GNN — maps to segger's --node-representation-dim + (its `in_channels` / cells_embedding_size, default 128). segger runs a PCA with this + many components on a genes x genes correlation matrix, so on small panels/crops (where + fewer genes survive segger's internal count filters) the PCA fails if this exceeds the + surviving-gene count. script.py auto-retries once at the ceiling segger reports, so this + is effectively an upper bound; leave at 128 for full panels. + default: 128 resources: - type: python_script diff --git a/src/methods_transcript_assignment/segger/script.py b/src/methods_transcript_assignment/segger/script.py index fe1c32ece..4293760d5 100644 --- a/src/methods_transcript_assignment/segger/script.py +++ b/src/methods_transcript_assignment/segger/script.py @@ -1,4 +1,5 @@ import os +import re import json import shutil import subprocess @@ -184,7 +185,7 @@ ################ SEGGER_OUT_DIR.mkdir(parents=True, exist_ok=True) -cmd = [ +base_cmd = [ "segger", "segment", "-i", str(XENIUM_DIR), "-o", str(SEGGER_OUT_DIR), @@ -198,8 +199,47 @@ env.setdefault("RAPIDS_NO_INITIALIZE", "1") env.setdefault("CUDF_NO_INITIALIZE", "1") env.setdefault("RMM_NO_INITIALIZE", "1") -print("Running segger:", " ".join(cmd), flush=True) -subprocess.run(cmd, check=True, env=env) + + +def run_segger(node_dim): + # --node-representation-dim is segger's `in_channels` (== cells_embedding_size): + # the PCA n_components for the gene/cell embeddings (segger's own default 128). + cmd = base_cmd + ["--node-representation-dim", str(int(node_dim))] + print("Running segger:", " ".join(cmd), flush=True) + # Stream segger's output live AND capture it, so we can recover the max valid + # embedding dim from the PCA error below without hiding the training logs. + proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, env=env + ) + captured = [] + for line in proc.stdout: + print(line, end="", flush=True) + captured.append(line) + proc.wait() + return proc.returncode, "".join(captured) + + +# segger runs PCA(n_components=node_dim) on a genes x genes correlation matrix in +# setup_anndata. On small panels/crops fewer genes survive segger's internal count +# filters (cells_min_counts etc.) than node_dim, so the PCA raises e.g. +# "n_components=128 must be between 0 and min(n_samples, n_features)=62". The surviving +# gene count depends on segger's own filtering (not cheaply predictable here), but segger +# reports the exact ceiling in that message, so retry once at that dim. On full panels +# (>=node_dim genes survive) this never triggers. +node_dim = int(par["node_representation_dim"]) +returncode, seg_out = run_segger(node_dim) +if returncode != 0: + m = re.search(r"must be between 0 and min\(n_samples, n_features\)=(\d+)", seg_out) + if m and 0 < int(m.group(1)) < node_dim: + max_dim = int(m.group(1)) + print( + f"segger's gene PCA rejected node_representation_dim={node_dim}: only {max_dim} " + f"genes survived its count filter. Retrying at {max_dim}.", + flush=True, + ) + returncode, seg_out = run_segger(max_dim) + if returncode != 0: + raise subprocess.CalledProcessError(returncode, base_cmd) seg_pq = SEGGER_OUT_DIR / "segger_segmentation.parquet" if not seg_pq.exists(): diff --git a/src/workflows/run_benchmark/config.vsh.yaml b/src/workflows/run_benchmark/config.vsh.yaml index a61489ffb..638799d3b 100644 --- a/src/workflows/run_benchmark/config.vsh.yaml +++ b/src/workflows/run_benchmark/config.vsh.yaml @@ -68,7 +68,7 @@ argument_groups: A list of transcript assignment methods to run. type: string multiple: true - default: "basic_transcript_assignment:baysor:clustermap:pciseq:comseg:proseg:segger" + default: "basic_transcript_assignment:baysor:clustermap:pciseq:comseg:proseg:segger:fastreseg" - name: "--count_aggregation_methods" description: | A list of count aggregation methods to run. @@ -167,6 +167,7 @@ dependencies: - name: methods_transcript_assignment/comseg - name: methods_transcript_assignment/proseg - name: methods_transcript_assignment/segger + - name: methods_transcript_assignment/fastreseg - name: methods_count_aggregation/basic_count_aggregation - name: methods_qc_filter/basic_qc_filter - name: methods_calculate_cell_volume/alpha_shapes diff --git a/src/workflows/run_benchmark/main.nf b/src/workflows/run_benchmark/main.nf index f1725397b..63a1b0e6c 100644 --- a/src/workflows/run_benchmark/main.nf +++ b/src/workflows/run_benchmark/main.nf @@ -139,7 +139,8 @@ workflow run_wf { pciseq, comseg, proseg, - segger + segger, + fastreseg ] segm_ass_ch = segm_ch