geneflowdecoder 0.1.6__tar.gz → 0.1.7__tar.gz

{geneflowdecoder-0.1.6/src/geneflowdecoder.egg-info → geneflowdecoder-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geneflowdecoder
-Version: 0.1.6
+Version: 0.1.7
 Summary: Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X.
 Author: GeneFlowDecoder Contributors
 License: MIT

{geneflowdecoder-0.1.6 → geneflowdecoder-0.1.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "geneflowdecoder"
-version = "0.1.6"
+version = "0.1.7"
 description = "Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X."
 readme = "README.md"
 requires-python = ">=3.9"

{geneflowdecoder-0.1.6 → geneflowdecoder-0.1.7}/src/geneflowdecoder/_core.py

@@ -1,83 +1,59 @@
 # -*- coding: utf-8 -*-
 """
-GeneFlowDecoder 0.1.2
-z (cell state, *frozen*) --> X (log-normalized expression, log1p space)
+GeneFlowDecoder  ——  z 为【固定表征】版本
+z -> X (log1p expression)
 
-Public API (stable):
+接口（稳定）：
     model = GeneFlowDecoder(n_genes=G, z_dim=D, ...)
     model.fit(X, z, X_val=None, z_val=None, ...)
-    model(z)                 # sample X from z  (flow, stochastic)
-    model.predict(z)         # deterministic X  (NB mean, log1p space)
-    model.save(path) / GeneFlowDecoder.load(path, **init_kwargs)
+    model(z)                 # 采样（Flow，随机）
+    model(z, deterministic=True) / model.predict(z)   # NB 均值，确定性
+    model.save(path) / GeneFlowDecoder.load(path, **kw)
 
-Requires: torch >= 2.1
+依赖: torch >= 2.1, numpy
 """
 
 from __future__ import annotations
 
 import math
 import warnings
-from dataclasses import dataclass, field
-from typing import Any, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import numpy as np
 
-# ----------------------------------------------------------------------
-# 可选的 torch 导入：缺 torch 时 import 不崩，构造模型时才明确报错
-# ----------------------------------------------------------------------
 try:
     import torch
     import torch.nn as nn
     import torch.nn.functional as F
-    from torch import Tensor
-
     _HAS_TORCH = True
 except Exception:  # pragma: no cover
     _HAS_TORCH = False
-    torch = None  # type: ignore[assignment]
-    nn = None  # type: ignore[assignment]
-    F = None  # type: ignore[assignment]
-
-    class Tensor:  # noqa: N801 -- placeholder type, never instantiated
-        """Stand-in so ``Array`` remains valid at annotation-check time."""
-
-        def __init_subclass__(cls, **_: Any) -> None:
-            pass
-
-    class _Noop:
-        def __getattr__(self, _name: str) -> "_Noop":
-            return self
-
-        def __call__(self, *__args: Any, **__kw: Any) -> None:
-            return None
 
-    nn = _Noop()  # type: ignore[assignment]
-    F = _Noop()  # type: ignore[assignment]
+Array = Union["torch.Tensor", np.ndarray]
 
 
-Array = Union[np.ndarray, "Tensor"]
-
-
-def _as_tensor(x: Array, device: "torch.device") -> "Tensor":
-    """numpy -> torch; 已在目标设备的 tensor 原样返回；cuda 输入不再拉回 cpu。"""
+def _as_tensor(x: Array, device) -> "torch.Tensor":
     if isinstance(x, np.ndarray):
         x = torch.from_numpy(x)
-    t: Tensor = x
-    if t.device != device:
-        t = t.to(device, non_blocking=True)
-    return t
+    if _HAS_TORCH and isinstance(x, torch.Tensor):
+        return x.to(device, non_blocking=True)
+    return x
 
 
-def _require_torch() -> None:
-    if not _HAS_TORCH:
+def _shape_ok(t: "torch.Tensor", ref: Tuple[int, ...], name: str) -> None:
+    """真实前向路径的形状断言：只在形状真的错时才报，且不猜维度。"""
+    if not _HAS_TORCH or not isinstance(t, torch.Tensor):
+        return
+    if tuple(t.shape) != tuple(ref):
         raise RuntimeError(
-            "GeneFlowDecoder requires PyTorch (>= 2.1). "
-            "Install it with: pip install 'torch>=2.1'"
+            f"[GeneFlowDecoder] 形状不一致: {name} = {tuple(t.shape)}, "
+            f"期望 {tuple(ref)}。通常是因为 X 的列数与 n_genes 不一致，"
+            f"或上游线性层输出维数不是 1。"
         )
 
 
 # ============================================================
-# 条件化模块
+# AdaLN 条件化
 # ============================================================
 
 class _AdaLNModulation(nn.Module):
@@ -85,12 +61,12 @@ class _AdaLNModulation(nn.Module):
         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.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: Tensor, z: Tensor) -> Tensor:
+    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
         g, b, gate = self.proj(z).chunk(3, dim=-1)
         g = g.unsqueeze(1)
         b = b.unsqueeze(1)
@@ -104,10 +80,10 @@ class _CondBlock(nn.Module):
         self.attn_norm = _AdaLNModulation(d, cond_dim)
         self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
         self.mlp_norm = _AdaLNModulation(d, cond_dim)
-        hd = int(d * mlp_ratio)
-        self.mlp = nn.Sequential(nn.Linear(d, hd), nn.GELU(), nn.Linear(hd, d))
+        h = int(d * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(d, h), nn.GELU(), nn.Linear(h, d))
 
-    def forward(self, x: Tensor, z: Tensor) -> Tensor:
+    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
@@ -115,36 +91,25 @@ class _CondBlock(nn.Module):
 
 
 # ============================================================
-# Gene Transformer backbone
-# 输出恒为 [B, G, d]；chunked 分支显式保证拼接后形状
+# Gene Transformer backbone: [B, G, d] -> [B, G, d]
 # ============================================================
 
 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,
-    ):
+    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 = max(1, chunk_size)
-        self._ckpt = ckpt
-
         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: Tensor, z: Tensor) -> Tensor:
+    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)
@@ -152,73 +117,84 @@ class _GeneTransformer(nn.Module):
                 x = blk(x, z)
         return self.head_norm(x)
 
-    def forward(self, gene_ids: Tensor, z: Tensor) -> Tensor:
-        B, G = gene_ids.shape[0], gene_ids.shape[1]
+    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)
-
-        chunks: list[Tensor] = []
+        outs: list[torch.Tensor] = []
         for s in range(0, G, self.chunk_size):
             e = min(s + self.chunk_size, G)
-            gids = gene_ids[:, s:e]
-            h = self.gene_embed(gids) + self.pos_embed[:, s:e, :]
-            chunks.append(self._run(h, z))
-        out = torch.cat(chunks, dim=1)
-        # 防御性断言：backbone 永远交出 [B, G, d]
-        if out.shape != (B, G, self.d):  # pragma: no cover - 仅开发期暴露
-            raise RuntimeError(
-                f"backbone shape invariant broken: got {tuple(out.shape)}, "
-                f"expected ({B},{G},{self.d})"
-            )
-        return out
+            h = self.gene_embed(gene_ids[..., s:e]) + self.pos_embed[:, s:e, :]
+            outs.append(self._run(h, z))
+        return torch.cat(outs, dim=1)
 
 
 # ============================================================
-# NB 输出头
+# 输出头：速度场 与 NB 均值，都只对【特征维 d】做线性 -> 1
 # ============================================================
 
+class _VelocityHead(nn.Module):
+    """h:[B,G,d] -> v:[B,G]。out_features=1，绝不对 n_genes 做映射。"""
+
+    def __init__(self, d: int):
+        super().__init__()
+        self.net = nn.Linear(d, 1)
+        nn.init.zeros_(self.net.weight)
+        nn.init.zeros_(self.net.bias)
+
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        # h 可能是 [B,G,d] 或退化 [G,d]/[d]
+        if h.ndim < 2:
+            raise RuntimeError(f"_VelocityHead 输入维度异常: {tuple(h.shape)}")
+        v = self.net(h).squeeze(-1)          # [B,G,1]->[B,G]
+        if v.ndim == 1:                      # [G] -> [1,G]
+            v = v.unsqueeze(0)
+        return v
+
+
 class _NBHead(nn.Module):
+    """mu 逐基因逐细胞一个标量；theta 每个基因一个 dispersion。"""
+
     def __init__(self, d: int, n_genes: int):
         super().__init__()
