preprocess_data_cell_level.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path
from typing import Any

import h5py
import numpy as np
import pandas as pd
import scanpy as sc
import tifffile
from scipy import sparse
import shapely.wkb

logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s] %(levelname)s - %(message)s',
    datefmt='%H:%M:%S',
)
log = logging.getLogger('prad_cell_prep')


# ---------------------------
# Basic cleaning / parsing
# ---------------------------

def clean_scalar_str(x: Any) -> str:
    if isinstance(x, (bytes, bytearray, np.bytes_)):
        return x.decode('utf-8')
    return str(x)


def clean_transcript_df(df_tx: pd.DataFrame) -> pd.DataFrame:
    if 'feature_name' not in df_tx.columns or 'cell_id' not in df_tx.columns:
        raise ValueError("Transcript parquet must contain 'feature_name' and 'cell_id' columns.")
    if len(df_tx) == 0:
        raise ValueError('Transcript parquet is empty.')

    df_tx = df_tx.copy()
    df_tx['feature_name'] = df_tx['feature_name'].map(clean_scalar_str)
    df_tx['cell_id'] = df_tx['cell_id'].map(clean_scalar_str)

    bad_ids = {'-1', 'UNASSIGNED', 'unassigned', 'nan', 'None'}
    df_tx = df_tx[~df_tx['cell_id'].isin(bad_ids)].copy()
    df_tx = df_tx[df_tx['feature_name'].str.len() > 0].copy()
    return df_tx


def detect_seg_id_column(df_seg: pd.DataFrame) -> str:
    candidates = ['cell_id', 'id', 'label', 'object_id']
    for c in candidates:
        if c in df_seg.columns:
            return c
    return '__index__'


def parse_segmentation(seg_path: Path) -> pd.DataFrame:
    df_seg = pd.read_parquet(seg_path).copy()
    id_col = detect_seg_id_column(df_seg)

    if id_col == '__index__':
        df_seg.index = df_seg.index.map(clean_scalar_str)
    else:
        df_seg[id_col] = df_seg[id_col].map(clean_scalar_str)
        df_seg = df_seg.set_index(id_col, drop=True)

    if 'geometry' not in df_seg.columns:
        raise ValueError("Segmentation parquet must contain a 'geometry' column.")

    def to_geom(g: Any):
        return g if hasattr(g, 'centroid') else shapely.wkb.loads(g)

    df_seg['geometry_obj'] = df_seg['geometry'].apply(to_geom)
    df_seg['he_x'] = df_seg['geometry_obj'].apply(lambda g: g.centroid.x)
    df_seg['he_y'] = df_seg['geometry_obj'].apply(lambda g: g.centroid.y)
    return df_seg


def build_sparse_counts(df_tx: pd.DataFrame):
    cell_codes, cell_uniques = pd.factorize(df_tx['cell_id'], sort=True)
    gene_codes, gene_uniques = pd.factorize(df_tx['feature_name'], sort=True)

    data = np.ones(len(df_tx), dtype=np.uint16)
    X = sparse.coo_matrix(
        (data, (cell_codes, gene_codes)),
        shape=(len(cell_uniques), len(gene_uniques)),
        dtype=np.uint32,
    ).tocsr()
    X.sum_duplicates()
    return X, np.array(cell_uniques, dtype=str), np.array(gene_uniques, dtype=str)


# ---------------------------
# Patch export
# ---------------------------

def crop_patch(wsi: np.ndarray, x: float, y: float, radius: int) -> np.ndarray:
    H, W = wsi.shape[:2]
    x = int(round(float(x)))
    y = int(round(float(y)))
    x0, x1 = x - radius, x + radius
    y0, y1 = y - radius, y + radius

    pad_x0 = max(0, -x0)
    pad_x1 = max(0, x1 - W)
    pad_y0 = max(0, -y0)
    pad_y1 = max(0, y1 - H)

    patch = np.pad(
        wsi[max(0, y0):min(H, y1), max(0, x0):min(W, x1)],
        ((pad_y0, pad_y1), (pad_x0, pad_x1), (0, 0)),
        mode='constant',
    )
    side = 2 * radius
    if patch.shape[:2] != (side, side):
        raise RuntimeError(f'Bad patch shape {patch.shape} for centroid ({x},{y}).')
    return patch


def write_patch_h5(path: Path, cell_ids: np.ndarray, coords: np.ndarray, images: np.ndarray):
    with h5py.File(path, 'w') as f:
        f.create_dataset('cell_ids', data=np.array(cell_ids, dtype=h5py.string_dtype('utf-8')))
        f.create_dataset('coords', data=coords.astype(np.float32))
        f.create_dataset('images', data=images, compression='gzip')


