gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/find_latent/gnn/distribution.py

@@ -0,0 +1,498 @@
+import warnings
+
+import torch
+import torch.nn.functional as F
+from torch.distributions import Distribution, Gamma, Poisson, constraints
+from torch.distributions.utils import (
+    broadcast_all,
+    lazy_property,
+    logits_to_probs,
+    probs_to_logits,
+)
+
+
+def log_zinb_positive(
+    x: torch.Tensor, mu: torch.Tensor, theta: torch.Tensor, pi: torch.Tensor, eps=1e-8
+):
+    """
+    From scVI.
+    Log likelihood (scalar) of a minibatch according to a zinb model.
+
+    Parameters
+    ----------
+    x
+        Data
+    mu
+        mean of the negative binomial (has to be positive support) (shape: minibatch x vars)
+    theta
+        inverse dispersion parameter (has to be positive support) (shape: minibatch x vars)
+    pi
+        logit of the dropout parameter (real support) (shape: minibatch x vars)
+    eps
+        numerical stability constant
+
+    Notes
+    -----
+    We parametrize the bernoulli using the logits, hence the softplus functions appearing.
+    """
+    # theta is the dispersion rate. If .ndimension() == 1, it is shared for all cells (regardless of batch or labels)
+    if theta.ndimension() == 1:
+        theta = theta.view(
+            1, theta.size(0)
+        )  # In this case, we reshape theta for broadcasting
+
+    softplus_pi = F.softplus(-pi)  # uses log(sigmoid(x)) = -softplus(-x)
+    log_theta_eps = torch.log(theta + eps)
+    log_theta_mu_eps = torch.log(theta + mu + eps)
+    pi_theta_log = -pi + theta * (log_theta_eps - log_theta_mu_eps)
+
+    case_zero = F.softplus(pi_theta_log) - softplus_pi
+    mul_case_zero = torch.mul((x < eps).type(torch.float32), case_zero)
+
+    case_non_zero = (
+        -softplus_pi
+        + pi_theta_log
+        + x * (torch.log(mu + eps) - log_theta_mu_eps)
+        + torch.lgamma(x + theta)
+        - torch.lgamma(theta)
+        - torch.lgamma(x + 1)
+    )
+    mul_case_non_zero = torch.mul((x > eps).type(torch.float32), case_non_zero)
+
+    res = mul_case_zero + mul_case_non_zero
+
+    return res
+
+
+def log_nb_positive(x: torch.Tensor, mu: torch.Tensor, theta: torch.Tensor, eps=1e-8):
+    """
+    From scVI.
+    Log likelihood (scalar) of a minibatch according to a nb model.
+
+    Parameters
+    ----------
+    x
+        data
+    mu
+        mean of the negative binomial (has to be positive support) (shape: minibatch x vars)
+    theta
+        inverse dispersion parameter (has to be positive support) (shape: minibatch x vars)
+    eps
+        numerical stability constant
+
+    Notes
+    -----
+    We parametrize the bernoulli using the logits, hence the softplus functions appearing.
+
+    """
+    if theta.ndimension() == 1:
+        theta = theta.view(
+            1, theta.size(0)
+        )  # In this case, we reshape theta for broadcasting
+
+    log_theta_mu_eps = torch.log(theta + mu + eps)
+
+    res = (
+        theta * (torch.log(theta + eps) - log_theta_mu_eps)
+        + x * (torch.log(mu + eps) - log_theta_mu_eps)
+        + torch.lgamma(x + theta)
+        - torch.lgamma(theta)
+        - torch.lgamma(x + 1)
+    )
+
+    return res
+
+
+def log_mixture_nb(
+    x: torch.Tensor,
+    mu_1: torch.Tensor,
+    mu_2: torch.Tensor,
+    theta_1: torch.Tensor,
+    theta_2: torch.Tensor,
+    pi_logits: torch.Tensor,
+    eps=1e-8,
+):
+    """
+    From scVI.
+    Log likelihood (scalar) of a minibatch according to a mixture nb model.
+
+    pi_logits is the probability (logits) to be in the first component.
+    For totalVI, the first component should be background.
+
+    Parameters
+    ----------
+    x
+        Observed data
+    mu_1
+        Mean of the first negative binomial component (has to be positive support) (shape: minibatch x features)
+    mu_2
+        Mean of the second negative binomial (has to be positive support) (shape: minibatch x features)
+    theta_1
+        First inverse dispersion parameter (has to be positive support) (shape: minibatch x features)
+    theta_2
+        Second inverse dispersion parameter (has to be positive support) (shape: minibatch x features)
+        If None, assume one shared inverse dispersion parameter.
+    pi_logits
+        Probability of belonging to mixture component 1 (logits scale)
+    eps
+        Numerical stability constant
+    """
+    if theta_2 is not None:
+        log_nb_1 = log_nb_positive(x, mu_1, theta_1)
+        log_nb_2 = log_nb_positive(x, mu_2, theta_2)
+    # this is intended to reduce repeated computations
+    else:
+        theta = theta_1
+        if theta.ndimension() == 1:
+            theta = theta.view(
+                1, theta.size(0)
+            )  # In this case, we reshape theta for broadcasting
+
+        log_theta_mu_1_eps = torch.log(theta + mu_1 + eps)
+        log_theta_mu_2_eps = torch.log(theta + mu_2 + eps)
+        lgamma_x_theta = torch.lgamma(x + theta)
+        lgamma_theta = torch.lgamma(theta)
+        lgamma_x_plus_1 = torch.lgamma(x + 1)
+
+        log_nb_1 = (
+            theta * (torch.log(theta + eps) - log_theta_mu_1_eps)
+            + x * (torch.log(mu_1 + eps) - log_theta_mu_1_eps)
+            + lgamma_x_theta
+            - lgamma_theta
+            - lgamma_x_plus_1
+        )
+        log_nb_2 = (
+            theta * (torch.log(theta + eps) - log_theta_mu_2_eps)
+            + x * (torch.log(mu_2 + eps) - log_theta_mu_2_eps)
+            + lgamma_x_theta
+            - lgamma_theta
+            - lgamma_x_plus_1
+        )
+
+    logsumexp = torch.logsumexp(torch.stack(
+        (log_nb_1, log_nb_2 - pi_logits)), dim=0)
+    softplus_pi = F.softplus(-pi_logits)
+
+    log_mixture_nb = logsumexp - softplus_pi
+
+    return log_mixture_nb
+
+
+def _convert_mean_disp_to_counts_logits(mu, theta, eps=1e-6):
+    r"""
+    From scVI.
+    NB parameterizations conversion.
+
+    Parameters
+    ----------
+    mu
+        mean of the NB distribution.
+    theta
+        inverse overdispersion.
+    eps
+        constant used for numerical log stability. (Default value = 1e-6)
+
+    Returns
+    -------
+    type
+        the number of failures until the experiment is stopped
+        and the success probability.
+    """
+    if not (mu is None) == (theta is None):
+        raise ValueError(
+            "If using the mu/theta NB parameterization, both parameters must be specified"
+        )
+    logits = (mu + eps).log() - (theta + eps).log()
+    total_count = theta
+    return total_count, logits
+
+
+def _convert_counts_logits_to_mean_disp(total_count, logits):
+    """
+    From scVI.
+    NB parameterizations conversion.
+
+    Parameters
+    ----------
+    total_count
+        Number of failures until the experiment is stopped.
+    logits
+        success logits.
+
+    Returns
+    -------
+    type
+        the mean and inverse overdispersion of the NB distribution.
+
+    """
+    theta = total_count
+    mu = logits.exp() * theta
+    return mu, theta
+
+
+def _gamma(theta, mu):
+    concentration = theta
+    rate = theta / mu
+    # Important remark: Gamma is parametrized by the rate = 1/scale!
+    gamma_d = Gamma(concentration=concentration, rate=rate)
+    return gamma_d
+
+
+class NegativeBinomial(Distribution):
+    r"""
+    From scVI.
+    Negative binomial distribution.
+
+    One of the following parameterizations must be provided:
+
+    (1), (`total_count`, `probs`) where `total_count` is the number of failures until
+    the experiment is stopped and `probs` the success probability. (2), (`mu`, `theta`)
+    parameterization, which is the one used by scvi-tools. These parameters respectively
+    control the mean and inverse dispersion of the distribution.
+
+    In the (`mu`, `theta`) parameterization, samples from the negative binomial are generated as follows:
+
+    1. :math:`w \sim \textrm{Gamma}(\underbrace{\theta}_{\text{shape}}, \underbrace{\theta/\mu}_{\text{rate}})`
+    2. :math:`x \sim \textrm{Poisson}(w)`
+
+    Parameters
+    ----------
+    total_count
+        Number of failures until the experiment is stopped.
+    probs
+        The success probability.
+    mu
+        Mean of the distribution.
+    theta
+        Inverse dispersion.
+    scale
+        Normalized mean expression of the distribution.
+    validate_args
+        Raise ValueError if arguments do not match constraints
+    """
+
+    arg_constraints = {
+        "mu": constraints.greater_than_eq(0),
+        "theta": constraints.greater_than_eq(0),
+    }
+    support = constraints.nonnegative_integer
+
+    def __init__(
+        self,
+        total_count: torch.Tensor | None = None,
+        probs: torch.Tensor | None = None,
+        logits: torch.Tensor | None = None,
+        mu: torch.Tensor | None = None,
+        theta: torch.Tensor | None = None,
+        scale: torch.Tensor | None = None,
+        validate_args: bool = False,
+    ):
+        self._eps = 1e-8
+        if (mu is None) == (total_count is None):
+            raise ValueError(
+                "Please use one of the two possible parameterizations. Refer to the documentation for more information."
+            )
+
+        using_param_1 = total_count is not None and (
+            logits is not None or probs is not None
+        )
+        if using_param_1:
+            logits = logits if logits is not None else probs_to_logits(probs)
+            total_count = total_count.type_as(logits)
+            total_count, logits = broadcast_all(total_count, logits)
+            mu, theta = _convert_counts_logits_to_mean_disp(
+                total_count, logits)
+        else:
+            mu, theta = broadcast_all(mu, theta)
+        self.mu = mu
+        self.theta = theta
+        self.scale = scale
+        super().__init__(validate_args=validate_args)
+
+    @property
+    def mean(self):
+        return self.mu
+
+    @property
+    def variance(self):
+        return self.mean + (self.mean**2) / self.theta
+
+    @torch.inference_mode()
+    def sample(
+        self,
+        sample_shape: torch.Size | tuple | None = None,
+    ) -> torch.Tensor:
+        """Sample from the distribution."""
+        sample_shape = sample_shape or torch.Size()
+        gamma_d = self._gamma()
+        p_means = gamma_d.sample(sample_shape)
+
+        # Clamping as distributions objects can have buggy behaviors when
+        # their parameters are too high
+        l_train = torch.clamp(p_means, max=1e8)
+        # Shape : (n_samples, n_cells_batch, n_vars)
+        counts = Poisson(l_train).sample()
+        return counts
+
+    @torch.inference_mode()
+    def rsample(self, sample_shape: torch.Size | tuple | None = None):
+        """Sample from the distribution."""
+        return self.sample(sample_shape=sample_shape)
+
+    def log_prob(self, value: torch.Tensor) -> torch.Tensor:
+        if self._validate_args:
+            try:
+                self._validate_sample(value)
+            except ValueError:
+                warnings.warn(
+                    "The value argument must be within the support of the distribution",
+                    UserWarning, stacklevel=2,
+                )
+
+        return log_nb_positive(value, mu=self.mu, theta=self.theta, eps=self._eps)
+
+    def _gamma(self):
+        return _gamma(self.theta, self.mu)
+
+    def pearson_residuals(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Compute the Pearson residuals.
+
+        Parameters
+        ----------
+        x
+            Observed data.
+
+        Returns
+        -------
+        type
+            Pearson residuals.
+        """
+        mean = self.mean
+        variance = self.variance
+        return (x - mean) / torch.sqrt(variance)
+
+
+class ZeroInflatedNegativeBinomial(NegativeBinomial):
+    r"""
+    From scVI.
+    Zero-inflated negative binomial distribution.
+
+    One of the following parameterizations must be provided:
+
+    (1), (`total_count`, `probs`) where `total_count` is the number of failures until
+    the experiment is stopped and `probs` the success probability. (2), (`mu`, `theta`)
+    parameterization, which is the one used by scvi-tools. These parameters respectively
+    control the mean and inverse dispersion of the distribution.
+
+    In the (`mu`, `theta`) parameterization, samples from the negative binomial are generated as follows:
+
+    1. :math:`w \sim \textrm{Gamma}(\underbrace{\theta}_{\text{shape}}, \underbrace{\theta/\mu}_{\text{rate}})`
+    2. :math:`x \sim \textrm{Poisson}(w)`
+
+    Parameters
+    ----------
+    total_count
+        Number of failures until the experiment is stopped.
+    probs
+        The success probability.
+    mu
+        Mean of the distribution.
+    theta
+        Inverse dispersion.
+    zi_logits
+        Logits scale of zero inflation probability.
+    scale
+        Normalized mean expression of the distribution.
+    validate_args
+        Raise ValueError if arguments do not match constraints
+    """
+
+    arg_constraints = {
+        "mu": constraints.greater_than_eq(0),
+        "theta": constraints.greater_than_eq(0),
+        "zi_probs": constraints.half_open_interval(0.0, 1.0),
+        "zi_logits": constraints.real,
+    }
+    support = constraints.nonnegative_integer
+
+    def __init__(
+        self,
+        total_count: torch.Tensor | None = None,
+        probs: torch.Tensor | None = None,
+        logits: torch.Tensor | None = None,
+        mu: torch.Tensor | None = None,
+        theta: torch.Tensor | None = None,
+        zi_logits: torch.Tensor | None = None,
+        scale: torch.Tensor | None = None,
+        validate_args: bool = False,
+    ):
+        super().__init__(
+            total_count=total_count,
+            probs=probs,
+            logits=logits,
+            mu=mu,
+            theta=theta,
+            scale=scale,
+            validate_args=validate_args,
+        )
+        self.zi_logits, self.mu, self.theta = broadcast_all(
+            zi_logits, self.mu, self.theta
+        )
+
+    @property
+    def mean(self):
+        pi = self.zi_probs
+        return (1 - pi) * self.mu
+
+    @property
+    def variance(self):
+        pi = self.zi_probs
+        return (1 - pi) * self.mu * (1 + self.mu * (pi + 1 / self.theta))
+
+    @lazy_property
+    def zi_logits(self) -> torch.Tensor:
+        """ZI logits."""
+        return probs_to_logits(self.zi_probs, is_binary=True)
+
+    @lazy_property
+    def zi_probs(self) -> torch.Tensor:
+        return logits_to_probs(self.zi_logits, is_binary=True)
+
+    @torch.inference_mode()
+    def sample(
+        self,
+        sample_shape: torch.Size | tuple | None = None,
+    ) -> torch.Tensor:
+        """Sample from the distribution."""
+        sample_shape = sample_shape or torch.Size()
+        samp = super().sample(sample_shape=sample_shape)
+        is_zero = torch.rand_like(samp) <= self.zi_probs
+        samp_ = torch.where(
+            is_zero, torch.tensor(0.0, dtype=torch.float32, device=samp.device), samp
+        )
+        return samp_
+
+    @torch.inference_mode()
+    def rsample(  # type: ignore
+        self,
+        sample_shape: torch.Size | tuple | None = None,
+    ) -> torch.Tensor:
+        """Sample from the distribution."""
+        sample_shape = sample_shape or torch.Size()
+        samp = super().rsample(sample_shape=sample_shape)
+        is_zero = torch.rand_like(samp) <= self.zi_probs
+        samp_ = torch.where(is_zero, torch.tensor(
+            0.0, dtype=torch.float32, device=samp.device), samp)
+        return samp_
+
+    def log_prob(self, value: torch.Tensor) -> torch.Tensor:
+        """Log probability."""
+        try:
+            self._validate_sample(value)
+        except ValueError:
+            warnings.warn(
+                "The value argument must be within the support of the distribution",
+                UserWarning, stacklevel=2,
+            )
+        return log_zinb_positive(value, self.mu, self.theta, self.zi_logits, eps=1e-08)

gsMap/find_latent/gnn/encoder_decoder.py

@@ -0,0 +1,186 @@
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+from torch.distributions import kl_divergence as kl
+
+from .gene_former import GeneModuleFormer
+
+
+def full_block(in_dim, out_dim, p_drop=0.1):
+    return nn.Sequential(
+        nn.Linear(in_dim, out_dim),
+        nn.BatchNorm1d(out_dim),
+        nn.ReLU(),
+        nn.Dropout(p=p_drop),
+    )
+
+class transform(nn.Module):
+    """
+    batch transform encoder
+    """
+    def __init__(self,
+                 input_size,
+                 hidden_size,
+                 batch_emb_size,
+                 module_dim,
+                 hidden_gmf,
+                 n_modules,
+                 nhead,
+                 n_enc_layer,
+                 use_tf):
+
+        super().__init__()
+        self.use_tf = use_tf
+
+        if self.use_tf:
+            self.input_size = hidden_gmf + batch_emb_size
+            self.gmf = GeneModuleFormer(input_dim=input_size,
+                                        module_dim=module_dim,
+                                        hidden_dim=hidden_gmf,
+                                        n_modules=n_modules,
+                                        nhead=nhead,
+                                        n_enc_layer=n_enc_layer
+                                        )
+            self.transform = full_block(self.input_size,hidden_size)
+        else:
+            self.input_size = input_size + batch_emb_size
+            self.transform = full_block(self.input_size,hidden_size)
+            self.norm = nn.LayerNorm(hidden_size)
+
+    def forward(self, x, batch):
+        if self.use_tf:
+            x = self.gmf(x)
+            x = self.transform(torch.cat([x,batch],dim=1))
+        else:
+            x = self.transform(torch.cat([x,batch],dim=1))
+        return self.norm(x)
+
+
+
+class Encoder(nn.Module):
+    """
+    GCN encoder
+    """
+    def __init__(self,
+                 input_size,
+                 hidden_size,
+                 emb_size,
+                 batch_emb_size,
+                 module_dim,
+                 hidden_gmf,
+                 n_modules,
+                 nhead,
+                 n_enc_layer,
+                 use_tf,
+                 variational=True):
+
+        super().__init__()
+        self.variational = variational
+
+        self.tf = transform(
+            input_size,
+            hidden_size,
+            batch_emb_size,
+            module_dim,
+            hidden_gmf,
+            n_modules,
+            nhead,
+            n_enc_layer,
+            use_tf
+        )
+
+        self.mlp = nn.Sequential(full_block(hidden_size, hidden_size),
+                                 full_block(hidden_size,hidden_size))
+
+
+        self.fc_mean = nn.Linear(hidden_size,emb_size)
+        self.fc_var = nn.Linear(hidden_size,emb_size)
+
+    def forward(self, x, batch):
+
+        xtf = self.tf(x,batch)
+        h = self.mlp(xtf)
+        if not self.variational:
+            mu = self.fc_mean(h)
+            return mu
+
+        mu = self.fc_mean(h)
+        logvar = self.fc_var(h)
+        self.mu = mu
+        self.sigma = logvar.exp().sqrt()
+        self.dist = Normal(self.mu, self.sigma)
+        return self.dist.rsample()
+
+    def kl_loss(self):
+        if not hasattr(self, "dist"):
+            return 0
+
+        mean = torch.zeros_like(self.mu)
+        scale = torch.ones_like(self.sigma)
+        kl_loss = kl(self.dist, Normal(mean, scale))
+        return kl_loss.mean()
+
+class Decoder(nn.Module):
+    """
+    Shared decoder
+    """
+    def __init__(self,
+                 out_put_size,
+                 hidden_size,
+                 emb_size,
+                 batch_emb_size,
+                 class_size,
+                 decoder_type,
+                 distribution,
+                 n_layers=3):
+        super().__init__()
+
+        self.decoder_type = decoder_type
+        self.mlp = nn.ModuleList()
+
+        # Set initial input size
+        if decoder_type == 'reconstruction':
+            input_size = emb_size + batch_emb_size
+        elif decoder_type == 'classification':
+            input_size = emb_size * 2 + batch_emb_size
+        else:
+            raise ValueError(f"Unknown decoder_type: {decoder_type}")
+
+        # Build MLP layers with batch embedding concat at each step
+        if isinstance(n_layers, int):
+            n_layers = [hidden_size] * n_layers
+
+        for hidden_size in n_layers:
+            self.mlp.append(full_block(input_size, hidden_size))
+            input_size = hidden_size + batch_emb_size  # update for next layer input
+
+        # Final output layer
+        if decoder_type == 'reconstruction':
+            self.zi_logit = nn.Linear(input_size, out_put_size)
+            self.fc_rec = nn.Linear(input_size, out_put_size)
+        elif decoder_type == 'classification':
+            self.fc_class = nn.Linear(input_size, class_size)
+
+        if distribution in ['nb','zinb']:
+            self.act = nn.Softmax(dim=-1)
+        else:
+            self.act = nn.Identity()
+
+    def forward(self, z, batch):
+        x = torch.cat([z, batch], dim=1)
+
+        for layer in self.mlp:
+            x = layer(x)
+            x = torch.cat([x, batch], dim=1)  # concat batch after each layer
+
+        if self.decoder_type == 'reconstruction':
+            x_hat = self.act(self.fc_rec(x))
+            zi_logit = self.zi_logit(x)
+            return x_hat, zi_logit
+
+        elif self.decoder_type == 'classification':
+            x_class = self.fc_class(x)
+            return x_class
+
+        else:
+            raise ValueError(f"Unknown decoder_type: {self.decoder_type}")