pg-sui 1.6.14.dev9__py3-none-any.whl → 1.7.0__py3-none-any.whl

pgsui/impute/unsupervised/loss_functions.py

@@ -1,261 +1,215 @@
-from typing import List, Literal
+from __future__ import annotations
+
+from typing import Literal, cast
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 
-class SafeFocalCELoss(nn.Module):
-    """Focal cross-entropy with ignore_index and numeric guards.
+class FocalCELoss(nn.Module):
+    """Focal cross-entropy with ignore_index and optional scaling.
+
+    Supports logits of shape (N, C) or (N, C, d1, d2, ...). Targets must be shape-compatible: (N) or (N, d1, d2, ...).
 
-    This class implements the focal loss function, which is designed to address class imbalance by down-weighting easy examples and focusing training on hard negatives. It also includes handling for ignored indices and numeric stability.
+    The optional `recon_scale` is useful in reconstruction settings (e.g., VAE) when your base reduction is "mean" over a sparse mask. Multiplying the final reduced loss by `recon_scale` makes the reconstruction term more "sum-like" per batch/sample, preventing KL from dominating.
     """
 
     def __init__(
         self,
-        gamma: float,
-        weight: torch.Tensor | None = None,
+        *,
+        alpha: torch.Tensor | None = None,
+        gamma: float = 2.0,
         ignore_index: int = -1,
-        eps: float = 1e-8,
-    ):
-        """Initialize the SafeFocalCELoss.
-
-        This class sets up the focal loss with specified focusing parameter, class weights, ignore index, and a small epsilon for numerical stability.
+        reduction: Literal["mean", "sum", "none"] = "mean",
+    ) -> None:
+        """Initialize the focal cross-entropy loss.
 
         Args:
-            gamma (float): Focusing parameter.
-            weight (torch.Tensor | None): A manual rescaling weight given to each class. If given, has to be a Tensor of size C (number of classes). Defaults to None.
-            ignore_index (int): Specifies a target value that is ignored and does not contribute to the input gradient. Default is -1.
-            eps (float): Small value to avoid numerical issues. Default is 1e-8.
+            alpha: Optional per-class weights of shape (C,).
+            gamma: Focusing parameter.
+            ignore_index: Target value to ignore.
+            reduction: "mean", "sum", or "none".
         """
         super().__init__()
-        self.gamma = gamma
-        self.weight = weight
-        self.ignore_index = ignore_index
-        self.eps = eps
-
-    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
-        """Calculates the focal loss on pre-flattened tensors.
-
-        Args:
-            logits (torch.Tensor): Logits from the model of shape (N, C) where N is the number of samples and C is the number of classes.
-            targets (torch.Tensor): Ground truth labels of shape (N,).
-
-        Returns:
-            torch.Tensor: The computed scalar loss value.
-        """
-        # logits: (N, C), targets: (N,)
-        valid = targets != self.ignore_index
-
-        if not valid.any():
-            return logits.new_tensor(0.0)
-
-        logits_v = logits[valid]
-        targets_v = targets[valid]
-
-        logp = F.log_softmax(logits_v, dim=-1)  # stable
-        ce = F.nll_loss(logp, targets_v, weight=self.weight, reduction="none")
-
-        # p_t = exp(logp[range, targets])
-        p_t = torch.exp(logp.gather(1, targets_v.unsqueeze(1)).squeeze(1))
-
-        # focal factor with clamp to avoid 0**gamma and NaNs
-        focal = (1.0 - p_t).clamp_min(self.eps).pow(self.gamma)
-
-        loss_vec = focal * ce
-
-        # guard remaining inf/nan if any slipped through
-        loss_vec = torch.nan_to_num(loss_vec, nan=0.0, posinf=1e6, neginf=0.0)
-        return loss_vec.mean()
-
-
-class WeightedMaskedCCELoss(nn.Module):
-    def __init__(
-        self,
-        alpha: float | List[float] | torch.Tensor | None = None,
-        reduction: Literal["mean", "sum"] = "mean",
-    ):
-        """A weighted, masked Categorical Cross-Entropy loss function.
-
-        This method computes the categorical cross-entropy loss while allowing for class weights and masking of invalid (missing) entries. It is particularly useful for sequence data where some positions may be missing or should not contribute to the loss calculation.
-
-        Args:
-            alpha (float | List | Tensor | None): A manual rescaling weight given to each class. If given, has to be a Tensor of size C (number of classes). Defaults to None.
-            reduction (str, optional): Specifies the reduction to apply to the output: 'mean' or 'sum'. Defaults to "mean".
-        """
-        super(WeightedMaskedCCELoss, self).__init__()
+        self._gamma = float(gamma)
+        self.ignore_index = int(ignore_index)
         self.reduction = reduction
