nest-st 0.1.0__tar.gz

nest_st-0.1.0/PKG-INFO

@@ -0,0 +1,109 @@
+Metadata-Version: 2.4
+Name: nest-st
+Version: 0.1.0
+Summary: NEST: Native Expression reconstruction for Spatial Transcriptomics
+Author-email: Tianyi Liu <1694414073@qq.com>
+License: MIT
+Project-URL: Homepage, https://github.com/lty10086/NEST
+Project-URL: Repository, https://github.com/lty10086/NEST
+Project-URL: Issues, https://github.com/lty10086/NEST/issues
+Keywords: spatial transcriptomics,single-cell,imputation,dropout,graph autoencoder
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.23
+Requires-Dist: pandas>=1.5
+Requires-Dist: scipy>=1.9
+Requires-Dist: anndata>=0.9
+Requires-Dist: scanpy>=1.9
+Requires-Dist: scikit-learn>=1.2
+Requires-Dist: matplotlib>=3.6
+Requires-Dist: torch>=2.0
+Requires-Dist: torch-geometric>=2.3
+
+# NEST
+
+NEST, Native Expression reconstruction for Spatial Transcriptomics, reconstructs biologically faithful spatial transcriptomes under native spatial constraints.
+
+It supports two tasks with one parameter:
+
+- `mode="missing_gene"`: imputes unmeasured genes in imaging-based spatial transcriptomics using a matched scRNA-seq reference.
+- `mode="dropout"`: recovers dropout-corrupted false zeros in sequencing-based spatial transcriptomics using ST data only.
+
+## Installation for development
+
+```bash
+git clone https://github.com/yourname/nest-st.git
+cd nest-st
+pip install -e .
+```
+
+For CUDA environments, install PyTorch and PyTorch Geometric using the wheel index recommended for your CUDA version before installing the package.
+
+## Python usage
+
+### Missing-gene imputation
+
+```python
+import scanpy as sc
+from nest_st import NEST, NESTConfig
+
+adata_st = sc.read_h5ad("data/st.h5ad")
+adata_sc = sc.read_h5ad("data/sc.h5ad")
+
+sc.pp.log1p(adata_st)
+sc.pp.log1p(adata_sc)
+
+config = NESTConfig(mode="missing_gene", ae_epochs=1000, gae_epochs=2000, device="cuda")
+adata_imputed = NEST(config).fit_transform(adata_st=adata_st, adata_sc=adata_sc)
+adata_imputed.write_h5ad("nest_missing_gene_result.h5ad")
+```
+
+### Dropout recovery
+
+```python
+import scanpy as sc
+from nest_st import NEST, NESTConfig
+
+adata_st = sc.read_h5ad("data/st.h5ad")
+sc.pp.log1p(adata_st)
+
+config = NESTConfig(mode="dropout", fill_zeros_only=True, ae_epochs=1000, gae_epochs=2000, device="cuda")
+adata_recovered = NEST(config).fit_transform(adata_st=adata_st)
+adata_recovered.write_h5ad("nest_dropout_result.h5ad")
+```
+
+## Command-line usage
+
+```bash
+python main.py --mode missing_gene --st data/st.h5ad --sc data/sc.h5ad --output result_gene.h5ad --log1p
+python main.py --mode dropout --st data/st.h5ad --output result_dropout.h5ad --log1p --fill-zeros-only
+```
+
+## Input requirements
+
+- Input files should be AnnData `.h5ad` files.
+- `adata_st.obsm["spatial"]` must contain spatial coordinates.
+- For `missing_gene` mode, `adata_sc` must contain genes to be transferred into the ST domain.
+- For `dropout` mode, no scRNA-seq reference is required.
+
+## Upload to PyPI
+
+```bash
+python -m pip install --upgrade build twine
+python -m build
+python -m twine upload dist/*
+```
+
+For TestPyPI:
+
+```bash
+python -m twine upload --repository testpypi dist/*
+```
+
+Update the package name, author information and GitHub URLs in `pyproject.toml` before publishing.

nest_st-0.1.0/README.md

