name: totalvi_scarches
namespace: "integrate"
scope: "public"
description: |
  Performs totalVI integration by mapping the query dataset to a reference dataset or model. 
info:
  links:
    documentation: https://docs.scvi-tools.org/en/1.3.2/tutorials/notebooks/multimodal/totalVI_reference_mapping.html
authors:
  - __merge__: /src/authors/vladimir_shitov.yaml
argument_groups:
  - name: Inputs
    arguments:
      - name: "--input"
        alternatives: ["-i"]
        type: file
        description: Input h5mu file with query data to integrate with reference.
        direction: input
        required: true
      - name: "--reference"
        alternatives: ["-r"]
        type: file
        description: Input h5mu file with reference data to train the TOTALVI model.
        direction: input
        required: true
      - name: "--force_retrain"
        alternatives: ["-f"]
        type: boolean_true
        description: If true, retrain the model and save it to reference_model_path
      - name: "--query_modality"
        type: string
        default: "rna"
        required: false
      - name: "--query_proteins_modality"
        type: string
        description: Name of the modality in the input (query) h5mu file containing protein data
        required: false
      - name: "--reference_modality"
        type: string
        default: "rna"
        required: false
      - name: "--reference_proteins_modality"
        type: string
        description: Name of the modality containing proteins in the reference
        default: "prot"
        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: "--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: "--var_input"
        type: string
        required: false
        description: ".var column containing highly variable genes. By default, do not subset genes."
  - name: Outputs
    arguments:
      - name: "--output"
        alternatives: ["-o"]
        type: file
        description: Output h5mu file.
        direction: output
        required: true
      - name: "--obsm_output"
        type: string
        default: "X_integrated_totalvi"
        required: false
        description: "In which .obsm slot to store the resulting integrated embedding."
      - name: "--obsm_normalized_rna_output"
        type: string
        default: "X_totalvi_normalized_rna"
        required: false
        description: "In which .obsm slot to store the normalized RNA from TOTALVI."
      - name: "--obsm_normalized_protein_output"
        type: string
        default: "X_totalvi_normalized_protein"
        required: false
        description: "In which .obsm slot to store the normalized protein data from TOTALVI."
      - name: "--reference_model_path"
        type: file
        description: Directory with the reference model. If not exists, trained model will be saved there
        required: false
        default: "totalvi_model_reference/"
        direction: output
      - name: "--query_model_path"
        type: file
        description: Directory, where the query model will be saved
        required: false
        default: "totalvi_model_query/"
        direction: output
  - name: "Learning parameters"
    arguments:
      - name: "--max_epochs"
        type: integer
        description: "Number of passes through the dataset"
        required: false
        default: 400
      - name: "--max_query_epochs"
        type: integer
        description: "Number of passes through the dataset, when fine-tuning model for query"
        required: false
        default: 200
      - name: "--weight_decay"
        type: double
        description: "Weight decay, when fine-tuning model for query"
        required: false
        default: 0.0
resources:
  - type: python_script
    path: script.py
  - path: /src/utils/setup_logger.py
test_resources:
  - type: python_script
    path: test.py
  - path: /resources_test/pbmc_1k_protein_v3/pbmc_1k_protein_v3_mms.h5mu

engines:
- type: docker
  image: pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
  setup:
    - type: python
      __merge__: [/src/base/requirements/anndata_mudata.yaml, .]
      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: [highmem, highcpu, highdisk, gpu]
