name: densmap
namespace: "dimred"
scope: "public"
description: |
  A modification of UMAP that adds an extra cost term in order to preserve information 
  about the relative local density of the data. It is performed on the same inputs as UMAP.
authors:
  - __merge__: /src/authors/jakub_majercik.yaml
    roles: [ maintainer ]
argument_groups:
  - name: Inputs
    arguments:
      - name: "--input"
        type: file
        description: Input h5mu file
        direction: input
        required: true
        example: input.h5mu

      - name: "--modality"
        description: |
          Which modality from the input MuData file to process.
        type: string
        default: "rna"
        required: false

      - name: "--uns_neighbors"
        type: string
        default: "neighbors"
        description: The `.uns` neighbors slot as output by the `find_neighbors` component.

      - name: "--obsm_pca"
        type: string
        description: |
          The slot in `.obsm` where the PCA results are stored.
        required: true
    
  - name: Outputs
    arguments:
      - name: "--output"
        alternatives: ["-o"]
        type: file
        description: Output h5mu file.
        direction: output
        required: true
        example: output.h5mu

      - name: "--obsm_output"
        type: string
        description: The .obsm key to use for storing the densMAP results..
        default: "X_densmap"
    __merge__: [., /src/base/h5_compression_argument.yaml]

    
  - name: Arguments UMAP
    arguments:
      - name: "--min_dist"
        type: double
        description: |
          The effective minimum distance between embedded points. Smaller values will result 
          in a more clustered/clumped embedding where nearby points on the manifold are drawn 
          closer together, while larger values will result on a more even dispersal of points. 
          The value should be set relative to the spread value, which determines the scale at 
          which embedded points will be spread out. 
        default: 0.5
        min: 0.0
        max: 10.0

      - name: "--spread"
        type: double
        description: | 
          The effective scale of embedded points. In combination with `min_dist` this 
          determines how clustered/clumped the embedded points are.
        default: 1.0
        min: 0.0
        max: 10.0

      - name: "--num_components"
        type: integer
        description: The number of dimensions of the embedding.
        default: 2
        min: 1

      - name: "--max_iter"
        type: integer
        description: | 
          The number of iterations (epochs) of the optimization. Called `n_epochs` 
          in the original UMAP. Default is set to 500 if 
          neighbors['connectivities'].shape[0] <= 10000, else 200.
        default: 0
        min: 0
        max: 1000

      - name: "--alpha"
        type: double
        description: The initial learning rate for the embedding optimization.
        default: 1.0

      - name: "--gamma"
        type: double
        description: | 
          Weighting applied to negative samples in low dimensional embedding optimization. 
          Values higher than one will result in greater weight being given to negative samples.
        default: 1.0

      - name: "--negative_sample_rate"
        type: integer
        description: |
          The number of negative samples to select per positive sample
          in the optimization process. Increasing this value will result
          in greater repulsive force being applied, greater optimization
          cost, but slightly more accuracy.
        default: 5
          
      - name: "--init_pos"
        type: string
        description: |
          How to initialize the low dimensional embedding. Called `init` in the original UMAP. Options are:
            
          * Any key from `.obsm`
          * `'paga'`: positions from `paga()`
          * `'spectral'`: use a spectral embedding of the graph
          * `'random'`: assign initial embedding positions at random.

        default: spectral
        choices: [ 'paga', 'spectral', 'random' ]

  - name: Arguments densMAP
    arguments:
      - name: "--lambda"
        type: double
        description: |
          Controls the regularization weight of the density correlation term in densMAP. 
          Higher values prioritize density preservation over the UMAP objective, and vice versa 
          for values closer to zero. Setting this parameter to zero is equivalent to running 
          the original UMAP algorithm.
        default: 2.0
        min: 0.01
        max: 10.0

      - name: "--fraction"
        type: double
        description: | 
          Controls the fraction of epochs (between 0 and 1) where the density-augmented objective 
          is used in densMAP. The first (1 - dens_frac) fraction of epochs optimize the original 
          UMAP objective before introducing the density correlation term.
        default: 0.3

      - name: "--var_shift"
        type: double
        description: | 
          A small constant added to the variance of local radii in the embedding when calculating 
          the density correlation objective to prevent numerical instability from dividing by a 
          small number.
        default: 0.1


resources:
  - type: python_script
    path: script.py
  - path: /src/utils/setup_logger.py
  - path: /src/utils/compress_h5mu.py
test_resources:
  - type: python_script
    path: test.py
  - path: /resources_test/pbmc_1k_protein_v3
engines:
  - type: docker
    image: python:3.12-slim
    setup:
      - type: apt
        packages:
          - procps
      - type: python
        packages:
          - umap-learn
        __merge__: [/src/base/requirements/anndata_mudata.yaml, .]
    __merge__: [/src/base/requirements/python_test_setup.yaml, .]
runners:
  - type: executable
  - type: nextflow
    directives:
      label: [highcpu, midmem, middisk]
