name: cellbender_remove_background
namespace: "correction"
scope: "public"
description: |
  Eliminating technical artifacts from high-throughput single-cell RNA sequencing data.

  This module removes counts due to ambient RNA molecules and random barcode swapping from (raw) UMI-based scRNA-seq count matrices. 
  At the moment, only the count matrices produced by the CellRanger count pipeline is supported. Support for additional tools and protocols 
  will be added in the future. A quick start tutorial can be found here.

  Fleming et al. 2022, bioRxiv.
argument_groups:
  - name: Inputs
    arguments:
      - name: "--input"
        alternatives: [-i]
        type: file
        description: "Input h5mu file. Data file on which to run tool. Data must be un-filtered: it should include empty droplets."
        direction: input
        required: true
        example: input.h5mu
      - name: "--modality"
        description: List of modalities to process.
        type: string
        default: "rna"
        required: false
  - name: Outputs
    arguments:
      - name: "--output"
        alternatives: [-o]
        type: file
        description: Full count matrix as an h5mu file, with background RNA removed. This file contains all the original droplet barcodes.
        direction: output
        required: true
        example: output.h5mu
      - name: "--layer_output"
        description: Output layer
        type: string
        default: "cellbender_corrected"
      - name: "--obs_background_fraction"
        type: string
        default: "cellbender_background_fraction"
      - name: "--obs_cell_probability"
        type: string
        default: "cellbender_cell_probability"
      - name: "--obs_cell_size"
        type: string
        default: "cellbender_cell_size"
      - name: "--obs_droplet_efficiency"
        description: |
          Name of the column in the .obs dataframe to store the droplet efficiencies in.
        type: string
        default: "cellbender_droplet_efficiency"
      - name: "--obs_latent_scale"
        type: string
        default: "cellbender_latent_scale"
      - name: "--var_ambient_expression"
        type: string
        default: "cellbender_ambient_expression"
      - name: "--obsm_gene_expression_encoding"
        type: string
        default: "cellbender_gene_expression_encoding"
    __merge__: [., /src/base/h5_compression_argument.yaml]

  - name: Arguments
    arguments:
      # custom arg
      - name: "--expected_cells_from_qc"
        type: boolean
        description: Will use the Cell Ranger QC to determine the estimated number of cells
        default: false
      - name: "--expected_cells"
        type: integer
        description: Number of cells expected in the dataset (a rough estimate within a factor of 2 is sufficient).
        example: 1000
      - name: "--total_droplets_included"
        type: integer
        description: |
          The number of droplets from the rank-ordered UMI plot
          that will have their cell probabilities inferred as an
          output. Include the droplets which might contain cells.
          Droplets beyond TOTAL_DROPLETS_INCLUDED should be
          'surely empty' droplets.
        example: 25000
      - name: "--force_cell_umi_prior"
        type: integer
        description: "Ignore CellBender's heuristic prior estimation, and use this prior for UMI counts in cells."
      - name: "--force_empty_umi_prior"
        type: integer
        description: "Ignore CellBender's heuristic prior estimation, and use this prior for UMI counts in empty droplets."
      - name: "--model"
        type: string
        choices: [naive, simple, ambient, swapping, full]
        description: |
          Which model is being used for count data.
          
          * 'naive' subtracts the estimated ambient profile.
          * 'simple' does not model either ambient RNA or random barcode swapping (for debugging purposes -- not recommended).
          * 'ambient' assumes background RNA is incorporated into droplets.
          * 'swapping' assumes background RNA comes from random barcode swapping (via PCR chimeras).
          * 'full' uses a combined ambient and swapping model.
        default: full
      - name: "--epochs"
        type: integer
        description: Number of epochs to train. 
        default: 150
      - name: "--low_count_threshold"
        type: integer
        description: |
          Droplets with UMI counts below this number are completely 
          excluded from the analysis. This can help identify the correct 
          prior for empty droplet counts in the rare case where empty 
          counts are extremely high (over 200).
        default: 5
      - name: "--z_dim"
        type: integer
        description: |
          Dimension of latent variable z.
        default: 64
      - name: "--z_layers"
        type: integer
        description: |
          Dimension of hidden layers in the encoder for z.
        multiple: true
        default: [512]
      - name: "--training_fraction"
        type: double
        description: |
          Training detail: the fraction of the data used for training.
          The rest is never seen by the inference algorithm. Speeds up learning.
        default: 0.9
      - name: "--empty_drop_training_fraction"
        type: double
        description: |
          Training detail: the fraction of the training data each epoch that 
          is drawn (randomly sampled) from surely empty droplets.
        default: 0.2
      - name: "--ignore_features"
        type: integer
        description: |
          Integer indices of features to ignore entirely. In the output
          count matrix, the counts for these features will be unchanged.
        multiple: true
      - name: "--fpr"
        type: double
        description: |
          Target 'delta' false positive rate in [0, 1). Use 0 for a cohort
          of samples which will be jointly analyzed for differential expression.
          A false positive is a true signal count that is erroneously removed.
          More background removal is accompanied by more signal removal at
          high values of FPR. You can specify multiple values, which will
          create multiple output files.
        multiple: true
        default: [0.01]
      - name: "--exclude_feature_types"
        type: string
        description: |
          Feature types to ignore during the analysis. These features will
          be left unchanged in the output file.
        multiple: true
      - name: "--projected_ambient_count_threshold"
        type: double
        description: |
          Controls how many features are included in the analysis, which
          can lead to a large speedup. If a feature is expected to have less
          than PROJECTED_AMBIENT_COUNT_THRESHOLD counts total in all cells
          (summed), then that gene is excluded, and it will be unchanged
          in the output count matrix. For example, 
          PROJECTED_AMBIENT_COUNT_THRESHOLD = 0 will include all features
          which have even a single count in any empty droplet.
        default: 0.1
      - name: "--learning_rate"
        type: double
        description: |
          Training detail: lower learning rate for inference.
          A OneCycle learning rate schedule is used, where the
          upper learning rate is ten times this value. (For this
          value, probably do not exceed 1e-3).
        default: 0.0001
      - name: "--final_elbo_fail_fraction"
        type: double
        description: |
          Training is considered to have failed if 
          (best_test_ELBO - final_test_ELBO)/(best_test_ELBO - initial_test_ELBO) > FINAL_ELBO_FAIL_FRACTION.
          Training will automatically re-run if --num-training-tries > 1.
          By default, will not fail training based on final_training_ELBO.
      - name: "--epoch_elbo_fail_fraction"
        type: double
        description: |
          Training is considered to have failed if 
          (previous_epoch_test_ELBO - current_epoch_test_ELBO)/(previous_epoch_test_ELBO - initial_train_ELBO) > EPOCH_ELBO_FAIL_FRACTION.
          Training will automatically re-run if --num-training-tries > 1.
          By default, will not fail training based on epoch_training_ELBO.
      - name: "--num_training_tries"
        type: integer
        description: |
          Number of times to attempt to train the model. At each subsequent attempt,
          the learning rate is multiplied by LEARNING_RATE_RETRY_MULT.
        default: 1
      - name: "--learning_rate_retry_mult"
        type: double
        description: |
          Learning rate is multiplied by this amount each time a new training
          attempt is made. (This parameter is only used if training fails based
          on EPOCH_ELBO_FAIL_FRACTION or FINAL_ELBO_FAIL_FRACTION and
          NUM_TRAINING_TRIES is > 1.) 
        default: 0.2
      - name: --posterior_batch_size
        type: integer
        description: |
          Training detail: size of batches when creating the posterior.
          Reduce this to avoid running out of GPU memory creating the posterior
          (will be slower).
        default: 128
      - name: --posterior_regulation
        type: string
        description: |
          Posterior regularization method. (For experts: not required for normal usage,
          see documentation). 
          
          * PRq is approximate quantile-targeting.
          * PRmu is approximate mean-targeting aggregated over genes (behavior of v0.2.0).
          * PRmu_gene is approximate mean-targeting per gene.
        choices: [PRq, PRmu, PRmu_gene]
      - name: "--alpha"
        type: double
        description: |
          Tunable parameter alpha for the PRq posterior regularization method
          (not normally used: see documentation).
      - name: "--q"
        type: double
        description: |
          Tunable parameter q for the CDF threshold estimation method (not
          normally used: see documentation).
      - name: "--estimator"
        type: string
        description: |
          Output denoised count estimation method. (For experts: not required
          for normal usage, see documentation).
        default: mckp
        choices: [map, mean, cdf, sample, mckp]
      - name: "--estimator_multiple_cpu"
        type: boolean_true
        description: |
          Including the flag --estimator-multiple-cpu will use more than one
          CPU to compute the MCKP output count estimator in parallel (does nothing
          for other estimators).
      - name: "--constant_learning_rate"
        type: boolean
        description: |
          Including the flag --constant-learning-rate will use the ClippedAdam
          optimizer instead of the OneCycleLR learning rate schedule, which is
          the default. Learning is faster with the OneCycleLR schedule.
          However, training can easily be continued from a checkpoint for more
          epochs than the initial command specified when using ClippedAdam. On
          the other hand, if using the OneCycleLR schedule with 150 epochs
          specified, it is not possible to pick up from that final checkpoint
          and continue training until 250 epochs.
      - name: "--debug"
        type: boolean_true
        description: |
          Including the flag --debug will log extra messages useful for debugging.
      - name: "--cuda"
        type: boolean_true
        description: |
          Including the flag --cuda will run the inference on a
          GPU.