-        self.mu_proj = nn.Linear(d, n_genes)
+        self.mu_proj = nn.Linear(d, 1)       # <- 关键：不是 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: Tensor) -> Tuple[Tensor, Tensor]:
-        # h: [B, G, d] -> log_mu: [B, G]
-        return self.mu_proj(h).squeeze(-1), F.softplus(self.theta) + 1e-3
+    def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        log_mu = self.mu_proj(h).squeeze(-1)
+        if log_mu.ndim == 1:
+            log_mu = log_mu.unsqueeze(0)
+        return log_mu, F.softplus(self.theta) + 1e-3
 
     @staticmethod
-    def nll(
-        log_mu: Tensor,
-        theta: Tensor,
-        x: Tensor,
-        mask: Optional[Tensor] = None,
-    ) -> Tensor:
+    def nll(log_mu: torch.Tensor, theta: torch.Tensor,
+            x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        _shape_ok(log_mu, x.shape, "log_mu")
         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())
-        )
+        # theta:[G] -> [1,G]，与 [B,G] 广播；不赌 x 的 stride/layout
+        th = theta.unsqueeze(0)
+        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:
-    """训练目标权重；均为非负标量。"""
+# ============================================================
+# 配置
+# ============================================================
 
-    sigma_min: float = 1e-4
-    nb_weight: float = 1.0
-    corr_weight: float = 0.1
+class FlowConfig:
+    def __init__(self, sigma_min: float = 1e-4,
+                 nb_weight: float = 1.0, corr_weight: float = 0.1):
+        self.sigma_min = sigma_min
+        self.nb_weight = nb_weight
+        self.corr_weight = corr_weight
 
 
 # ============================================================
@@ -227,18 +203,12 @@ class FlowConfig:
 
 class GeneFlowDecoder(nn.Module):
     """
-    z -> X (log1p expression).
-
-    ``z`` 默认视为**固定表征**（freeze_z=True）：不参与反传、不被优化器更新，
-    fit() 开头还会对其做一次 detach + clone 快照，训练中途外部 encoder
-    若继续更新也不会污染本 decoder 的输入。
+    z (固定表征) -> X (log1p 表达)
 
-    Examples
-    --------
-    >>> model = GeneFlowDecoder(n_genes=2048, z_dim=256, d=512, depth=8)
+    >>> model = GeneFlowDecoder(n_genes=2048, z_dim=256, freeze_z=True)
     >>> model.fit(X_train, z_train, epochs=200, batch_size=64)
-    >>> X_hat = model(z_test)          # 采样 (flow, stochastic)
-    >>> X_rec = model.predict(z_test)  # 确定性还原: NB mean (log1p 空间)
+    >>> X_hat = model(z_test)          # 采样
+    >>> X_rec  = model.predict(z_test) # 确定性 NB 均值
     """
 
     def __init__(
@@ -254,9 +224,12 @@ class GeneFlowDecoder(nn.Module):
         freeze_z: bool = True,
         device: str = "auto",
     ):
-        _require_torch()
+        if not _HAS_TORCH:
+            raise RuntimeError(
+                "GeneFlowDecoder requires PyTorch (torch >= 2.1). "
+                "当前环境未找到 torch，请先 `pip install 'torch>=2.1'`。"
+            )
         super().__init__()
-
         if device == "auto":
             device = "cuda" if torch.cuda.is_available() else "cpu"
         self.device = torch.device(device)
@@ -266,22 +239,27 @@ class GeneFlowDecoder(nn.Module):
         self._cfg = cfg if cfg is not None else FlowConfig()
 
         self.backbone = _GeneTransformer(
-            self.n_genes, d, depth, n_heads, self.z_dim, chunk_size, gradient_checkpoint
+            self.n_genes, d, depth, n_heads, self.z_dim, chunk_size,
+            gradient_checkpoint,
         )
         self.nb_head = _NBHead(d, self.n_genes)
-
-        # ★ 关键：每个 (cell, gene) 位置预测一个标量速度，out_features 必须是 1
-        self.vel_proj = nn.Linear(d, 1)
-        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.vel_head = _VelocityHead(d)          # <- 新增独立速度头，d->1
+        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)
 
     # --------------------------------------------------------