def export_patches(
    adata_raw: sc.AnnData,
    wsi_path: Path,
    patch_path: Path,
    patch_radius: int = 64,
    use_page: int = 0,
) -> dict[str, Any]:
    log.info('Loading WSI page %d: %s', use_page, wsi_path)
    with tifffile.TiffFile(str(wsi_path)) as tif:
        page = tif.pages[use_page]
        wsi = page.asarray()

    if wsi.ndim == 2:
        wsi = np.stack([wsi, wsi, wsi], axis=-1)

    coords = np.asarray(adata_raw.obsm['spatial_he'])
    log.info('Cropping %d patches with radius=%d', len(coords), patch_radius)

    images = np.empty(
        (adata_raw.n_obs, patch_radius * 2, patch_radius * 2, wsi.shape[2]),
        dtype=wsi.dtype,
    )

    for i, (x, y) in enumerate(coords):
        if i % 5000 == 0 and i > 0:
            log.info('Cropped %d / %d patches', i, adata_raw.n_obs)
        images[i] = crop_patch(wsi, x, y, patch_radius)

    write_patch_h5(
        patch_path,
        adata_raw.obs_names.to_numpy(dtype=str),
        coords,
        images,
    )

    return {
        'path': str(patch_path),
        'n_patches': int(images.shape[0]),
        'patch_shape': list(images.shape[1:]),
        'patch_dtype': str(images.dtype),
        'patch_radius': int(patch_radius),
        'wsi_page': int(use_page),
    }


# ---------------------------
# AnnData builders
# ---------------------------

def make_raw_adata(
    X,
    cell_ids: np.ndarray,
    gene_names: np.ndarray,
    df_seg: pd.DataFrame,
    sample_id: str,
) -> sc.AnnData:
    seg_use = df_seg.loc[cell_ids].copy()

    adata = sc.AnnData(X=X)
    adata.obs_names = cell_ids.astype(str)
    adata.var_names = gene_names.astype(str)

    adata.obs['cell_id'] = adata.obs_names.astype(str)
    adata.obs['sample_id'] = sample_id
    adata.obs['barcode'] = adata.obs_names.astype(str)

    adata.obsm['spatial_he'] = seg_use[['he_x', 'he_y']].to_numpy()
    if all(c in seg_use.columns for c in ['x', 'y']):
        adata.obsm['spatial'] = seg_use[['x', 'y']].to_numpy()
    else:
        adata.obsm['spatial'] = adata.obsm['spatial_he'].copy()

    total_counts = np.asarray(adata.X.sum(axis=1)).ravel()
    n_genes = np.asarray((adata.X > 0).sum(axis=1)).ravel()
    adata.obs['total_counts'] = total_counts
    adata.obs['n_genes_by_counts'] = n_genes

    mito_mask = np.array([g.upper().startswith(('MT-', 'MT.')) for g in adata.var_names], dtype=bool)
    adata.var['mito'] = mito_mask
    if mito_mask.any():
        mito_counts = np.asarray(adata.X[:, mito_mask].sum(axis=1)).ravel()
    else:
        mito_counts = np.zeros(adata.n_obs, dtype=np.float32)
    adata.obs['total_counts_mito'] = mito_counts
    adata.obs['pct_counts_mito'] = np.where(total_counts > 0, mito_counts / total_counts * 100.0, 0.0)

    adata.uns['preprocess_info'] = {
        'level': 'cell',
        'sample_id': sample_id,
        'obs_unit': 'xenium cell',
        'data_state': 'raw_counts_master',
    }
    return adata


def build_raw_cell_adata(root: Path, sample_id: str):
    transcript_path = root / 'transcripts' / f'{sample_id}_transcripts.parquet'
    seg_path = root / 'xenium_seg' / f'{sample_id}_xenium_cell_seg.parquet'

    for p in [transcript_path, seg_path]:
        if not p.exists():
            raise FileNotFoundError(p)

    log.info('Loading transcripts: %s', transcript_path)
    df_tx = pd.read_parquet(transcript_path)
    df_tx = clean_transcript_df(df_tx)

    log.info('Loading segmentation: %s', seg_path)
    df_seg = parse_segmentation(seg_path)

    log.info('Building sparse cell-by-gene matrix')
    X, cell_ids, gene_names = build_sparse_counts(df_tx)

    common = [cid for cid in cell_ids if cid in df_seg.index]
    if len(common) == 0:
        raise RuntimeError('No overlap between transcript cell IDs and segmentation IDs.')
    if len(common) < len(cell_ids):
        log.warning('Dropping %d transcript cells missing from segmentation.', len(cell_ids) - len(common))

    keep_idx = pd.Index(cell_ids).get_indexer(common)
    X = X[keep_idx]
    cell_ids = np.array(common, dtype=str)

    log.info('Creating raw AnnData for %d cells x %d genes', X.shape[0], X.shape[1])
    adata = make_raw_adata(X, cell_ids, gene_names, df_seg, sample_id)

    summary = {
        'transcripts': str(transcript_path),
        'xenium_seg': str(seg_path),
        'n_cells': int(adata.n_obs),
        'n_genes': int(adata.n_vars),
        'x_dtype': str(adata.X.dtype),
        'x_sparse': bool(sparse.issparse(adata.X)),
    }
    return adata, summary


