name: scvi
namespace: "integrate"
scope: "public"
description: "Performs scvi integration as done in the human lung cell atlas https://github.com/LungCellAtlas/HLCA"
authors:
  - __merge__: /src/authors/malte_luecken.yaml
    roles: [ author ]
  - __merge__: /src/authors/dries_schaumont.yaml
    roles: [ maintainer ]
  - __merge__: /src/authors/matthias_beyens.yaml
    roles: [ contributor ]
argument_groups:
  - name: Inputs
    arguments:
      - name: "--input"
        alternatives: ["-i"]
        type: file
        description: Input h5mu file
        direction: input
        required: true
      - name: "--modality"
        description: |
          Which modality from the input MuData file to process.
        type: string
        default: "rna"
        required: false
      - name: "--input_layer"
        type: string
        required: false
        description: "Input layer to use. If None, X is used"
      - name: "--obs_batch"
        type: string
        default: "sample_id"
        required: false
        description: Column name discriminating between your batches.
      - name: "--var_gene_names"
        type: string
        required: false
        description: ".var column containing gene names. By default, use the index."
      - name: "--var_input"
        type: string
        required: false
        description: ".var column containing highly variable genes. By default, do not subset genes."
      - name: "--obs_labels"
        type: string
        required: false
        description: |
          Key in adata.obs for label information. Categories will automatically be 
          converted into integer categories and saved to adata.obs['_scvi_labels'].
          If None, assigns the same label to all the data.
      - name: "--obs_size_factor"
        type: string
        required: false
        description: |
          Key in adata.obs for size factor information. Instead of using library size as a size factor,
          the provided size factor column will be used as offset in the mean of the likelihood.
          Assumed to be on linear scale.
      - name: "--obs_categorical_covariate"
        type: string
        required: false
        multiple: true
        description: |
          Keys in adata.obs that correspond to categorical data. These covariates can be added in
          addition to the batch covariate and are also treated as nuisance factors
          (i.e., the model tries to minimize their effects on the latent space).
          Thus, these should not be used for biologically-relevant factors that you do _not_ want to correct for.
      - name: "--obs_continuous_covariate"
        type: string
        required: false
        multiple: true
        description: |
          Keys in adata.obs that correspond to continuous data. These covariates can be added in
          addition to the batch covariate and are also treated as nuisance factors
          (i.e., the model tries to minimize their effects on the latent space). Thus, these should not be
          used for biologically-relevant factors that you do _not_ want to correct for.
      - name: "--sanitize_ensembl_ids"
        type: boolean
        description: Whether to sanitize ensembl ids by removing version numbers.
        default: true
  - name: Outputs
    arguments:
      - name: "--output"
        alternatives: ["-o"]
        type: file
        description: Output h5mu file.
        direction: output
        required: true
      - name: "--output_model"
        type: file
        description: Folder where the state of the trained model will be saved to.
        required: false
        direction: output
      - name: "--obsm_output"
        type: string
        default: "X_scvi_integrated"
        required: false
        description: "In which .obsm slot to store the resulting integrated embedding."
    __merge__: [., /src/base/h5_compression_argument.yaml]

  - name: "SCVI options"
    arguments:
      - name: --n_hidden_nodes
        type: integer
        required: false
        default: 128 # this is not the default in SCVI
        description: "Number of nodes per hidden layer."
      - name: --n_dimensions_latent_space
        type: integer
        required: false
        default: 30 # this is not the default in SCVI
        description: "Dimensionality of the latent space."
      - name: "--n_hidden_layers"
        required: false
        type: integer
        default: 2 # this is not the default in SCVI
        description:  Number of hidden layers used for encoder and decoder neural-networks.
      - name: --dropout_rate
        type: double
        required: false
        default: 0.1 
        description: "Dropout rate for the neural networks."
      - name: "--dispersion"
        type: string
        required: false
        choices: ["gene", "gene-batch", "gene-label", "gene-cell"]
        default: "gene"
        description: |
          Set the behavior for the dispersion for negative binomial distributions:
          - gene: dispersion parameter of negative binomial is constant per gene across cells
          - gene-batch: dispersion can differ between different batches
          - gene-label: dispersion can differ between different labels
          - gene-cell:  dispersion can differ for every gene in every cell
      - name: "--gene_likelihood"
        type: string
        default: "nb" # This is not the default in SCVI
        choices: ["nb", "zinb", "poisson"]
        description: |
          Model used to generate the expression data from a count-based likelihood distribution.
          - nb: Negative binomial distribution
          - zinb: Zero-inflated negative binomial distribution
          - poisson: Poisson distribution
  - name: "Variational auto-encoder model options"
    arguments:
      - name: "--use_layer_normalization"
        description: |
          Neural networks for which to enable layer normalization. 
        type: string
        choices: ["encoder", "decoder", "none", "both"]
        default: "both" # This is not the default in SCVI
      - name: "--use_batch_normalization"
        description: |
          Neural networks for which to enable batch normalization. 
        type: string
        choices: ["encoder", "decoder", "none", "both"]
        default: "none" # This is not the default in SCVI
      - name: "--encode_covariates"
        type: boolean_false # This is not the default in SCVI
        description: "Whether to concatenate covariates to expression in encoder"
      - name: "--deeply_inject_covariates"
        type: boolean_true
        description: |
          Whether to concatenate covariates into output of hidden layers in encoder/decoder. 
          This option only applies when n_layers > 1. The covariates are concatenated to
          the input of subsequent hidden layers.
      - name: "--use_observed_lib_size"
        type: boolean_true # This is not the default in SCVI
        description: |
          Use observed library size for RNA as scaling factor in mean of conditional distribution.
  - name: "Early stopping arguments"
    arguments:
      - name: "--early_stopping"
        required: false
        type: boolean
        description: "Whether to perform early stopping with respect to the validation set."
      - name: "--early_stopping_monitor"
        choices: ["elbo_validation", "reconstruction_loss_validation", "kl_local_validation"]
        default: "elbo_validation"
        type: string
        description: "Metric logged during validation set epoch."
      - name: "--early_stopping_patience"
        type: integer
        min: 1
        default: 45
        description: "Number of validation epochs with no improvement after which training will be stopped."
      - name: "--early_stopping_min_delta"
        min: 0
        type: double
        default: 0.0
        description: "Minimum change in the monitored quantity to qualify as an improvement, 
                      i.e. an absolute change of less than min_delta, will count as no improvement."
  - name: "Learning parameters"
    arguments:
      - name: "--max_epochs"
        type: integer
        description: "Number of passes through the dataset, defaults to (20000 / number of cells) * 400 or 400; whichever is smallest."
        required: false
      - name: "--reduce_lr_on_plateau"
        description: "Whether to monitor validation loss and reduce learning rate when validation set `lr_scheduler_metric` plateaus."
        type: boolean
        default: True
      - name: "--lr_factor"
        description: "Factor to reduce learning rate."
        type: double
        default: 0.6
        min: 0
      - name: "--lr_patience"
        description: "Number of epochs with no improvement after which learning rate will be reduced."
        type: double
        default: 30
        min: 0
  - name: "Data validition"
    arguments:
      - name: "--n_obs_min_count"
        type: integer
        description: "Minimum number of cells threshold ensuring that every obs_batch category has sufficient observations (cells) for model training."
        required: false
        default: 0
      - name: "--n_var_min_count"
        type: integer
        description: "Minimum number of genes threshold ensuring that every var_input filter has sufficient observations (genes) for model training."
        required: false
        default: 0
resources:
  - type: python_script
    path: script.py
  - path: /src/utils/subset_vars.py
  - path: /src/utils/compress_h5mu.py
  - path: /src/utils/set_var_index.py
test_resources:
  - type: python_script
    path: test.py
  - path: /resources_test/pbmc_1k_protein_v3/pbmc_1k_protein_v3_mms.h5mu
  - path: /resources_test/annotation_test_data/TS_Blood_filtered.h5mu

engines:
- type: docker
  image: pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
  setup:
    - type: python
      __merge__: [/src/base/requirements/anndata_mudata.yaml, /src/base/requirements/scanpy.yaml, .]
    - type: python
      packages:
        - scvi-tools[cuda]~=1.4.1
  test_setup:
    - type: python
      __merge__: [ /src/base/requirements/viashpy.yaml, .]

runners:
- type: executable
  # docker_run_args: ["--gpus all"]
- type: nextflow
  directives:
    label: [midcpu, midmem, gpu, highdisk]
