A small, reproducible demo that uses Kernel Spectral Joint Embeddings (KSJE) — Ding & Ma, "Joint embedding of high-dimensional datasets via duo-landmark integral operators", JASA 2025 (code) — to remove a batch effect between two single-cell RNA-seq datasets, and compares it against a no-correction baseline.
New to single-cell? Read Background first.
A single-cell dataset is a big table: one row per cell, one column per gene, each entry a count of how many mRNA molecules of that gene were captured in that cell (tens of thousands of genes, thousands of cells, mostly zeros).
A batch effect is the problem we are fighting. When two datasets are produced in separate experiments — different days, machines, chemistry versions, donors, or labs ("batches") — systematic technical differences get added on top of the real biology. Pool the two naively and the cells often group by batch instead of by cell type: a T cell from batch A looks more like a B cell from batch A than like a T cell from batch B. That buries the biology you care about (which cell types exist, how they differ) under an artifact. Step 3 below shows this.
An embedding is a small set of coordinates (here, ~30 numbers, or 2 for a plot) per cell, arranged so that distance ≈ dissimilarity: similar cells land near each other.
A joint embedding puts cells from both datasets into one shared coordinate system, built to capture the structure common to both. Done well, the same cell type from either batch lands in the same region — the batch artifact is suppressed while genuine biological differences are kept. That is exactly the goal of dataset integration.
How KSJE does it (intuition). Instead of embedding each dataset on its own, KSJE measures every cell in dataset 1 against every cell in dataset 2 (a Gaussian kernel on their cross-distances — each dataset's cells act as landmarks for the other, the "duo-landmark" idea). The singular value decomposition of that cross-kernel yields directions of shared variation, and both datasets are embedded along them. Because the embedding is defined only through cross-batch comparisons, batch-specific quirks that don't help match the two datasets are down-weighted.
| Step | Script | Output |
|---|---|---|
| 2. Download + preprocess two PBMC datasets with a known batch effect | scripts/01_preprocess.py |
data/combined.h5ad, data/ksje_{X,Y}.csv |
| 3. Baseline: concat → PCA → UMAP, coloured by batch | scripts/02_baseline.py |
figures/01_baseline_umap.png |
| 4. KSJE genuine joint embedding | scripts/03_run_ksje.R |
data/ksje_embed_b{1,2}.csv |
| 4+5. KSJE UMAP + batch-mixing comparison | scripts/04_compare.py |
figures/02_baseline_vs_ksje_umap.png, results/comparison.txt |
| 6. Biology check: label cells by type, recolour baseline vs KSJE UMAPs | scripts/05_celltypes.py |
figures/03_celltypes.png |
Biology check (step 6). The batch-coloured UMAP shows the batches blend, but
not whether real cell types survived. Step 6 labels each cell by canonical PBMC
markers (CD3D→T, MS4A1→B, LYZ→monocyte, …), scored per cell with scanpy's
score_genes, then recolours both UMAPs by type. Under KSJE each cell type forms
one cluster containing both batches — batches mixed, biology kept. Note the
markers are reloaded from the full gene set: classic ones like CD3D are too
uniform to survive the highly-variable-gene filter, so they are absent from
combined.h5ad. Labels are automated (not expert-curated), but the recovered
type proportions (T cells the majority, etc.) match known blood composition.
Datasets. Two public 10x Genomics PBMC datasets — pbmc3k (v1 chemistry, hg19) and pbmc8k (v2 chemistry, GRCh38). Different donors, chemistries and reference genomes give a clear batch effect over the same underlying immune cell types. Each batch is subsampled to 1,200 cells.
Preprocessing (standard scanpy): basic QC filtering, CP10k + log1p normalization, 1,000 highly-variable genes shared across batches, then per-gene scaling. The scaled HVG matrix is the common feature space both PCA and KSJE use.
Metric. scripts/metrics.py implements two neighbourhood batch-mixing scores
(no R/kBET dependency): a kBET-style acceptance rate (chi-square test that a
cell's k-NN batch mix matches the global mix) and a k-NN batch entropy. For
both, higher = better mixed = batch effect better removed. They are computed
in the embedding space (baseline 30 PCs vs. KSJE 30 dims), so the comparison is
apples-to-apples.
KSJE ships as a single R function, DL.embed, in ksje_src/main_function.R
(fetched verbatim from the upstream repo). We run that genuine function for
the embedding (step 4), and use Python/scanpy for data handling, plotting and
metrics. Everything lives in one conda environment (environment.yml) that
provides both Python and R, so there is nothing to install system-wide.
Two transparent adaptations are made in scripts/03_run_ksje.R, documented
in-file: the two library() calls DL.embed never uses are stripped, and its
internal data.Y1/data.Y2 variables are bound to the inputs. The numerical
body of DL.embed is left untouched.
Prerequisites: macOS/Linux, internet. No admin rights needed — the environment is
created in-project with a bundled micromamba.
# 1. one-time environment (Python + R, ~2 GB) via the bundled micromamba
export MAMBA_ROOT_PREFIX="$PWD/.micromamba"
export CURL_CA_BUNDLE="$PWD/.micromamba_cacert.pem" # CA bundle for the static binary
./bin/micromamba create -y -f environment.yml -r "$MAMBA_ROOT_PREFIX"
# 2. run the whole pipeline
bash run_all.shThen look at figures/02_baseline_vs_ksje_umap.png and results/comparison.txt
for the batch-mixing result, and figures/03_celltypes.png for the biology check.
ksje-integration/
├── environment.yml # one env: python (scanpy ...) + R (RSpectra, Rfast)
├── run_all.sh # end-to-end driver
├── ksje_src/main_function.R # genuine KSJE DL.embed (upstream, unmodified)
├── scripts/
│ ├── config.py # dataset URLs + all parameters
│ ├── 01_preprocess.py # download + QC + normalize + export
│ ├── 02_baseline.py # naive concat + PCA + UMAP
│ ├── 03_run_ksje.R # run genuine DL.embed
│ ├── 04_compare.py # joint UMAP + metric comparison
│ ├── 05_celltypes.py # biology check: cell-type colouring
│ └── metrics.py # kBET-style acceptance + kNN batch entropy
├── figures/ # generated plots
└── results/ # metrics + comparison table
Ding, X. & Ma, R. (2025). Joint embedding of high-dimensional datasets via duo-landmark integral operators. Journal of the American Statistical Association. Code: https://github.com/rongstat/KSJE.