PathwayVI 0.1.0__py3-none-any.whl

PathwayVI/PathwayVI.py

@@ -0,0 +1,1310 @@
+import pyro
+import pyro.distributions as dist
+from pyro.optim import ExponentialLR
+from pyro.infer import SVI, JitTraceEnum_ELBO, TraceEnum_ELBO, config_enumerate
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.distributions.utils import logits_to_probs, probs_to_logits, clamp_probs
+from torch.distributions import constraints
+from torch.distributions.transforms import SoftmaxTransform
+from torch.func import jacrev, vmap
+
+from .utils.custom_mlp import MLP, Exp, PathwayDecoder
+from .utils.utils import CustomDataset, tensor_to_numpy, convert_to_tensor
+
+
+from pyro.distributions import LKJCorrCholesky
+
+import zuko 
+from pyro.contrib.zuko import ZukoToPyro
+
+import os
+import argparse
+import random
+import numpy as np
+import datatable as dt
+from tqdm import tqdm
+from scipy import sparse
+from sklearn.cluster import KMeans
+
+import scanpy as sc
+from .atac import binarize
+
+from typing import Literal
+
+import warnings
+warnings.filterwarnings("ignore")
+
+import dill as pickle
+import gzip
+from packaging.version import Version
+torch_version = torch.__version__
+
+
+def set_random_seed(seed):
+    # Set seed for PyTorch
+    torch.manual_seed(seed)
+    
+    # If using CUDA, set the seed for CUDA
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)  # For multi-GPU setups.
+    
+    # Set seed for NumPy
+    np.random.seed(seed)
+    
+    # Set seed for Python's random module
+    random.seed(seed)
+
+    # Set seed for Pyro
+    pyro.set_rng_seed(seed)
+
+    
+class PathwayVI(nn.Module):
+    """SUccinct REpresentation of single-omics cells
+
+    Parameters
+    ----------
+    inpute_size
+        Number of features (e.g., genes, peaks, proteins, etc.) per cell.
+    codebook_size
+        Number of metacells.
+    context_sizes
+        Number of context factors. 
+    transforms
+        Number of neural spline flows
+    z_dim
+        Dimensionality of latent states and metacells. 
+    hidden_layers
+        A list gives the numbers of neurons for each hidden layer.
+    flow_hidden_layers
+        A list gives the numbers of neurons for each hidden layer in neural spline flows.
+        
+        One of the following: 
+        * ``'negbinomial'`` -  negative binomial distribution (default)
+        * ``'poisson'`` - poisson distribution
+        * ``'multinomial'`` - multinomial distribution
+    use_cuda
+        A boolean option for switching on cuda device. 
+
+    Examples
+    --------
+    >>>
+    >>>
+    >>>
+
+    """
+    def __init__(self,
+                 input_dim: int,
+                 n_pathways: int,
+                 masks: torch.Tensor,
+                 n_latent: int = 20,
+                 codebook_size: int = 50,
+                 context_sizes: list = [0],
+                 perturb_size: int = 0,
+                 covariate_sizes: list = [0],
+                 total_count: int = 10000,
+                 dispersion: float = 10.0,
+                 hidden_layers: list = [500],
+                 hidden_layer_activation: Literal['relu','softplus','leakyrelu','linear'] = 'relu',
+                 flow_transforms: int = 5,
+                 flow_hidden_layers: list = [256,256],
+                 nn_dropout: float = 0.1,
+                 post_layer_fct: list = ['layernorm'],
+                 post_act_fct: list = None,
+                 config_enum: str = 'parallel',
+                 use_cuda: bool = True,
+                 seed: int = 42,
+                 dtype = torch.float32, # type: ignore
+                 ):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.n_pathways = n_pathways
+        self.n_latent = n_latent
+        self.codebook_size = codebook_size
+        self.context_sizes = context_sizes
+        self.perturb_size = perturb_size
+        self.covariate_sizes = covariate_sizes
+        self.covariate_dim = n_pathways + n_latent #if covariate_dim is None else covariate_dim
+        self.dispersion = dispersion
+        self.latent_dim = n_pathways + n_latent
+        self.hidden_layers = hidden_layers
+        self.decoder_hidden_layers = hidden_layers[::-1]
+        self.allow_broadcast = config_enum == 'parallel'
+        self.use_cuda = use_cuda
+        self.options = None
+        self.dtype = dtype
+        self.nn_dropout = nn_dropout
+        self.post_layer_fct = post_layer_fct
+        self.post_act_fct = post_act_fct
+        self.hidden_layer_activation = hidden_layer_activation
+        self.total_count = int(total_count)
+        
+        self.codebook_weights = None
+        self.centroids = None
+        
+        self.masks = masks
+        
+        self.transforms = flow_transforms
+        self.flow_hidden_layers = flow_hidden_layers
+        
+        self.seed = seed
+        
+        self.lkj_eta = 5.0
+        self.sparse_scale = 1.0
+
+        set_random_seed(seed)
+        self.setup_networks()
+
+    def setup_networks(self):
+        latent_dim = self.latent_dim
+        hidden_sizes = self.hidden_layers
+
+        nn_layer_norm, nn_batch_norm, nn_layer_dropout = False, False, False
+        na_layer_norm, na_batch_norm, na_layer_dropout = False, False, False
+
+        if self.post_layer_fct is not None:
+            nn_layer_norm=True if ('layernorm' in self.post_layer_fct) or ('layer_norm' in self.post_layer_fct) else False
+            nn_batch_norm=True if ('batchnorm' in self.post_layer_fct) or ('batch_norm' in self.post_layer_fct) else False
+            nn_layer_dropout=True if 'dropout' in self.post_layer_fct else False
+
+        if self.post_act_fct is not None:
+            na_layer_norm=True if ('layernorm' in self.post_act_fct) or ('layer_norm' in self.post_act_fct) else False
+            na_batch_norm=True if ('batchnorm' in self.post_act_fct) or ('batch_norm' in self.post_act_fct) else False
+            na_layer_dropout=True if 'dropout' in self.post_act_fct else False
+
+        if nn_layer_norm and nn_batch_norm and nn_layer_dropout:
+            post_layer_fct = lambda layer_ix, total_layers, layer: nn.Sequential(nn.Dropout(self.nn_dropout),nn.BatchNorm1d(layer.module.out_features), nn.LayerNorm(layer.module.out_features))
+        elif nn_layer_norm and nn_layer_dropout:
+            post_layer_fct = lambda layer_ix, total_layers, layer: nn.Sequential(nn.Dropout(self.nn_dropout), nn.LayerNorm(layer.module.out_features))
+        elif nn_batch_norm and nn_layer_dropout:
+            post_layer_fct = lambda layer_ix, total_layers, layer: nn.Sequential(nn.Dropout(self.nn_dropout), nn.BatchNorm1d(layer.module.out_features))
+        elif nn_layer_norm and nn_batch_norm:
+            post_layer_fct = lambda layer_ix, total_layers, layer: nn.Sequential(nn.BatchNorm1d(layer.module.out_features), nn.LayerNorm(layer.module.out_features))
+        elif nn_layer_norm:
+            post_layer_fct = lambda layer_ix, total_layers, layer: nn.LayerNorm(layer.module.out_features)
+        elif nn_batch_norm:
+            post_layer_fct = lambda layer_ix, total_layers, layer:nn.BatchNorm1d(layer.module.out_features)
+        elif nn_layer_dropout:
+            post_layer_fct = lambda layer_ix, total_layers, layer: nn.Dropout(self.nn_dropout)
+        else:
+            post_layer_fct = lambda layer_ix, total_layers, layer: None
+
+        if na_layer_norm and na_batch_norm and na_layer_dropout:
+            post_act_fct = lambda layer_ix, total_layers, layer: nn.Sequential(nn.Dropout(self.nn_dropout),nn.BatchNorm1d(layer.module.out_features), nn.LayerNorm(layer.module.out_features))
+        elif na_layer_norm and na_layer_dropout:
+            post_act_fct = lambda layer_ix, total_layers, layer: nn.Sequential(nn.Dropout(self.nn_dropout), nn.LayerNorm(layer.module.out_features))
+        elif na_batch_norm and na_layer_dropout:
+            post_act_fct = lambda layer_ix, total_layers, layer: nn.Sequential(nn.Dropout(self.nn_dropout), nn.BatchNorm1d(layer.module.out_features))
+        elif na_layer_norm and na_batch_norm:
+            post_act_fct = lambda layer_ix, total_layers, layer: nn.Sequential(nn.BatchNorm1d(layer.module.out_features), nn.LayerNorm(layer.module.out_features))
+        elif na_layer_norm:
+            post_act_fct = lambda layer_ix, total_layers, layer: nn.LayerNorm(layer.module.out_features)
+        elif na_batch_norm:
+            post_act_fct = lambda layer_ix, total_layers, layer:nn.BatchNorm1d(layer.module.out_features)
+        elif na_layer_dropout:
+            post_act_fct = lambda layer_ix, total_layers, layer: nn.Dropout(self.nn_dropout)
+        else:
+            post_act_fct = lambda layer_ix, total_layers, layer: None
+
+        if self.hidden_layer_activation == 'relu':
+            activate_fct = nn.ReLU
+        elif self.hidden_layer_activation == 'softplus':
+            activate_fct = nn.Softplus
+        elif self.hidden_layer_activation == 'leakyrelu':
+            activate_fct = nn.LeakyReLU
+        elif self.hidden_layer_activation == 'linear':
+            activate_fct = nn.Identity
+
+        self.encoder_alpha = MLP(
+                [self.latent_dim] + hidden_sizes + [self.codebook_size],
+                activation=activate_fct,
+                output_activation=None,
+                post_layer_fct=post_layer_fct,
+                post_act_fct=post_act_fct,
+                allow_broadcast=self.allow_broadcast,
+                use_cuda=self.use_cuda,
+            )
+        
+        self.encoder_coord = MLP(
+            [self.input_dim] + hidden_sizes + [[latent_dim, latent_dim]],
+            activation=activate_fct,
+            output_activation=[None, Exp],
+            post_layer_fct=post_layer_fct,
+            post_act_fct=post_act_fct,
+            allow_broadcast=self.allow_broadcast,
+            use_cuda=self.use_cuda,
+        )
+        
+        self.latent_decoder = NormalizeFlow(z_dim=self.latent_dim, num_flows=self.transforms, flow_type='planar')
+        if self.use_cuda:
+            self.latent_decoder.to('cuda')
+        
+        if np.sum(self.context_sizes)>0:
+            self.context_tokens = nn.ModuleList()
+            for context_size in self.context_sizes:
+                self.context_tokens.append(MLP(
+                        [context_size] + self.decoder_hidden_layers + [self.latent_dim],
+                        activation=activate_fct,
+                        output_activation=None,
+                        post_layer_fct=post_layer_fct,
+                        post_act_fct=post_act_fct,
+                        allow_broadcast=self.allow_broadcast,
+                        use_cuda=self.use_cuda,
+                    )
+                )
+            self.context_effect = MLP(
+                        [self.latent_dim+self.latent_dim] + self.decoder_hidden_layers + [self.latent_dim],
+                        activation=activate_fct,
+                        output_activation=None,
+                        post_layer_fct=post_layer_fct,
+                        post_act_fct=post_act_fct,
+                        allow_broadcast=self.allow_broadcast,
+                        use_cuda=self.use_cuda,
+                    )
+                
+        if self.perturb_size>0:
+            self.perturb_token = MLP(
+                [self.perturb_size] + self.decoder_hidden_layers + [self.latent_dim],
+                activation=activate_fct,
+                output_activation=None,
+                post_layer_fct=post_layer_fct,
+                post_act_fct=post_act_fct,
+                allow_broadcast=self.allow_broadcast,
+                use_cuda=self.use_cuda,
+            )
+            self.perturb_effect = MLP(
+                [self.latent_dim+self.latent_dim] + self.decoder_hidden_layers + [self.latent_dim],
+                activation=activate_fct,
+                output_activation=None,
+                post_layer_fct=post_layer_fct,
+                post_act_fct=post_act_fct,
+                allow_broadcast=self.allow_broadcast,
+                use_cuda=self.use_cuda,
+            )
+                
+        if np.sum(self.covariate_sizes)>0:
+            self.covariate_tokens = nn.ModuleList()
+            for covariate_size in self.covariate_sizes:
+                self.covariate_tokens.append(MLP(
+                        [covariate_size] + self.decoder_hidden_layers + [self.covariate_dim],
+                        activation=activate_fct,
+                        output_activation=None,
+                        post_layer_fct=post_layer_fct,
+                        post_act_fct=post_act_fct,
+                        allow_broadcast=self.allow_broadcast,
+                        use_cuda=self.use_cuda,
+                    )
+                )
+            self.covariate_effect = MLP(
+                        [self.covariate_dim] + self.decoder_hidden_layers + [self.latent_dim],
+                        activation=activate_fct,
+                        output_activation=None,
+                        post_layer_fct=post_layer_fct,
+                        post_act_fct=post_act_fct,
+                        allow_broadcast=self.allow_broadcast,
+                        use_cuda=self.use_cuda,
+                    )
+        
+        self.decoder_log_mu = nn.ModuleList()
+        for _ in range(self.n_latent):
+            self.decoder_log_mu.append(MLP(
+                    [1] + hidden_sizes + [self.input_dim],
+                    activation=nn.Identity,    # linear
+                    output_activation=None,
+                    allow_broadcast=self.allow_broadcast,
+                    use_cuda=self.use_cuda,
+                ))
+            
+        self.pathway_decoder_log_mu = PathwayDecoder(n_latent=self.n_pathways,
+                                                     n_genes=self.input_dim,
+                                                     n_pathways=self.n_pathways,
+                                                     masks=self.masks,
+                                                     sub_hidden=hidden_sizes,
+                                                     activation=activate_fct,
+                                                     output_activation=None,
+                                                     post_layer_fct=post_layer_fct,
+                                                     post_act_fct=post_act_fct,
+                                                     allow_broadcast=self.allow_broadcast,
+                                                     use_cuda=self.use_cuda)
+        
+        self.codebook = MLP(
+            [self.codebook_size] + hidden_sizes + [latent_dim],
+            activation=activate_fct,
+            output_activation=None,
+            post_layer_fct=post_layer_fct,
+            post_act_fct=post_act_fct,
+            allow_broadcast=self.allow_broadcast,
+            use_cuda=self.use_cuda,
+        )
+            
+        if self.use_cuda:
+            self.cuda()
+
+    def get_device(self):
+        return next(self.parameters()).device
+
+    def lse_pool(self, split_logits: torch.Tensor) -> torch.Tensor:
+        """LSE pooling over the pathway split dimension."""
+        return torch.logsumexp(split_logits, dim=1)
+    
+    def cutoff(self, xs, thresh=None):
+        eps = torch.finfo(xs.dtype).eps
+        
+        if not thresh is None:
+            if eps < thresh:
+                eps = thresh
+
+        xs = xs.clamp(min=eps)
+
+        if torch.any(torch.isnan(xs)):
+            xs[torch.isnan(xs)] = eps
+
+        return xs
+
+    def softmax(self, xs):
+        #xs = SoftmaxTransform()(xs)
+        #xs = dist.Multinomial(total_count=1, logits=xs).mean
+        logits = xs - xs.logsumexp(-1, keepdim=True)  # log softmax
+        return logits.exp()
+    
+    def sigmoid(self, xs):
+        #sigm_enc = nn.Sigmoid()
+        #xs = sigm_enc(xs)
+        #xs = clamp_probs(xs)
+        xs = dist.Bernoulli(logits=xs).mean
+        return xs
+
+    def softmax_logit(self, xs):
+        eps = torch.finfo(xs.dtype).eps
+        xs = self.softmax(xs)
+        xs = torch.logit(xs, eps=eps)
+        return xs
+
+    def logit(self, xs):
+        eps = torch.finfo(xs.dtype).eps
+        xs = torch.logit(xs, eps=eps)
+        return xs
+
+    def dirimulti_param(self, xs):
+        xs = self.dirimulti_mass * self.sigmoid(xs)
+        return xs
+
+    def multi_param(self, xs):
+        xs = self.softmax(xs)
+        return xs
+
+    def model(self, xs, cs=None, ps=None, fs=None):
+        pyro.module('sure', self)
+
+        eps = torch.finfo(xs.dtype).eps
+        batch_size = xs.size(0)
+        self.options = dict(dtype=xs.dtype, device=xs.device)
+        
+        # latent modeling and data likelihood        
+        acs_scale = pyro.param("codebook_scale", xs.new_ones(self.latent_dim), constraint=constraints.positive)
+
+        I = torch.eye(self.codebook_size)
+        acs_loc = self.codebook(I)
+            
+        with pyro.plate('data'):
+            prior = torch.zeros(batch_size, self.codebook_size, **self.options)
+            ns = pyro.sample('n', dist.OneHotCategorical(logits=prior))
+            _, ind = torch.topk(ns, 1)
+            
+            zn_loc = acs_loc[ind.squeeze()]
+            zn_scale = acs_scale
+            zns = pyro.sample('coord', dist.Normal(zn_loc, zn_scale).to_event(1))
+            
+            zs_ = torch.zeros_like(zns)
+            zs_ += zns
+            if (np.sum(self.context_sizes)>0) and (cs is not None):
+                zct = torch.zeros_like(zs_)
+                shift = 0
+                for i, context_size in enumerate(self.context_sizes):
+                    cs_i = cs[:,shift:(shift+context_size)]
+                    zct += self.context_tokens[i](cs_i)
+                    shift += context_size
+                zcs = self.context_effect([zs_, zct])
+                zs_ += zcs
+            
+            if (self.perturb_size>0) and (ps is not None):
+                zpt = self.perturb_token(ps)
+                zps = self.perturb_effect([zs_, zpt])
+                zs_ += zps
+                
+            if (np.sum(self.covariate_sizes)>0) and (fs is not None):
+                zft = torch.zeros(batch_size, self.covariate_dim).to(self.get_device())
+                shift = 0
+                for i, covariate_size in enumerate(self.covariate_sizes):
+                    fs_i = fs[:,shift:(shift+covariate_size)]
+                    zft += self.covariate_tokens[i](fs_i)
+                    shift += covariate_size
+                zfs = self.covariate_effect(zft)
+            else:
+                zfs = torch.zeros(batch_size, self.latent_dim).to(self.get_device())
+                
+            zs = self.latent_decoder(zs_)
+            zs += zfs
+            
+            pathway_chunk,novel_chunk = torch.chunk(zs, chunks=[self.n_pathways,self.n_latent], dim=1)
+            
+            chunks = torch.chunk(novel_chunk, chunks=len(self.decoder_log_mu), dim=1)
+            log_mu_parts = [self.decoder_log_mu[i](chunk) for i, chunk in enumerate(chunks)]
+            novel_log_mu = torch.stack(log_mu_parts, dim=1)  # [N, E, H] (假设输出带专家维度)
+    
+            pathway_log_mu = self.pathway_decoder_log_mu(pathway_chunk)
+            
+            total_log_mu = torch.stack([pathway_log_mu, novel_log_mu], dim=1)
+            log_mu = self.lse_pool(total_log_mu)
+            logits = log_mu - log_mu.logsumexp(-1, keepdim=True)  # log softmax
+            pyro.sample('x', dist.Multinomial(total_count=self.total_count, logits=logits), obs=xs)
+            
+    def guide(self, xs, cs=None, ps=None, fs=None):
+        if self.options is None:
+            self.options = dict(dtype=xs.dtype, device=xs.device)
+            
+        with pyro.plate('data'):
+            zn_loc, zn_scale = self.encoder_coord(xs)
+            zns = pyro.sample('coord', dist.Normal(zn_loc, zn_scale).to_event(1))
+            
+            alpha = self.encoder_alpha(zns)
+            ns = pyro.sample('n', dist.OneHotCategorical(logits=alpha))
+
+    def log_prob(self, xs, zs):
+        mu = self._decode_latent(zs)
+        logp = dist.Multinomial(probs=mu).log_prob(xs).sum()
+        return logp
+    
+    def _decode_latent(self, zs):
+        zs = self.latent_decoder(zs)
+        
+        pathway_chunk,novel_chunk = torch.chunk(zs, chunks=[self.n_pathways,self.n_latent], dim=1)
+        
+        chunks = torch.chunk(novel_chunk, chunks=len(self.decoder_log_mu), dim=1)
+        log_mu_parts = [self.decoder_log_mu[i](chunk) for i, chunk in enumerate(chunks)]
+        novel_log_mu = torch.stack(log_mu_parts, dim=1)  # [N, E, H] (假设输出带专家维度)
+                
+        pathway_log_mu = self.pathway_decoder_log_mu(pathway_chunk)
+            
+        total_log_mu = torch.stack([pathway_log_mu,novel_log_mu], dim=1)
+        log_mu = self.lse_pool(total_log_mu)
+        logits = log_mu - log_mu.logsumexp(-1, keepdim=True)  # log softmax
+        mu = logits.exp()
+            
+        return mu
+    
+    def manifold_representation(self, zs, batch_size: int = 1024, show_progress: bool = True):
+        """
+        """
+        ##################
+        zs = convert_to_tensor(zs, device='cpu')
+        dataset = CustomDataset(zs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        zs = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for Z_batch, _ in dataloader:
+                Z_batch = Z_batch.to(self.get_device())
+                z = self.latent_decoder(Z_batch)
+                zs.append(tensor_to_numpy(z))
+                pbar.update(1)
+
+        zs = np.concatenate(zs)
+        return zs
+    
+    def decode(self, zs, batch_size: int = 1024, show_progress: bool = True):
+        """
+        Return cells' latent representations
+
+        Parameters
+        ----------
+        xs
+            Single-cell expression matrix. It should be a Numpy array or a Pytorch Tensor.
+        batch_size
+            Size of batch processing.
+        show_progress
+            Verbose on or off
+        """
+        ##################
+        zs = convert_to_tensor(zs, device='cpu')
+        dataset = CustomDataset(zs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        mu = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for Z_batch, _ in dataloader:
+                Z_batch = Z_batch.to(self.get_device())
+                mu_ = self._decode_latent(Z_batch)
+                mu.append(tensor_to_numpy(mu_))
+                pbar.update(1)
+
+        mu = np.concatenate(mu)
+        return mu
+    
+    def _get_basal_embedding(self, xs):           
+        zns, _ = self.encoder_coord(xs)
+        return zns 
+    
+    def get_cell_embedding(self, 
+                             xs, 
+                             css:list = None,
+                             ps: np.array = None,
+                             batch_size: int = 1024,
+                             show_progress: bool = True):
+        """
+        Return cells' latent representations
+
+        Parameters
+        ----------
+        xs
+            Single-cell expression matrix. It should be a Numpy array or a Pytorch Tensor.
+        batch_size
+            Size of batch processing.
+        show_progress
+            Verbose on or off
+        """
+        ##################
+        xs = self.preprocess(xs)
+        xs = convert_to_tensor(xs, device='cpu')
+        dataset = CustomDataset(xs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        Z = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for X_batch, _ in dataloader:
+                X_batch = X_batch.to(self.get_device())
+                zns = self._get_basal_embedding(X_batch)
+                Z.append(tensor_to_numpy(zns))
+                pbar.update(1)
+
+        Z = np.concatenate(Z)
+        
+        #######################
+        CT = np.zeros_like(Z)
+        if css is not None:
+            for i in range(len(css)):
+                CT += self.get_context_token(css[i], i, batch_size=batch_size, show_progress=show_progress)
+            Z += self.get_context_effect(Z, CT, batch_size=batch_size, show_progress=show_progress)
+                
+        #######################
+        PT = np.zeros_like(Z)
+        if ps is not None:
+            PT = self.get_perturb_token(ps, batch_size=batch_size, show_progress=show_progress)
+            Z += self.get_perturb_effect(Z, PT)
+            
+        #####################
+        Z = convert_to_tensor(Z, device='cpu')
+        dataset = CustomDataset(Z)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+        
+        latent = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for ZC_batch, _ in dataloader:
+                ZC_batch = ZC_batch.to(self.get_device())
+                zns = self.latent_decoder(ZC_batch)
+                latent.append(tensor_to_numpy(zns))
+                pbar.update(1)
+
+        latent = np.concatenate(latent)
+        return latent
+    
+    def get_context_token(self, cs, i, batch_size=1024, show_progress=True):
+        cs = convert_to_tensor(cs, dtype=self.dtype, device='cpu')
+        
+        dataset = CustomDataset(cs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        A = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for C_batch, idx in dataloader:
+                C_batch = C_batch.to(self.get_device())
+                
+                dzs = self.context_tokens[i](C_batch)
+                
+                A.append(tensor_to_numpy(dzs))
+                pbar.update(1)
+
+        A = np.concatenate(A)
+        return A
+    
+    def get_context_effect(self, zs, zcs, batch_size=1024, show_progress=True):
+        zs = convert_to_tensor(zs, dtype=self.dtype, device='cpu')
+        zcs = convert_to_tensor(zcs, dtype=self.dtype, device='cpu')
+        
+        dataset = CustomDataset(zcs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        A = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for C_batch, idx in dataloader:
+                C_batch = C_batch.to(self.get_device())
+                Z_batch = zs[idx].to(self.get_device())
+                
+                dzs = self.context_effect([Z_batch,C_batch])
+                
+                A.append(tensor_to_numpy(dzs))
+                pbar.update(1)
+
+        A = np.concatenate(A)
+        return A
+    
+    def get_perturb_token(self, ps, batch_size=1024, show_progress=True):
+        ps = convert_to_tensor(ps, dtype=self.dtype, device='cpu')
+        
+        dataset = CustomDataset(ps)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        A = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for P_batch, idx in dataloader:
+                P_batch = P_batch.to(self.get_device())
+                
+                dzs = self.perturb_token(P_batch)
+                
+                A.append(tensor_to_numpy(dzs))
+                pbar.update(1)
+
+        A = np.concatenate(A)
+        return A
+    
+    def get_perturb_effect(self, zs, zps, batch_size=1024, show_progress=True):
+        zs = convert_to_tensor(zs, dtype=self.dtype, device='cpu')
+        zps = convert_to_tensor(zps, dtype=self.dtype, device='cpu')
+        
+        dataset = CustomDataset(zs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        A = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for Z_batch, idx in dataloader:
+                Z_batch = Z_batch.to(self.get_device())
+                P_batch = zps[idx].to(self.get_device())
+                
+                dzs = self.perturb_effect([Z_batch,P_batch])
+                
+                A.append(tensor_to_numpy(dzs))
+                pbar.update(1)
+
+        A = np.concatenate(A)
+        return A
+    
+    def _codebook(self):
+        I = torch.eye(self.codebook_size)
+        acs_loc = self.codebook(I)
+        return acs_loc
+    
+    def get_codebook(self):
+        return tensor_to_numpy(self._codebook())
+    
+    def _code(self, xs):
+        #zns,_ = self.encoder_zn(xs)
+        zns = self._get_basal_embedding(xs)
+        alpha = self.encoder_alpha(zns)
+        return alpha
+    
+    def _hard_assignments(self, xs):
+        alpha = self._code(xs)
+        res, ind = torch.topk(alpha, 1)
+        ns = torch.zeros_like(alpha).scatter_(1, ind, 1.0)
+        return ns,ind
+    
+    def hard_assignments(self, xs, batch_size=1024, show_progress=True):
+        """
+        Map cells to metacells and return the assigned metacell identities.
+        """
+        xs = self.preprocess(xs)
+        xs = convert_to_tensor(xs, device='cpu')
+        dataset = CustomDataset(xs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        A = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for X_batch, _ in dataloader:
+                X_batch = X_batch.to(self.get_device())
+                a,_ = self._hard_assignments(X_batch)
+                A.append(tensor_to_numpy(a))
+                pbar.update(1)
+
+        A = np.concatenate(A)
+        return A
+    
+    def _soft_assignments(self, xs):
+        alpha = self._code(xs)
+        alpha = self.softmax(alpha)
+        return alpha
+    
+    def soft_assignments(self, xs, batch_size=1024, show_progress=True):
+        """
+        Map cells to metacells and return the probabilistic values of metacell assignments
+        """
+        xs = self.preprocess(xs)
+        xs = convert_to_tensor(xs, device='cpu')
+        dataset = CustomDataset(xs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        A = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for X_batch, _ in dataloader:
+                X_batch = X_batch.to(self.get_device())
+                a = self._soft_assignments(X_batch)
+                A.append(tensor_to_numpy(a))
+                pbar.update(1)
+
+        A = np.concatenate(A)
+        return A
+    
+    def module_expression(self):
+        """
+        Generate gene expression prediction from given cell data and covariates.
+        This function can be used for simulating cells' transcription profiles at new conditions.
+        
+        :param self: SURE model
+        :param xs: Cell data at the source condition
+        :param cs: Covariates specifying the target condition for generation
+        :param batch_size: Data size per batch
+        :param show_progress: Toggle on or off message output
+        """
+        A = []
+        
+        z = np.array([[1]])
+        z = convert_to_tensor(z, device=self.get_device())
+            
+        for i in range(self.latent_dim):
+            log_mu = self.decoder_log_mu[i](z)
+            #logits = log_mu - log_mu.logsumexp(-1, keepdim=True)
+            #mu = logits.exp()
+            A.append(tensor_to_numpy(log_mu).squeeze())
+            
+        return np.vstack(A)
+    
+    def get_full_representation(self, xs:np.array, cs_list:list=None, ps:np.array=None, batch_size=1024, show_progress=True):
+        """
+        """
+        xs = self.preprocess(xs)
+        xs = convert_to_tensor(xs, dtype=self.dtype, device='cpu')
+        if cs_list is not None:
+            cs = np.hstack(cs_list)
+            cs = convert_to_tensor(cs, dtype=self.dtype, device='cpu')
+        if ps is not None:
+            ps = convert_to_tensor(ps, dtype=self.dtype, device='cpu')
+        
+        dataset = CustomDataset(xs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        A = []
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for X_batch, idx in dataloader:
+                X_batch = X_batch.to(self.get_device())
+                
+                z_basal = self._get_basal_embedding(X_batch)
+                
+                zct = torch.zeros_like(z_basal)
+                if cs_list is not None:
+                    C_batch = cs[idx].to(self.get_device())
+                    shift = 0
+                    for i, context_size in enumerate(self.context_sizes):
+                        C_batch_i = C_batch[:, shift:(shift+context_size)]
+                        zct += self.context_tokens[i](C_batch_i)
+                        shift += context_size
+                    zcs = self.context_effect([z_basal,zct])
+                    z_basal += zcs
+                    
+                if ps is not None:
+                    P_batch = ps[idx].to(self.get_device())
+                    zpt = self.perturb_token(P_batch)
+                    zps = self.perturb_effect([z_basal,zpt])
+                    z_basal += zps 
+                                   
+                A.append(tensor_to_numpy(z_basal))
+                pbar.update(1)
+
+        A = np.concatenate(A)
+        return A
+    
+    def predict(self, xs:np.array, cs_list:list=None, ps:np.array=None, fs_list:list=None, codebook:bool=False, batch_size=1024, show_progress=True):
+        """
+        Generate gene expression prediction from given cell data and covariates.
+        This function can be used for simulating cells' transcription profiles at new conditions.
+        
+        :param self: DRVI model
+        :param xs: Cell data at the source condition
+        :param cs: Covariates specifying the target condition for generation
+        :param batch_size: Data size per batch
+        :param show_progress: Toggle on or off message output
+        """
+        xs = self.preprocess(xs)
+        xs = convert_to_tensor(xs, dtype=self.dtype, device='cpu')
+        if cs_list is not None:
+            cs = np.hstack(cs_list)
+            cs = convert_to_tensor(cs, dtype=self.dtype, device='cpu')
+        if ps is not None:
+            ps = convert_to_tensor(ps, dtype=self.dtype, device='cpu')
+        if fs_list is not None:
+            fs = np.hstack(fs_list)
+            fs = convert_to_tensor(fs, dtype=self.dtype, device='cpu')
+        
+        dataset = CustomDataset(xs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        A = []
+        loc = self._codebook()
+        with tqdm(total=len(dataloader), disable=not show_progress, desc='', unit='batch') as pbar:
+            for X_batch, idx in dataloader:
+                X_batch = X_batch.to(self.get_device())
+                library_size = torch.sum(X_batch, 1, keepdim=True)
+                
+                if not codebook:
+                    z_basal = self._get_basal_embedding(X_batch)
+                else:
+                    ns = self._soft_assignments(X_batch)
+                    z_basal = ns @ loc
+                
+                zct = torch.zeros_like(z_basal)
+                if cs_list is not None:
+                    C_batch = cs[idx].to(self.get_device())
+                    shift = 0
+                    for i, context_size in enumerate(self.context_sizes):
+                        C_batch_i = C_batch[:, shift:(shift+context_size)]
+                        zct += self.context_tokens[i](C_batch_i)
+                        shift += context_size
+                    zcs = self.context_effect([z_basal,zct])
+                    z_basal += zcs
+                    
+                if ps is not None:
+                    P_batch = ps[idx].to(self.get_device())
+                    zpt = self.perturb_token(P_batch)
+                    zps = self.perturb_effect([z_basal,zpt])
+                    z_basal += zps 
+                    
+                zfs = torch.zeros(X_batch.shape[0], self.latent_dim).to(self.get_device())
+                if fs_list is not None:
+                    zfs_ = torch.zeros(X_batch.shape[0], self.covariate_dim).to(self.get_device())
+                    F_batch = fs[idx].to(self.get_device())
+                    shift = 0
+                    for i, covariate_size in enumerate(self.covariate_sizes):
+                        F_batch_i = F_batch[:, shift:(shift+covariate_size)]
+                        zfs_ += self.covariate_tokens[i](F_batch_i)#.repeat(1,self.latent_dim)
+                        shift += covariate_size
+                    zfs = self.covariate_effect(zfs_)
+                             
+                zas = self.latent_decoder(z_basal)
+                
+                pathway_chunk, novel_chunk = torch.chunk(zas, chunks=[self.n_pathways, self.n_latent], dim=1)
+                
+                pathway_log_mu = self.pathway_decoder_log_mu(pathway_chunk)
+                
+                chunks = torch.chunk(novel_chunk, chunks=len(self.decoder_log_mu), dim=1)
+                log_mu_parts = [self.decoder_log_mu[i](chunk) for i, chunk in enumerate(chunks)]
+                log_mu = torch.stack(log_mu_parts, dim=1)  # [N, E, H] (假设输出带专家维度)
+            
+                total_log_mu = torch.stack([pathway_log_mu, log_mu], dim=1)
+                log_mu = self.lse_pool(total_log_mu)
+                logits = log_mu - log_mu.logsumexp(-1, keepdim=True)  # log softmax
+                mu = logits.exp()
+                counts = mu * library_size
+            
+                A.append(tensor_to_numpy(counts))
+                pbar.update(1)
+
+        A = np.concatenate(A)
+        return A
+    
+    def preprocess(self, xs):
+        xs = np.round(xs)
+            
+        if sparse.issparse(xs):
+            xs = xs.toarray()
+        return xs 
+
+    def fit(self, xs: np.array, 
+            css: list = None, 
+            ps: np.array = None,
+            fss: list = None,
+            num_epochs: int = 100, 
+            learning_rate: float = 0.0001, 
+            use_mask: bool = False,
+            mask_ratio: float = 0.15,
+            batch_size: int = 2000, 
+            algo: Literal['adam','rmsprop','adamw'] = 'adam', 
+            beta_1: float = 0.9, 
+            weight_decay: float = 0.005, 
+            decay_rate: float = 0.9,
+            config_enum: str = 'parallel',
+            use_jax: bool = False,
+            show_progress: bool = True,
+            # Early stopping 相关参数
+            patience: int = 10,
+            min_delta: float = 1e-4,
+            restore_best_weights: bool = True,
+            monitor: str = 'loss'):
+        """
+        Train the SURENF model.
+
+        Parameters
+        ----------
+        xs
+            Single-cell expression matrix. It should be a Numpy array or a Pytorch Tensor. 
+            Rows are cells and columns are features.
+        cs
+            cell-level factor matrix. 
+        num_epochs
+            Number of training epochs.
+        learning_rate
+            Parameter for training.
+        batch_size
+            Size of batch processing.
+        algo
+            Optimization algorithm.
+        beta_1
+            Parameter for optimization.
+        weight_decay
+            Parameter for optimization.
+        decay_rate 
+            Parameter for optimization.
+        patience
+            Number of epochs with no improvement after which training will be stopped.
+        min_delta
+            Minimum change in the monitored quantity to qualify as an improvement.
+        restore_best_weights
+            Whether to restore model weights from the epoch with the best value of the monitored quantity.
+        monitor
+            Quantity to be monitored. Currently supports 'loss'.
+        use_jax
+            If toggled on, Jax will be used for speeding up.
+        show_progress
+            Whether to show training progress bar.
+        """
+        self.use_mask = use_mask
+        self.mask_ratio = mask_ratio
+        
+        xs = self.preprocess(xs)
+        xs = convert_to_tensor(xs, dtype=self.dtype, device='cpu')
+        if css is not None:
+            cs = np.hstack(css)
+            cs = convert_to_tensor(cs, dtype=self.dtype, device='cpu')
+        if ps is not None:
+            ps = convert_to_tensor(ps, dtype=self.dtype, device='cpu')
+        if fss is not None:
+            fs = np.hstack(fss)
+            fs = convert_to_tensor(fs, dtype=self.dtype, device='cpu')
+
+        dataset = CustomDataset(xs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+        # setup the optimizer
+        optim_params = {'lr': learning_rate, 'betas': (beta_1, 0.999), 'weight_decay': weight_decay}
+
+        if algo.lower()=='rmsprop':
+            optimizer_class = torch.optim.RMSprop
+        elif algo.lower()=='adam':
+            optimizer_class = torch.optim.Adam
+        elif algo.lower() == 'adamw':
+            optimizer_class = torch.optim.AdamW
+        else:
+            raise ValueError("An optimization algorithm must be specified.")
+        
+        # 初始化优化器
+        scheduler_config = {
+            'optimizer': optimizer_class, 
+            'optim_args': optim_params, 
+            'gamma': decay_rate
+        }
+        
+        pyro.clear_param_store()
+
+        # ------------- Stage 1 -----------------------------
+        # set up the loss(es) for inference
+        Elbo = JitTraceEnum_ELBO if use_jax else TraceEnum_ELBO
+        elbo = Elbo(max_plate_nesting=1, strict_enumeration_warning=False)
+        #Elbo = JitTraceEnum_ELBO if use_jax else RegTraceEnumELBO
+        #elbo = Elbo(max_plate_nesting=1, strict_enumeration_warning=False, ortho_scale=lambda_orth, sparse_scale=lambda_sparse)
+        guide = config_enumerate(self.guide, config_enum, expand=True)
+        
+        # 创建 scheduler
+        scheduler = ExponentialLR(scheduler_config)
+        loss_basic = SVI(self.model, guide, scheduler, loss=elbo)
+
+        # build a list of all losses considered
+        losses = [loss_basic]
+        num_losses = len(losses)
+        
+        # Early stopping 初始化
+        best_loss = float('inf')
+        best_epoch = 0
+        patience_counter = 0
+        
+        # 保存最佳状态
+        best_state = None
+        
+        with tqdm(total=num_epochs, disable=not show_progress, desc='Training', unit='epoch') as pbar:
+            for epoch in range(num_epochs):
+                epoch_losses = [0.0] * num_losses
+                for batch_x, idx in dataloader:
+                    batch_x = batch_x.to(self.get_device())
+                    for loss_id in range(num_losses):
+                        if css is None:
+                            batch_c = None
+                        else:
+                            batch_c = cs[idx].to(self.get_device())
+                        if ps is None:
+                            batch_p = None
+                        else:
+                            batch_p = ps[idx].to(self.get_device())
+                        if fss is None:
+                            batch_f = None
+                        else:
+                            batch_f = fs[idx].to(self.get_device())
+                            
+                        new_loss = losses[loss_id].step(xs=batch_x, cs=batch_c, ps=batch_p, fs=batch_f)
+                        epoch_losses[loss_id] += new_loss
+
+                avg_epoch_losses_ = list(map(lambda v: v / len(dataloader), epoch_losses))
+                avg_epoch_losses_str = list(map(lambda v: "{:.4f}".format(v), avg_epoch_losses_))
+                
+                # 计算当前 epoch 的总损失
+                current_loss = sum(avg_epoch_losses_)
+                
+                # Early stopping 逻辑
+                if current_loss < best_loss - min_delta:
+                    # 有显著改进
+                    best_loss = current_loss
+                    best_epoch = epoch
+                    patience_counter = 0
+                    
+                    # 保存最佳状态（模型参数 + 优化器状态）
+                    if restore_best_weights:
+                        best_state = self.get_model_and_optimizer_state(scheduler)
+                        
+                    pbar.set_postfix({
+                        'loss': ' '.join(avg_epoch_losses_str),
+                        'best': f"{best_loss:.4f}",
+                        'patience': f"{patience_counter}/{patience}"
+                    })
+                else:
+                    # 没有改进
+                    patience_counter += 1
+                    pbar.set_postfix({
+                        'loss': ' '.join(avg_epoch_losses_str),
+                        'best': f"{best_loss:.4f}",
+                        'patience': f"{patience_counter}/{patience}"
+                    })
+                    
+                    # 检查是否应该提前停止
+                    if patience_counter >= patience:
+                        print(f"\nEarly stopping triggered at epoch {epoch + 1}")
+                        print(f"Best loss {best_loss:.4f} achieved at epoch {best_epoch + 1}")
+                        
+                        # 恢复最佳状态
+                        if restore_best_weights and best_state is not None:
+                            self.load_model_and_optimizer_state(best_state, scheduler)
+                            print("Restored model and optimizer states from best epoch.")
+                        
+                        break
+                
+                pbar.update(1)
+            
+            # 训练完成（达到最大 epoch 数）
+            if epoch == num_epochs - 1:
+                print(f"\nTraining completed. Best loss {best_loss:.4f} at epoch {best_epoch + 1}")
+                if restore_best_weights and best_state is not None:
+                    self.load_model_and_optimizer_state(best_state, scheduler)
+                    print("Restored model and optimizer states from best epoch.")      
+                    
+    def get_model_and_optimizer_state(self, scheduler=None):
+        """获取模型和优化器的完整状态"""
+        state_dict = {
+            'model_params': dict(pyro.get_param_store()),  # 保存所有参数
+        }
+        
+        # 保存优化器状态
+        if scheduler is not None and hasattr(scheduler, 'optimizer'):
+            # Pyro scheduler 通常包含 optimizer
+            optimizer = scheduler.optimizer
+            if optimizer is not None:
+                state_dict['optimizer_state'] = optimizer.state_dict()
+        
+        # 保存 scheduler 状态
+        if scheduler is not None and hasattr(scheduler, 'scheduler'):
+            scheduler_obj = scheduler.scheduler
+            if scheduler_obj is not None and hasattr(scheduler_obj, 'state_dict'):
+                state_dict['scheduler_state'] = scheduler_obj.state_dict()
+        
+        return state_dict
+
+    def load_model_and_optimizer_state(self, state_dict, scheduler=None):
+        """加载模型和优化器的完整状态"""
+        
+        # 1. 加载模型参数
+        if 'model_params' in state_dict:
+            param_store = pyro.get_param_store()
+            for name, value in state_dict['model_params'].items():
+                # 确保参数存在
+                if name in param_store:
+                    # 创建可训练的参数
+                    param_store[name] = value.clone().detach().requires_grad_(True)
+                else:
+                    # 如果参数不存在，则创建它
+                    param_store[name] = torch.nn.Parameter(
+                        value.clone().detach().requires_grad_(True)
+                    )
+        
+        # 2. 加载优化器状态
+        if scheduler is not None and 'optimizer_state' in state_dict:
+            optimizer = scheduler.optimizer
+            if optimizer is not None:
+                # 确保优化器的参数与当前模型参数匹配
+                current_params = dict(pyro.get_param_store())
+                
+                # 创建一个新的状态字典，只包含当前存在的参数
+                filtered_optimizer_state = {
+                    'state': {},
+                    'param_groups': state_dict['optimizer_state']['param_groups']
+                }
+                
+                # 过滤状态，只保留当前存在的参数
+                for param_idx, param_state in state_dict['optimizer_state']['state'].items():
+                    # 检查参数是否仍然存在
+                    param_key = f'Param_{param_idx}'
+                    if param_idx < len(optimizer.param_groups[0]['params']):
+                        filtered_optimizer_state['state'][param_idx] = param_state
+                
+                optimizer.load_state_dict(filtered_optimizer_state)
+        
+        # 3. 加载 scheduler 状态
+        if (scheduler is not None and 
+            'scheduler_state' in state_dict and 
+            hasattr(scheduler, 'scheduler')):
+            scheduler_obj = scheduler.scheduler
+            if (scheduler_obj is not None and 
+                hasattr(scheduler_obj, 'load_state_dict')):
+                scheduler_obj.load_state_dict(state_dict['scheduler_state'])
+
+    def get_optimizer_param_mapping(self, optimizer):
+        """获取优化器参数到模型参数的映射关系"""
+        if optimizer is None:
+            return {}
+        
+        mapping = {}
+        param_store = dict(pyro.get_param_store())
+        
+        for group_idx, param_group in enumerate(optimizer.param_groups):
+            for param_idx, param in enumerate(param_group['params']):
+                # 找到对应的模型参数名
+                for name, model_param in param_store.items():
+                    if param is model_param or torch.equal(param, model_param):
+                        mapping[f'group{group_idx}_param{param_idx}'] = name
+                        break
+        
+        return mapping
+
+    # 使用自动微分计算 Jacobian
+    def compute_jacobian(self, zs):
+        """
+        zs: torch.Tensor of shape (B, latent_dim)
+        returns: Jacobian of shape (B, input_dim, latent_dim)
+        """
+        zs = convert_to_tensor(zs, dtype=self.dtype, device=self.get_device())
+        zs.requires_grad_(True)
+        mu = self._decode_latent(zs)
+        
+        jacobians = []
+        for i in range(mu.shape[-1]):
+            grad = torch.autograd.grad(
+                mu[:, i].sum(), zs,
+                create_graph=True,
+                retain_graph=True
+            )[0]  # (B, d)
+            jacobians.append(grad.unsqueeze(1))
+    
+        return tensor_to_numpy(torch.cat(jacobians, dim=1))  # (B, input_dim, d)
+
+    def delta_z(self, J, delta_x, rcond=1e-6):
+        """
+        J: (input_dim, latent_dim)
+        delta_x: (input_dim,)
+        return: delta_z (latent_dim,)
+        """
+        J = convert_to_tensor(J, dtype=self.dtype, device=self.get_device())
+        delta_x = convert_to_tensor(delta_x, dtype=self.dtype, device=self.get_device())
+        JtJ = J.T @ J            # (latent_dim, latent_dim)
+        JtJ_inv = torch.linalg.pinv(JtJ, rcond=rcond)
+        delta_z = JtJ_inv @ J.T @ delta_x
+        return tensor_to_numpy(delta_z)
+
+    @classmethod
+    def save_model(cls, model, file_path, compression=False):
+        """Save the model to the specified file path."""
+        file_path = os.path.abspath(file_path)
+
+        model.eval()
+        if compression:
+            with gzip.open(file_path, 'wb') as pickle_file:
+                pickle.dump(model, pickle_file)
+        else:
+            with open(file_path, 'wb') as pickle_file:
+                pickle.dump(model, pickle_file)
+
+        print(f'Model saved to {file_path}')
+
+    @classmethod
+    def load_model(cls, file_path):
+        """Load the model from the specified file path and return an instance."""
+        print(f'Model loaded from {file_path}')
+
+        file_path = os.path.abspath(file_path)
+        if file_path.endswith('gz'):
+            with gzip.open(file_path, 'rb') as pickle_file:
+                model = pickle.load(pickle_file)
+        else:
+            with open(file_path, 'rb') as pickle_file:
+                model = pickle.load(pickle_file)
+        
+        return model
+
+
+
+    
+# decoder_flow
+
+from pyro.distributions.transforms import Planar as PlanarFlow
+from pyro.distributions.transforms import Radial as RadialFlow
+from pyro.nn import PyroModule
+from pyro.distributions.transforms import ComposeTransform
+
+
+class NormalizeFlow(PyroModule):
+    def __init__(
+        self,
+        z_dim: int,
+        num_flows: int = 2,
+        flow_type: Literal['planar','radial'] = 'planar',
+    ):
+        super().__init__()
+        self.z_dim = z_dim
+        if flow_type == 'planar':
+            self.flows = torch.nn.ModuleList(
+                [PlanarFlow(z_dim) for _ in range(num_flows)]
+            )
+        elif flow_type == 'radial':
+            self.flows = torch.nn.ModuleList(
+                [RadialFlow(z_dim) for _ in range(num_flows)]
+            )
+        else:
+            raise ValueError(f"Unknown flow_type: {flow_type}")
+        
+    def transforms(self):
+        """返回 ComposeTransform，供 TransformedDistribution 使用"""
+        return ComposeTransform(list(self.flows))
+
+    def forward(self, z):
+        """
+        z: (batch, z_dim)
+        returns z': same shape, after flow transforms
+        """
+        for flow in self.flows:
+            z = flow(z)
+        return z