pg-sui 1.0.2.1__py3-none-any.whl → 1.6.8__py3-none-any.whl

pgsui/impute/unsupervised/models/vae_model.py

@@ -1,707 +1,348 @@
-import logging
-import os
-import sys
-import warnings
+from typing import List, Literal, Tuple
 
-# Import tensorflow with reduced warnings.
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-logging.getLogger("tensorflow").disabled = True
-warnings.filterwarnings("ignore", category=UserWarning)
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from snpio.utils.logging import LoggerManager
+from torch.distributions import Normal
 
-import tensorflow as tf
+from pgsui.impute.unsupervised.loss_functions import MaskedFocalLoss
 
-# Disable can't find cuda .dll errors. Also turns of GPU support.
-tf.config.set_visible_devices([], "GPU")
 
-from tensorflow.python.util import deprecation
+class Sampling(nn.Module):
+    """A layer that samples from a latent distribution using the reparameterization trick.
 
-# Disable warnings and info logs.
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-tf.get_logger().setLevel(logging.ERROR)
-
-
-# Monkey patching deprecation utils to supress warnings.
-# noinspection PyUnusedLocal
-def deprecated(
-    date, instructions, warn_once=True
-):  # pylint: disable=unused-argument
-    def deprecated_wrapper(func):
-        return func
-
-    return deprecated_wrapper
-
-
-deprecation.deprecated = deprecated
-
-from tensorflow.keras.layers import (
-    Dropout,
-    Dense,
-    Reshape,
-    Activation,
-    LeakyReLU,
-    PReLU,
-)
-
-from tensorflow.keras.regularizers import l1_l2
-from tensorflow.keras import backend as K
-
-# Custom Modules
-try:
-    from ..neural_network_methods import NeuralNetworkMethods
-except (ModuleNotFoundError, ValueError, ImportError):
-    from impute.unsupervised.neural_network_methods import NeuralNetworkMethods
+    This layer is a core component of a Variational Autoencoder (VAE). It takes the mean and log-variance of a latent distribution as input and generates a sample from that distribution. By using the reparameterization trick ($z = \mu + \sigma \cdot \epsilon$), it allows gradients to be backpropagated through the random sampling process, making the VAE trainable.
+    """
 
+    def forward(self, z_mean: torch.Tensor, z_log_var: torch.Tensor) -> torch.Tensor:
+        """Performs the forward pass to generate a latent sample.
 
-class Sampling(tf.keras.layers.Layer):
-    """Layer to calculate Z to sample from latent dimension."""
+        Args:
+            z_mean (torch.Tensor): The mean of the latent normal distribution.
+            z_log_var (torch.Tensor): The log of the variance of the latent normal distribution.
 
-    def __init__(self, *args, **kwargs):
-        self.is_placeholder = True
-        super(Sampling, self).__init__(*args, **kwargs)
+        Returns:
+            torch.Tensor: A sampled vector from the latent space.
+        """
+        z_sigma = torch.exp(0.5 * z_log_var)  # Precompute outside
 
-    def call(self, inputs):
-        """Sampling during forward pass."""
-        z_mean, z_log_var = inputs
-        z_sigma = tf.math.exp(0.5 * z_log_var)
-        batch = tf.shape(z_mean)[0]
-        dim = tf.shape(z_mean)[1]
-        epsilon = tf.random.normal(shape=(batch, dim))
+        # Ensure on GPU
+        # rand_like takes random values from a normal distribution
+        # of the same shape as z_mean.
+        epsilon = torch.randn_like(z_mean, device=z_mean.device)
         return z_mean + z_sigma * epsilon
 
 
-class Encoder(tf.keras.layers.Layer):
-    """VAE encoder to Encode genotypes to (z_mean, z_log_var, z).
-
-    Args:
-        n_features (int): Number of featuresi in input dataset.
-
-        num_classes (int): Number of classes in target data.
-
-        latent_dim (int): Number of latent dimensions to use.
-
-        hidden_layer_sizes (list of int): List of hidden layer sizes to use.
-
-        dropout_rate (float): Dropout rate for Dropout layer.
-
-        activation (str): Hidden activation function to use.
-
-        kernel_initializer (str): Initializer to use for weights.
-
-        kernel_regularizer (tf.keras.regularizers): L1 and/ or L2 objects.
-
-        beta (float, optional): KL divergence beta to use. Defualts to 1.0.
-
-        name (str): Name of model. Defaults to Encoder.
+class Encoder(nn.Module):
+    """The Encoder module of a Variational Autoencoder (VAE).
 