-        self.alpha = alpha
+
+        if alpha is not None:
+            if alpha.dim() != 1:
+                raise ValueError("alpha must be a 1D tensor of shape (C,).")
+            # Register as buffer so it moves with the module across devices.
+            self.register_buffer("alpha", alpha)
+        else:
+            self.alpha = None
 
     def forward(
         self,
         logits: torch.Tensor,
         targets: torch.Tensor,
-        valid_mask: torch.Tensor | None = None,
+        *,
+        recon_scale: torch.Tensor | float | None = None,
     ) -> torch.Tensor:
-        """Compute the masked categorical cross-entropy loss.
+        """Compute focal cross-entropy loss.
 
         Args:
-            logits (torch.Tensor): Logits from the model of shape
-                (batch_size, seq_len, num_classes).
-            targets (torch.Tensor): Ground truth labels of shape (batch_size, seq_len).
-            valid_mask (torch.Tensor, optional): Boolean mask of shape (batch_size, seq_len) where True indicates a valid (observed) value to include in the loss.
-                Defaults to None, in which case all values are considered valid.
+            logits: Tensor of shape (N, C) or (N, C, d1, d2, ...).
+            targets: Tensor of shape (N) or (N, d1, d2, ...).
+            recon_scale: Optional scalar multiplier applied to the final loss.
+                - If reduction is "mean" or "sum", multiplies the scalar loss.
+                - If reduction is "none", multiplies elementwise.
 
         Returns:
-            torch.Tensor: The computed scalar loss value.
+            Loss tensor:
+                - Scalar if reduction in {"mean","sum"}
+                - Tensor shaped like `targets` if reduction == "none"
         """
-        # Automatically detect the device from the input tensor
-        device = logits.device
-        num_classes = logits.shape[-1]
-
-        # Ensure targets are on the correct device and are Long type
-        targets = targets.to(device).long()
-
-        # Prepare weights and pass them directly to the loss function
-        class_weights = None
-        if self.alpha is not None:
-            if not isinstance(self.alpha, torch.Tensor):
-                class_weights = torch.tensor(
-                    self.alpha, dtype=torch.float, device=device
-                )
-            else:
-                class_weights = self.alpha.to(device)
-
-        loss = F.cross_entropy(
-            logits.reshape(-1, num_classes),
-            targets.reshape(-1),
-            weight=class_weights,
-            reduction="none",
-            ignore_index=-1,  # Ignore all targets with the value -1
-        )
-
-        # If a mask is provided, filter the losses for the training set
-        if valid_mask is not None:
-            loss = loss[valid_mask.reshape(-1)]
-
-        # If after masking no valid losses remain, return 0
-        if loss.numel() == 0:
-            return torch.tensor(0.0, device=device)
-
-        # Apply the final reduction
-        if self.reduction == "mean":
-            return loss.mean()
-        elif self.reduction == "sum":
-            return loss.sum()
-        else:
-            msg = f"Reduction mode '{self.reduction}' not supported."
-            raise ValueError(msg)
+        # Move C (dim 1) to the last position for flattening:
+        # (N, C, d1, ...) -> (N, d1, ..., C)
+        if logits.dim() > 2:
+            logits = logits.permute(0, *range(2, logits.dim()), 1)
 
+        logits_flat = logits.reshape(-1, logits.size(-1))
+        targets_flat = targets.reshape(-1).long()
 
-class MaskedFocalLoss(nn.Module):
-    """Focal loss (gamma > 0) with optional class weights and a boolean valid mask.
+        valid_mask = targets_flat != self.ignore_index
 
-    This method implements the focal loss function, which is designed to address class imbalance by down-weighting easy examples and focusing training on hard negatives. It also supports masking of invalid (missing) entries, making it suitable for sequence data with missing values.
-    """
+        # Early exit if everything is ignored
+        if not bool(valid_mask.any()):
+            out = torch.tensor(0.0, device=logits.device, dtype=logits.dtype)
+            # preserve grad path behavior if caller expects it
+            out = out.requires_grad_(True)
+            return out
 
-    def __init__(
-        self,
-        gamma: float = 2.0,
-        alpha: torch.Tensor | None = None,
-        reduction: Literal["mean", "sum"] = "mean",
-    ):
-        """Initialize the MaskedFocalLoss.
+        logits_v = logits_flat[valid_mask]
+        targets_v = targets_flat[valid_mask]
 
