{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bring your own data to NicheDECODE\n",
    "Train a new model using an annotated spatial reference and predict niche proportions\n",
    "for unlabelled bulk RNA-seq samples. Start with the [input and parameter guide](../docs/bring-your-own-data.html).\n",
    "This template requires raw RNA counts, a reference-derived SpatialDE table and GO embeddings.\n",
    "It supports single-cell spatial references and cell-number-weighted spot references.\n",
    "Outputs describe relative composition of the supplied niches, not cell coordinates.\n",
    "Run from this repository using the NicheDECODE environment and a CUDA GPU.\n",
    "External tools are run separately; their installation is not included in environment.yml."
   ],
   "id": "byod-00"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import sys, json\n",
    "ROOT = next(p for p in [Path.cwd(), *Path.cwd().parents] if (p / 'model/nicheDeconv.py').exists())\n",
    "sys.path.insert(0, str(ROOT))\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import anndata as ad\n",
    "import torch\n",
    "from sklearn.model_selection import train_test_split\n",
    "from model.nicheDeconv import nicheDeconv\n",
    "from model.utils import auto_log_then_minmax, set_global_seed\n",
    "from tutorials.byod_utils import (\n",
    "    SPATIAL_FEATURES, validate_counts, attach_cell_counts, make_pseudotissues,\n",
    "    pack_model_input, functional_graph, predict_unlabelled, dense,\n",
    ")\n",
    "set_global_seed(2021)"
   ],
   "id": "byod-01"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 Set paths and sampling parameters\n",
    "Replace the paths below. Use a new RUN_NAME for each configuration: the existing\n",
    "training API automatically loads save_models/RUN_NAME/best_model.pt if it exists.\n",
    "The guard below prevents accidental reuse. One reference file must contain one section.\n",
    "WIDTH_FRACTION is relative to its x-coordinate span; inspect the resulting windows.\n",
    "These are starting settings, not optimized settings for every tissue."
   ],
   "id": "byod-02"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "INPUT = ROOT / 'data/custom'\n",
    "REFERENCE = INPUT / 'reference.h5ad'\n",
    "BULK = INPUT / 'bulk_counts.csv'  # gene rows, sample columns, first column = gene ID\n",
    "SPATIAL = INPUT / 'spatialde.csv'  # column g + the ten SPATIAL_FEATURES\n",
    "EMBEDDINGS = INPUT / 'go_embeddings.csv'  # header; first column = matching gene ID\n",
    "NICHE_KEY = 'niche'\n",
    "COUNTS_LAYER = 'counts'  # set None only when reference.X holds raw counts\n",
    "SPOT_RESOLUTION = False\n",
    "CELL_COUNTS = INPUT / 'spot_cell_counts.csv'  # spot_id,cell_count; required for spots\n",
    "RUN_NAME = 'custom_run_001'\n",
    "WIDTH_FRACTION = 0.10\n",
    "OVERLAP = 0.5\n",
    "ANGLES = [0, 45, 90]\n",
    "MIN_OBSERVATIONS = 10\n",
    "MIN_NICHES = 1\n",
    "RANDOM_SAMPLES = 1000\n",
    "RANDOM_OBSERVATIONS = None  # None: mean retained window observation count\n",
    "GRAPH_K = 5\n",
    "EPOCHS, BATCH_SIZE, LEARNING_RATE, PATIENCE = 200, 32, 1e-4, 10\n",
    "OUT = ROOT / 'res' / RUN_NAME\n",
    "CHECKPOINT = ROOT / 'save_models' / RUN_NAME / 'best_model.pt'\n",
    "if CHECKPOINT.exists() or OUT.exists():\n",
    "    raise FileExistsError('Choose a new RUN_NAME to protect an existing run.')"
   ],
   "id": "byod-03"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 Load the reference and attach niche annotations\n",
    "reference.obs_names are cell/spot IDs; reference.var_names are gene IDs;\n",
    "reference.obsm['spatial'] contains x,y coordinates. The string/categorical column\n",
    "reference.obs[NICHE_KEY] assigns one niche to each observation.\n",
    "If annotations are in a separate CSV, join by ID as illustrated below. NicheDECODE\n",
    "does not discover niches: obtain spatial domains using a suitable spatial annotation\n",
    "method or expert histology, inspect their spatial distribution and markers, and\n",
    "use consistent biological labels. Cell-type labels alone are not niche annotations."
   ],
   "id": "byod-04"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "reference = ad.read_h5ad(REFERENCE)\n",
    "if COUNTS_LAYER is not None:\n",
    "    reference.X = reference.layers[COUNTS_LAYER].copy()\n",
    "# Optional import; uncomment if labels are not already present:\n",
    "# annotations = pd.read_csv(INPUT / 'niche_annotations.csv', index_col=0)\n",
    "# assert annotations.index.is_unique\n",
    "# reference.obs[NICHE_KEY] = annotations['niche'].reindex(reference.obs_names)\n",
    "validate_counts(reference, spatial=True, niche_key=NICHE_KEY)\n",
    "print(reference)\n",
    "print(reference.obs[NICHE_KEY].value_counts())\n",
    "# If the file contains multiple sections, subset ONE before proceeding:\n",
    "# reference = reference[reference.obs['section'] == 'section_A'].copy()"
   ],
   "id": "byod-05"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 Spot-resolution branch with CytoSPACE\n",
    "For single-cell spatial data, skip the external step and leave SPOT_RESOLUTION=False.\n",
    "For spots, run [CytoSPACE](https://github.com/digitalcytometry/cytospace) with a compatible,\n",
    "cell-type-annotated scRNA reference and the spatial raw counts/coordinates (see guide).\n",
    "Convert assigned_locations.csv to the two-column file expected below:\n",
    "\n",
    "```python\n",
    "assignments = pd.read_csv('cytospace_output/assigned_locations.csv')\n",
    "assert assignments['UniqueCID'].is_unique\n",
    "counts = assignments.groupby('SpotID').size().rename('cell_count')\n",
    "counts.rename_axis('spot_id').to_csv(INPUT / 'spot_cell_counts.csv')\n",
    "```\n",
    "\n",
    "Join using exact spot IDs, including any section prefixes. Missing or zero-assignment\n",
    "spots require explicit QC; they are not filled with one cell. The helper sums observed\n",
    "spot expression once per selected spot, and uses estimated cell counts only for labels:\n",
    "`p(niche c) = sum(cell_count of selected spots labelled c) / sum(cell_count of all selected spots)`.\n",
    "Do not multiply spot expression by cell counts or substitute CytoSPACE cell-type fractions\n",
    "for niche labels."
   ],
   "id": "byod-06"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "if SPOT_RESOLUTION:\n",
    "    cell_counts = pd.read_csv(CELL_COUNTS, index_col=0)['cell_count']\n",
    "    attach_cell_counts(reference, cell_counts)\n",
    "    print(reference.obs['cell_count'].describe())"
   ],
   "id": "byod-07"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 Prepare unlabelled bulk targets and align genes\n",
    "CSV orientation is genes x samples; AnnData orientation is samples x genes. No niche\n",
    "labels or spatial coordinates are needed for bulk samples. Use unique IDs in the same\n",
    "species and identifier system as the reference and embeddings. Resolve duplicate or\n",
    "ambiguous mappings upstream; do not append suffixes to biological gene IDs.\n",
    "For bulk .h5ad, replace the CSV conversion with ad.read_h5ad and set X to raw counts.\n",
    "Gene intersection uses availability only, not target outcomes. For a large reference,\n",
    "select informative genes using the reference alone before this step and record the list."
   ],
   "id": "byod-08"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "bulk_table = pd.read_csv(BULK, index_col=0)\n",
    "bulk = ad.AnnData(bulk_table.T.astype(np.float32))\n",
    "validate_counts(bulk)\n",
    "spatial_table = pd.read_csv(SPATIAL).set_index('g')\n",
    "if not spatial_table.index.is_unique:\n",
    "    raise ValueError('SpatialDE must contain one row per gene.')\n",
    "genes = [g for g in reference.var_names if g in bulk.var_names and g in spatial_table.index]\n",
    "if len(genes) <= GRAPH_K:\n",
    "    raise ValueError('Too few shared genes; check identifiers and matrix orientation.')\n",
    "print(f'Shared genes: {len(genes)}/{reference.n_vars} reference genes; {bulk.n_vars} bulk genes')\n",
    "reference = reference[:, genes].copy()\n",
    "bulk = bulk[:, genes].copy()\n",
    "validate_counts(reference, spatial=True, niche_key=NICHE_KEY)\n",
    "validate_counts(bulk)\n",
    "prior_raw = spatial_table.loc[genes, SPATIAL_FEATURES].astype(float)\n",
    "if not np.isfinite(prior_raw.to_numpy()).all():\n",
    "    raise ValueError('Resolve failed/nonfinite SpatialDE fits before training.')\n",
    "prior_scaled, log_columns = auto_log_then_minmax(prior_raw)\n",
    "prior = prior_scaled.to_numpy(dtype=np.float32)\n",
    "embeddings = pd.read_csv(EMBEDDINGS, index_col=0)\n",
    "embeddings.index = embeddings.index.astype(str)\n",
    "graph_genes, edge_index, edge_weight = functional_graph(embeddings, genes, k=GRAPH_K)\n",
    "print(f'Graph coverage: {len(graph_genes)}/{len(genes)} genes')"
   ],
   "id": "byod-09"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing the spatial features and functional embeddings\n",
    "Run SpatialDE on the reference only, using its raw counts and coordinates. See the\n",
    "guide for executable preprocessing code and the exact ten output columns. Do not\n",
    "reuse the NSCLC spatial prior for another tissue. GO embeddings may be reused when\n",
    "species and gene identifiers agree. The guide includes the FRoGS source and a CSV\n",
    "conversion example; GO-term membership tables are not embedding matrices."
   ],
   "id": "byod-10"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 Construct pseudo-tissues\n",
    "Use this template to train a new model for your reference and target cohort.\n",
    "Random mixing samples observations with replacement and recomputes labels from actual\n",
    "draws. For spots, requested sampling fractions need not equal cell-weighted labels.\n",
    "No niche is automatically removed. Review rare categories before choosing the output\n",
    "label set; the model cannot estimate an unrepresented niche."
   ],
   "id": "byod-11"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "width = np.ptp(reference.obsm['spatial'][:, 0]) * WIDTH_FRACTION\n",
    "pseudo = make_pseudotissues(\n",
    "    reference, niche_key=NICHE_KEY, width=width, overlap=OVERLAP, angles=ANGLES,\n",
    "    min_observations=MIN_OBSERVATIONS, min_niches=MIN_NICHES,\n",
    "    random_samples=RANDOM_SAMPLES, random_observations=RANDOM_OBSERVATIONS,\n",
    "    cell_count_key='cell_count' if SPOT_RESOLUTION else None,\n",
    ")\n",
    "print(pseudo)\n",
    "print(pseudo.obsm['sampling'].groupby('kind')['n_observations'].describe())\n",
    "print(pseudo.obs.describe())\n",
    "if pseudo.n_obs < 10:\n",
    "    raise ValueError('Too few pseudo-tissues; adjust sampling before training.')"
   ],
   "id": "byod-12"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6 Train using reference-derived validation samples\n",
    "The random split below is for early stopping only. Overlapping windows and random\n",
    "mixtures share reference observations, so this is not independent accuracy evidence.\n",
    "For a benchmark, hold out complete donors/sections before constructing pseudo-tissues\n",
    "and fitting spatial priors. Do not tune on clinical outcomes of the bulk cohort.\n",
    "The legacy training API requests a labelled target_data argument, but training uses\n",
    "source_data and early stopping uses valid_data. Here valid_data fills that interface\n",
    "slot; real bulk samples remain unlabelled and are predicted separately."
   ],
   "id": "byod-13"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "if not torch.cuda.is_available():\n",
    "    raise RuntimeError('The current NicheDECODE model requires CUDA; select a CUDA-enabled environment.')\n",
    "train_ids, valid_ids = train_test_split(np.arange(pseudo.n_obs), test_size=0.2, random_state=2021)\n",
    "source_data = pack_model_input(pseudo[train_ids].copy(), genes, prior, graph_genes)\n",
    "valid_data = pack_model_input(pseudo[valid_ids].copy(), genes, prior, graph_genes)\n",
    "target_data = pack_model_input(bulk, genes, prior, graph_genes)\n",
    "model = nicheDeconv(EPOCHS, BATCH_SIZE, LEARNING_RATE)\n",
    "model.register_graph(edge_index, edge_weight, len(graph_genes))\n",
    "loss, best_weights = model.train(source_data, valid_data, valid_data, PATIENCE, RUN_NAME)\n",
    "if best_weights is None:\n",
    "    raise RuntimeError('Training produced no valid checkpoint; inspect loss and inputs.')\n",
    "for name in ['encoder', 'gnn', 'predictor']:\n",
    "    getattr(model, name).load_state_dict(best_weights[name])"
   ],
   "id": "byod-14"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7 Predict and save the model contract\n",
    "The output CSV has sample IDs as rows and niche labels as columns. Values are\n",
    "nonnegative and sum to one per sample. Accuracy cannot be computed without measured\n",
    "ground truth. The frozen gene order, prior, graph and label order must accompany the\n",
    "weights for reuse; changing the panel or reference requires a new run."
   ],
   "id": "byod-15"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "predictions = predict_unlabelled(model, target_data, batch_size=BATCH_SIZE)\n",
    "assert np.isfinite(predictions.to_numpy()).all()\n",
    "assert np.allclose(predictions.sum(axis=1), 1, atol=1e-5)\n",
    "OUT.mkdir(parents=True)\n",
    "CHECKPOINT.parent.mkdir(parents=True, exist_ok=True)\n",
    "predictions.to_csv(OUT / 'niche_proportions.csv', index_label='sample_id')\n",
    "torch.save(best_weights, CHECKPOINT)\n",
    "np.savez_compressed(OUT / 'model_inputs.npz', spatial_prior=prior,\n",
    "                    edge_index=edge_index.numpy(), edge_weight=edge_weight.numpy())\n",
    "manifest = dict(genes=genes, graph_genes=graph_genes, niches=list(model.labels),\n",
    "    spatial_features=SPATIAL_FEATURES, reference=str(REFERENCE), bulk=str(BULK),\n",
    "    spatial_table=str(SPATIAL), embeddings=str(EMBEDDINGS),\n",
    "    normalization='raw counts -> total 10000 -> log1p -> row maximum',\n",
    "    spot_resolution=SPOT_RESOLUTION, cell_counts=str(CELL_COUNTS) if SPOT_RESOLUTION else None,\n",
    "    width=float(width), overlap=OVERLAP, angles=ANGLES, min_observations=MIN_OBSERVATIONS,\n",
    "    min_niches=MIN_NICHES, random_samples=RANDOM_SAMPLES, random_observations=RANDOM_OBSERVATIONS,\n",
    "    graph_k=GRAPH_K, seed=2021, epochs=EPOCHS, batch_size=BATCH_SIZE,\n",
    "    learning_rate=LEARNING_RATE, patience=PATIENCE)\n",
    "(OUT / 'manifest.json').write_text(json.dumps(manifest, indent=2), encoding='utf-8')\n",
    "print(predictions.head())\n",
    "print('Saved predictions and preprocessing contract to', OUT)"
   ],
   "id": "byod-16"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}