geneflowdecoder 0.1.0__tar.gz

geneflowdecoder-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 GeneFlowDecoder Contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

geneflowdecoder-0.1.0/PKG-INFO

@@ -0,0 +1,154 @@
+Metadata-Version: 2.4
+Name: geneflowdecoder
+Version: 0.1.0
+Summary: Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X.
+Author: GeneFlowDecoder Contributors
+License: MIT
+Project-URL: Homepage, https://github.com/ZengFLab/geneflowdecoder
+Project-URL: Documentation, https://github.com/ZengFLab/geneflowdecoder#readme
+Project-URL: Repository, https://github.com/ZengFLab/geneflowdecoder
+Project-URL: Issues, https://github.com/ZengFLab/geneflowdecoder/issues
+Keywords: single-cell,transcriptomics,decoder,flow-matching,gene-transformer,scvi,representation-learning
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1
+Requires-Dist: numpy>=1.24
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Dynamic: license-file
+
+# GeneFlowDecoder
+
+`z → X` 的细胞状态解码器：**Gene Transformer + Conditional Rectified Flow (CFM)**，
+把固定的细胞表征 `z` 解码成 log1p 归一化表达谱 `X`。
+
+> 设计目标：接口极简、可复现、单卡可跑 2 万+ 基因面板；`z` 默认冻结，
+> 不参与反传，也不被优化器更新。
+
+```python
+from geneflowdecoder import GeneFlowDecoder
+
+model = GeneFlowDecoder(n_genes=2048, z_dim=256, d=512, depth=8)
+
+model.fit(X_train, z_train,               # X: [N, G] log1p; z: [N, D]
+          X_val=X_val, z_val=z_val,       # 可选，开启早停
+          epochs=200, batch_size=64)
+
+X_hat = model(z_test)           # 采样（生成式，log1p 空间）
+X_hat = model(z_test, steps=8, deterministic=False)
+
+X_rec = model.predict(z_test)  # 确定性还原：NB 均值
+```
+
+---
+
+## 安装
+
+```bash
+# 从源码（开发模式）
+git clone https://github.com/your-org/geneflowdecoder.git
+cd geneflowdecoder
+pip install -e .[dev]
+
+# 或构建 wheel 后安装
+python -m build
+pip install dist/geneflowdecoder-0.1.0-py3-none-any.whl
+```
+
+依赖：`torch >= 2.1`、`numpy >= 1.24`。没有 GPU 也能在 CPU 上跑（只是慢）。
+
+---
+
+## API
+
+### `GeneFlowDecoder(...)`
+
+| 参数 | 默认 | 说明 |
+|---|---|---|
+| `n_genes` | — | 基因数，即输出维度 `G` |
+| `z_dim` | — | 细胞状态 `z` 的维度 |
+| `d` | 768 | Transformer 宽度 |
+| `depth` | 16 | Transformer 层数 |
+| `n_heads` | 12 | 注意力头数 |
+| `chunk_size` | 2048 | 基因轴分块，控制显存峰值 |
+| `gradient_checkpoint` | `True` | 激活检查点，省显存 |
+| `cfg` | `FlowConfig()` | flow / loss 超参 |
+| `freeze_z` | `True` | **`z` 全程 `detach`，只读条件** |
+| `device` | `"auto"` | `"auto"` / `"cuda"` / `"cpu"` |
+
+### `model.fit(X, z, *, X_val=None, z_val=None, epochs=200, batch_size=64, lr=2e-4, weight_decay=0.02, patience=20, num_workers=0, verbose=1) -> self`
+
+- `X` / `X_val`：`[N, G]` log1p 表达，numpy 或 torch 都行
+- `z` / `z_val`：`[N, z_dim]`；`freeze_z=True` 时绝不会被修改或反传
+- 给了验证集 → 早停 + 自动回滚最优权重
+
+### `model(z, *, steps=4, method="midpoint", deterministic=False) -> Tensor`
+
+Flow 采样，`steps=4` 通常够用。
+
+### `model.predict(z) -> Tensor`
+
+确定性输出 `E[X]` 的 log1p（NB mean）。转伪计数：`(log_mu.exp() - 1).clamp(min=0)`。
+
+### `model.save(path)` / `GeneFlowDecoder.load(path, device="auto")`
+
+---
+
+## `FlowConfig`
+
+```python
+from geneflowdecoder import FlowConfig
+
+cfg = FlowConfig(
+    sigma_min=1e-4,     # t 的截断，防 t≈0 数值爆炸
+    nb_weight=1.0,      # NB 似然辅助损失权重
+    corr_weight=0.1,    # cell-wise Pearson 相关性正则
+)
+model = GeneFlowDecoder(n_genes=G, z_dim=D, cfg=cfg)
+```
+
+---
+
+## 发布
+
+```bash
+python -m build          # 生成 sdist + wheel 到 dist/
+twine check dist/*       # 校验元数据
+twine upload dist/*      # 发布到 PyPI（需账号 token）
+```
+
+---
+
+## 目录结构
+
+```
+geneflowdecoder/
+├── pyproject.toml
+├── README.md
+├── LICENSE
+├── src/geneflowdecoder/
+│   ├── __init__.py        # 公开 API
+│   ├── _core.py           # 模型 + fit/forward/predict
+│   └── _version.py
+└── tests/
+    └── test_api.py
+```
+
+## License
+
+MIT