resources:
  - type: python_script
    path: script.py
  - path: /src/utils/setup_logger.py
test_resources:
  - type: python_script
    path: test.py
  # normally cellbender should only be run on unfiltered data, but for the purposes of the unit test the filtered data will do.
  - path: /resources_test/pbmc_1k_protein_v3/pbmc_1k_protein_v3_filtered_feature_bc_matrix.h5mu

engines:
- type: docker
  image: nvcr.io/nvidia/cuda:11.8.0-devel-ubuntu22.04
  # run_args: ["--gpus all"]
  setup:
    - type: docker
      env:
        - PYENV_ROOT="/root/.pyenv"
        - PATH="$PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH"
        - PYTHON_VERSION=3.7.16
      run: |
        apt update && DEBIAN_FRONTEND=noninteractive apt install -y make build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget ca-certificates curl llvm libncurses5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev mecab-ipadic-utf8 git \
        && curl https://pyenv.run | bash \
        && pyenv update \
        && pyenv install $PYTHON_VERSION \
        && pyenv global $PYTHON_VERSION \
        && apt-get clean
    - type: python
      packages:
        # Use this version of mudata because it is the last one that supports python 3.7
        - lxml~=4.8.0
        - mudata~=0.2.1
        - cellbender~=0.3.0

runners:
- type: executable
- type: nextflow
  directives:
    label: [midcpu, midmem, gpu]
