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

.gitignore

@@ -0,0 +1,34 @@
+# Byte-compiled / cache
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Build / packaging artifacts
+build/
+dist/
+*.egg-info/
+.eggs/
+*.egg
+
+# Test / coverage
+.pytest_cache/
+.coverage
+.coverage.*
+htmlcov/
+
+# Type/lint caches
+.mypy_cache/
+.ruff_cache/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Virtual environments
+venv/
+env/
+.venv/
+
+# Editors / OS
+.vscode/
+.idea/
+.DS_Store

CHANGELOG.md

@@ -0,0 +1,34 @@
+# Changelog
+
+## Breaking changes
+
+- `setup_reference()` no longer runs QC filtering or HVG selection automatically. It raises an error and tells you which `scanpy` function to run first.
+- `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. MT/HLA/RPL genes now get a warning instead of automatic removal.
+- `DOT.fit()` no longer accepts `mode`, `ratios_weight`, or `max_spot_size`. Pass these to `DOT()` when you construct it instead. `max_spot_size` was renamed to `max_size`.
+- `DOT.get_weights()` returns a `pandas.DataFrame` instead of a numpy array.
+- `kmeans_define_subtypes()` now returns a `DataFrame`.
+
+## New features
+
+- `setup_reference()` now supports predefined subtype annotations. This works alongside the existing k-means discovery. New functions: `predefined_subtypes()`, `check_subtype_consistency()`, `summarize_subtypes()`, `select_de_genes()`, `validate_reference_input()`.
+- `DOT()` accepts `gene_weight`, `spot_weight`, `spatial_weight`, `ratios_weight`, `sparsity_coef`, and `cluster_weight` directly. You can also override any lambda value directly, or pass your own `weights_to_lambdas` function.
+- `cluster_weight` implements `l_c`. This is the cluster-wise cosine term from the R package. It matches R's behaviour and is off by default.
+- `setup_reference()` exposes the k-means subtype-count heuristic and per-type cell cap through `kws_kmeans`. Defaults match the previous hardcoded values. It also adds `min_frac` to filter out small clusters as noise.
+- `setup_reference()` can further sub-cluster predefined subtypes that are under-resolved relative to the k-means heuristic. Set `refine_undersized=True` and pass `kws_refine` to control it. New functions: `plan_subtype_refinement()`, `refine_predefined_subtypes()`.
+- `setup_reference()` stores a diagnostic table in `uns['subtype_plan']` when using predefined sub-clusters, with per-label cell counts, expected-vs-current cluster counts, and (when `refine_undersized=True`) how many sub-clusters were actually produced.
+- `kmeans_define_subtypes()` exposes `max_genes` to control the gene-narrowing threshold, previously hardcoded at 500.
+- `DOT.print_config()` prints the resolved optimisation configuration.
+- `get_weights()` and `get_cell_types()` accept a `level` parameter. This aggregates results at any level of the subtype hierarchy.
+- `plot_spatial_weights()` and `plot_cell_type_proportions()` accept a `DataFrame` directly.
+- Added `dotpy.ds.toy_reference()` for generating synthetic test data.
+
+## Fixes
+
+- Fixed `_safe_log2`. It returned different values than R's `safelog2` for zero and negative inputs. This affects the abundance-matching and spatial-coherence loss terms.
+
+## Other
+
+- Added a `pytest` test suite with about 349 tests and 97% line coverage.
+- Updated `run_dot_cli.py`, `example.py`, `README.md`, and `QUICKSTART.md` for the new API.

QUICKSTART.md

@@ -77,6 +77,14 @@ from dotpy import setup_reference, setup_spatial, DOT, plot_spatial_weights
 ref_adata = sc.read_h5ad('your_reference.h5ad')
 spatial_adata = sc.read_h5ad('your_spatial.h5ad')
 
+# setup_reference() expects pre-filtered (positive) raw counts
+# QC, MT/HLA/RPL removal, and HVG selection are standard scanpy steps left to you.
+sc.pp.filter_cells(ref_adata, min_counts=1)
+sc.pp.filter_genes(ref_adata, min_cells=1)
+sc.pp.highly_variable_genes(adata, batch_key='sample')
+
+adata = adata[:, adata.var['highly_variable_intersection']]
+
 # 2. Process data
 print("Processing reference...")
 ref_processed = setup_reference(
@@ -102,17 +110,15 @@ print(f"Running on: {device}")
 dot = DOT(
     spatial_processed, 
     ref_processed,
+    mode='highres',  # For Xenium, MERFISH, CosMx, etc.
+    # mode='lowres', max_size=20,  # OR for Visium, ST, etc.
     batch_size=500,  # Adjust based on GPU memory
     device=device
 )
 
 # 4. Run deconvolution
 print("Running DOT...")
-# For high-resolution data (Xenium, MERFISH, CosMx)
-dot.fit(mode='highres', iterations=100, verbose=True)
-
-# OR for low-resolution data (Visium, ST)
-# dot.fit(mode='lowres', max_spot_size=20, iterations=100, verbose=True)
+dot.fit(iterations=100, verbose=True)
 
 # 5. Get results
 weights = dot.get_weights(normalize=True)
@@ -132,8 +138,8 @@ plot_spatial_weights(
 
 # 7. Save results
 spatial_adata.obsm['dot_weights'] = weights
-for i, ct in enumerate(cell_types):
-    spatial_adata.obs[f'dot_{ct}'] = weights[:, i]
+for ct in cell_types:
+    spatial_adata.obs[f'dot_{ct}'] = weights[ct]
 
 spatial_adata.write('spatial_deconvolved.h5ad')
 print("Results saved!")
@@ -143,24 +149,37 @@ print("Results saved!")
 
 ### Reference Processing
 
+`setup_reference()` only does sub-clustering + DE gene selection -- it expects
+`adata` to already be QC-filtered (no empty cells/genes) with raw (positive) counts in
+`.X`. It validates this instead of fixing it for you: it raises on empty
+cells/genes, negative values, or non-integer counts (unless `check_counts=False`,
+e.g. for background-corrected data), and on more than `max_input_genes` genes.
+MT-/HLA-/RPL-prefixed genes are never removed automatically -- if present,
+you'll get a warning with per-prefix counts (silence with `warn_mt=False`).
+
 ```python
 ref_processed = setup_reference(
     adata,
     cell_type_key='cell_type',  # Column with cell type labels
     subcluster_size=10,         # Max subclusters per cell type (higher = more granular)
-    max_genes=5000,             # Number of genes to use (higher = more info, slower)
-    remove_mt=True,             # Remove mitochondrial genes
+    max_genes=5000,             # Genes to keep in the final DE panel
+    max_input_genes=5000,       # Genes adata is allowed to have on input (independent of max_genes)
     th_inner_logfold=0.75,      # Log-fold threshold for gene selection in subclustering
     random_state=42,            # Random seed for reproducibility
-    verbose=True,               # Print progress
-    copy=True                   # Copy adata before processing
+    warn_mt=True,               # Warn (don't remove) if MT/HLA/RPL genes are present
+    check_counts=True,          # Verify adata.X looks like raw counts
+    verbose=True                # Print progress
 )
 ```
 
 **When to adjust:**
 - Increase `subcluster_size` for more heterogeneous cell types
 - Increase `max_genes` if you have many similar cell types
 - Adjust `th_inner_logfold` to control gene selection stringency
+- Raise `max_input_genes` if your reference has more genes than the default
+  5000 and you'd rather not (re-)run your own HVG selection first
+- Set `check_counts=False` if your counts aren't strictly integer (e.g.
+  ambient-RNA/background-corrected data)
 
 ### Spatial Processing
 
@@ -171,7 +190,7 @@ spatial_processed = setup_spatial(
     th_spatial=0.84,             # Similarity threshold for spatial neighbors
     th_gene_low=0.01,            # Min expression frequency
     th_gene_high=0.99,           # Max expression frequency
-    remove_mt=True,              # Remove mitochondrial genes
+    warn_mt=True,                # Warn if MT/HLA/RPL genes are present
     radius='auto',               # Spatial neighborhood radius
     verbose=True,
     copy=True
@@ -206,11 +225,17 @@ dot = DOT(
 
 ### DOT Fitting
 
+`mode`, `ratios_weight`, and `max_size` are configured when constructing `DOT`
+(see above), not at `fit()` time:
+
 ```python
 # High-resolution (subcellular)