-        This class sets up the focal loss with specified focusing parameter, class weights, and reduction method. It is designed to handle missing data through a valid mask, ensuring that only relevant entries contribute to the loss calculation.
+        # Numerically stable log-softmax
+        log_probs = F.log_softmax(logits_v, dim=-1)
+        log_pt = log_probs.gather(1, targets_v.unsqueeze(1)).squeeze(1)
+        pt = log_pt.exp()
 
-        Args:
-            gamma (float): Focusing parameter.
-            alpha (torch.Tensor | None): Class weights.
-            reduction (Literal["mean", "sum"]): Reduction mode ('mean' or 'sum').
-        """
-        super().__init__()
-        self.gamma = gamma
-        self.alpha = alpha
-        self.reduction = reduction
-
-    def forward(
-        self,
-        logits: torch.Tensor,  # Expects (N, C) where N = batch*features
-        targets: torch.Tensor,  # Expects (N,)
-        valid_mask: torch.Tensor,  # Expects (N,)
-    ) -> torch.Tensor:
-        """Calculates the focal loss on pre-flattened tensors.
+        focal_term = (1.0 - pt).pow(self.gamma)
+        loss_vec = -focal_term * log_pt
 
-        Args:
-            logits (torch.Tensor): Logits from the model of shape (N, C) where N is the number of samples (batch_size * seq_len) and C is the number of classes.
-            targets (torch.Tensor): Ground truth labels of shape (N,).
-            valid_mask (torch.Tensor): Boolean mask of shape (N,) where True indicates a valid (observed) value to include in the loss.
-
-        Returns:
-            torch.Tensor: The computed scalar loss value.
-        """
-        device = logits.device
+        if self.alpha is not None:
+            loss_vec = loss_vec * self.alpha[targets_v]
 
-        # Calculate standard cross-entropy loss per-token (no reduction)
-        ce = F.cross_entropy(
-            logits,
-            targets,
-            weight=(self.alpha.to(device) if self.alpha is not None else None),
-            reduction="none",
-            ignore_index=-1,
-        )
+        # Apply reduction
+        if self.reduction == "mean":
+            out = loss_vec.mean()
+        elif self.reduction == "sum":
+            out = loss_vec.sum()
+        else:  # "none"
+            out_flat = torch.zeros_like(targets_flat, dtype=loss_vec.dtype)
+            out_flat[valid_mask] = loss_vec
+            out = out_flat.view(targets.shape)
+
+        # Optional scaling (useful for VAE recon term)
+        if recon_scale is not None:
+            if not isinstance(recon_scale, torch.Tensor):
+                recon_scale = torch.tensor(
+                    float(recon_scale), device=out.device, dtype=out.dtype
+                )
+            else:
+                recon_scale = recon_scale.to(device=out.device, dtype=out.dtype)
 
-        # Calculate p_t from the cross-entropy loss
-        pt = torch.exp(-ce)
-        focal = ((1 - pt) ** self.gamma) * ce
+            out = out * recon_scale
 
-        # Apply the valid mask. We select only the elements that should contribute to the loss.
-        focal = focal[valid_mask]
+        return out
 
-        # Return early if no valid elements exist to avoid NaN results
-        if focal.numel() == 0:
-            return torch.tensor(0.0, device=device)
+    @property
+    def gamma(self) -> float:
+        return self._gamma
 
-        # Apply reduction
-        if self.reduction == "mean":
-            return focal.mean()
-        elif self.reduction == "sum":
-            return focal.sum()
-        else:
-            msg = f"Reduction mode '{self.reduction}' not supported."
-            raise ValueError(msg)
+    @gamma.setter
+    def gamma(self, value: torch.Tensor | float) -> None:
+        if isinstance(value, torch.Tensor):
+            value = float(value.item())
+        self._gamma = float(value)
 
 
 def safe_kl_gauss_unit(
     mu: torch.Tensor, logvar: torch.Tensor, reduction: str = "mean"
 ) -> torch.Tensor:
-    """KL divergence between N(mu, exp(logvar)) and N(0, I) with guards."""
-    logvar = logvar.clamp(min=-30.0, max=20.0)
-    kl = -0.5 * (1.0 + logvar - mu.pow(2) - logvar.exp())
+    """Compute KL divergence between N(mu, var) and N(0, I) with numeric guards.
+
+    Args:
+        mu (torch.Tensor): Latent mean (shape: [B, D]).
+        logvar (torch.Tensor): Latent log-variance (shape: [B, D]).
+        reduction (str): Reduction method ('mean' or 'sum').
+
+    Returns:
+        torch.Tensor: KL divergence (scalar).
+    """
+    kl = -0.5 * (1.0 + logvar - mu.pow(2) - logvar.exp())  # (B, D)
+    kl = kl.sum(dim=-1)  # (B,)
+
     if reduction == "sum":
         kl = kl.sum()
     elif reduction == "mean":
         kl = kl.mean()
+    else:
+        raise ValueError(f"Invalid reduction: {reduction}")
+
     return torch.nan_to_num(kl, nan=0.0, posinf=1e6, neginf=0.0)
 
 
 def compute_vae_loss(
+    criterion: nn.Module,
     recon_logits: torch.Tensor,
     targets: torch.Tensor,
     *,
     mu: torch.Tensor,
     logvar: torch.Tensor,
-    class_weights: torch.Tensor | None,
-    gamma: float,
-    beta: float,
-    ignore_index: int = -1,
+    kl_beta: torch.Tensor | float,
+    reduction: str = "mean",
+    recon_scale: torch.Tensor | float | None = None,
 ) -> torch.Tensor:
-    """Focal reconstruction + beta * KL with normalized class weights."""
-    cw = None
-    if class_weights is not None:
-        cw = class_weights / class_weights.mean().clamp_min(1e-8)
-
-    criterion = SafeFocalCELoss(
-        gamma=gamma,
-        weight=cw,
-        ignore_index=ignore_index,
-    )
-    rec = criterion(recon_logits.view(-1, recon_logits.size(-1)), targets.view(-1))
-    kl = safe_kl_gauss_unit(mu, logvar, reduction="mean")
-    return rec + beta * kl
+    """Compute VAE loss: reconstruction + KL divergence, with optional recon scaling.
+
+    Args:
+        criterion: Reconstruction loss module (e.g., FocalCELoss / CrossEntropyLoss).
+            Must accept (logits_2d, targets_1d). If it supports `recon_scale`, it will
+            be passed through; otherwise it will be called without it.
+        recon_logits: Reconstruction logits from decoder. Shape: (N, L, C) or (N_eval, C).
+        targets: Ground truth targets. Shape: (N, L) or (N_eval,).
+        mu: Latent mean. Shape: (B, D) (or compatible with safe_kl_gauss_unit).
+        logvar: Latent log-variance. Shape: (B, D).
+        kl_beta: Scalar KL weight.
+        reduction: KL reduction: "mean" or "sum".
+        recon_scale: Optional scalar multiplier applied to reconstruction term.
+            Use this to make reconstruction more "sum-like" for high-dimensional data.
+
+    Returns:
+        Scalar loss tensor.
+    """
+    # Flatten logits/targets to (N_total, C) and (N_total,)
+    if recon_logits.dim() == 3:
+        logits_2d = recon_logits.reshape(-1, recon_logits.size(-1))
+    elif recon_logits.dim() == 2:
+        logits_2d = recon_logits
+    else:
+        msg = f"recon_logits must be 2D or 3D; got shape {tuple(recon_logits.shape)}"
+        raise ValueError(msg)
+
+    tgt_1d = targets.reshape(-1) if targets.dim() > 1 else targets
+
+    # Reconstruction loss (criterion may ignore_index internally)
+    try:
+        rec = criterion(logits_2d, tgt_1d, recon_scale=recon_scale)
+    except TypeError:
+        # Criterion doesn't accept recon_scale (e.g., torch.nn.CrossEntropyLoss)
+        rec = criterion(logits_2d, tgt_1d)
+        if recon_scale is not None:
+            if isinstance(recon_scale, torch.Tensor):
+                rec = rec * recon_scale.to(device=rec.device, dtype=rec.dtype)
+            else:
+                rec = rec * float(recon_scale)
+
+    # KL term
+    kl = safe_kl_gauss_unit(mu, logvar, reduction=reduction)
+    loss = rec + (kl_beta * kl)
+
+    return torch.nan_to_num(loss, nan=1e6, posinf=1e6, neginf=1e6)

pgsui/impute/unsupervised/models/autoencoder_model.py

@@ -5,7 +5,6 @@ import torch
 import torch.nn as nn
 from snpio.utils.logging import LoggerManager
 
-from pgsui.impute.unsupervised.loss_functions import MaskedFocalLoss
 from pgsui.utils.logging_utils import configure_logger
 
 
@@ -44,8 +43,8 @@ class Encoder(nn.Module):
         for hidden_size in hidden_layer_sizes:
             layers.append(nn.Linear(input_dim, hidden_size))
             layers.append(nn.BatchNorm1d(hidden_size))
-            layers.append(nn.Dropout(dropout_rate))
             layers.append(activation)
+            layers.append(nn.Dropout(dropout_rate))
             input_dim = hidden_size
 
         self.hidden_layers = nn.Sequential(*layers)
@@ -98,8 +97,8 @@ class Decoder(nn.Module):
         for hidden_size in hidden_layer_sizes:
             layers.append(nn.Linear(input_dim, hidden_size))
             layers.append(nn.BatchNorm1d(hidden_size))
-            layers.append(nn.Dropout(dropout_rate))
             layers.append(activation)
+            layers.append(nn.Dropout(dropout_rate))
             input_dim = hidden_size
 
         self.hidden_layers = nn.Sequential(*layers)
@@ -128,17 +127,17 @@ class AutoencoderModel(nn.Module):
 
     **Model Architecture and Objective:**
 
-    The autoencoder consists of two parts: an encoder, $f_{\theta}$, and a decoder, $g_{\phi}$.
+    The autoencoder consists of two parts: an encoder, $f_{\\theta}$, and a decoder, $g_{\\phi}$.
         1.  The **encoder** maps the input data $x$ to a latent representation $z$:
             $$
             z = f_{\theta}(x)
             $$
-        2.  The **decoder** reconstructs the data $\hat{x}$ from the latent representation:
+        2.  The **decoder** reconstructs the data $\\hat{x}$ from the latent representation:
             $$
-            \hat{x} = g_{\phi}(z)
+            \\hat{x} = g_{\\phi}(z)
             $$
 
-    The model is trained by minimizing a reconstruction loss, $L(x, \hat{x})$, which measures the dissimilarity between the original input and the reconstructed output. This implementation uses a `MaskedFocalLoss` to handle missing values and class imbalance effectively.
+    The model is trained by minimizing a reconstruction loss, $L(x, \\hat{x})$, which measures the dissimilarity between the original input and the reconstructed output. This implementation uses a ``FocalCELoss`` to handle missing values and class imbalance effectively.
     """
 
     def __init__(
@@ -151,7 +150,7 @@ class AutoencoderModel(nn.Module):
         latent_dim: int = 2,
         dropout_rate: float = 0.2,
         activation: Literal["relu", "elu", "selu", "leaky_relu"] = "relu",
-        gamma: float = 2.0,
+        gamma: torch.Tensor = torch.tensor(2.0),
         device: Literal["cpu", "gpu", "mps"] = "cpu",
         verbose: bool = False,
         debug: bool = False,
@@ -222,47 +221,6 @@ class AutoencoderModel(nn.Module):
         reconstruction = self.decoder(z)
         return reconstruction
 
-    def compute_loss(
-        self,
-        reconstruction: torch.Tensor,
-        y: torch.Tensor,
-        mask: torch.Tensor | None = None,
-        class_weights: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """Computes the reconstruction loss for the Autoencoder model.
-
-        This method calculates the reconstruction loss using a masked focal loss, which is suitable for categorical data with missing values and class imbalance.
-
-        Args:
-            reconstruction (torch.Tensor): The reconstructed output (logits) from the model's forward pass.
-            y (torch.Tensor): The target data tensor, expected to be one-hot encoded. It is converted to class indices internally for the loss calculation.
-            mask (torch.Tensor | None): A boolean mask to exclude missing values from the loss calculation.
-            class_weights (torch.Tensor | None): Weights to apply to each class in the loss to handle imbalance.
-
-        Returns:
-            torch.Tensor: The computed scalar loss value.
-        """
-        if class_weights is None:
-            class_weights = torch.ones(self.num_classes, device=y.device)
-
-        logits_flat = reconstruction.view(-1, self.num_classes)
-        targets_flat = torch.argmax(y, dim=-1).view(-1)
-
-        if mask is None:
-            mask_flat = torch.ones_like(targets_flat, dtype=torch.bool)
-        else:
-            mask_flat = mask.view(-1)
-
-        criterion = MaskedFocalLoss(alpha=class_weights, gamma=self.gamma)
-
-        reconstruction_loss = criterion(
-            logits_flat.to(self.device),
-            targets_flat.to(self.device),
-            valid_mask=mask_flat.to(self.device),
-        )
-
-        return reconstruction_loss
-
     def _resolve_activation(
         self, activation: Literal["relu", "elu", "leaky_relu", "selu"]
     ) -> torch.nn.Module: