Data Preparation

import scanpy as sc
import scvi
import numpy as np

# Load data
adata = sc.read_h5ad("data.h5ad")

# Check what's in adata.X
print(f"Shape: {adata.shape}")
print(f"X dtype: {adata.X.dtype}")
print(f"X contains integers: {np.allclose(adata.X.data, adata.X.data.astype(int))}")
print(f"X min: {adata.X.min()}, max: {adata.X.max()}")

# scvi-tools needs INTEGER counts
# If X appears normalized, check for raw counts

if hasattr(adata, 'raw') and adata.raw is not None:
    print("Found adata.raw")
    # Use raw counts
    adata = adata.raw.to_adata()
    
# Or check layers
if 'counts' in adata.layers:
    print("Found counts layer")
    # Will specify layer in setup_anndata

# Filter cells (standard QC)
sc.pp.filter_cells(adata, min_genes=200)
sc.pp.filter_cells(adata, max_genes=5000)

# Calculate mito percent if not present
# Handle both human (MT-) and mouse (mt-, Mt-) mitochondrial genes
adata.var['mt'] = (
    adata.var_names.str.startswith('MT-') |
    adata.var_names.str.startswith('mt-') |
    adata.var_names.str.startswith('Mt-')
)
sc.pp.calculate_qc_metrics(adata, qc_vars=['mt'], inplace=True)
adata = adata[adata.obs['pct_counts_mt'] < 20].copy()

# Filter genes
sc.pp.filter_genes(adata, min_cells=3)

print(f"After filtering: {adata.shape}")

# Store raw counts in a layer
adata.layers["counts"] = adata.X.copy()

# Now you can normalize for other purposes (HVG selection)
# But scvi will use the counts layer

# Normalize for HVG selection only
adata_hvg = adata.copy()
sc.pp.normalize_total(adata_hvg, target_sum=1e4)
sc.pp.log1p(adata_hvg)

# Select HVGs
sc.pp.highly_variable_genes(
    adata_hvg,
    n_top_genes=2000,
    flavor="seurat"  # or "cell_ranger"
)

# Transfer HVG annotation
adata.var['highly_variable'] = adata_hvg.var['highly_variable']

# Use seurat_v3 flavor with batch_key
# This selects genes variable across batches
sc.pp.highly_variable_genes(
    adata,
    n_top_genes=2000,
    flavor="seurat_v3",
    batch_key="batch",  # Your batch column
    layer="counts"      # Use raw counts
)

# Subset to highly variable genes
adata = adata[:, adata.var['highly_variable']].copy()
print(f"After HVG selection: {adata.shape}")

scvi.model.SCVI.setup_anndata(
    adata,
    layer="counts"  # Specify layer with raw counts
)

scvi.model.SCVI.setup_anndata(
    adata,
    layer="counts",
    batch_key="batch"  # Column in adata.obs
)

scvi.model.SCANVI.setup_anndata(
    adata,
    layer="counts",
    batch_key="batch",
    labels_key="cell_type"  # Column with cell type labels
)

scvi.model.SCVI.setup_anndata(
    adata,
    layer="counts",
    batch_key="batch",
    continuous_covariate_keys=["percent_mito", "n_genes"]
)

scvi.model.SCVI.setup_anndata(
    adata,
    layer="counts",
    batch_key="batch",
    categorical_covariate_keys=["donor", "technology"]
)

# Protein data in adata.obsm
# RNA in adata.X, protein in separate matrix

# Add protein data
adata.obsm["protein_expression"] = protein_counts  # numpy array

# Setup for totalVI
scvi.model.TOTALVI.setup_anndata(
    adata,
    layer="counts",
    batch_key="batch",
    protein_expression_obsm_key="protein_expression"
)

# RNA and ATAC in separate AnnData objects or MuData

import mudata as md

# If using MuData
mdata = md.read("multiome.h5mu")

scvi.model.MULTIVI.setup_mudata(
    mdata,
    rna_layer="counts",
    protein_layer=None,
    batch_key="batch",
    modalities={"rna": "rna", "accessibility": "atac"}
)

from model_utils import prepare_adata

# Prepare data with QC, HVG selection, and layer setup
adata = prepare_adata(
    adata,
    batch_key="batch",
    n_top_genes=2000,
    min_genes=200,
    max_mito_pct=20
)

# Then setup for your model
import scvi
scvi.model.SCVI.setup_anndata(adata, layer="counts", batch_key="batch")

# View registered data
print(adata.uns['_scvi_manager_uuid'])
print(adata.uns['_scvi_adata_minify_type'])

# For scVI
scvi.model.SCVI.view_anndata_setup(adata)