-    # 时间编码
-    # --------------------------------------------------------
-    def _sincos(self, t: Tensor) -> Tensor:
+    def _maybe_detach(self, z: torch.Tensor, *, snapshot: bool = False):
+        if not self.freeze_z or not isinstance(z, torch.Tensor):
+            return z
+        z = z.detach()
+        return z.clone() if snapshot else z
+
+    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)
@@ -292,55 +270,38 @@ class GeneFlowDecoder(nn.Module):
             e = F.pad(e, (0, 128 - e.shape[-1]))
         return e
 
-    # --------------------------------------------------------
-    # z 隔离：freeze_z 时 detach；fit 入口额外 clone 成快照
-    # --------------------------------------------------------
-    def _maybe_detach(self, z: Tensor, *, snapshot: bool = False) -> Tensor:
-        if not self.freeze_z:
-            return z
-        z = z.detach()
-        return z.clone() if snapshot else z
+    def _gene_ids_batch(self, B: int, G: int) -> torch.Tensor:
+        return self._gene_ids[:G].unsqueeze(0).expand(B, -1).contiguous()
 
     # --------------------------------------------------------
-    # 速度场 v(x_t, t, z) -> [B, G]
+    # 速度场 v(x_t, t, z)  ->  [B, G]
     # --------------------------------------------------------
-    def _velocity(self, x_t: Tensor, t: Tensor, z: Tensor) -> Tensor:
+    def _velocity(self, x_t: torch.Tensor, t: torch.Tensor,
+                  z: torch.Tensor) -> torch.Tensor:
         z = self._maybe_detach(z)
         B, G = x_t.shape
-        t_emb = self.time_mlp(self._sincos(t)).unsqueeze(1)  # [B,1,d]
-
-        # 明确 [B, G] 的 gene id 网格，避免广播维歧义
-        gids = self._gene_ids[:G].unsqueeze(0).expand(B, -1)
-        h = self.backbone(gids, z) + t_emb              # [B,G,d]
-
-        v = self.vel_proj(h).squeeze(-1)                # [B,G,1]->[B,G]
-
-        if v.ndim != 2 or v.shape != (B, G):            # pragma: no cover
-            raise RuntimeError(
-                f"_velocity shape bug: got {tuple(v.shape)}, expected ({B},{G}). "
-                f"vel_proj 必须是 nn.Linear(d, 1)，请重装 >=0.1.2 版本。"
-            )
+        t_emb = self.time_mlp(self._sincos(t)).unsqueeze(1)   # [B,1,d]
+        gids = self._gene_ids_batch(B, G)
+        h = self.backbone(gids, z) + t_emb                    # [B,G,d]
+        v = self.vel_head(h)                                  # [B,G]
+        _shape_ok(v, x_t.shape, "_velocity")
         return v
 
     # --------------------------------------------------------
     # 损失
     # --------------------------------------------------------
