Predefined subtype support, configurable optimisation weights, AnnData-based I/O, testing suite

[robinfallegger]

Sorry this became a much bigger PR than I had originally thought. The main goal was to increase tuning and make storing reference centroids easier in python. More details below and in CHANGELOG.md

Closes #4
Closes #5

Summary

This PR adds the option to use a predefined cell/cluster subtype hierarchy instead of running kmeans subclustering. The preprocessing functions setup_reference and setup_spatial now return AnnData objects, allowing you to easily store reference clusters for future use. The optimisation weights are now exposed and can be tuned by the user. It also adds tests for main components.

In detail

setup_reference changed how it handles preprocessing. It no longer runs QC filtering or HVG selection automatically. It validates your input instead and raises an error telling you which scanpy function to run first. It also exposes the k-means subtype-count heuristic and per-type cell cap as parameters instead of hardcoding them. The defaults match the previous behaviour. This would also allow to drop scanpy from the dependencies.

Predefined subtype support lets you pass your own hierarchical cell subtype annotations into setup_reference. Cluster centroids can be directly computed from these annotations, increasing preprocessing speed by entirely skipping kmeans clustering. A new select_de_genes() step supports pluggable DE gene scoring.

When running with fine-grained subtype annotations, setup_reference can also further sub-cluster predefined subtypes that are under-resolved relative to the k-means heuristic (that defines the number of expected clusters based on the number of cells). This prevents major cell types with few sub-clusters from being under-represented in the optimisation just because they were coarsely pre-labelled. In setup_reference, set refine_undersized=True and pass kws_refine to control it. This refinement runs kmeans clustering only on major cell types with too few subtypes compared to their number of cells. The new function plan_subtype_refinement() returns a dataframe with cell counts per major cell type, and expected-vs-current cluster counts.

setup_reference and setup_spatial now return AnnData objects instead of dicts. DOT() still accepts the old dict format but raises a FutureWarning.

DOT() now accepts gene_weight, spot_weight, spatial_weight, ratios_weight, sparsity_coef, and cluster_weight directly. These parameters are transformed to the lambdas used in the loss-function used in the optimisation after normalisation with respect to the dimensionality of the feature space, number of spatial units or the number of clusters as described in the original DOT publication (see the new _default_weights_to_lambdas function). Alternatively, you can set any lambda value directly, or pass your own weights_to_lambdas function. cluster_weight now implements l_c. This is the cluster-wise cosine term from the R package. It was missing from the Python implementation until now. It is off by default (same as in the R version), so existing behaviour is unchanged.

If running the optimisation with a hierarchical cell type annotation, the get_weights() method can now return the aggregated weights at any of the annotation levels. For example, if your reference dataset has three levels of annotations (e.g. level_keys = ['major_type', 'subset', 'cell_state']):

├── Fibroblasts
├── Epithelial cells
└── Immune cells
    ├── Myeloid
    └── Lymphoid
        ├── Naive B cells
        └── Memory B cells

get_weights(level="major_type") returns weights aggregated over Fibroblasts, Epithelial cells and Immune cells. get_weights(level="subset") returns weights aggregated for Myeloid and Lymphoid in the immune compartment (and for other subsets of Fibroblasts and Epithelial cells if they exist). get_weights(level="subtype") returns weights at the finest resolution, for separating Naive and Memory B cells.

Breaking changes

• setup_reference no longer runs QC filtering or HVG selection. It validates and raises instead.
• setup_reference dropped the copy parameter.
• setup_reference and setup_spatial return AnnData instead of dict.
• remove_mt was renamed to warn_mt in both functions.
• DOT.fit() no longer accepts mode, ratios_weight, or max_spot_size. Set these on DOT() instead. max_spot_size was renamed to max_size.
• DOT.get_weights() returns a DataFrame instead of a numpy array.
• kmeans_define_subtypes() returns a DataFrame.

Full list in the changelog.

Also fixed

_safe_log2 returned different values than R's safelog2 for zero and negative inputs. This affects the abundance-matching and spatial-coherence loss terms.

Test plan

• Full test suite passes (pytest tests/, about 349 tests)
• 97% line coverage
• run_dot_cli.py and example.py updated for the new API
• Verified default weight and lambda values match R for every mode, including l_c

[robinfallegger]

Just to address your comments in issues #4:

• I agree the pp tl etc is not required yet as still relatively few functions
• I basically switched to using AnnData instead of dicts. Dicts are still accepted by DOT()
• Yes I saw that many samples could be run, but I am usually running this in parallel on HPC (e.g. with snakemake/nextflow)

For changes proposed in #5:

1. I basically removed any preprocessing from setup_reference as I think this forces decisions on the user (I was always annoyed by this in R too). For example, doing batch aware HVG selection was basically impossible. The limitation on the number of HVGs (was always capped at 5000) might make it impossible to capture rarer cell types well. Therefore, I would recommend for DOT(py) to actually focus on the centroid definition and optimisation and leave the preprocessing entirely to scanpy (could be removed as a dependency).
2. I made an alternative to the kmeans clustering. In some of my projects, biologists often want to resolve their single-cell data very granularly (and sometimes I really think it is over-clustered already) and keep matching cell type labels. The new predefined_subtypes function takes cell type labels at face-value and warns you if there are over-/under-clustered cell types (i.e. too few or too many clusters with respect to the number of cells.)
   • If you have a hierarchical structure of cell types, you can aggregate the solution at any level after optimisation
   • If some clusters are too small you can run kmeans only for these clusters.
3. I split the cluster definition (k-means or predefined) from the aggregation (summarize_subtypes)
4. I exposed the optimisation weights and lambdas from the loss-function. The user can now set the mode (highres/lowres defaults), the weights, the function that transforms weights to lambdas, or set lambdas directly (or combination of these).
5. The optimisation did not contain the l_c parameter (was always 0 in R too). Can be used now if desired.

I also added a test suite to check that this actually all works. I also checked that the new setup reference function returns the same results as the original function.

Happy to discuss changes you think are necessary. At some point I wanted to add some docs with simple vignettes using examples from the original paper.

[earmingol]

Thanks @robinfallegger .

These changes look very promising overall. As soon I find some time I'll do a code review, and propose some restructuring changes to keep the essence of how DOT was built, while leaving it flexible as you did here.

[robinfallegger]

Sounds good! Let me know if anything is unclear or needs restructuring