# ---------------------------
# Expression preprocessing for scGPT input
# ---------------------------

def preprocess_adata_for_scgpt_input(
    adata_raw: sc.AnnData,
    n_hvg: int = 1200,
    min_gene_counts: int = 3,
    target_sum: float = 1e4,
    convert_dense: bool = True,
) -> sc.AnnData:
    '''
    Build an scGPT-input AnnData while preserving the exact raw/master cell set.

    Important design choice:
      - NO cell filtering here.
      - NO scGPT vocab filtering here.
      - Keep row-wise pairing with raw image patches valid.
      - Let the downstream scGPT encoder / extraction code handle vocab mapping.
    '''
    adata = adata_raw.copy()

    # preserve raw counts
    adata.layers['counts'] = adata.X.copy()

    # gene filtering only
    if min_gene_counts and min_gene_counts > 0:
        sc.pp.filter_genes(adata, min_counts=min_gene_counts)

    # normalize + log1p
    sc.pp.normalize_total(adata, target_sum=target_sum)
    sc.pp.log1p(adata)

    # HVG selection on the same cell set
    if n_hvg and adata.n_vars > int(n_hvg):
        try:
            sc.pp.highly_variable_genes(
                adata,
                n_top_genes=int(n_hvg),
                flavor='seurat_v3',
                layer='counts',
                subset=True,
            )
        except Exception:
            sc.pp.highly_variable_genes(
                adata,
                n_top_genes=int(n_hvg),
                flavor='seurat_v3',
                subset=True,
            )

    if convert_dense and sparse.issparse(adata.X):
        adata.X = adata.X.toarray()
    adata.X = np.asarray(adata.X, dtype=np.float32)

    adata.uns['scgpt_input_preprocess_info'] = {
        'n_hvg': int(n_hvg),
        'min_gene_counts': int(min_gene_counts),
        'target_sum': float(target_sum),
        'cell_filtering_applied_here': False,
        'vocab_filtering_applied_here': False,
        'note': 'Cell set intentionally preserved to keep transcript/image row alignment. scGPT vocab mapping is deferred to embedding extraction.',
    }
    return adata


# ---------------------------
# Save helpers
# ---------------------------

def save_manifest(path: Path, payload: dict[str, Any]) -> None:
    with open(path, 'w') as f:
        json.dump(payload, f, indent=2)


# ---------------------------
# Main orchestration
# ---------------------------

def run_raw_mode(root: Path, sample_id: str, sample_out: Path):
    adata_raw, raw_summary = build_raw_cell_adata(root, sample_id)

    raw_path = sample_out / f'{sample_id}_cells_raw.h5ad'
    adata_raw.write_h5ad(raw_path)
    log.info('Saved raw AnnData: %s', raw_path)

    return adata_raw, {
        'adata_raw': str(raw_path),
        'raw_summary': raw_summary,
    }


def run_scgpt_input_mode(
    sample_id: str,
    sample_out: Path,
    n_hvg: int,
    min_gene_counts: int,
    target_sum: float,
    input_raw_adata: Path | None = None,
):
    raw_path = input_raw_adata or (sample_out / f'{sample_id}_cells_raw.h5ad')
    if not raw_path.exists():
        raise FileNotFoundError(f'Raw h5ad not found: {raw_path}')

    log.info('Loading raw master AnnData for scGPT-input preprocessing: %s', raw_path)
    adata_raw = sc.read_h5ad(raw_path)

    adata_scgpt = preprocess_adata_for_scgpt_input(
        adata_raw=adata_raw,
        n_hvg=n_hvg,
        min_gene_counts=min_gene_counts,
        target_sum=target_sum,
        convert_dense=True,
    )

    if not np.array_equal(adata_raw.obs_names.astype(str), adata_scgpt.obs_names.astype(str)):
        raise RuntimeError('scGPT-input preprocessing changed cell order or cell set unexpectedly.')

    scgpt_path = sample_out / f'{sample_id}_cells_scgpt_input.h5ad'
    adata_scgpt.write_h5ad(scgpt_path)
    log.info('Saved scGPT-input AnnData: %s', scgpt_path)

    return {
        'adata_scgpt_input': str(scgpt_path),
        'scgpt_input_summary': {
            'n_cells': int(adata_scgpt.n_obs),
            'n_genes': int(adata_scgpt.n_vars),
            'x_dtype': str(adata_scgpt.X.dtype),
            'x_dense': True,
        },
    }