-    def _loss(
-        self,
-        x: Tensor,
-        z: Tensor,
-        mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        z = self._maybe_detach(z)
-        B = x.shape[0]
+    def _loss(self, x: torch.Tensor, z: torch.Tensor,
+              mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        z = self._maybe_detach(z, snapshot=False)
+        B, G = x.shape
         t = torch.rand(B, device=self.device).clamp(
-            self._cfg.sigma_min, 1.0 - self._cfg.sigma_min
-        )
+            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)      # [B,G]
-        v_target = eps - x                       # [B,G]，rectified-flow 目标
+        v_pred = self._velocity(x_t, t, z)     # [B,G]
+        v_target = eps - x                      # [B,G]
+        _shape_ok(v_pred, v_target.shape, "flow v_pred")
         mse = F.mse_loss(v_pred, v_target, reduction="none")
         if mask is not None:
             mse = mse * mask
@@ -348,37 +309,34 @@ class GeneFlowDecoder(nn.Module):
         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[: x.shape[1]].unsqueeze(0).expand(B, -1)
+            gids = self._gene_ids_batch(B, G)
             log_mu, theta = self.nb_head(self.backbone(gids, z))
+            _shape_ok(log_mu, x.shape, "NB log_mu")
             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:
+            x0_hat = x_t.lerp(v_pred.detach(), t.unsqueeze(1))
             corr = self._corr(x0_hat, x)
 
         return mse + self._cfg.nb_weight * nb + self._cfg.corr_weight * corr
 
     @staticmethod
-    def _corr(p: Tensor, q: Tensor) -> Tensor:
+    def _corr(p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
         p = p - p.mean(1, keepdim=True)
         q = q - q.mean(1, keepdim=True)
         d = p.norm(dim=1) * q.norm(dim=1) + 1e-8
-        return (1 - (p * q).sum(dim=1) / d).mean()
+        return (1 - (p * q).sum(1) / d).mean()
 
     # --------------------------------------------------------
     # fit
     # --------------------------------------------------------
     @staticmethod
-    def _want_pin_memory(X: Tensor, z: Tensor) -> bool:
-        """pin_memory 只对 CPU tensor 有意义；cuda 输入直接关掉。"""
-
-        def on_gpu(a: Tensor) -> bool:
+    def _want_pin_memory(X: torch.Tensor, z: torch.Tensor) -> bool:
+        def on_gpu(a):
             return isinstance(a, torch.Tensor) and a.device.type != "cpu"
-
         return not (on_gpu(X) or on_gpu(z))
 
     def fit(
@@ -397,42 +355,30 @@ class GeneFlowDecoder(nn.Module):
     ) -> "GeneFlowDecoder":
         X = _as_tensor(X, self.device)
         z = _as_tensor(z, self.device)
-        # ★ 快照：freeze_z 时拷出一份独立副本，训练期与上游 z 彻底解耦
-        z = self._maybe_detach(z, snapshot=True)
-
-        if X.ndim != 2 or z.ndim != 2:
-            raise ValueError(f"X/z 必须是 2D: got {X.ndim}D / {z.ndim}D")
-        if X.shape[0] != z.shape[0]:
-            raise ValueError(f"X 与 z 样本数不一致: {X.shape[0]} vs {z.shape[0]}")
-        if z.shape[1] != self.z_dim:
-            raise ValueError(f"z 特征维 {z.shape[1]} 与 z_dim={self.z_dim} 不符")
-        if X.shape[1] != self.n_genes:
-            raise ValueError(
-                f"X 基因数 {X.shape[1]} 与 n_genes={self.n_genes} 不符；"
-                f"请用实际 X.shape[1] 构造模型"
+        z = self._maybe_detach(z, snapshot=True)     # 冻结：每轮用同一份快照
+
+        if X.ndim != 2 or X.shape[1] != self.n_genes:
+            raise RuntimeError(
+                f"X 形状应为 [N, n_genes={self.n_genes}]，实际 {tuple(X.shape)}"
+            )
+        if z.ndim != 2 or z.shape[1] != self.z_dim:
+            raise RuntimeError(
+                f"z 形状应为 [N, z_dim={self.z_dim}]，实际 {tuple(z.shape)}"
             )
 
-        # ★ 先算 pin_memory，再断言（0.1.1 的顺序 bug 已修正）
         pin_memory = self._want_pin_memory(X, z)
-        assert (
-            X.device.type == "cpu" or z.device.type == "cpu" or not pin_memory
-        ), "cuda 输入不应开启 pin_memory"
-
         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=pin_memory,
+            ds, batch_size=int(batch_size), shuffle=True,
+            num_workers=int(num_workers), pin_memory=pin_memory,
         )
 
         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)
-            if Xv.shape[1] != self.n_genes or zv.shape[1] != self.z_dim:
-                raise ValueError("验证集维度与训练集不一致")
+            zv = self._as_z_val(z_val)
+            if Xv.shape[1] != self.n_genes:
+                raise RuntimeError(f"X_val 列数 {Xv.shape[1]} != n_genes {self.n_genes}")
             has_val = True
 
         params = [p for p in self.parameters() if p.requires_grad]
@@ -440,7 +386,7 @@ class GeneFlowDecoder(nn.Module):
         sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=int(epochs))
 
         best_val = float("inf")
-        best_state: Optional[dict[str, Tensor]] = None
+        best_state: Optional[dict] = None
         no_improve = 0
 
         for ep in range(1, int(epochs) + 1):
@@ -448,6 +394,7 @@ class GeneFlowDecoder(nn.Module):
             tot = 0.0
             n_seen = 0
             for xb, zb in dl:
+                zb = self._maybe_detach(zb)
                 opt.zero_grad()
                 loss = self._loss(xb, zb)
                 loss.backward()
@@ -465,8 +412,7 @@ class GeneFlowDecoder(nn.Module):
                         vl = self._loss(Xv, zv).item()
                     msg += f" | val={vl:.4f}"
                     if vl < best_val - 1e-5:
-                        best_val = vl
-                        no_improve = 0
+                        best_val, no_improve = vl, 0
                         best_state = {
                             k: v.detach().cpu().clone()
                             for k, v in self.state_dict().items()
@@ -485,6 +431,10 @@ class GeneFlowDecoder(nn.Module):
                 print(f"[fit] restored best val={best_val:.4f}")
         return self
 
+    def _as_z_val(self, z: Array) -> torch.Tensor:
+        z = _as_tensor(z, self.device)
+        return self._maybe_detach(z, snapshot=True)
+
     # --------------------------------------------------------
     # 采样 / 确定性还原
     # --------------------------------------------------------
@@ -495,15 +445,11 @@ class GeneFlowDecoder(nn.Module):
         steps: int = 4,
         method: str = "midpoint",
         deterministic: bool = False,
-    ) -> Tensor:
-        """
-        给定 z 采样 X（log1p 空间）。
-        deterministic=True 等价于 predict()；steps=4 通常够用。
-        """
-        z = self._maybe_detach(_as_tensor(z, self.device))
+    ) -> torch.Tensor:
         if deterministic:
             return self.predict(z)
-
+        z = _as_tensor(z, self.device)
+        z = self._maybe_detach(z)
         self.eval()
         B = z.shape[0]
         G = self.n_genes
@@ -521,29 +467,40 @@ class GeneFlowDecoder(nn.Module):
         return x
 
     @torch.no_grad()
-    def predict(self, z: Array) -> Tensor:
-        """确定性还原：NB 均值（log1p 空间）。伪计数用 (exp(log_mu)-1).clamp(min=0)。"""
-        z = self._maybe_detach(_as_tensor(z, self.device))
+    def predict(self, z: Array) -> torch.Tensor:
+        """确定性还原：NB 均值（log1p 空间）。"""
+        z = _as_tensor(z, self.device)
+        z = self._maybe_detach(z)
         self.eval()
         B = z.shape[0]
-        gids = self._gene_ids[:self.n_genes].unsqueeze(0).expand(B, -1)
-        log_mu, _ = self.nb_head(self.backbone(gids, z))
+        gids = self._gene_ids_batch(B, self.n_genes)
+        log_mu, _ = self.nb_head(self.backbone(gids, z))   # [B,G]
+        _shape_ok(log_mu, (B, self.n_genes), "predict")
         return log_mu
 
     # --------------------------------------------------------
-    # 持久化
+    # 序列化
     # --------------------------------------------------------
     def save(self, path: str) -> None:
-        torch.save(
-            {"state": self.state_dict(), "cfg": self._cfg, "freeze_z": self.freeze_z},
-            path,
-        )
+        torch.save({
+            "state": self.state_dict(),
+            "cfg": {
+                "sigma_min": self._cfg.sigma_min,
+                "nb_weight": self._cfg.nb_weight,
+                "corr_weight": self._cfg.corr_weight,
+            },
+            "freeze_z": self.freeze_z,
+            "n_genes": self.n_genes,
+            "z_dim": self.z_dim,
+        }, path)
 
     @classmethod
-    def load(cls, path: str, **init_kwargs: Any) -> "GeneFlowDecoder":
+    def load(cls, path: str, **init_kwargs) -> "GeneFlowDecoder":
         ckpt = torch.load(path, map_location="cpu")
-        init_kwargs.setdefault("freeze_z", bool(ckpt.get("freeze_z", True)))
-        model = cls(**init_kwargs)
-        model._cfg = ckpt.get("cfg", model._cfg)
+        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))
+        cfg = FlowConfig(**(ckpt.get("cfg") or {}))
+        model = cls(cfg=cfg, **init_kwargs)
         model.load_state_dict(ckpt["state"])
         return model

{geneflowdecoder-0.1.6 → geneflowdecoder-0.1.7}/src/geneflowdecoder/_version.py

@@ -1,2 +1,2 @@
 # -*- coding: utf-8 -*-
-__version__ = "0.1.6"
+__version__ = "0.1.7"

{geneflowdecoder-0.1.6 → geneflowdecoder-0.1.7/src/geneflowdecoder.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: geneflowdecoder
-Version: 0.1.6
+Version: 0.1.7
 Summary: Gene Transformer + Conditional Rectified Flow decoder from cell-state z to log1p expression X.
 Author: GeneFlowDecoder Contributors
 License: MIT