+    This module defines the encoder network, which takes high-dimensional input data and maps it to the parameters of a lower-dimensional latent distribution. The architecture consists of a series of fully-connected hidden layers that process the flattened input. The network culminates in two separate linear layers that output the mean (`z_mean`) and log-variance (`z_log_var`) of the approximate posterior distribution, $q(z|x)$.
     """
 
     def __init__(
         self,
-        n_features,
-        num_classes,
-        latent_dim,
-        hidden_layer_sizes,
-        dropout_rate,
-        activation,
-        kernel_initializer,
-        kernel_regularizer,
-        beta=1.0,
-        name="Encoder",
-        **kwargs,
+        n_features: int,
+        num_classes: int,
+        latent_dim: int,
+        hidden_layer_sizes: List[int],
+        dropout_rate: float,
+        activation: torch.nn.Module,
     ):
-        super(Encoder, self).__init__(name=name, **kwargs)
+        """Initializes the Encoder module.
+
+        Args:
+            n_features (int): The number of features in the input data (e.g., SNPs).
+            num_classes (int): The number of possible classes for each input element (e.g., 4 alleles).
+            latent_dim (int): The dimensionality of the latent space.
+            hidden_layer_sizes (List[int]): A list of integers specifying the size of each hidden layer.
+            dropout_rate (float): The dropout rate for regularization in the hidden layers.
+            activation (torch.nn.Module): An instantiated activation function module (e.g., `nn.ReLU()`) for the hidden layers.
+        """
+        super(Encoder, self).__init__()
+        self.flatten = nn.Flatten()
+        self.activation = (
+            getattr(F, activation) if isinstance(activation, str) else activation
+        )
 
-        self.beta = beta * latent_dim
+        layers = []
+        # The input dimension accounts for channels
+        input_dim = n_features * num_classes
+        for hidden_size in hidden_layer_sizes:
+            layers.append(nn.Linear(input_dim, hidden_size))
 
-        self.dense2 = None
-        self.dense3 = None
-        self.dense4 = None
-        self.dense5 = None
+            # BatchNorm can lead to faster convergence.
+            layers.append(nn.BatchNorm1d(hidden_size))
 
-        # # n_features * num_classes.
-        self.flatten = tf.keras.layers.Flatten()
+            layers.append(nn.Dropout(dropout_rate))
+            layers.append(activation)
+            input_dim = hidden_size
 
-        self.dense1 = Dense(
-            hidden_layer_sizes[0],
-            input_shape=(n_features * num_classes,),
-            activation=activation,
-            kernel_initializer=kernel_initializer,
-            kernel_regularizer=kernel_regularizer,
-            name="Encoder1",
-        )
+        self.hidden_layers = nn.Sequential(*layers)
+        self.dense_z_mean = nn.Linear(input_dim, latent_dim)
+        self.dense_z_log_var = nn.Linear(input_dim, latent_dim)
+        self.sampling = Sampling()
 
-        if len(hidden_layer_sizes) >= 2:
-            self.dense2 = Dense(
-                hidden_layer_sizes[1],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-                name="Encoder2",
-            )
-
-        if len(hidden_layer_sizes) >= 3:
-            self.dense3 = Dense(
-                hidden_layer_sizes[2],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-                name="Encoder3",
-            )
-
-        if len(hidden_layer_sizes) >= 4:
-            self.dense4 = Dense(
-                hidden_layer_sizes[3],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-                name="Encoder4",
-            )
-
-        if len(hidden_layer_sizes) == 5:
-            self.dense5 = Dense(
-                hidden_layer_sizes[4],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-                name="Encoder5",
-            )
-
-        self.dense_z_mean = Dense(
-            latent_dim,
-            name="z_mean",
-        )
-        self.dense_z_log_var = Dense(
-            latent_dim,
-            name="z_log_var",
-        )
-        # z_mean and z_log_var are inputs.
-        self.sampling = Sampling(
-            name="z",
-        )
+    def forward(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Performs the forward pass through the encoder.
 
-        self.dense_latent = Dense(
-            latent_dim,
-            activation=activation,
-            kernel_initializer=kernel_initializer,
-            kernel_regularizer=kernel_regularizer,
-            name="Encoder5",
-        )
+        Args:
+            x (torch.Tensor): The input data tensor of shape `(batch_size, n_features, num_classes)`.
 
-        self.dropout_layer = Dropout(dropout_rate)
-
-    def call(self, inputs, training=None):
-        """Forward pass for model."""
-        x = self.flatten(inputs)
-        x = self.dense1(x)
-        x = self.dropout_layer(x, training=training)
-        if self.dense2 is not None:
-            x = self.dense2(x)
-            x = self.dropout_layer(x, training=training)
-        if self.dense3 is not None:
-            x = self.dense3(x)
-            x = self.dropout_layer(x, training=training)
-        if self.dense4 is not None:
-            x = self.dense4(x)
-            x = self.dropout_layer(x, training=training)
-        if self.dense5 is not None:
-            x = self.dense5(x)
-            x = self.dropout_layer(x, training=training)
-
-        x = self.dense_latent(x)
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing the latent mean (`z_mean`), latent log-variance (`z_log_var`), and a sample from the latent distribution (`z`).
+        """
+        x = self.flatten(x)
+        x = self.hidden_layers(x)
         z_mean = self.dense_z_mean(x)
         z_log_var = self.dense_z_log_var(x)