@@ -0,0 +1,81 @@
+# NEST
+
+NEST, Native Expression reconstruction for Spatial Transcriptomics, reconstructs biologically faithful spatial transcriptomes under native spatial constraints.
+
+It supports two tasks with one parameter:
+
+- `mode="missing_gene"`: imputes unmeasured genes in imaging-based spatial transcriptomics using a matched scRNA-seq reference.
+- `mode="dropout"`: recovers dropout-corrupted false zeros in sequencing-based spatial transcriptomics using ST data only.
+
+## Installation for development
+
+```bash
+git clone https://github.com/yourname/nest-st.git
+cd nest-st
+pip install -e .
+```
+
+For CUDA environments, install PyTorch and PyTorch Geometric using the wheel index recommended for your CUDA version before installing the package.
+
+## Python usage
+
+### Missing-gene imputation
+
+```python
+import scanpy as sc
+from nest_st import NEST, NESTConfig
+
+adata_st = sc.read_h5ad("data/st.h5ad")
+adata_sc = sc.read_h5ad("data/sc.h5ad")
+
+sc.pp.log1p(adata_st)
+sc.pp.log1p(adata_sc)
+
+config = NESTConfig(mode="missing_gene", ae_epochs=1000, gae_epochs=2000, device="cuda")
+adata_imputed = NEST(config).fit_transform(adata_st=adata_st, adata_sc=adata_sc)
+adata_imputed.write_h5ad("nest_missing_gene_result.h5ad")
+```
+
+### Dropout recovery
+
+```python
+import scanpy as sc
+from nest_st import NEST, NESTConfig
+
+adata_st = sc.read_h5ad("data/st.h5ad")
+sc.pp.log1p(adata_st)
+
+config = NESTConfig(mode="dropout", fill_zeros_only=True, ae_epochs=1000, gae_epochs=2000, device="cuda")
+adata_recovered = NEST(config).fit_transform(adata_st=adata_st)
+adata_recovered.write_h5ad("nest_dropout_result.h5ad")
+```
+
+## Command-line usage
+
+```bash
+python main.py --mode missing_gene --st data/st.h5ad --sc data/sc.h5ad --output result_gene.h5ad --log1p
+python main.py --mode dropout --st data/st.h5ad --output result_dropout.h5ad --log1p --fill-zeros-only
+```
+
+## Input requirements
+
+- Input files should be AnnData `.h5ad` files.
+- `adata_st.obsm["spatial"]` must contain spatial coordinates.
+- For `missing_gene` mode, `adata_sc` must contain genes to be transferred into the ST domain.
+- For `dropout` mode, no scRNA-seq reference is required.
+
+## Upload to PyPI
+
+```bash
+python -m pip install --upgrade build twine
+python -m build
+python -m twine upload dist/*
+```
+
+For TestPyPI:
+
+```bash
+python -m twine upload --repository testpypi dist/*
+```
+
+Update the package name, author information and GitHub URLs in `pyproject.toml` before publishing.

nest_st-0.1.0/nest_st/__init__.py

@@ -0,0 +1,13 @@
+"""NEST: Native Expression reconstruction for Spatial Transcriptomics."""
+
+from .core import NEST, NESTConfig, impute_missing_genes, recover_dropout, run_nest
+
+__all__ = [
+    "NEST",
+    "NESTConfig",
+    "run_nest",
+    "impute_missing_genes",
+    "recover_dropout",
+]
+
+__version__ = "0.1.0"

nest_st-0.1.0/nest_st/core.py

