scGPT integration workflow

CHANGELOG.md

@@ -36,6 +36,8 @@
 
 * `scgpt/binning` component: Added a scGPT pre-processing binning component (PR #765).
 
+* `workflows/integration/scgpt_leiden` workflow integrate with scGPT followed by Leiden clustering (PR #794).
+
 * `transform/clr` component: Added the option to set the `axis` along which to apply CLR. Possible to override
   on workflow level as well (PR #767).
 

src/scgpt/embedding/config.vsh.yaml

@@ -25,27 +25,6 @@ functionality:
           type: string
           default: "rna"
           required: false
-        - name: "--model"
-          type: file
-          direction: input
-          required: true
-          example: best_model.pt
-          description: |
-            Path to scGPT model file.
-        - name: "--model_vocab"
-          type: file
-          direction: input
-          required: true
-          example: vocab.json
-          description: |
-            Path to scGPT model vocabulary file.
-        - name: "--model_config"
-          type: file
-          direction: input
-          required: true
-          example: args.json
-          description: |
-            Path to scGPT model config file.
         - name: "--obsm_gene_tokens"
           type: string
           default: "gene_id_tokens"
@@ -70,6 +49,35 @@ functionality:
           type: string
           description: |
             The name of the adata.obs column containing the batch labels. Must be provided when 'dsbn' is set to True.
+    - name: Model
+      arguments:
+        - name: "--model"
+          type: file
+          direction: input
+          required: true
+          example: best_model.pt
+          description: |
+            Path to scGPT model file.
+        - name: "--model_vocab"
+          type: file
+          direction: input
+          required: true
+          example: vocab.json
+          description: |
+            Path to scGPT model vocabulary file.
+        - name: "--model_config"
+          type: file
+          direction: input
+          required: true
+          example: args.json
+          description: |
+            Path to scGPT model config file.
+        - name: "--finetuned_checkpoints_key"
+          type: string
+          required: false
+          example: model_state_dict
+          description: |
+            Key in the model file containing the pretrained checkpoints. Only relevant for fine-tuned models.
 
     - name: Outputs
       arguments:

src/scgpt/embedding/script.py

@@ -145,9 +145,22 @@ def setup_logger():
     pre_norm=False  #TODO: Parametrize when GPU-based machine types are supported
     )
 
+
+logger.info("Loading model")
+model_file = par["model"]
+model_dict = torch.load(model_file, map_location=device)
+
+# Ensure the provided model has the correct architecture
+if par["finetuned_checkpoints_key"]:
+    if par["finetuned_checkpoints_key"] not in model_dict.keys():
+        finetuned_checkpoints_key = par["finetuned_checkpoints_key"]
+        raise KeyError(f"The key '{finetuned_checkpoints_key}' provided for '--finetuned_checkpoints_key' could not be found in the provided --model file. The finetuned model file for cell type annotation requires valid keys for the checkpoints and the label mapper.")
+    model_dict = model_dict[par["finetuned_checkpoints_key"]]
+
+# Load model
 load_pretrained(
     model,
-    torch.load(model_file, map_location=device),
+    model_dict,
     verbose=False
     )
 

src/scgpt/embedding/test.py

@@ -1,5 +1,6 @@
 import pytest
 import subprocess
+import torch
 import re
 import sys
 import mudata as mu
@@ -20,10 +21,23 @@
 
 input = f"{meta['resources_dir']}/Kim2020_Lung_subset.h5mu"
 model_file = f"{meta['resources_dir']}/source/best_model.pt"
+ft_model = f'{meta["resources_dir"]}/ft_best_model.pt'
 vocab_file = f"{meta['resources_dir']}/source/vocab.json"
 model_config_file = f"{meta['resources_dir']}/source/args.json"
 input_file = mu.read(input)
 
+def scgpt_to_ft_scgpt(scgpt_path, ft_scgpt_path, state_dict_key, mapper_key):
+    f_model_dict = torch.load(scgpt_path, map_location="cpu")
+    model_dict = {}
+    model_dict[state_dict_key] = f_model_dict
+    model_dict[mapper_key] = {k: str(k) for k in range(15)}
+    torch.save(model_dict, ft_scgpt_path)
+
+# Convert foundation model into fine-tuned model architecture:
+# To be able to do a cell type label mapping, the model architecture needs to contain a class to label mapper dictionary
+scgpt_to_ft_scgpt(model_file, ft_model, "model_state_dict", "id_to_class")
+
+
 ## START TEMPORARY WORKAROUND DATA PREPROCESSING
 #TODO: Remove this workaround once full scGPT preprocessing workflow is implemented
 # Read in data
@@ -259,5 +273,63 @@ def test_integration_embedding_non_existing_keys(run_component, tmp_path):
         err.value.stdout.decode('utf-8'))
 
 
+def test_finetuned_model(run_component, tmp_path):
+    output_embedding_file = tmp_path / "Kim2020_Lung_subset_embedded.h5mu"
+
+    run_component([
+        "--input", tokenized_data_path,
+        "--modality", "rna",
+        "--model", ft_model,
+        "--model_vocab", vocab_file,
+        "--model_config", model_config_file,
+        "--dsbn", "True",
+        "--obs_batch_label", "sample",
+        "--obsm_gene_tokens", "gene_id_tokens",
+        "--obsm_tokenized_values", "values_tokenized",
+        "--obsm_padding_mask", "padding_mask",
+        "--finetuned_checkpoints_key", "model_state_dict",
+        "--output", output_embedding_file
+    ])
+
+    # Read output file
+    output_mdata = mu.read(output_embedding_file)
+    output_adata = output_mdata.mod["rna"]
+
+    # check that embedding obs is present
+    assert 'X_scGPT' in output_adata.obsm.keys(), "X_scGPT is not present in anndata obsm keys"
+
+    # check embedding size
+    assert output_adata.obsm["X_scGPT"].shape[1] == 512, "Embedding size does not equal 512"
+
+    # check embedding value range
+    assert not all(np.isnan(output_adata.obsm["X_scGPT"][0])), "Embedding values are nan"
+    assert all([all(i > -1) & all(i < 1) for i in output_adata.obsm["X_scGPT"]]), "Range of embedding values is outside of [-1, 1]"
+
+
+def test_finetuned_model_architecture(run_component, tmp_path):
+    output_embedding_file = tmp_path / "Kim2020_Lung_subset_embedded.h5mu"
+
+    args = [
+        "--input", tokenized_data_path,
+        "--modality", "rna",
+        "--model", ft_model,
+        "--model_vocab", vocab_file,
+        "--model_config", model_config_file,
+        "--dsbn", "True",
+        "--obs_batch_label", "sample",
+        "--obsm_gene_tokens", "gene_id_tokens",
+        "--obsm_tokenized_values", "values_tokenized",
+        "--obsm_padding_mask", "padding_mask",
+        "--finetuned_checkpoints_key", "dummy_checkpoints_key",
+        "--output", output_embedding_file
+    ]
+
+    with pytest.raises(subprocess.CalledProcessError) as err:
+        run_component(args)
+    assert re.search(
+        r'KeyError: "The key \'dummy_checkpoints_key\' provided for \'--finetuned_checkpoints_key\' could not be found in the provided --model file. The finetuned model file for cell type annotation requires valid keys for the checkpoints and the label mapper."',
+        err.value.stdout.decode('utf-8'))
+
+
 if __name__ == '__main__':
     sys.exit(pytest.main([__file__]))

src/workflows/integration/scgpt_leiden/config.vsh.yaml

@@ -0,0 +1,180 @@
+ffunctionality:
+  name: "scgpt_leiden"
+  namespace: "workflows/integration"
+  description: "Run scGPT integration (cell embedding generation) followed by neighbour calculations, leiden clustering and run umap on the result."
+  authors:
+    - __merge__: /src/authors/dorien_roosen.yaml
+      roles: [ author, maintainer ]
+  argument_groups:
+    - name: "Inputs"
+      arguments:
+        - name: "--id"
+          required: true
+          type: string
+          description: ID of the sample.
+          example: foo
+        - name: "--input"
+          type: file
+          required: true
+          description: Path to the input file.
+          example: input.h5mu
+        - name: "--modality"
+          type: string
+          default: "rna"
+          required: false
+        - name: "--input_layer"
+          type: string
+          required: False
+          description: |
+            Mudata layer (key from layers) to use as input data for hvg subsetting and binning; if not specified, X is used.
+        - name: "--var_gene_names"
+          type: string
+          required: false
+          description: |
+            The name of the adata var column containing gene names; when no gene_name_layer is provided, the var index will be used.
+        - name: "--obs_batch_label"
+          type: string
+          description: |
+            The name of the adata obs column containing the batch labels.
+    - name: Model
+      arguments:
+        - name: "--model"
+          type: file
+          required: true
+          example: resources_test/scgpt/best_model.pt
+          description: |
+            Path to scGPT model file.
+        - name: "--model_vocab"
+          type: file
+          direction: input
+          required: true
+          example: resources_test/scgpt/vocab.json
+          description: |
+            Path to scGPT model vocabulary file.
+        - name: "--model_config"
+          type: file
+          direction: input
+          required: true
+          example: args.json
+          description: |
+            Path to scGPT model config file.
+        - name: "--finetuned_checkpoints_key"
+          type: string
+          required: false
+          example: model_state_dict
+          description: |
+            Key in the model file containing the pretrained checkpoints. Only relevant for fine-tuned models.
+    - name: "Outputs"
+      arguments:
+        - name: "--output"
+          type: file
+          required: true
+          direction: output
+          description: Output file path
+          example: output.h5mu
+        - name: "--obsm_integrated"
+          type: string
+          default: "X_scgpt"
+          required: false
+          description: "In which .obsm slot to store the resulting integrated embedding."
+        - name: "--output_compression"
+          type: string
+          example: "gzip"
+          required: false
+          choices: ["gzip", "lzf"]
+          description: |
+            The compression algorithm to use for the output h5mu file.
+
+    - name: "Padding arguments"
+      arguments:
+        - name: "--pad_token"
+          type: string
+          default: "<pad>"
+          required: false
+          description: |
+            Token used for padding.
+        - name: "--pad_value"
+          type: integer
+          default: -2
+          required: false
+          description: |
+            The value of the padding token.
+    
+    - name: "HVG subset arguments"
+      arguments:
+        - name: "--n_hvg"
+          type: integer
+          default: 1200
+          description: |
+            Number of highly variable genes to subset for.
+    - name: "Tokenization arguments"
+      arguments:
+        - name: "--max_seq_len"
+          type: integer
+          required: false
+          description: |
+            The maximum sequence length of the tokenized data.
+    - name: "Embedding arguments"
+      arguments:
+        - name: --dsbn
+          type: boolean
+          default: true
+          description: |
+            Apply domain-specific batch normalization
+        - name: "--batch_size"
+          type: integer
+          default: 64
+          description: |
+            The batch size to be used for embedding inference.
+
+    - name: "Binning arguments"
+      arguments:
+        - name: "--n_input_bins"
+          type: integer
+          default: 51
+          required: False
+          min: 1
+          description: |
+            The number of bins to discretize the data into; When no value is provided, data won't be binned.
+        - name: "--seed"
+          type: integer
+          required: false
+          description: |
+            Seed for random number generation used for binning. If not set, no seed is used.
+  
+    - name: "Clustering arguments"
+      arguments:
+        - name: "--leiden_resolution"
+          type: double
+          description: Control the coarseness of the clustering. Higher values lead to more clusters.
+          default: [1]
+          multiple: true
+
+  resources:
+    - type: nextflow_script
+      path: main.nf
+      entrypoint: run_wf
+
+  dependencies:
+    - name: scgpt/cross_check_genes
+    - name: scgpt/binning
+    - name: feature_annotation/highly_variable_features_scanpy
+    - name: filter/do_filter
+    - name: scgpt/pad_tokenize
+    - name: scgpt/embedding
+    - name: dimred/umap
+    - name: neighbors/find_neighbors
+    - name: cluster/leiden
+    - name: metadata/move_obsm_to_obs
+
+  test_resources:
+    - type: nextflow_script
+      path: test.nf
+      entrypoint: test_wf
+    - type: nextflow_script
+      path: test.nf
+      entrypoint: test_wf2
+    - path: /resources_test/scgpt
+
+platforms:
+  - type: nextflow