-
-        # Compute the KL divergence
-        kl_loss = -0.5 * tf.reduce_sum(
-            1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var), axis=-1
-        )
-        # Add the KL divergence to the model's total loss
-        self.add_loss(self.beta * tf.reduce_mean(kl_loss))
-
-        z = self.sampling([z_mean, z_log_var])
-
+        z = self.sampling(z_mean, z_log_var)
         return z_mean, z_log_var, z
 
 
-class Decoder(tf.keras.layers.Layer):
-    """Converts z, the encoded vector, back into the reconstructed output.
-
-    Args:
-        n_features (int): Number of features in input dataset.
-
-        num_classes (int): Number of classes in input dataset.
+class Decoder(nn.Module):
+    """The Decoder module of a Variational Autoencoder (VAE).
 
-        latent_dim (int): Number of latent dimensions to use.
-
-        hidden_layer_sizes (list of int): List of hidden layer sizes to use.
-
-        dropout_rate (float): Dropout rate for Dropout layer.
-
-        activation (str): Hidden activation function to use.
-
-        kernel initializer (str): Function for initilizing weights.
-
-        kernel_regularizer (tf.keras.regularizer): Initialized L1 and/ or L2 regularizer.
-
-        name (str): Name of model. Defaults to "Decoder".
+    This module defines the decoder network, which takes a sample from the low-dimensional latent space and maps it back to the high-dimensional data space. It aims to reconstruct the original input data. The architecture consists of a series of fully-connected hidden layers followed by a final linear layer that produces the reconstructed data, which is then reshaped to match the original input's dimensions.
     """
 
     def __init__(
         self,
-        n_features,
-        num_classes,
-        latent_dim,
-        hidden_layer_sizes,
-        dropout_rate,
-        activation,
-        kernel_initializer,
-        kernel_regularizer,
-        name="Decoder",
-        **kwargs,
-    ):
-        super(Decoder, self).__init__(name=name, **kwargs)
-
-        self.dense2 = None
-        self.dense3 = None
-        self.dense4 = None
-        self.dense5 = None
-
-        self.dense1 = Dense(
-            hidden_layer_sizes[0],
-            input_shape=(latent_dim,),
-            activation=activation,
-            kernel_initializer=kernel_initializer,
-            kernel_regularizer=kernel_regularizer,
-            name="Decoder1",
-        )
-
-        if len(hidden_layer_sizes) >= 2:
-            self.dense2 = Dense(
-                hidden_layer_sizes[1],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-                name="Decoder2",
-            )
-
-        if len(hidden_layer_sizes) >= 3:
-            self.dense3 = Dense(
-                hidden_layer_sizes[2],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-                name="Decoder3",
-            )
-
-        if len(hidden_layer_sizes) >= 4:
-            self.dense4 = Dense(
-                hidden_layer_sizes[3],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-                name="Decoder4",
-            )
-
-        if len(hidden_layer_sizes) == 5:
-            self.dense5 = Dense(
-                hidden_layer_sizes[4],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-                name="Decoder5",
-            )
-
-        # No activation for final layer.
-        self.dense_output = Dense(
-            n_features * num_classes,
-            kernel_initializer=kernel_initializer,
-            kernel_regularizer=kernel_regularizer,
-            name="DecoderExpanded",
-        )
-
-        self.rshp = Reshape((n_features, num_classes))
-
-        self.dropout_layer = Dropout(dropout_rate)
-
-    def call(self, inputs, training=None):
-        """Forward pass for model."""
-        x = self.dense1(inputs)
-        x = self.dropout_layer(x, training=training)
-        if self.dense2 is not None:
-            x = self.dense2(x)
-            x = self.dropout_layer(x, training=training)
-        if self.dense3 is not None:
-            x = self.dense3(x)
-            x = self.dropout_layer(x, training=training)
-        if self.dense4 is not None:
-            x = self.dense4(x)
-            x = self.dropout_layer(x, training=training)
-        if self.dense5 is not None:
-            x = self.dense5(x)
-            x = self.dropout_layer(x, training=training)
-
-        x = self.dense_output(x)
-        return self.rshp(x)
-
-
-class VAEModel(tf.keras.Model):
-    """Variational Autoencoder model. Runs the encdoer and decoder and outputs the reconsruction.
-
-    Args:
-        output_shape (tuple): Shape of output. Defaults to None.
-
-        n_components (int): Number of latent dimensions to use. Defaults to 3.
-
-        weights_initializer (str, optional): kernel initializer to use. Defaults to "glorot_normal".
-
-        hidden_layer_sizes (str or list, optional): List of hidden layer sizes to use, or can use "midpoint", "log2", or "sqrt". Defaults to "midpoint".
-
-        num_hidden_layers (int, optional): Number of hidden layers to use. Defaults to 1.
-
-        hidden_activation (str, optional): Hidden activation function to use. Defaults to "elu".
+        n_features: int,
+        num_classes: int,
+        latent_dim: int,
+        hidden_layer_sizes: List[int],
+        dropout_rate: float,
+        activation: torch.nn.Module,
+    ) -> None:
+        """Initializes the Decoder module.
+
+        Args:
+            n_features (int): The number of features in the output data (e.g., SNPs).
+            num_classes (int): The number of possible classes for each output element (e.g., 4 alleles).
+            latent_dim (int): The dimensionality of the input latent space.
+            hidden_layer_sizes (List[int]): A list of integers specifying the size of each hidden layer.
+            dropout_rate (float): The dropout rate for regularization in the hidden layers.
+            activation (torch.nn.Module): An instantiated activation function module (e.g., `nn.ReLU()`) for the hidden layers.
+        """
+        super(Decoder, self).__init__()
 