@@ -0,0 +1,627 @@
+"""Core implementation of NEST.
+
+NEST supports two tasks:
+1. missing_gene: impute genes absent from spatial transcriptomics data using a matched
+   scRNA-seq reference, followed by native-domain correction and graph-guided recovery.
+2. dropout: recover dropout-corrupted false zeros in spatial transcriptomics data without
+   requiring a scRNA-seq reference.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Literal, Optional, Sequence, Tuple
+
+import anndata as ad
+import numpy as np
+import scanpy as sc
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from scipy import sparse
+from sklearn.cluster import KMeans
+from sklearn.metrics.pairwise import cosine_similarity as sklearn_cosine_similarity
+from sklearn.neighbors import LocalOutlierFactor, NearestNeighbors
+from torch.nn.functional import cosine_similarity
+from torch_geometric.data import Data
+from torch_geometric.nn import GCNConv
+from torch_geometric.utils import add_self_loops, dense_to_sparse
+
+NESTMode = Literal["missing_gene", "dropout"]
+
+
+@dataclass
+class NESTConfig:
+    """Configuration for NEST.
+
+    Parameters
+    ----------
+    mode
+        ``"missing_gene"`` imputes unmeasured genes using a scRNA-seq reference.
+        ``"dropout"`` recovers dropout events using the ST matrix only.
+    k
+        Number of cell neighbors used for graph construction. If ``None``, NEST uses
+        ``max(10, min(30, n_cells // 100))``, bounded by the number of cells.
+    kmap
+        Number of nearest scRNA-seq neighbors used for cross-modal mapping.
+        Used only when ``mode="missing_gene"``.
+    kg
+        Number of shared gene neighbors used for native-domain correction. If ``None``,
+        NEST uses ``max(5, n_common_genes // 20)``.
+    latent_size
+        Latent dimension for AE and GAE. If ``None``, NEST uses one quarter of the
+        input dimension, with a minimum of 2.
+    lr
+        Learning rate for Adam.
+    ae_epochs
+        Number of AE training epochs.
+    gae_epochs
+        Number of GAE training epochs.
+    device
+        ``"cuda"`` or ``"cpu"``. If ``None``, CUDA is used when available.
+    fill_zeros_only
+        In dropout mode, if ``True``, observed non-zero values are kept and only zero
+        entries are replaced by reconstructed values. If ``False``, the full GAE
+        reconstruction is returned.
+    random_state
+        Random seed for K-means and Scanpy-related operations.
+    verbose
+        Whether to print progress messages.
+    """
+
+    mode: NESTMode = "missing_gene"
+    k: Optional[int] = None
+    kmap: int = 50
+    kg: Optional[int] = None
+    latent_size: Optional[int] = None
+    lr: float = 1e-3
+    ae_epochs: int = 1000
+    gae_epochs: int = 2000
+    device: Optional[str] = None
+    fill_zeros_only: bool = True
+    random_state: int = 42
+    verbose: bool = True
+
+
+class AutoEncoder(nn.Module):
+    """Two-layer fully connected autoencoder used to learn the shared ST latent space."""
+
+    def __init__(self, input_dim: int, latent_size: int):
+        super().__init__()
+        hidden_size = max(2, input_dim // 2)
+        latent_size = max(2, latent_size)
+
+        self.fc1 = nn.Linear(input_dim, hidden_size)
+        self.layer_norm1 = nn.LayerNorm(hidden_size)
+        self.fc2 = nn.Linear(hidden_size, latent_size)
+        self.layer_norm2 = nn.LayerNorm(latent_size)
+        self.fc3 = nn.Linear(latent_size, hidden_size)
+        self.fc4 = nn.Linear(hidden_size, input_dim)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        encoded = self.fc1(x)
+        encoded = self.layer_norm1(encoded)
+        latent = self.fc2(encoded)
+        latent = self.layer_norm2(latent)
+        decoded = F.relu(self.fc3(latent))
+        reconstructed = F.relu(self.fc4(decoded))
+        return latent, reconstructed
+
+
+class GraphAutoEncoder(nn.Module):
+    """Graph autoencoder with two GCN layers and an auxiliary cluster classifier."""
+
+    def __init__(self, input_dim: int, latent_size: int, num_clusters: int):
+        super().__init__()
+        hidden_size = max(2, input_dim // 2)
+        latent_size = max(2, latent_size)
+        num_clusters = max(1, num_clusters)
+
+        self.gc1 = GCNConv(input_dim, hidden_size)
+        self.layer_norm1 = nn.LayerNorm(hidden_size)
+        self.gc2 = GCNConv(hidden_size, latent_size)
+        self.layer_norm2 = nn.LayerNorm(latent_size)
+        self.fc1 = nn.Linear(latent_size, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, input_dim)
+        self.classifier = nn.Linear(latent_size, num_clusters)
+
+    def forward(
+        self, x: torch.Tensor, edge_index: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden = self.gc1(x, edge_index)
+        hidden = self.layer_norm1(hidden)
+        latent = self.gc2(hidden, edge_index)
+        latent = self.layer_norm2(latent)
+        decoded = F.relu(self.fc1(latent))
+        reconstructed = F.relu(self.fc2(decoded))
+        logits = self.classifier(latent)
+        return latent, reconstructed, logits
+
+
+def _as_numpy(x) -> np.ndarray:
+    if sparse.issparse(x):
+        return x.toarray()
+    return np.asarray(x)
+
+
+def _copy_adata_with_numpy(adata: ad.AnnData) -> ad.AnnData:
+    out = adata.copy()
+    out.X = _as_numpy(out.X).astype(np.float32, copy=False)
+    return out
+
+
+def _resolve_device(device: Optional[str]) -> torch.device:
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    return torch.device(device)
+
+
+def _resolve_k(n_cells: int, k: Optional[int]) -> int:
+    if n_cells < 3:
+        raise ValueError("NEST requires at least three spatial cells/spots.")
+    if k is None:
+        k = max(10, min(30, n_cells // 100))
+    return int(max(2, min(k, n_cells - 1)))
+
+
+def _resolve_latent_size(input_dim: int, latent_size: Optional[int]) -> int:
+    if input_dim < 2:
+        raise ValueError("NEST requires at least two input genes.")
+    if latent_size is None:
+        latent_size = input_dim // 4
+    return int(max(2, min(latent_size, input_dim)))
+
+
+def _require_spatial(adata_st: ad.AnnData) -> np.ndarray:
+    if "spatial" not in adata_st.obsm:
+        raise ValueError("adata_st.obsm['spatial'] is required for spatial graph construction.")
+    return np.asarray(adata_st.obsm["spatial"])
+
+
+def _log(message: str, verbose: bool = True) -> None:
+    if verbose:
+        print(message)
+
+
+def _pre_clustering(
+    adata: ad.AnnData,
+    k: int,
+    random_state: int,
+    scale: bool,
+    verbose: bool,
+) -> Tuple[np.ndarray, np.ndarray]:
+    _log("Preliminary cell clustering...", verbose)
+    adata_pre = adata.copy()
+    if scale:
+        sc.pp.scale(adata_pre, max_value=10)
+    n_components = min(adata_pre.n_obs, adata_pre.n_vars) - 1
+    n_components = max(1, min(40, n_components))
+    sc.pp.pca(adata_pre, n_comps=n_components, random_state=random_state)
+    sc.pp.neighbors(adata_pre, n_neighbors=k, n_pcs=n_components, random_state=random_state)
+    sc.tl.leiden(adata_pre, resolution=1.0, random_state=random_state)
+    sc.tl.umap(adata_pre, random_state=random_state)
+    labels = adata_pre.obs["leiden"].astype("category").cat.codes.values
+    return np.asarray(adata_pre.obsm["X_umap"]), labels
+
+
+def _create_cell_network(
+    adata_base: ad.AnnData,
+    adata_expr: ad.AnnData,
+    k: int,
+    random_state: int,
+    scale_base: bool,
+    scale_expr: bool,
+    verbose: bool,
+) -> Tuple[ad.AnnData, np.ndarray, np.ndarray]:
+    """Build the consensus cell graph used by GAE."""
+
+    spatial_coordinates = _require_spatial(adata_base)
+    z_umap_base, cell_cluster_labels = _pre_clustering(
+        adata_base, k=k, random_state=random_state, scale=scale_base, verbose=verbose
+    )
+    adata_expr.obs["leiden"] = cell_cluster_labels.astype(str)
+
+    if adata_expr is adata_base:
+        z_umap_expr = z_umap_base.copy()
+    else:
+        z_umap_expr, _ = _pre_clustering(
+            adata_expr, k=k, random_state=random_state, scale=scale_expr, verbose=verbose
+        )
+
+    _log("Searching for scattered cells...", verbose)
+    outlier_indices_all = []
+    for cluster in np.unique(cell_cluster_labels):
+        cluster_cells = np.where(cell_cluster_labels == cluster)[0]
+        if len(cluster_cells) < 5:
+            continue
+        cluster_data = z_umap_expr[cluster_cells]
+        n_neighbors = max(2, min(len(cluster_cells) - 1, len(cluster_cells) // 3))
+        lof = LocalOutlierFactor(n_neighbors=n_neighbors, contamination=0.3)
+        y_pred = lof.fit_predict(cluster_data)
+        local_outliers = np.where(y_pred == -1)[0]
+        if local_outliers.size > 0:
+            outlier_indices_all.append(cluster_cells[local_outliers])
+    if outlier_indices_all:
+        outlier_indices_all = np.concatenate(outlier_indices_all)
+    else:
+        outlier_indices_all = np.array([], dtype=int)
+
+    _log("Creating cell network...", verbose)
+    nn_expr = NearestNeighbors(n_neighbors=min(k, z_umap_base.shape[0]))
+    nn_expr.fit(z_umap_base)
+    _, expr_neighbors = nn_expr.kneighbors(z_umap_base)
+
+    g1 = np.zeros((z_umap_base.shape[0], z_umap_base.shape[0]), dtype=np.float32)
+    for i in range(z_umap_base.shape[0]):
+        for neigh in expr_neighbors[i]:
+            if i < neigh and cell_cluster_labels[i] == cell_cluster_labels[neigh]:
+                g1[i, neigh] = 1.0
+                g1[neigh, i] = 1.0
+
+    g2 = g1.copy()
+    nn_spatial = NearestNeighbors(n_neighbors=min(k, spatial_coordinates.shape[0]))
+    nn_spatial.fit(spatial_coordinates)
+    _, spatial_neighbors = nn_spatial.kneighbors(spatial_coordinates)
+
+    for i in outlier_indices_all:
+        neigh = spatial_neighbors[i]
+        neigh = neigh[neigh != i]
+        same_cluster = neigh[cell_cluster_labels[neigh] == cell_cluster_labels[i]]
+        for j in same_cluster:
+            g2[i, j] = 1.0
+            g2[j, i] = 1.0
+
+    return adata_expr, g1, g2
+
+
+def _train_autoencoder_missing_gene(
+    adata_sc: ad.AnnData,
+    adata_st: ad.AnnData,
+    common_genes: Sequence[str],
+    latent_size: int,
+    lr: float,
+    kmap: int,
+    kg: int,
+    epochs: int,
+    device: torch.device,
+    verbose: bool,
+) -> ad.AnnData:
+    criterion = nn.MSELoss()
+    _log("ST autoencoder begins training for missing-gene imputation...", verbose)
+
+    sc_norm = _as_numpy(adata_sc[:, common_genes].X).astype(np.float32)
+    st_norm = _as_numpy(adata_st[:, common_genes].X).astype(np.float32)
+    x_sc = torch.tensor(sc_norm, dtype=torch.float32, device=device)
+    x_st = torch.tensor(st_norm, dtype=torch.float32, device=device)
+
+    model = AutoEncoder(x_st.shape[1], latent_size).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999))
+
+    for epoch in range(epochs):
+        model.train()
+        optimizer.zero_grad()
+        _, xhat_st = model(x_st)
+        loss_rec = criterion(xhat_st, x_st)
+        loss_cos = 1 - cosine_similarity(xhat_st, x_st, dim=1).mean()
+        loss = loss_rec + loss_cos
+        loss.backward()
+        optimizer.step()
+        if verbose and (epoch + 1) % 500 == 0:
+            print(
+                f"[AE {epoch + 1}/{epochs}] loss={loss.item():.4f} "
+                f"mse={loss_rec.item():.4f} cos={loss_cos.item():.4f}"
+            )
+
+    _log("Cell-domain mapping...", verbose)
+    model.eval()
+    with torch.no_grad():
+        st_embedding, _ = model(x_st)
+        sc_embedding, _ = model(x_sc)
+
+    sc_embedding_np = sc_embedding.detach().cpu().numpy()
+    st_embedding_np = st_embedding.detach().cpu().numpy()
+    kmap = min(kmap, sc_embedding_np.shape[0])
+    nbrs = NearestNeighbors(n_neighbors=kmap, metric="euclidean").fit(sc_embedding_np)
+    distances, indices = nbrs.kneighbors(st_embedding_np)
+    weights = 1.0 / (distances + 1e-8)
+    weights /= weights.sum(axis=1, keepdims=True)
+
+    sc_all = _as_numpy(adata_sc.X).astype(np.float32)
+    mapped_x = (weights[..., None] * sc_all[indices]).sum(axis=1)
+
+    adata_map = ad.AnnData(X=mapped_x)
+    adata_map.obs_names = adata_st.obs_names.copy()
+    adata_map.var_names = adata_sc.var_names.copy()
+    adata_map.obs = adata_st.obs.copy()
+    adata_map.var = adata_sc.var.copy()
+    if "spatial" in adata_st.obsm:
+        adata_map.obsm["spatial"] = np.asarray(adata_st.obsm["spatial"]).copy()
+    adata_map.obsm["X_latent"] = st_embedding_np
+
+    _log("Native gene-domain correction...", verbose)
+    non_common_genes = [g for g in adata_map.var_names if g not in set(common_genes)]
+    if len(non_common_genes) == 0:
+        _log("No non-common genes found. Gene-domain correction skipped.", verbose)
+        return adata_map
+
+    valid_common = [g for g in common_genes if g in adata_st.var_names and g in adata_map.var_names]
+    if len(valid_common) == 0:
+        return adata_map
+
+    x_corrected = _as_numpy(adata_map.X).astype(np.float32).copy()
+    idx_common_map = [adata_map.var_names.get_loc(g) for g in valid_common]
+    idx_common_st = [adata_st.var_names.get_loc(g) for g in valid_common]
+    expr_map_full = _as_numpy(adata_map.X).astype(np.float32)
+    expr_st_common = _as_numpy(adata_st.X[:, idx_common_st]).astype(np.float32)
+    expr_common_in_map = expr_map_full[:, idx_common_map].T
+
+    for gene in non_common_genes:
+        idx_gene = adata_map.var_names.get_loc(gene)
+        expr_gene = expr_map_full[:, idx_gene].reshape(1, -1)
+        sim = sklearn_cosine_similarity(expr_gene, expr_common_in_map).flatten()
+        top_k = min(kg, len(sim))
+        neighbor_indices = np.argsort(-sim)[:top_k]
+        neighbor_expr = expr_st_common[:, neighbor_indices]
+        gene_weights = sim[neighbor_indices]
+        if np.allclose(gene_weights.sum(), 0):
+            gene_weights = np.ones_like(gene_weights) / len(gene_weights)
+        else:
+            gene_weights = gene_weights / gene_weights.sum()
+        x_corrected[:, idx_gene] = neighbor_expr @ gene_weights
+
+    adata_map.X = x_corrected
+    for gene in valid_common:
+        adata_map[:, gene].X = _as_numpy(adata_st[:, gene].X).copy()
+
+    return adata_map
+
+
+def _train_autoencoder_dropout(
+    adata_st: ad.AnnData,
+    genes: Sequence[str],
+    latent_size: int,
+    lr: float,
+    epochs: int,
+    device: torch.device,
+    verbose: bool,
+) -> ad.AnnData:
+    criterion = nn.MSELoss()
+    _log("ST autoencoder begins training for dropout recovery...", verbose)
+    st_norm = _as_numpy(adata_st[:, genes].X).astype(np.float32)
+    x_st = torch.tensor(st_norm, dtype=torch.float32, device=device)
+
+    model = AutoEncoder(x_st.shape[1], latent_size).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999))
+
+    for epoch in range(epochs):
+        model.train()
+        optimizer.zero_grad()
+        _, xhat_st = model(x_st)
+        mask = x_st != 0
+        loss_rec = criterion(xhat_st[mask], x_st[mask])
+        loss_cos = 1 - cosine_similarity(xhat_st * mask, x_st * mask, dim=1).mean()
+        loss = loss_rec + loss_cos
+        loss.backward()
+        optimizer.step()
+        if verbose and (epoch + 1) % 500 == 0:
+            print(
+                f"[AE {epoch + 1}/{epochs}] loss={loss.item():.4f} "
+                f"mse={loss_rec.item():.4f} cos={loss_cos.item():.4f}"
+            )
+
+    model.eval()
+    with torch.no_grad():
+        st_embedding, _ = model(x_st)
+    adata_out = adata_st.copy()
+    adata_out.obsm["X_latent"] = st_embedding.detach().cpu().numpy()
+    return adata_out
+
+
+def _train_graph_autoencoder(
+    adata_expr: ad.AnnData,
+    graph: np.ndarray,
+    latent_size: int,
+    lr: float,
+    epochs: int,
+    device: torch.device,
+    observed_only_loss: bool,
+    random_state: int,
+    verbose: bool,
+) -> np.ndarray:
+    _log("Graph autoencoder begins training...", verbose)
+    edge_index, _ = dense_to_sparse(torch.tensor(graph, dtype=torch.float32))
+    edge_index, _ = add_self_loops(edge_index, num_nodes=graph.shape[0])
+    edge_index = edge_index.to(device)
+
+    x = torch.tensor(_as_numpy(adata_expr.X).astype(np.float32), dtype=torch.float32, device=device)
+    data = Data(x=x, edge_index=edge_index)
+
+    latent_np = np.asarray(adata_expr.obsm["X_latent"])
+    n_clusters = int(max(1, len(np.unique(adata_expr.obs["leiden"]))))
+    if n_clusters == 1:
+        labels = np.zeros(latent_np.shape[0], dtype=int)
+    else:
+        labels = KMeans(n_clusters=n_clusters, n_init=10, random_state=random_state).fit_predict(latent_np)
+    pseudo_labels = torch.tensor(labels, dtype=torch.long, device=device)
+
+    model = GraphAutoEncoder(x.size(1), latent_size, n_clusters).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999))
+    criterion_rec = nn.MSELoss()
+    criterion_cls = nn.CrossEntropyLoss()
+
+    for epoch in range(epochs):
+        model.train()
+        optimizer.zero_grad()
+        latent, reconstructed, logits = model(data.x, data.edge_index)
+        if observed_only_loss:
+            mask = data.x != 0
+            loss_rec = criterion_rec(reconstructed[mask], data.x[mask])
+            loss_cos = 1 - cosine_similarity(reconstructed * mask, data.x * mask, dim=1).mean()
+        else:
+            loss_rec = criterion_rec(reconstructed, data.x)
+            loss_cos = 1 - cosine_similarity(reconstructed, data.x, dim=1).mean()
+        loss_cls = criterion_cls(logits, pseudo_labels)
+        loss = loss_rec + loss_cos + loss_cls
+        loss.backward()
+        optimizer.step()
+        if verbose and (epoch + 1) % 500 == 0:
+            print(
+                f"[GAE {epoch + 1}/{epochs}] loss={loss.item():.4f} "
+                f"mse={loss_rec.item():.4f} cos={loss_cos.item():.4f} label={loss_cls.item():.4f}"
+            )
+
+    model.eval()
+    with torch.no_grad():
+        _, imputed, _ = model(data.x, data.edge_index)
+    return imputed.detach().cpu().numpy()
+
+
+class NEST:
+    """High-level NEST runner."""
+
+    def __init__(self, config: Optional[NESTConfig] = None):
+        self.config = config or NESTConfig()
+
+    def fit_transform(
+        self,
+        adata_st: ad.AnnData,
+        adata_sc: Optional[ad.AnnData] = None,
+        common_genes: Optional[Iterable[str]] = None,
+    ) -> ad.AnnData:
+        """Run NEST and return an AnnData object containing reconstructed expression."""
+
+        config = self.config
+        torch.manual_seed(config.random_state)
+        np.random.seed(config.random_state)
+        device = _resolve_device(config.device)
+
+        adata_st_work = _copy_adata_with_numpy(adata_st)
+        k = _resolve_k(adata_st_work.n_obs, config.k)
+
+        if config.mode == "missing_gene":
+            if adata_sc is None:
+                raise ValueError("adata_sc is required when mode='missing_gene'.")
+            adata_sc_work = _copy_adata_with_numpy(adata_sc)
+            if common_genes is None:
+                common_genes = list(adata_st_work.var_names.intersection(adata_sc_work.var_names))
+            else:
+                common_genes = list(common_genes)
+            if len(common_genes) < 2:
+                raise ValueError("At least two common genes are required for missing-gene imputation.")
+            kg = config.kg if config.kg is not None else max(5, len(common_genes) // 20)
+            latent_size = _resolve_latent_size(len(common_genes), config.latent_size)
+            adata_expr = _train_autoencoder_missing_gene(
+                adata_sc=adata_sc_work,
+                adata_st=adata_st_work,
+                common_genes=common_genes,
+                latent_size=latent_size,
+                lr=config.lr,
+                kmap=config.kmap,
+                kg=kg,
+                epochs=config.ae_epochs,
+                device=device,
+                verbose=config.verbose,
+            )
+            # Match the original NEST.py logic: cell clusters are learned from the
+            # original measured ST matrix, while outlier detection is performed on
+            # the mapped and gene-domain-corrected expression matrix.
+            adata_expr, _, graph = _create_cell_network(
+                adata_base=adata_st_work.copy(),
+                adata_expr=adata_expr,
+                k=k,
+                random_state=config.random_state,
+                scale_base=False,
+                scale_expr=False,
+                verbose=config.verbose,
+            )
+            latent_size_gae = _resolve_latent_size(adata_expr.n_vars, config.latent_size)
+            reconstructed = _train_graph_autoencoder(
+                adata_expr=adata_expr,
+                graph=graph,
+                latent_size=latent_size_gae,
+                lr=config.lr,
+                epochs=config.gae_epochs,
+                device=device,
+                observed_only_loss=False,
+                random_state=config.random_state,
+                verbose=config.verbose,
+            )
+            adata_out = adata_expr.copy()
+            adata_out.X = reconstructed.astype(np.float32)
+            for gene in common_genes:
+                if gene in adata_out.var_names and gene in adata_st_work.var_names:
+                    adata_out[:, gene].X = _as_numpy(adata_st_work[:, gene].X).copy()
+            return adata_out
+
+        if config.mode == "dropout":
+            genes = list(adata_st_work.var_names)
+            latent_size = _resolve_latent_size(len(genes), config.latent_size)
+            adata_expr = _train_autoencoder_dropout(
+                adata_st=adata_st_work,
+                genes=genes,
+                latent_size=latent_size,
+                lr=config.lr,
+                epochs=config.ae_epochs,
+                device=device,
+                verbose=config.verbose,
+            )
+            adata_expr, _, graph = _create_cell_network(
+                adata_base=adata_expr,
+                adata_expr=adata_expr,
+                k=k,
+                random_state=config.random_state,
+                scale_base=True,
+                scale_expr=True,
+                verbose=config.verbose,
+            )
+            reconstructed = _train_graph_autoencoder(
+                adata_expr=adata_expr,
+                graph=graph,
+                latent_size=latent_size,
+                lr=config.lr,
+                epochs=config.gae_epochs,
+                device=device,
+                observed_only_loss=True,
+                random_state=config.random_state,
+                verbose=config.verbose,
+            )
+            original = _as_numpy(adata_st_work.X).astype(np.float32)
+            if config.fill_zeros_only:
+                out_x = original.copy()
+                zero_mask = original == 0
+                out_x[zero_mask] = reconstructed[zero_mask]
+            else:
+                out_x = reconstructed
+            adata_out = adata_st_work.copy()
+            adata_out.X = out_x.astype(np.float32)
+            return adata_out
+
+        raise ValueError("config.mode must be either 'missing_gene' or 'dropout'.")
+
+
+def run_nest(
+    adata_st: ad.AnnData,
+    adata_sc: Optional[ad.AnnData] = None,
+    mode: NESTMode = "missing_gene",
+    **kwargs,
+) -> ad.AnnData:
+    """Convenience function for running NEST in one line."""
+    config = NESTConfig(mode=mode, **kwargs)
+    return NEST(config).fit_transform(adata_st=adata_st, adata_sc=adata_sc)
+
+
+def impute_missing_genes(
+    adata_st: ad.AnnData,
+    adata_sc: ad.AnnData,
+    **kwargs,
+) -> ad.AnnData:
+    """Impute unmeasured genes in ST data using a matched scRNA-seq reference."""
+    return run_nest(adata_st=adata_st, adata_sc=adata_sc, mode="missing_gene", **kwargs)
+
+
+def recover_dropout(
+    adata_st: ad.AnnData,
+    **kwargs,
+) -> ad.AnnData:
+    """Recover dropout-corrupted zeros in ST data without a scRNA-seq reference."""
+    return run_nest(adata_st=adata_st, mode="dropout", **kwargs)

nest_st-0.1.0/nest_st.egg-info/PKG-INFO

@@ -0,0 +1,109 @@
+Metadata-Version: 2.4
+Name: nest-st
+Version: 0.1.0
+Summary: NEST: Native Expression reconstruction for Spatial Transcriptomics
+Author-email: Tianyi Liu <1694414073@qq.com>
+License: MIT
+Project-URL: Homepage, https://github.com/lty10086/NEST
+Project-URL: Repository, https://github.com/lty10086/NEST
+Project-URL: Issues, https://github.com/lty10086/NEST/issues
+Keywords: spatial transcriptomics,single-cell,imputation,dropout,graph autoencoder
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.23
+Requires-Dist: pandas>=1.5
+Requires-Dist: scipy>=1.9
+Requires-Dist: anndata>=0.9
+Requires-Dist: scanpy>=1.9
+Requires-Dist: scikit-learn>=1.2
+Requires-Dist: matplotlib>=3.6
+Requires-Dist: torch>=2.0
+Requires-Dist: torch-geometric>=2.3
+
+# NEST
+
+NEST, Native Expression reconstruction for Spatial Transcriptomics, reconstructs biologically faithful spatial transcriptomes under native spatial constraints.
+
+It supports two tasks with one parameter:
+
+- `mode="missing_gene"`: imputes unmeasured genes in imaging-based spatial transcriptomics using a matched scRNA-seq reference.
+- `mode="dropout"`: recovers dropout-corrupted false zeros in sequencing-based spatial transcriptomics using ST data only.
+
+## Installation for development
+
+```bash
+git clone https://github.com/yourname/nest-st.git
+cd nest-st
+pip install -e .
+```
+
+For CUDA environments, install PyTorch and PyTorch Geometric using the wheel index recommended for your CUDA version before installing the package.
+
+## Python usage
+
+### Missing-gene imputation
+
+```python
+import scanpy as sc
+from nest_st import NEST, NESTConfig
+
+adata_st = sc.read_h5ad("data/st.h5ad")
+adata_sc = sc.read_h5ad("data/sc.h5ad")
+
+sc.pp.log1p(adata_st)
+sc.pp.log1p(adata_sc)
+
+config = NESTConfig(mode="missing_gene", ae_epochs=1000, gae_epochs=2000, device="cuda")
+adata_imputed = NEST(config).fit_transform(adata_st=adata_st, adata_sc=adata_sc)
+adata_imputed.write_h5ad("nest_missing_gene_result.h5ad")
+```
+
+### Dropout recovery
+
+```python
+import scanpy as sc
+from nest_st import NEST, NESTConfig
+
+adata_st = sc.read_h5ad("data/st.h5ad")
+sc.pp.log1p(adata_st)
+
+config = NESTConfig(mode="dropout", fill_zeros_only=True, ae_epochs=1000, gae_epochs=2000, device="cuda")
+adata_recovered = NEST(config).fit_transform(adata_st=adata_st)
+adata_recovered.write_h5ad("nest_dropout_result.h5ad")
+```
+
+## Command-line usage
+
+```bash
+python main.py --mode missing_gene --st data/st.h5ad --sc data/sc.h5ad --output result_gene.h5ad --log1p
+python main.py --mode dropout --st data/st.h5ad --output result_dropout.h5ad --log1p --fill-zeros-only
+```
+
+## Input requirements
+
+- Input files should be AnnData `.h5ad` files.
+- `adata_st.obsm["spatial"]` must contain spatial coordinates.
+- For `missing_gene` mode, `adata_sc` must contain genes to be transferred into the ST domain.
+- For `dropout` mode, no scRNA-seq reference is required.
+
+## Upload to PyPI
+
+```bash
+python -m pip install --upgrade build twine
+python -m build
+python -m twine upload dist/*
+```
+
+For TestPyPI:
+
+```bash
+python -m twine upload --repository testpypi dist/*
+```
+
+Update the package name, author information and GitHub URLs in `pyproject.toml` before publishing.

nest_st-0.1.0/nest_st.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+README.md
+pyproject.toml
+nest_st/__init__.py
+nest_st/core.py
+nest_st.egg-info/PKG-INFO
+nest_st.egg-info/SOURCES.txt
+nest_st.egg-info/dependency_links.txt
+nest_st.egg-info/requires.txt
+nest_st.egg-info/top_level.txt

nest_st-0.1.0/nest_st.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

nest_st-0.1.0/nest_st.egg-info/requires.txt

@@ -0,0 +1,9 @@
+numpy>=1.23
+pandas>=1.5
+scipy>=1.9
+anndata>=0.9
+scanpy>=1.9
+scikit-learn>=1.2
+matplotlib>=3.6
+torch>=2.0
+torch-geometric>=2.3

nest_st-0.1.0/nest_st.egg-info/top_level.txt

@@ -0,0 +1 @@
+nest_st

nest_st-0.1.0/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "nest-st"
+version = "0.1.0"
+description = "NEST: Native Expression reconstruction for Spatial Transcriptomics"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "MIT"}
+authors = [
+  {name = "Tianyi Liu", email = "1694414073@qq.com"}
+]
+keywords = ["spatial transcriptomics", "single-cell", "imputation", "dropout", "graph autoencoder"]
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+  "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+dependencies = [
+  "numpy>=1.23",
+  "pandas>=1.5",
+  "scipy>=1.9",
+  "anndata>=0.9",
+  "scanpy>=1.9",
+  "scikit-learn>=1.2",
+  "matplotlib>=3.6",
+  "torch>=2.0",
+  "torch-geometric>=2.3",
+]
+
+[project.urls]
+Homepage = "https://github.com/lty10086/NEST"
+Repository = "https://github.com/lty10086/NEST"
+Issues = "https://github.com/lty10086/NEST/issues"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["nest_st*"]

nest_st-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+