def preprocess_sample(
    root: Path,
    sample_id: str,
    out_root: Path,
    mode: str,
    patch_radius: int = 64,
    use_page: int = 0,
    export_patches_flag: bool = False,
    scgpt_n_hvg: int = 1200,
    scgpt_min_gene_counts: int = 3,
    scgpt_target_sum: float = 1e4,
    input_raw_adata: Path | None = None,
):
    sample_out = out_root / sample_id
    sample_out.mkdir(parents=True, exist_ok=True)

    manifest = {
        'sample_id': sample_id,
        'mode': mode,
        'input': {'root': str(root)},
        'output': {'sample_out': str(sample_out)},
    }

    adata_raw = None

    if mode in {'raw', 'full'}:
        adata_raw, raw_info = run_raw_mode(root, sample_id, sample_out)
        manifest['output'].update(raw_info)
    elif mode == 'scgpt_input':
        raw_path = input_raw_adata or (sample_out / f'{sample_id}_cells_raw.h5ad')
        manifest['input']['raw_adata_for_scgpt_input'] = str(raw_path)

    if export_patches_flag:
        if adata_raw is None:
            raw_path = input_raw_adata or (sample_out / f'{sample_id}_cells_raw.h5ad')
            if not raw_path.exists():
                raise FileNotFoundError(f'Cannot export patches without raw master AnnData: {raw_path}')
            adata_raw = sc.read_h5ad(raw_path)

        wsi_path = root / 'wsis' / f'{sample_id}.tif'
        if not wsi_path.exists():
            raise FileNotFoundError(wsi_path)

        patch_path = sample_out / f'{sample_id}_cell_patches.h5'
        patch_info = export_patches(
            adata_raw=adata_raw,
            wsi_path=wsi_path,
            patch_path=patch_path,
            patch_radius=patch_radius,
            use_page=use_page,
        )
        manifest['input']['wsi'] = str(wsi_path)
        manifest['output']['patches'] = patch_info['path']
        manifest['patches_summary'] = patch_info

    if mode in {'scgpt_input', 'full'}:
        scgpt_info = run_scgpt_input_mode(
            sample_id=sample_id,
            sample_out=sample_out,
            n_hvg=scgpt_n_hvg,
            min_gene_counts=scgpt_min_gene_counts,
            target_sum=scgpt_target_sum,
            input_raw_adata=input_raw_adata,
        )
        manifest['output'].update(scgpt_info)

    manifest_path = sample_out / 'manifest.json'
    save_manifest(manifest_path, manifest)
    log.info('Saved manifest: %s', manifest_path)
    return manifest


# ---------------------------
# CLI
# ---------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description='Build raw cell-level PRAD data and/or an scGPT-input copy.'
    )
    p.add_argument('--root', type=Path, required=True,
                   help='Path to PRAD directory containing transcripts/, xenium_seg/, wsis/.')
    p.add_argument('--sample-id', type=str, required=True)
    p.add_argument('--output-root', type=Path, default=Path('data/PRAD_cell'))

    p.add_argument('--mode', type=str, choices=['raw', 'scgpt_input', 'full'], default='full',
                   help='raw: only raw master adata; scgpt_input: only build scGPT-input copy from existing raw adata; full: raw + scgpt_input.')
    p.add_argument('--input-raw-adata', type=Path, default=None,
                   help='Optional explicit path to an existing *_cells_raw.h5ad (useful for mode=scgpt_input).')

    p.add_argument('--export-patches', action='store_true',
                   help='Also export raw cell patches from the WSI. Patch export always follows the raw/master cell set.')
    p.add_argument('--patch-radius', type=int, default=64,
                   help='Patch half-size in pixels; output patch is 2*radius square.')
    p.add_argument('--wsi-page', type=int, default=0,
                   help='TIFF page index to read when exporting patches.')

    p.add_argument('--scgpt-n-hvg', type=int, default=1200)
    p.add_argument('--scgpt-min-gene-counts', type=int, default=3)
    p.add_argument('--scgpt-target-sum', type=float, default=1e4)

    return p.parse_args()


def main():
    args = parse_args()
    manifest = preprocess_sample(
        root=args.root,
        sample_id=args.sample_id,
        out_root=args.output_root,
        mode=args.mode,
        patch_radius=args.patch_radius,
        use_page=args.wsi_page,
        export_patches_flag=args.export_patches,
        scgpt_n_hvg=args.scgpt_n_hvg,
        scgpt_min_gene_counts=args.scgpt_min_gene_counts,
        scgpt_target_sum=args.scgpt_target_sum,
        input_raw_adata=args.input_raw_adata,
    )
    print(json.dumps(manifest, indent=2))


if __name__ == '__main__':
    main()