-        l1_penalty (float, optional): L1 regularization penalty to use. Defaults to 1e-06.
+        layers = []
+        input_dim = latent_dim
+        for hidden_size in hidden_layer_sizes:
+            layers.append(nn.Linear(input_dim, hidden_size))
 
-        l2_penalty (float, optional): L2 regularization penalty. Defaults to 1e-06.
+            # BatchNorm can lead to faster convergence.
+            layers.append(nn.BatchNorm1d(hidden_size))
 
-        dropout_rate (float, optional): Dropout rate to use for Dropout layers.
+            layers.append(nn.Dropout(dropout_rate))
+            layers.append(activation)
+            input_dim = hidden_size
 
-        kl_beta (float, optional): Beta to use for KL divergence loss. The KL divergence gets multiplied by this value. Defaults to 1.0 (no scaling).
+        self.hidden_layers = nn.Sequential(*layers)
+        # UPDATED: Output dimension must account for channels
+        output_dim = n_features * num_classes
+        self.dense_output = nn.Linear(input_dim, output_dim)
+        # UPDATED: Reshape must account for channels
+        self.reshape = (n_features, num_classes)
 
-        num_classes (int, optional): Number of classes in input data. Defaults to 10.
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Performs the forward pass through the decoder.
 
-        sample_weight (list of float, optional): sample weights to use. Defaults to None.
+        Args:
+            x (torch.Tensor): The input latent tensor of shape `(batch_size, latent_dim)`.
 
