Metadata-Version: 2.4
Name: geneflowdecoder
Version: 0.2.10
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 均值
```

---

## 快速基线解码器：`FastDecoder`

当**不需要**流匹配的迭代积分、只想“z 直接还原 log1p 表达”时（插值、z 空间检索、
基线对比、sanity check），用 `FastDecoder`：单次 MLP 前向，无 ODE、无时间嵌入、
无梯度检查点，训练与推理都比 flow 模型快 1–2 个数量级。

> 它与 `GeneFlowDecoder` **完全独立、互不影响**，只是同一个包里的另一条解码路径。
> 输出尺度一致：log1p 表达谱。

```python
from geneflowdecoder import FastDecoder, FastDecoderConfig
import numpy as np, scipy.sparse as sp

# 一次性转稠密 + 连续内存（稀疏矩阵直接喂会慢）
Xtr = np.ascontiguousarray(
    adata.X.toarray() if sp.isspmatrix(adata.X) else adata.X, dtype=np.float32)
ztr = np.ascontiguousarray(adata.obsm['X_latent'], dtype=np.float32)

cfg = FastDecoderConfig(
    z_dim=ztr.shape[1], n_genes=Xtr.shape[1],
    hidden_dim=256, depth=2,       # 通常 1~2 层就够
    
)
dec = FastDecoder(cfg, device='cuda' if __import__('torch').cuda.is_available() else 'cpu')

# 训练：z -> log1p(X)。
#   loss='mse'        （默认）对数域高斯 NLL，200 HVG log1p 面板就用它，稳又快
#   loss='poisson'    计数/伪计数尺度且想 var=mean 时
#   loss='multinomial' 200 基因是同一文库深度下的"份额"时用（构成型 CE）
dec.fit(ztr, Xtr, epochs=200, batch_size=1024, lr=1e-3, loss='mse', verbose=1)

# 想试构成型解码（要求 X 非负、各基因尺度可比）：
# dec.fit(ztr, np.clip(Xtr,0,None), epochs=200, batch_size=1024, loss='multinomial')


# 推理：任意 z 秒出 log-normalized 表达（自动分 chunk，不爆显存）
X_hat = dec.predict(some_z)        # numpy [N, G]
```

### `FastDecoderConfig`

| 参数 | 默认 | 说明 |
|---|---|---|
| `z_dim` | — | 细胞状态 `z` 的维度 |
| `n_genes` | — | 输出基因数 `G` |
| `hidden_dim` | 256 | MLP 隐层宽度（逐基因独立缩放/偏置另计） |
| `depth` | 2 | MLP 层数 |
| `dropout` | 0.0 |  dropout，训练样本少时可开 |
| `nb_dispersion` | `False` | **已弃用（0.2.3）**：传 `True` 只会触发一条警告，模型不再建 dispersion 参数。NB 头已移除，改用 `loss='mse'`（默认）、`'poisson'` 或 `'multinomial'` |
| `gene_wise_scale` | `True` | 每个基因独立缩放+偏置，log1p 数据几乎必开 |
| `device` | `"auto"` | `"auto"` / `"cuda"` / `"cpu"` |

### 方法速览

| 方法 | 说明 |
|---|---|
| `dec.fit(z, X, *, epochs=200, batch_size=256, lr=1e-3, loss='mse', patience=20, verbose=1, **nan_guard=True, rollback_on_nan=True, ckpt=None**)` | 拟合。`loss='mse'`（默认，log1p 域）\| `'poisson'` \| `'multinomial'`。出现 inf/nan 自动丢弃坏步，连续损坏则停止并回滚到首个坏轮前的干净权重，落盘 `<ckpt>.nan_rollback.pt` |
| `dec(z) -> Tensor` | 前向，返回 log1p 预测 `[N, G]` |
| `dec.predict(z, *, batch_size=512) -> ndarray` | 分块推理，返回 numpy，不爆显存 |
| `dec.poisson_nll(z, X) -> Tensor` / `dec.multinomial_ce(z, X) -> Tensor` | 另两种损失的手动调用版本（fit 里已自动选用） |
| 权重存取 | `torch.save(dec.state_dict(), path)` / `dec.load_state_dict(torch.load(path))`（标准 PyTorch 写法；0.2.2 之前的 NB 权重会因缺 `theta` 被 `strict=False` 安全跳过） |

### 什么时候用哪个

| 需求 | 选哪个 |
|---|---|
| 只要 z→表达的快速还原 / 基线 / 检索 | **`FastDecoder`** |
| 想验证“z 里到底有没有表达信息” | 先跑 `FastDecoder`，MSE 都降不动就是 z 或数据的问题 |
| 需要生成式采样、流形插值、精细表达分布 | `GeneFlowDecoder` |


---

## 安装

```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           # GeneFlowDecoder：flow 模型 + fit/forward/predict
│   ├── _fastdec.py        # FastDecoder：确定性 MLP 解码器
│   └── _version.py
└── tests/
    ├── test_api.py
    └── test_fastdec.py    # FastDecoder 形状契约与收敛测试
```

## License

MIT