geneflowdecoder-0.1.0/README.md

@@ -0,0 +1,120 @@
+# GeneFlowDecoder
+
+`z → X` 的细胞状态解码器：**Gene Transformer + Conditional Rectified Flow (CFM)**，
+把固定的细胞表征 `z` 解码成 log1p 归一化表达谱 `X`。
+
+> 设计目标：接口极简、可复现、单卡可跑 2 万+ 基因面板；`z` 默认冻结，
+> 不参与反传，也不被优化器更新。
+
+```python
+from geneflowdecoder import GeneFlowDecoder
+
+model = GeneFlowDecoder(n_genes=2048, z_dim=256, d=512, depth=8)
+
+model.fit(X_train, z_train,               # X: [N, G] log1p; z: [N, D]
+          X_val=X_val, z_val=z_val,       # 可选，开启早停
+          epochs=200, batch_size=64)
+
+X_hat = model(z_test)           # 采样（生成式，log1p 空间）
+X_hat = model(z_test, steps=8, deterministic=False)
+
+X_rec = model.predict(z_test)  # 确定性还原：NB 均值
+```
+
+---
+
+## 安装
+
+```bash
+# 从源码（开发模式）
+git clone https://github.com/your-org/geneflowdecoder.git
+cd geneflowdecoder
+pip install -e .[dev]
+
+# 或构建 wheel 后安装
+python -m build
+pip install dist/geneflowdecoder-0.1.0-py3-none-any.whl
+```
+
+依赖：`torch >= 2.1`、`numpy >= 1.24`。没有 GPU 也能在 CPU 上跑（只是慢）。
+
+---
+
+## API
+
+### `GeneFlowDecoder(...)`
+
+| 参数 | 默认 | 说明 |
+|---|---|---|
+| `n_genes` | — | 基因数，即输出维度 `G` |
+| `z_dim` | — | 细胞状态 `z` 的维度 |
+| `d` | 768 | Transformer 宽度 |
+| `depth` | 16 | Transformer 层数 |
+| `n_heads` | 12 | 注意力头数 |
+| `chunk_size` | 2048 | 基因轴分块，控制显存峰值 |
+| `gradient_checkpoint` | `True` | 激活检查点，省显存 |
+| `cfg` | `FlowConfig()` | flow / loss 超参 |
+| `freeze_z` | `True` | **`z` 全程 `detach`，只读条件** |
+| `device` | `"auto"` | `"auto"` / `"cuda"` / `"cpu"` |
+
+### `model.fit(X, z, *, X_val=None, z_val=None, epochs=200, batch_size=64, lr=2e-4, weight_decay=0.02, patience=20, num_workers=0, verbose=1) -> self`
+
+- `X` / `X_val`：`[N, G]` log1p 表达，numpy 或 torch 都行
+- `z` / `z_val`：`[N, z_dim]`；`freeze_z=True` 时绝不会被修改或反传
+- 给了验证集 → 早停 + 自动回滚最优权重
+
+### `model(z, *, steps=4, method="midpoint", deterministic=False) -> Tensor`
+
+Flow 采样，`steps=4` 通常够用。
+
+### `model.predict(z) -> Tensor`
+
+确定性输出 `E[X]` 的 log1p（NB mean）。转伪计数：`(log_mu.exp() - 1).clamp(min=0)`。
+
+### `model.save(path)` / `GeneFlowDecoder.load(path, device="auto")`
+
+---
+
+## `FlowConfig`
+
+```python
+from geneflowdecoder import FlowConfig
+
+cfg = FlowConfig(
+    sigma_min=1e-4,     # t 的截断，防 t≈0 数值爆炸
+    nb_weight=1.0,      # NB 似然辅助损失权重
+    corr_weight=0.1,    # cell-wise Pearson 相关性正则
+)
+model = GeneFlowDecoder(n_genes=G, z_dim=D, cfg=cfg)
+```
+
+---
+
+## 发布
+
+```bash
+python -m build          # 生成 sdist + wheel 到 dist/
+twine check dist/*       # 校验元数据
+twine upload dist/*      # 发布到 PyPI（需账号 token）
+```
+
+---
+
+## 目录结构
+
+```
+geneflowdecoder/
+├── pyproject.toml
+├── README.md
+├── LICENSE
+├── src/geneflowdecoder/
+│   ├── __init__.py        # 公开 API
+│   ├── _core.py           # 模型 + fit/forward/predict
+│   └── _version.py
+└── tests/
+    └── test_api.py
+```
+
+## License
+
+MIT

geneflowdecoder-0.1.0/pyproject.toml

@@ -0,0 +1,64 @@
+[project]
+name = "geneflowdecoder"
+version = "0.1.0"
+description = "Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X."
+readme = "README.md"
+requires-python = ">=3.9"
+license = {text = "MIT"}
+authors = [
+    {name = "GeneFlowDecoder Contributors"}
+]
+keywords = [
+    "single-cell", "transcriptomics", "decoder", "flow-matching",
+    "gene-transformer", "scvi", "representation-learning",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "torch >= 2.1",
+    "numpy >= 1.24",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest >= 7",
+    "build",
+    "twine",
+]
+test = ["pytest >= 7"]
+
+[project.urls]
+Homepage = "https://github.com/ZengFLab/geneflowdecoder"
+Documentation = "https://github.com/ZengFLab/geneflowdecoder#readme"
+Repository = "https://github.com/ZengFLab/geneflowdecoder"
+Issues = "https://github.com/ZengFLab/geneflowdecoder/issues"
+
+[build-system]
+requires = ["setuptools >= 61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["geneflowdecoder*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q --import-mode=importlib"
+
+[tool.ruff]
+line-length = 88
+target-version = "py39"
+
+[tool.mypy]
+python_version = "3.9"
+ignore_missing_imports = true

geneflowdecoder-0.1.0/setup.cfg

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

geneflowdecoder-0.1.0/src/geneflowdecoder/__init__.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+"""GeneFlowDecoder: z -> X decoder (Gene Transformer + Rectified Flow)."""
+
+from ._core import GeneFlowDecoder, FlowConfig
+from ._version import __version__
+
+__all__ = [
+    "GeneFlowDecoder",
+    "FlowConfig",
+    "__version__",
+]
+
+__doc__ = GeneFlowDecoder.__doc__ or ""

geneflowdecoder-0.1.0/src/geneflowdecoder/_core.py

@@ -0,0 +1,567 @@
+# -*- coding: utf-8 -*-
+"""
+GeneFlowDecoder
+===============
+
+z (fixed cell-state representation) --> X (log-normalized expression, log1p space)
+
+Public API
+----------
+    model = GeneFlowDecoder(n_genes=2048, z_dim=256)
+    model.fit(X, z)                      # train
+    X_hat = model(z)                     # sample (generative)
+    X_hat = model.predict(z)             # deterministic (NB mean)
+
+This module is framework-agnostic beyond ``torch`` and ``numpy``.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Optional, Tuple, Union
+
+import numpy as np
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    _no_grad = torch.no_grad
+except Exception:  # pragma: no cover - soft dep for import-only environments
+    torch = None  # type: ignore
+    nn = object   # type: ignore
+    F = object    # type: ignore
+
+    def _no_grad(fn):  # identity decorator when torch is absent
+        return fn
+
+    Tensor = object  # placeholder for type annotations only
+
+try:
+    from torch import Tensor as _TorchTensor  # noqa: F401
+    Tensor = _TorchTensor  # type: ignore
+except Exception:
+    pass
+
+Array = Union[np.ndarray, torch.Tensor]
+
+
+def _as_tensor(x: Array, device: torch.device) -> torch.Tensor:
+    if isinstance(x, np.ndarray):
+        x = torch.from_numpy(x)
+    return x.to(device, non_blocking=True)
+
+
+# ----------------------------------------------------------------------
+# 1. Conditional modulation & transformer blocks
+# ----------------------------------------------------------------------
+
+class _AdaLNModulation(nn.Module):
+    def __init__(self, d: int, cond_dim: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
+        self.proj = nn.Sequential(
+            nn.Linear(cond_dim, d), nn.SiLU(), nn.Linear(d, 3 * d)
+        )
+        nn.init.zeros_(self.proj[-1].weight)
+        nn.init.zeros_(self.proj[-1].bias)
+
+    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        g, b, gate = self.proj(z).chunk(3, dim=-1)
+        g, b, gate = g.unsqueeze(1), b.unsqueeze(1), gate.unsqueeze(1).sigmoid()
+        return gate * (self.norm(x) * (1 + g) + b)
+
+
+class _CondBlock(nn.Module):
+    def __init__(self, d: int, cond_dim: int, n_heads: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.attn_norm = _AdaLNModulation(d, cond_dim)
+        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
+        self.mlp_norm = _AdaLNModulation(d, cond_dim)
+        h = int(d * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(d, h), nn.GELU(), nn.Linear(h, d))
+
+    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        h = self.attn_norm(x, z)
+        h, _ = self.attn(h, h, h, need_weights=False)
+        x = x + h
+        return x + self.mlp(self.mlp_norm(x, z))
+
+
+class _GeneTransformer(nn.Module):
+    def __init__(
+        self,
+        n_genes: int,
+        d: int,
+        depth: int,
+        n_heads: int,
+        cond_dim: int,
+        chunk_size: int,
+        ckpt: bool,
+    ):
+        super().__init__()
+        self.d = d
+        self.n_genes = n_genes
+        self.chunk_size = chunk_size
+        self.gene_embed = nn.Embedding(n_genes, d)
+        nn.init.normal_(self.gene_embed.weight, std=0.02)
+        self.pos_embed = nn.Parameter(torch.zeros(1, n_genes, d))
+        nn.init.normal_(self.pos_embed, std=0.02)
+
+        self.blocks = nn.ModuleList(
+            [_CondBlock(d, cond_dim, n_heads) for _ in range(depth)]
+        )
+        self.head_norm = nn.LayerNorm(d, eps=1e-6)
+        self._ckpt = ckpt
+
+    def _run(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        for blk in self.blocks:
+            if self._ckpt and self.training:
+                x = torch.utils.checkpoint.checkpoint(
+                    blk, x, z, use_reentrant=False
+                )
+            else:
+                x = blk(x, z)
+        return self.head_norm(x)
+
+    def forward(self, gene_ids: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        B, G = gene_ids.shape[0], gene_ids.shape[-1]
+        if G <= self.chunk_size:
+            h = self.gene_embed(gene_ids) + self.pos_embed[:, :G, :]
+            return self._run(h, z)
+
+        outs: list[torch.Tensor] = []
+        s = 0
+        while s < G:
+            e = min(s + self.chunk_size, G)
+            h = self.gene_embed(gene_ids[..., s:e]) + self.pos_embed[:, s:e, :]
+            outs.append(self._run(h, z))
+            s = e
+        return torch.cat(outs, dim=1)
+
+
+# ----------------------------------------------------------------------
+# 2. Negative-binomial output head
+# ----------------------------------------------------------------------
+
+class _NBHead(nn.Module):
+    def __init__(self, d: int, n_genes: int):
+        super().__init__()
+        self.mu_proj = nn.Linear(d, n_genes)
+        self.theta = nn.Parameter(torch.zeros(n_genes))
+        nn.init.zeros_(self.mu_proj.weight)
+        nn.init.zeros_(self.mu_proj.bias)
+
+    def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self.mu_proj(h), F.softplus(self.theta) + 1e-3
+
+    @staticmethod
+    def nll(
+        log_mu: torch.Tensor,
+        theta: torch.Tensor,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        mu = log_mu.exp().clamp(min=1e-6)
+        th = theta.unsqueeze(0).expand_as(x)
+        ll = (
+            torch.lgamma(x + th)
+            - torch.lgamma(th)
+            - torch.lgamma(x + 1)
+            + th * (th.log() - (th + mu).log())
+            + x * (mu.log() - (th + mu).log())
+        )
+        if mask is not None:
+            ll = ll * mask
+            return -ll.sum() / mask.sum().clamp(min=1.0)
+        return -ll.mean()
+
+
+@dataclass
+class FlowConfig:
+    """Hyper-parameters of the conditional flow matching objective."""
+
+    sigma_min: float = 1e-4
+    nb_weight: float = 1.0
+    corr_weight: float = 0.1
+
+
+# ----------------------------------------------------------------------
+# 3. GeneFlowDecoder
+# ----------------------------------------------------------------------
+
+class GeneFlowDecoder(nn.Module):
+    """
+    Conditional Rectified-Flow decoder from a **fixed** cell-state vector *z*
+    to a log1p expression profile *X*.
+
+    Parameters
+    ----------
+    n_genes : int
+        Number of genes (output dimension).
+    z_dim : int
+        Dimensionality of the cell-state representation.
+    d, depth, n_heads : int
+        Transformer width, number of layers, number of attention heads.
+    chunk_size : int
+        Chunk size for the gene axis (enables >10k gene panels on one GPU).
+    gradient_checkpoint : bool
+        Enable activation checkpointing inside the transformer.
+    cfg : FlowConfig | None
+        Flow / loss hyper-parameters.
+    freeze_z : bool, default ``True``
+        When ``True`` (recommended), *z* is detached everywhere: it is treated
+        as a fixed, non-trainable conditioning signal and is never updated.
+    device : ``"auto"`` | ``str`` | ``torch.device``
+        Where to place the model. ``"auto"`` picks CUDA if available.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from geneflowdecoder import GeneFlowDecoder
+    >>> G, D = 2048, 128
+    >>> model = GeneFlowDecoder(n_genes=G, z_dim=D, d=512, depth=8)
+    >>> X = np.random.randn(1024, G).astype("float32")
+    >>> z = np.random.randn(1024, D).astype("float32")
+    >>> model.fit(X, z, epochs=10, batch_size=64)
+    >>> X_hat = model(z[:4])           # sampled, log1p space
+    >>> X_rec  = model.predict(z[:4])  # deterministic NB mean
+    """
+
+    def __init__(
+        self,
+        n_genes: int,
+        z_dim: int,
+        *,
+        d: int = 768,
+        depth: int = 16,
+        n_heads: int = 12,
+        chunk_size: int = 2048,
+        gradient_checkpoint: bool = True,
+        cfg: Optional[FlowConfig] = None,
+        freeze_z: bool = True,
+        device: str = "auto",
+    ):
+        super().__init__()
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.n_genes = int(n_genes)
+        self.z_dim = int(z_dim)
+        self.freeze_z = bool(freeze_z)
+        self._cfg = cfg or FlowConfig()
+
+        self.backbone = _GeneTransformer(
+            n_genes, d, depth, n_heads, z_dim, chunk_size, gradient_checkpoint
+        )
+        self.nb_head = _NBHead(d, n_genes)
+        self.vel_proj = nn.Linear(d, n_genes)
+        self.time_mlp = nn.Sequential(
+            nn.Linear(128, d), nn.SiLU(), nn.Linear(d, d)
+        )
+        self._gene_ids = nn.Parameter(
+            torch.arange(self.n_genes), requires_grad=False
+        )
+        self.to(self.device)
+
+    # ----------------------------------------------------------------
+    # internals
+    # ----------------------------------------------------------------
+    def _sincos(self, t: torch.Tensor) -> torch.Tensor:
+        half = 64
+        freqs = torch.exp(
+            -math.log(10000.0) * torch.arange(half, device=t.device) / (half - 1)
+        )
+        a = t[:, None] * freqs[None, :]
+        e = torch.cat([a.sin(), a.cos()], dim=-1)
+        if e.shape[-1] < 128:
+            e = F.pad(e, (0, 128 - e.shape[-1]))
+        return e
+
+    def _maybe_detach(
+        self, z: torch.Tensor, *, snapshot: bool = False
+    ) -> torch.Tensor:
+        """冻结 z 时截断计算图；``snapshot=True`` 再额外拷一份独立副本。
+
+        snapshot 只在 fit 入口用：把当轮 z 固化下来，之后外部 encoder
+        若继续 backward / 原地更新，训练集 z 的数值也不会跟着变。
+        单次推理（forward / predict）不开 snapshot，省一次显存拷贝。
+        """
+        if not self.freeze_z:
+            return z
+        z = z.detach()
+        return z.clone() if snapshot else z
+
+    def _velocity(
+        self, x_t: torch.Tensor, t: torch.Tensor, z: torch.Tensor
+    ) -> torch.Tensor:
+        z = self._maybe_detach(z)
+        B = x_t.shape[0]
+        t_emb = self.time_mlp(self._sincos(t)).unsqueeze(1)
+        gids = self._gene_ids.unsqueeze(0).expand(B, -1)
+        h = self.backbone(gids, z) + t_emb
+        return self.vel_proj(h)
+
+    # ----------------------------------------------------------------
+    # objective
+    # ----------------------------------------------------------------
+    def _loss(
+        self,
+        x: torch.Tensor,
+        z: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        z = self._maybe_detach(z)
+        B = x.shape[0]
+        t = torch.rand(B, device=self.device).clamp(
+            self._cfg.sigma_min, 1.0 - self._cfg.sigma_min
+        )
+        eps = torch.randn_like(x)
+        x_t = x.lerp(eps, t.unsqueeze(1))
+
+        v_pred = self._velocity(x_t, t, z)
+        v_target = eps - x  # rectified-flow target
+        mse = F.mse_loss(v_pred, v_target, reduction="none")
+        if mask is not None:
+            mse = mse * mask
+            mse = mse.sum() / mask.sum().clamp(min=1.0)
+        else:
+            mse = mse.mean()
+
+        x0_hat = x_t.lerp(v_pred.detach(), t.unsqueeze(1))
+
+        nb = torch.tensor(0.0, device=self.device)
+        if self._cfg.nb_weight > 0:
+            gids = self._gene_ids.unsqueeze(0).expand(B, -1)
+            log_mu, theta = self.nb_head(self.backbone(gids, z))
+            nb = self.nb_head.nll(log_mu, theta, x, mask)
+
+        corr = torch.tensor(0.0, device=self.device)
+        if self._cfg.corr_weight > 0 and mask is None:
+            corr = self._corr(x0_hat, x)
+
+        return (
+            mse
+            + self._cfg.nb_weight * nb
+            + self._cfg.corr_weight * corr
+        )
+
+    @staticmethod
+    def _corr(p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
+        p = p - p.mean(1, keepdim=True)
+        q = q - q.mean(1, keepdim=True)
+        denom = p.norm(dim=1) * q.norm(dim=1) + 1e-8
+        return (1 - (p * q).sum(dim=1) / denom).mean()
+
+    # ----------------------------------------------------------------
+    # fit
+    # ----------------------------------------------------------------
+    def fit(
+        self,
+        X: Array,
+        z: Array,
+        *,
+        X_val: Optional[Array] = None,
+        z_val: Optional[Array] = None,
+        epochs: int = 200,
+        batch_size: int = 64,
+        lr: float = 2e-4,
+        weight_decay: float = 0.02,
+        patience: int = 20,
+        num_workers: int = 0,
+        verbose: int = 1,
+    ) -> "GeneFlowDecoder":
+        """
+        Train the decoder.
+
+        Parameters
+        ----------
+        X : array, shape ``(N, n_genes)``
+            Log1p-normalized expression matrix.
+        z : array, shape ``(N, z_dim)``
+            Cell-state vectors. **Never updated** when ``freeze_z=True``.
+        X_val, z_val : arrays, optional
+            Validation split. If given, early stopping is enabled and the
+            best weights (lowest validation loss) are restored at the end.
+        epochs, batch_size, lr, weight_decay, patience
+            Standard optimisation knobs.
+        num_workers : int
+            Workers for the training DataLoader.
+        verbose : int
+            0 = silent, 1 = epoch logs.
+
+        Returns
+        -------
+        self
+        """
+        X = _as_tensor(X, self.device)
+        z = self._maybe_detach(_as_tensor(z, self.device), snapshot=True)
+
+        ds = torch.utils.data.TensorDataset(X, z)
+        dl = torch.utils.data.DataLoader(
+            ds,
+            batch_size=int(batch_size),
+            shuffle=True,
+            num_workers=int(num_workers),
+            pin_memory=True,
+        )
+
+        has_val = False
+        if X_val is not None and z_val is not None:
+            Xv = _as_tensor(X_val, self.device)
+            zv = self._maybe_detach(_as_tensor(z_val, self.device), snapshot=True)
+            has_val = True
+
+        # Only decoder-owned parameters are optimised.
+        params = [p for p in self.parameters() if p.requires_grad]
+        opt = torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay)
+        sched = torch.optim.lr_scheduler.CosineAnnealingLR(
+            opt, T_max=int(epochs)
+        )
+
+        best_val = float("inf")
+        best_state: Optional[dict[str, torch.Tensor]] = None
+        no_improve = 0
+
+        for ep in range(1, int(epochs) + 1):
+            self.train()
+            tot = 0.0
+            for xb, zb in dl:
+                opt.zero_grad()
+                loss = self._loss(xb, zb)
+                loss.backward()
+                nn.utils.clip_grad_norm_(params, 1.0)
+                opt.step()
+                tot += loss.item() * xb.shape[0]
+            sched.step()
+
+            if verbose:
+                msg = (
+                    f"[fit] epoch {ep:3d}/{epochs} "
+                    f"| train={tot / len(ds):.4f}"
+                )
+                if has_val:
+                    self.eval()
+                    with torch.no_grad():
+                        vl = self._loss(Xv, zv).item()
+                    msg += f" | val={vl:.4f}"
+                    if vl < best_val - 1e-5:
+                        best_val, no_improve = vl, 0
+                        best_state = {
+                            k: v.detach().cpu().clone()
+                            for k, v in self.state_dict().items()
+                        }
+                    else:
+                        no_improve += 1
+                    if no_improve >= int(patience):
+                        if verbose:
+                            print(f"[fit] early stop at epoch {ep}")
+                        break
+                print(msg)
+
+        if best_state is not None:
+            self.load_state_dict(best_state)
+            if verbose:
+                print(f"[fit] restored best val={best_val:.4f}")
+        return self
+
+    # ----------------------------------------------------------------
+    # inference
+    # ----------------------------------------------------------------
+    @torch.no_grad()
+    def forward(
+        self,
+        z: Array,
+        *,
+        steps: int = 4,
+        method: str = "midpoint",
+        deterministic: bool = False,
+    ) -> torch.Tensor:
+        """
+        Sample an expression profile from *z* (log1p space).
+
+        Parameters
+        ----------
+        z : array, shape ``(B, z_dim)``
+        steps : int
+            Number of flow-integration steps (4 is usually enough).
+        method : ``"midpoint"`` | ``"euler"``
+        deterministic : bool
+            If ``True``, returns the deterministic NB mean
+            (equivalent to ``self.predict(z)``).
+        """
+        z = self._maybe_detach(_as_tensor(z, self.device))
+        if deterministic:
+            return self.predict(z)
+
+        self.eval()
+        B = z.shape[0]
+        x = torch.randn(B, self.n_genes, device=self.device)
+        ts = torch.linspace(0.0, 1.0, int(steps) + 1, device=self.device)
+        for i in range(int(steps)):
+            t0, t1 = ts[i], ts[i + 1]
+            tmid = (t0 + t1) / 2
+            dt = (t1 - t0).item()
+            v = self._velocity(x, torch.full((B,), t0, device=self.device), z)
+            if method == "midpoint":
+                xm = x + 0.5 * dt * v
+                v = self._velocity(
+                    xm, torch.full((B,), tmid, device=self.device), z
+                )
+            x = x + dt * v
+        return x
+
+    @torch.no_grad()
+    def predict(self, z: Array) -> torch.Tensor:
+        """
+        Deterministic reconstruction: the NB mean in log1p space.
+
+        Convert to pseudo-counts with ``(log_mu.exp() - 1).clamp(min=0)``.
+        """
+        z = self._maybe_detach(_as_tensor(z, self.device))
+        self.eval()
+        B = z.shape[0]
+        gids = self._gene_ids.unsqueeze(0).expand(B, -1)
+        log_mu, _ = self.nb_head(self.backbone(gids, z))
+        return log_mu
+
+    # ----------------------------------------------------------------
+    # persistence
+    # ----------------------------------------------------------------
+    def save(self, path: str) -> None:
+        """Save model weights + config to ``path``."""
+        torch.save(
+            {
+                "state": self.state_dict(),
+                "cfg": self._cfg,
+                "freeze_z": self.freeze_z,
+                "n_genes": self.n_genes,
+                "z_dim": self.z_dim,
+            },
+            path,
+        )
+
+    @classmethod
+    def load(
+        cls,
+        path: str,
+        *,
+        device: str = "auto",
+        **init_kwargs,
+    ) -> "GeneFlowDecoder":
+        """
+        Load a saved model.
+
+        ``n_genes`` / ``z_dim`` / ``freeze_z`` are read from the checkpoint
+        unless explicitly overridden in ``init_kwargs``.
+        """
+        ckpt = torch.load(path, map_location="cpu")
+        init_kwargs.setdefault("n_genes", ckpt.get("n_genes"))
+        init_kwargs.setdefault("z_dim", ckpt.get("z_dim"))
+        init_kwargs.setdefault("freeze_z", ckpt.get("freeze_z", True))
+        init_kwargs["device"] = device
+        model = cls(**init_kwargs)
+        model._cfg = ckpt.get("cfg", model._cfg)
+        model.load_state_dict(ckpt["state"])
+        return model

geneflowdecoder-0.1.0/src/geneflowdecoder/_version.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+__version__ = "0.1.0"

geneflowdecoder-0.1.0/src/geneflowdecoder.egg-info/PKG-INFO

@@ -0,0 +1,154 @@
+Metadata-Version: 2.4
+Name: geneflowdecoder
+Version: 0.1.0
+Summary: Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X.
+Author: GeneFlowDecoder Contributors
+License: MIT
+Project-URL: Homepage, https://github.com/ZengFLab/geneflowdecoder
+Project-URL: Documentation, https://github.com/ZengFLab/geneflowdecoder#readme
+Project-URL: Repository, https://github.com/ZengFLab/geneflowdecoder
+Project-URL: Issues, https://github.com/ZengFLab/geneflowdecoder/issues
+Keywords: single-cell,transcriptomics,decoder,flow-matching,gene-transformer,scvi,representation-learning
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1
+Requires-Dist: numpy>=1.24
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Dynamic: license-file
+
+# GeneFlowDecoder
+
+`z → X` 的细胞状态解码器：**Gene Transformer + Conditional Rectified Flow (CFM)**，
+把固定的细胞表征 `z` 解码成 log1p 归一化表达谱 `X`。
+
+> 设计目标：接口极简、可复现、单卡可跑 2 万+ 基因面板；`z` 默认冻结，
+> 不参与反传，也不被优化器更新。
+
+```python
+from geneflowdecoder import GeneFlowDecoder
+
+model = GeneFlowDecoder(n_genes=2048, z_dim=256, d=512, depth=8)
+
+model.fit(X_train, z_train,               # X: [N, G] log1p; z: [N, D]
+          X_val=X_val, z_val=z_val,       # 可选，开启早停
+          epochs=200, batch_size=64)
+
+X_hat = model(z_test)           # 采样（生成式，log1p 空间）
+X_hat = model(z_test, steps=8, deterministic=False)
+
+X_rec = model.predict(z_test)  # 确定性还原：NB 均值
+```
+
+---
+
+## 安装
+
+```bash
+# 从源码（开发模式）
+git clone https://github.com/your-org/geneflowdecoder.git
+cd geneflowdecoder
+pip install -e .[dev]
+
+# 或构建 wheel 后安装
+python -m build
+pip install dist/geneflowdecoder-0.1.0-py3-none-any.whl
+```
+
+依赖：`torch >= 2.1`、`numpy >= 1.24`。没有 GPU 也能在 CPU 上跑（只是慢）。
+
+---
+
+## API
+
+### `GeneFlowDecoder(...)`
+
+| 参数 | 默认 | 说明 |
+|---|---|---|
+| `n_genes` | — | 基因数，即输出维度 `G` |
+| `z_dim` | — | 细胞状态 `z` 的维度 |
+| `d` | 768 | Transformer 宽度 |
+| `depth` | 16 | Transformer 层数 |
+| `n_heads` | 12 | 注意力头数 |
+| `chunk_size` | 2048 | 基因轴分块，控制显存峰值 |
+| `gradient_checkpoint` | `True` | 激活检查点，省显存 |
+| `cfg` | `FlowConfig()` | flow / loss 超参 |
+| `freeze_z` | `True` | **`z` 全程 `detach`，只读条件** |
+| `device` | `"auto"` | `"auto"` / `"cuda"` / `"cpu"` |
+
+### `model.fit(X, z, *, X_val=None, z_val=None, epochs=200, batch_size=64, lr=2e-4, weight_decay=0.02, patience=20, num_workers=0, verbose=1) -> self`
+
+- `X` / `X_val`：`[N, G]` log1p 表达，numpy 或 torch 都行
+- `z` / `z_val`：`[N, z_dim]`；`freeze_z=True` 时绝不会被修改或反传
+- 给了验证集 → 早停 + 自动回滚最优权重
+
+### `model(z, *, steps=4, method="midpoint", deterministic=False) -> Tensor`
+
+Flow 采样，`steps=4` 通常够用。
+
+### `model.predict(z) -> Tensor`
+
+确定性输出 `E[X]` 的 log1p（NB mean）。转伪计数：`(log_mu.exp() - 1).clamp(min=0)`。
+
+### `model.save(path)` / `GeneFlowDecoder.load(path, device="auto")`
+
+---
+
+## `FlowConfig`
+
+```python
+from geneflowdecoder import FlowConfig
+
+cfg = FlowConfig(
+    sigma_min=1e-4,     # t 的截断，防 t≈0 数值爆炸
+    nb_weight=1.0,      # NB 似然辅助损失权重
+    corr_weight=0.1,    # cell-wise Pearson 相关性正则
+)
+model = GeneFlowDecoder(n_genes=G, z_dim=D, cfg=cfg)
+```
+
+---
+
+## 发布
+
+```bash
+python -m build          # 生成 sdist + wheel 到 dist/
+twine check dist/*       # 校验元数据
+twine upload dist/*      # 发布到 PyPI（需账号 token）
+```
+
+---
+
+## 目录结构
+
+```
+geneflowdecoder/
+├── pyproject.toml
+├── README.md
+├── LICENSE
+├── src/geneflowdecoder/
+│   ├── __init__.py        # 公开 API
+│   ├── _core.py           # 模型 + fit/forward/predict
+│   └── _version.py
+└── tests/
+    └── test_api.py
+```
+
+## License
+
+MIT

geneflowdecoder-0.1.0/src/geneflowdecoder.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+LICENSE
+README.md
+pyproject.toml
+src/geneflowdecoder/__init__.py
+src/geneflowdecoder/_core.py
+src/geneflowdecoder/_version.py
+src/geneflowdecoder.egg-info/PKG-INFO
+src/geneflowdecoder.egg-info/SOURCES.txt
+src/geneflowdecoder.egg-info/dependency_links.txt
+src/geneflowdecoder.egg-info/requires.txt
+src/geneflowdecoder.egg-info/top_level.txt
+tests/test_api.py

geneflowdecoder-0.1.0/src/geneflowdecoder.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

geneflowdecoder-0.1.0/src/geneflowdecoder.egg-info/requires.txt

@@ -0,0 +1,10 @@
+torch>=2.1
+numpy>=1.24
+
+[dev]
+pytest>=7
+build
+twine
+
+[test]
+pytest>=7

geneflowdecoder-0.1.0/src/geneflowdecoder.egg-info/top_level.txt

@@ -0,0 +1 @@
+geneflowdecoder

geneflowdecoder-0.1.0/tests/test_api.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+"""Smoke tests for the public API of geneflowdecoder."""
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch", reason="torch is an optional test dependency")
+
+from geneflowdecoder import GeneFlowDecoder, FlowConfig
+
+
+def _make(n_genes: int = 256, n_cells: int = 32, z_dim: int = 32):
+    rng = np.random.default_rng(0)
+    X = rng.normal(0, 1, (n_cells, n_genes)).astype(np.float32)
+    z = rng.normal(0, 1, (n_cells, z_dim)).astype(np.float32)
+    return X, z
+
+
+def test_import_and_init():
+    model = GeneFlowDecoder(n_genes=128, z_dim=16, d=64, depth=2, freeze_z=True)
+    assert model.n_genes == 128
+    assert model.z_dim == 16
+    assert model.freeze_z is True
+
+
+def test_fit_runs_and_z_unchanged():
+    X, z = _make()
+    z_before = z.copy()
+
+    model = GeneFlowDecoder(
+        n_genes=X.shape[1], z_dim=z.shape[1], d=64, depth=2,
+        freeze_z=True, device="cpu",
+    )
+    model.fit(X, z, epochs=3, batch_size=16, verbose=0)
+
+    # z is not a model parameter and must be numerically untouched
+    assert np.array_equal(z, z_before)
+
+
+def test_forward_and_predict_shapes():
+    X, z = _make(n_genes=256, n_cells=16, z_dim=32)
+    model = GeneFlowDecoder(
+        n_genes=X.shape[1], z_dim=z.shape[1], d=64, depth=2,
+        freeze_z=True, device="cpu",
+    )
+    model.fit(X, z, epochs=3, batch_size=8, verbose=0)
+
+    samp = model(z[:5])
+    assert samp.shape == (5, X.shape[1])
+    assert torch.isfinite(samp).all()
+
+    rec = model.predict(z[:5])
+    assert rec.shape == (5, X.shape[1])
+    assert torch.isfinite(rec).all()
+
+
+def test_forward_deterministic_flag():
+    X, z = _make(n_genes=128, n_cells=8, z_dim=16)
+    model = GeneFlowDecoder(
+        n_genes=X.shape[1], z_dim=z.shape[1], d=64, depth=2,
+        freeze_z=True, device="cpu",
+    )
+    model.fit(X, z, epochs=2, batch_size=8, verbose=0)
+
+    a = model(z, deterministic=True)
+    b = model.predict(z)
+    assert torch.allclose(a, b)
+
+
+def test_save_load_roundtrip(tmp_path):
+    X, z = _make(n_genes=128, n_cells=8, z_dim=16)
+    model = GeneFlowDecoder(
+        n_genes=X.shape[1], z_dim=z.shape[1], d=64, depth=2,
+        freeze_z=True, device="cpu",
+    )
+    model.fit(X, z, epochs=2, batch_size=8, verbose=0)
+
+    p = tmp_path / "model.pt"
+    model.save(str(p))
+
+    loaded = GeneFlowDecoder.load(str(p), device="cpu")
+    assert loaded.n_genes == model.n_genes
+    assert loaded.freeze_z == model.freeze_z
+
+    a = model.predict(z)
+    b = loaded.predict(z)
+    assert torch.allclose(a, b, atol=1e-6)
+
+
+def test_flow_config_passed_through():
+    cfg = FlowConfig(nb_weight=2.0, corr_weight=0.0, sigma_min=1e-5)
+    model = GeneFlowDecoder(
+        n_genes=64, z_dim=8, d=32, depth=1, cfg=cfg, device="cpu",
+    )
+    assert model._cfg.nb_weight == 2.0
+    assert model._cfg.sigma_min == 1e-5