+dot = DOT(
+    spatial_processed, ref_processed,
+    mode='highres',               # For Xenium, MERFISH, CosMx, etc.
+    ratios_weight=0.0,            # Weight for matching reference abundances (0-1)
+)
 dot.fit(
-    mode='highres',              # For Xenium, MERFISH, CosMx, etc.
-    ratios_weight=0.0,           # Weight for matching reference abundances (0-1)
     iterations=100,              # Number of optimization iterations
     gap_threshold=0.01,          # Convergence threshold
     use_mixed_precision=False,   # Use float16 on GPU (saves memory)
@@ -221,10 +246,13 @@ dot.fit(
 )
 
 # Low-resolution (spot-based)
+dot = DOT(
+    spatial_processed, ref_processed,
+    mode='lowres',                # For Visium, ST, etc.
+    ratios_weight=0.3,            # Higher weight to match reference proportions
+    max_size=20,                  # Max cells per spot
+)
 dot.fit(
-    mode='lowres',               # For Visium, ST, etc.
-    ratios_weight=0.3,           # Higher weight to match reference proportions
-    max_spot_size=20,            # Max cells per spot
     iterations=100,
     gap_threshold=0.01,
     verbose=True
@@ -260,7 +288,7 @@ spatial_adata.var_names = spatial_adata.var_names.str.upper()
 dot = DOT(spatial_processed, ref_processed, batch_size=100)
 
 # Option 2: Use mixed precision
-dot.fit(mode='highres', use_mixed_precision=True)
+dot.fit(use_mixed_precision=True)
 
 # Option 3: Reduce genes
 ref_processed = setup_reference(
@@ -354,7 +382,6 @@ if torch.cuda.is_available():
 ```python
 # Save checkpoints during long runs
 dot.fit(
-    mode='highres',
     iterations=200,
     checkpoint_dir='./checkpoints',
     checkpoint_freq=20,  # Save every 20 iterations
@@ -363,7 +390,6 @@ dot.fit(
 
 # Resume from checkpoint
 dot.fit(
-    mode='highres',
     iterations=300,  # Continue to 300 total
     resume_from='./checkpoints/checkpoint_iter_200.pkl',
     verbose=True

README.md

@@ -47,6 +47,13 @@ from dotpy import DOT, setup_reference, setup_spatial, plot_spatial_weights
 ref_adata = sc.read_h5ad('reference.h5ad')
 spatial_adata = sc.read_h5ad('spatial.h5ad')
 
+# setup_reference() expects pre-filtered raw (positive) counts
+# QC, MT/HLA/RPL removal, and HVG selection are standard scanpy steps left to you.
+sc.pp.filter_cells(ref_adata, min_counts=1)
+sc.pp.filter_genes(ref_adata, min_cells=1)
+sc.pp.highly_variable_genes(adata, batch_key='sample')
+adata = adata[:,adata.var['highly_variable_intersection']]
+
 # Process reference and spatial data
 ref_processed = setup_reference(
     ref_adata,
@@ -67,11 +74,11 @@ spatial_processed = setup_spatial(
 dot = DOT(
     spatial_processed, 
     ref_processed,
+    mode='highres',
     batch_size=500  # Adjust for your GPU memory
 )
 
 dot.fit(
-    mode='highres',
     iterations=100,
     checkpoint_dir='./checkpoints',  # Save checkpoints
     checkpoint_freq=10,
@@ -96,7 +103,6 @@ plot_spatial_weights(
 
 ```python
 dot.fit(
-    mode='highres',
     iterations=100,
     resume_from='./checkpoints/checkpoint_iter_50.pkl',
     verbose=True
@@ -245,26 +251,28 @@ results/
 For subcellular resolution data where each spot typically contains 1 cell:
 
 ```python
-dot.fit(
+dot = DOT(
+    spatial_processed,
+    ref_processed,
     mode='highres',
     ratios_weight=0.0,
-    iterations=100,
-    verbose=True
 )
+dot.fit(iterations=100, verbose=True)
 ```
 
 ### Low-Resolution Data (Visium, ST)
 
 For spot-based technologies where spots contain multiple cells:
 
 ```python
-dot.fit(
+dot = DOT(
+    spatial_processed,
+    ref_processed,
     mode='lowres',
-    max_spot_size=20,  # Maximum cells per spot
+    max_size=20,  # Maximum cells per spot
     ratios_weight=0.3,  # Weight for matching cell type proportions
-    iterations=100,
-    verbose=True
 )
+dot.fit(iterations=100, verbose=True)
 ```
 
 ## Algorithm Overview
@@ -284,14 +292,17 @@ The optimization is performed using the Frank-Wolfe algorithm, which efficiently
 
 ```python
 # Setup reference with custom parameters
+# (ref_adata must already be QC-filtered, subset to highly variable genes, with raw counts in .X)
 ref_processed = setup_reference(
     ref_adata,
     cell_type_key='cell_type',
     subcluster_size=15,      # More subclusters per cell type
-    max_genes=10000,         # Use more genes
-    remove_mt=True,          # Remove mitochondrial genes
+    max_genes=5000,         # Genes to keep in the final DE panel
+    max_input_genes=5000,   # Genes ref_adata is allowed to have on input
     th_inner_logfold=0.75,   # Log-fold threshold for gene selection
     random_state=42,         # For reproducibility
+    warn_mt=True,            # Warn (don't remove) if MT/HLA/RPL genes are present
+    check_counts=True,       # Verify adata.X looks like raw counts
     verbose=True
 )
 
@@ -303,7 +314,7 @@ spatial_processed = setup_spatial(
     th_gene_low=0.01,        # Minimum gene expression frequency
     th_gene_high=0.99,       # Maximum gene expression frequency
     radius='auto',           # Or specify numeric value
-    remove_mt=True,          # Remove mitochondrial genes
+    warn_mt=True,            # Warn if MT/HLA/RPL genes are present
     verbose=True
 )
 
@@ -314,14 +325,14 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 dot = DOT(
     spatial_processed,
     ref_processed,
+    mode='highres',
+    ratios_weight=0.2,       # Weight for abundance matching
     batch_size=500,          # Adjust for GPU memory
     device=device            # Explicitly set device
 )
 
 # Fine-tune optimization
 dot.fit(
-    mode='highres',
-    ratios_weight=0.2,       # Weight for abundance matching
     iterations=200,          # More iterations
     gap_threshold=0.001,     # Tighter convergence
     use_mixed_precision=True,  # Use float16 on GPU
@@ -348,12 +359,12 @@ dot = DOT(
 ### Saving Results
 
 ```python
-# Add results to spatial AnnData
-spatial_adata.obsm['dot_weights'] = weights
+# Add results to spatial AnnData (weights is a DataFrame indexed by spot)
+spatial_adata.obsm['dot_weights'] = weights.loc[spatial_adata.obs.index,:]
 
 # Add individual cell type columns
-for i, ct in enumerate(cell_types):
-    spatial_adata.obs[f'dot_{ct}'] = weights[:, i]
+for ct in cell_types:
+    spatial_adata.obs[f'dot_{ct}'] = weights.loc[spatial_adata.obs.index,ct]
 
 # Save
 spatial_adata.write('spatial_with_deconvolution.h5ad')
@@ -436,11 +447,10 @@ ref_processed = setup_reference(
 )
 
 # Use smaller batch size
-dot = DOT(spatial, ref, batch_size=100)
+dot = DOT(spatial, ref, mode='highres', batch_size=100)
 
 # Enable mixed precision on GPU
 dot.fit(
-    mode='highres',
     use_mixed_precision=True,
     iterations=100
 )
@@ -450,11 +460,10 @@ dot.fit(
 
 ```python
 # Faster (fewer iterations)
-dot.fit(mode='highres', iterations=50)
+dot.fit(iterations=50)
 
 # More accurate (more iterations, tighter convergence)
 dot.fit(
-    mode='highres',
     iterations=200,
     gap_threshold=0.001
 )
@@ -492,7 +501,7 @@ Nat Commun 15, 4994 (2024). https://doi.org/10.1038/s41467-024-48868-z
 dot = DOT(spatial, ref, batch_size=100)
 
 # Solution 2: Enable mixed precision
-dot.fit(mode='highres', use_mixed_precision=True)
+dot.fit(use_mixed_precision=True)
 
 # Solution 3: Use CPU
 dot = DOT(spatial, ref, device='cpu')