-        missing_mask (np.ndarray, optional): Missing mask to use in model for masking missing values. Defaults to None.
+        Returns:
+            torch.Tensor: The reconstructed output data of shape `(batch_size, n_features, num_classes)`.
+        """
+        x = self.hidden_layers(x)
+        x = self.dense_output(x)
+        return x.view(-1, *self.reshape)
 
-        batch_size (int, optional): Batch size to use for training. Defaults to 32.
 
-        final_activation (str, optional): Final activation function to use. Should be either "sigmoid" (if doing multilabel classification) or "softmax" (if doing multiclass). If left None, then activation is not performed. Defaults to None.
+class VAEModel(nn.Module):
+    """A Variational Autoencoder (VAE) model for imputation.
 
-        y (np.ndarray, optional): Input dataset y. Should be the full dataset. It will get subset by a callback to get batch_size samples. Default to None.
+    This class combines an `Encoder` and a `Decoder` to form a VAE, a generative model for learning complex data distributions. It is designed for imputing missing values in categorical data, such as genomic SNPs. The model is trained by maximizing the Evidence Lower Bound (ELBO), which is a lower bound on the log-likelihood of the data.
 
+    **Objective Function (ELBO):**
+    The VAE loss function is derived from the ELBO and consists of two main components: a reconstruction term and a regularization term.
+    $$
+    \\mathcal{L}(\\theta, \\phi; x) = \\underbrace{\\mathbb{E}_{q_{\\phi}(z|x)}[\\log p_{\\theta}(x|z)]}_{\\text{Reconstruction Loss}} - \\underbrace{D_{KL}(q_{\\phi}(z|x) || p(z))}_{\\text{KL Divergence}}
+    $$
+    -   The **Reconstruction Loss** encourages the decoder to accurately reconstruct the input data from its latent representation. This implementation uses a `MaskedFocalLoss`.
+    -   The **KL Divergence** acts as a regularizer, forcing the approximate posterior distribution $q_{\\phi}(z|x)$ learned by the encoder to be close to a prior distribution $p(z)$ (typically a standard normal distribution).
     """
 
     def __init__(
         self,
-        output_shape=None,
-        n_components=3,
-        weights_initializer="glorot_normal",
-        hidden_layer_sizes="midpoint",
-        num_hidden_layers=1,
-        hidden_activation="elu",
-        l1_penalty=1e-6,
-        l2_penalty=1e-6,
-        dropout_rate=0.2,
-        kl_beta=1.0,
-        num_classes=10,
-        sample_weight=None,
-        missing_mask=None,
-        batch_size=32,
-        final_activation=None,
-        y=None,
+        n_features: int,
+        prefix: str,
+        *,
+        num_classes: int = 4,
+        hidden_layer_sizes: List[int] | np.ndarray = [128, 64],
+        latent_dim: int = 2,
+        dropout_rate: float = 0.2,
+        activation: Literal["relu", "elu", "selu", "leaky_relu"] = "relu",
+        gamma: float = 2.0,
+        beta: float = 1.0,
+        device: Literal["cpu", "gpu", "mps"] = "cpu",
+        verbose: bool = False,
+        debug: bool = False,
     ):
+        """Initializes the VAEModel.
+
+        Args:
+            n_features (int): The number of features in the input data (e.g., SNPs).
+            prefix (str): A prefix used for logging.
+            num_classes (int): The number of possible classes for each input element. Defaults to 4.
+            hidden_layer_sizes (List[int] | np.ndarray): A list of integers specifying the size of each hidden layer in the encoder and decoder. Defaults to [128, 64].
+            latent_dim (int): The dimensionality of the latent space. Defaults to 2.
+            dropout_rate (float): The dropout rate for regularization in the hidden layers. Defaults to 0.2.
+            activation (str): The name of the activation function to use in hidden layers. Defaults to "relu".
+            gamma (float): The focusing parameter for the focal loss component. Defaults to 2.0.
+            beta (float): A weighting factor for the KL divergence term in the total loss ($\beta$-VAE). Defaults to 1.0.
+            device (Literal["cpu", "gpu", "mps"]): The device to run the model on.
+            verbose (bool): If True, enables detailed logging. Defaults to False.
+            debug (bool): If True, enables debug mode. Defaults to False.
+        """
         super(VAEModel, self).__init__()
-
-        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
-        self.binary_accuracy_tracker = tf.keras.metrics.Mean(
-            name="binary_accuracy"
-        )
-
-        self.kl_beta = K.variable(0.0)
-        self.kl_beta._trainable = False
-
-        self._sample_weight = sample_weight
-        self._missing_mask = missing_mask
-        self._batch_idx = 0
-        self._batch_size = batch_size
-        self._y = y
-
-        self._final_activation = final_activation
-        if num_classes == 10 or num_classes == 3:
-            self.acc_func = tf.keras.metrics.categorical_accuracy
-        elif num_classes == 4:
-            self.acc_func = tf.keras.metrics.binary_accuracy
-
-        self.nn_ = NeuralNetworkMethods()
-
-        # y_train[1] dimension.
-        self.n_features = output_shape
-
-        self.n_components = n_components
-        self.weights_initializer = weights_initializer
-        self.hidden_layer_sizes = hidden_layer_sizes
-        self.num_hidden_layers = num_hidden_layers
-        self.hidden_activation = hidden_activation
-        self.l1_penalty = l1_penalty
-        self.l2_penalty = l2_penalty
-        self.dropout_rate = dropout_rate
         self.num_classes = num_classes
+        self.gamma = gamma
+        self.beta = beta
+        self.device = device
 
-        nn = NeuralNetworkMethods()
-
-        hidden_layer_sizes = nn.validate_hidden_layers(
-            self.hidden_layer_sizes, self.num_hidden_layers
-        )
-
-        hidden_layer_sizes = nn.get_hidden_layer_sizes(
-            self.n_features, self.n_components, hidden_layer_sizes, vae=True
+        logman = LoggerManager(
+            name=__name__, prefix=prefix, verbose=verbose, debug=debug
         )
+        self.logger = logman.get_logger()
 
-        hidden_layer_sizes = [h * self.num_classes for h in hidden_layer_sizes]
-
-        if self.l1_penalty == 0.0 and self.l2_penalty == 0.0:
-            kernel_regularizer = None
-        else:
-            kernel_regularizer = l1_l2(self.l1_penalty, self.l2_penalty)
-
-        kernel_initializer = self.weights_initializer
-
-        if self.hidden_activation.lower() == "leaky_relu":
-            activation = LeakyReLU(alpha=0.01)
-
-        elif self.hidden_activation.lower() == "prelu":
-            activation = PReLU()
-
-        elif self.hidden_activation.lower() == "selu":
-            activation = "selu"
-            kernel_initializer = "lecun_normal"
-
-        else:
-            activation = self.hidden_activation
-
-        if num_hidden_layers > 5:
-            raise ValueError(
-                f"The maximum number of hidden layers is 5, but got "
-                f"{num_hidden_layers}"
-            )
+        activation = self._resolve_activation(activation)
 
         self.encoder = Encoder(
-            self.n_features,
+            n_features,
             self.num_classes,
-            self.n_components,
+            latent_dim,
             hidden_layer_sizes,
-            self.dropout_rate,
+            dropout_rate,
             activation,
-            kernel_initializer,
-            kernel_regularizer,
-            beta=self.kl_beta,
         )
 
-        hidden_layer_sizes.reverse()
+        decoder_layer_sizes = list(reversed(hidden_layer_sizes))
 
         self.decoder = Decoder(
-            self.n_features,
+            n_features,
             self.num_classes,
-            self.n_components,
-            hidden_layer_sizes,
-            self.dropout_rate,
+            latent_dim,
+            decoder_layer_sizes,
+            dropout_rate,
             activation,
-            kernel_initializer,
-            kernel_regularizer,
         )
 
-        self.activation = Activation("sigmoid")
+    def forward(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Performs the forward pass through the full VAE model.
 
-    def call(self, inputs, training=None):
-        """Forward pass for model."""
-        z_mean, z_log_var, z = self.encoder(inputs)
-        reconstruction = self.decoder(z)
-        return self.activation(reconstruction)
+        Args:
+            x (torch.Tensor): The input data tensor of shape `(batch_size, n_features, num_classes)`.
 
-    def model(self):
-        """Here so that mymodel.model().summary() can be called for debugging.
-        :noindex:
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing the reconstructed output, the latent mean (`z_mean`), and the latent log-variance (`z_log_var`).
         """
-        x = tf.keras.Input(shape=(self.n_features, self.num_classes))
-        return tf.keras.Model(inputs=[x], outputs=self.call(x))
+        z_mean, z_log_var, z = self.encoder(x)
+        reconstruction = self.decoder(z)
+        return reconstruction, z_mean, z_log_var
 
-    def set_model_outputs(self):
-        """Set model output dimensions for building model.
-        :noindex:
+    def compute_loss(
+        self,
+        outputs: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+        y: torch.Tensor,
+        mask: torch.Tensor | None = None,
+        class_weights: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Computes the VAE loss function (negative ELBO).
+
+        The loss is the sum of a reconstruction term and a regularizing KL divergence term. The reconstruction loss is calculated using a masked focal loss, and the KL divergence measures the difference between the learned latent distribution and a standard normal prior.
+
+        Args:
+            outputs (Tuple[torch.Tensor, torch.Tensor, torch.Tensor]): The tuple of (reconstruction, z_mean, z_log_var) from the model's forward pass.
+            y (torch.Tensor): The target data tensor, expected to be one-hot encoded. This is converted to class indices internally for the loss function.
+            mask (torch.Tensor | None): A boolean mask to exclude missing values from the reconstruction loss.
+            class_weights (torch.Tensor | None): Weights to apply to each class in the reconstruction loss to handle imbalance.
+
+        Returns:
+            torch.Tensor: The computed scalar loss value.
         """
-        x = tf.keras.Input(shape=(self.n_features, self.num_classes))
-        model = tf.keras.Model(inputs=[x], outputs=self.call(x))
-        self.outputs = model.outputs
-
-    @property
-    def metrics(self):
-        return [
-            self.total_loss_tracker,
-            self.binary_accuracy_tracker,
-        ]
-
-    @tf.function
-    def train_step(self, data):
-        y = self._y
-
-        (
-            y_true,
-            sample_weight,
-            missing_mask,
-        ) = self.nn_.prepare_training_batches(
-            y,
-            y,
-            self._batch_size,
-            self._batch_idx,
-            True,
-            self.n_components,
-            self._sample_weight,
-            self._missing_mask,
-            ubp=False,
-        )
-
-        if sample_weight is not None:
-            sample_weight_masked = tf.convert_to_tensor(
-                sample_weight[~missing_mask], dtype=tf.float32
-            )
-        else:
-            sample_weight_masked = None
+        reconstruction, z_mean, z_log_var = outputs
 
-        y_true_masked = tf.boolean_mask(
-            tf.convert_to_tensor(y_true, dtype=tf.float32),
-            tf.reduce_any(tf.not_equal(y_true, -1), axis=-1),
+        # 1. KL Divergence Calculation
+        prior = Normal(torch.zeros_like(z_mean), torch.ones_like(z_log_var))
+        posterior = Normal(z_mean, torch.exp(0.5 * z_log_var))
+        kl_loss = (
+            torch.distributions.kl.kl_divergence(posterior, prior).sum(dim=1).mean()
         )
 
-        with tf.GradientTape() as tape:
-            reconstruction = self(y_true, training=True)
+        if class_weights is None:
+            class_weights = torch.ones(self.num_classes, device=y.device)
 
-            y_pred_masked = tf.boolean_mask(
-                reconstruction,
-                tf.reduce_any(tf.not_equal(y_true, -1), axis=-1),
-            )
+        # 2. Reconstruction Loss Calculation
+        # Reverting to the robust method of flattening tensors and using the
+        # custom loss function.
+        n_classes = reconstruction.shape[-1]
+        logits_flat = reconstruction.reshape(-1, n_classes)
 
-            # Returns binary crossentropy loss.
-            reconstruction_loss = self.compiled_loss(
-                y_true_masked,
-                y_pred_masked,
-                sample_weight=sample_weight_masked,
-                regularization_losses=self.losses,
-            )
+        # Convert one-hot `y` to class indices for the loss function.
+        targets_flat = torch.argmax(y, dim=-1).reshape(-1)
 
-            # Doesn't include KL Divergence Loss.
-            regularization_loss = sum(self.losses)
-
-            total_loss = reconstruction_loss + regularization_loss
-
-        grads = tape.gradient(total_loss, self.trainable_variables)
-        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
-
-        self.total_loss_tracker.update_state(total_loss)
-        self.binary_accuracy_tracker.update_state(
-            tf.keras.metrics.binary_accuracy(y_true_masked, y_pred_masked)
-        )
-
-        return {
-            "loss": self.total_loss_tracker.result(),
-            "binary_accuracy": self.binary_accuracy_tracker.result(),
-        }
-
-    @tf.function
-    def test_step(self, data):
-        y = self._y
-
-        (
-            y_true,
-            sample_weight,
-            missing_mask,
-        ) = self.nn_.prepare_training_batches(
-            y,
-            y,
-            self._batch_size,
-            self._batch_idx,
-            True,
-            self.n_components,
-            self._sample_weight,
-            self._missing_mask,
-            ubp=False,
-        )
-
-        if sample_weight is not None:
-            sample_weight_masked = tf.convert_to_tensor(
-                sample_weight[~missing_mask], dtype=tf.float32
-            )
+        if mask is None:
+            # If no mask is provided, all targets are considered valid.
+            mask_flat = torch.ones_like(targets_flat, dtype=torch.bool)
         else:
-            sample_weight_masked = None
-
-        y_true_masked = tf.boolean_mask(
-            tf.convert_to_tensor(y_true, dtype=tf.float32),
-            tf.reduce_any(tf.not_equal(y_true, -1), axis=-1),
-        )
+            # The mask needs to be reshaped to match the flattened targets.
+            mask_flat = mask.reshape(-1)
 
-        reconstruction = self(y_true, training=False)
+        # Logits, class-index targets, and the valid mask.
+        criterion = MaskedFocalLoss(alpha=class_weights, gamma=self.gamma)
 
-        y_pred_masked = tf.boolean_mask(
-            reconstruction,
-            tf.reduce_any(tf.not_equal(y_true, -1), axis=-1),
+        reconstruction_loss = criterion(
+            logits_flat.to(self.device),
+            targets_flat.to(self.device),
+            valid_mask=mask_flat.to(self.device),
         )
 
-        reconstruction_loss = self.compiled_loss(
-            y_true_masked,
-            y_pred_masked,
-            sample_weight=sample_weight_masked,
-            regularization_losses=self.losses,
-        )
+        return reconstruction_loss + self.beta * kl_loss
 
-        # Doesn't include KL Divergence Loss.
-        regularization_loss = sum(self.losses)
+    def _resolve_activation(
+        self, activation: Literal["relu", "elu", "leaky_relu", "selu"]
+    ) -> torch.nn.Module:
+        """Resolves an activation function module from a string name.
 
-        total_loss = reconstruction_loss + regularization_loss
+        Args:
+            activation (Literal["relu", "elu", "leaky_relu", "selu"]): The name of the activation function.
 
-        ### NOTE: If you get the error, "'tuple' object has no attribute
-        ### 'rank', then convert y_true to a tensor object."
-        self.total_loss_tracker.update_state(total_loss)
-        self.binary_accuracy_tracker.update_state(
-            tf.keras.metrics.binary_accuracy(y_true_masked, y_pred_masked)
-        )
+        Returns:
+            torch.nn.Module: The corresponding instantiated PyTorch activation function module.
 
-        return {
-            "loss": self.total_loss_tracker.result(),
-            "binary_accuracy": self.binary_accuracy_tracker.result(),
-        }
-
-    @property
-    def batch_size(self):
-        """Batch (=step) size per epoch."""
-        return self._batch_size
-
-    @property
-    def batch_idx(self):
-        """Current batch (=step) index."""
-        return self._batch_idx
-
-    @property
-    def y(self):
-        return self._y
-
-    @property
-    def missing_mask(self):
-        return self._missing_mask
-
-    @property
-    def sample_weight(self):
-        return self._sample_weight
-
-    @batch_size.setter
-    def batch_size(self, value):
-        """Set batch_size parameter."""
-        self._batch_size = int(value)
-
-    @batch_idx.setter
-    def batch_idx(self, value):
-        """Set current batch (=step) index."""
-        self._batch_idx = int(value)
-
-    @y.setter
-    def y(self, value):
-        """Set y after each epoch."""
-        self._y = value
-
-    @missing_mask.setter
-    def missing_mask(self, value):
-        """Set y after each epoch."""
-        self._missing_mask = value
-
-    @sample_weight.setter
-    def sample_weight(self, value):
-        self._sample_weight = value
+        Raises:
+            ValueError: If the provided activation name is not supported.
+        """
+        if isinstance(activation, str):
+            activation = activation.lower()
+            if activation == "relu":
+                return nn.ReLU()
+            elif activation == "elu":
+                return nn.ELU()
+            elif activation in ["leaky_relu", "leakyrelu"]:
+                return nn.LeakyReLU()
+            elif activation == "selu":
+                return nn.SELU()
+            else:
+                msg = f"Activation {activation} not supported."
+                self.logger.error(msg)
+                raise ValueError(msg)
+
+        return activation