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

pgsui/impute/unsupervised/models/ubp_model.py

@@ -1,1340 +1,221 @@
-import logging
-import os
-import sys
-import warnings
+from typing import List, Literal
 
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"  # or any {'0', '1', '2', '3'}
-logging.getLogger("tensorflow").disabled = True
+import numpy as np
+import torch
+import torch.nn as nn
+from snpio.utils.logging import LoggerManager
 
-# Import tensorflow with reduced warnings.
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-logging.getLogger("tensorflow").disabled = True
-warnings.filterwarnings("ignore", category=UserWarning)
+from pgsui.impute.unsupervised.loss_functions import MaskedFocalLoss
 
-# noinspection PyPackageRequirements
-import tensorflow as tf
 
-# Disable can't find cuda .dll errors. Also turns of GPU support.
-tf.config.set_visible_devices([], "GPU")
+class UBPModel(nn.Module):
+    """An Unsupervised Backpropagation (UBP) model with a multi-phase decoder.
 
-from tensorflow.python.util import deprecation
+    This class implements a deep neural network that serves as the decoder component in an unsupervised imputation pipeline. It's designed to reconstruct high-dimensional genomic data from a low-dimensional latent representation. The model features a unique multi-phase architecture with two distinct decoding paths:
 
-# Disable warnings and info logs.
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-tf.get_logger().setLevel(logging.ERROR)
+    1.  **Phase 1 Decoder:** A simple, shallow linear network.
+    2.  **Phase 2 & 3 Decoder:** A deeper, multi-layered, fully-connected network with batch normalization and dropout for regularization.
 
+    This phased approach allows for progressive training strategies. The model is tailored for two-channel allele data, where it learns to predict allele probabilities for each of the two channels at every SNP locus.
 
-# 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
+    **Model Architecture:**
 
-    return deprecated_wrapper
+    The model's forward pass maps a latent vector, $z$, to a reconstructed output, $\hat{x}$, via one of two paths.
 
+    -   **Phase 1 Path (Shallow Decoder):**
+        $$
+        \hat{x}_{p1} = W_{p1} z + b_{p1}
+        $$
 
-deprecation.deprecated = deprecated
+    -   **Phase 2/3 Path (Deep Decoder):**
+        This path uses a series of hidden layers with non-linear activations, $f(\cdot)$:
+        $$
+        h_1 = f(W_1 z + b_1)
+        $$
+        $$
+        \dots
+        $$
+        $$
+        h_L = f(W_L h_{L-1} + b_L)
+        $$
+        $$
+        \hat{x}_{p23} = W_{L+1} h_L + b_{L+1}
+        $$
 
-from tensorflow.keras.layers import (
-    Dropout,
-    Dense,
-    Reshape,
-    LeakyReLU,
-    PReLU,
-    Activation,
-)
-
-from tensorflow.keras.regularizers import l1_l2
-
-# Custom Modules
-try:
-    from ..neural_network_methods import NeuralNetworkMethods
-except (ModuleNotFoundError, ValueError, ImportError):
-    from impute.unsupervised.neural_network_methods import NeuralNetworkMethods
-
-
-class UBPPhase1(tf.keras.Model):
-    """UBP Phase 1 single layer perceptron model to train predict imputations.
-
-    This model is subclassed from the tensorflow/ Keras framework.
-
-    UBPPhase1 subclasses the tf.keras.Model and overrides the train_step function, which does training and evalutation for each batch in each epoch.
-
-    UBPPhase1 is a single-layer perceptron model used to initially refine V. After Phase 1 the Phase 1 weights are discarded.
-
-    Args:
-        V (numpy.ndarray(float)): V should have been randomly initialized and will be used as the input data that gets refined during training. Defaults to None.
-
-        y (numpy.ndarray): Target values to predict. Actual input data. Defaults to None.
-
-        batch_size (int, optional): Batch size per epoch. Defaults to 32.
-
-        missing_mask (numpy.ndarray): Missing data mask for y. Defaults to None.
-
-        output_shape (int): Output units for n_features dimension. Output will be of shape (batch_size, n_features). Defaults to None.
-
-        n_components (int, optional): Number of features in input V to use. Defaults to 3.
-
-        weights_initializer (str, optional): Kernel initializer to use for initializing model weights. Defaults to "glorot_normal".
-
-        hidden_layer_sizes (NoneType, optional): Output units for each hidden layer. List should be of same length as the number of hidden layers. Not used for UBP Phase 1, but is here for compatibility. Defaults to "midpoint".
-
-        num_hidden_layers (int, optional): Number of hidden layers to use. Not used in UBP Phase 1, but is here for compatibility. Defaults to 1.
-
-        hidden_activation (str, optional): Activation function to use for hidden layers. Defaults to "elu".
-
-        l1_penalty (float, optional): L1 regularization penalty to use to reduce overfitting. Defaults to 0.01.
-
-        l2_penalty (float, optional): L2 regularization penalty to use to reduce overfitting. Defaults to 0.01.
-
-        dropout_rate (float, optional): Dropout rate during training to reduce overfitting. Must be a float between 0 and 1. Defaults to 0.2.
-
-        num_classes (int, optional): Number of classes in output. Corresponds to the 3rd dimension of the output shape (batch_size, n_features, num_classes). Defaults to 3.
-
-        phase (int, optional): Current phase if doing UBP model. Defaults to 1.
-
-        sample_weight (numpy.ndarray, optional): 2D sample weights of shape (n_samples, n_features). Should have values for each class weighted. Defaults to None.
-
-    Example:
-        >>>model = UBPPhase1(V=V, y=y, batch_size=32, missing_mask=missing_mask, output_shape=y_train.shape[1], 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, num_classes=3, phase=3)
-        >>>model.compile(optimizer=optimizer, loss=loss_func, metrics=[my_metrics], run_eagerly=True)
-        >>>history = model.fit(X, y, batch_size=batch_size, epochs=epochs, callbacks=[MyCallback()], validation_split=validation_split, shuffle=False)
-
-    Raises:
-        TypeError: V, y, missing_mask, output_shape must not be NoneType.
-        ValueError: Maximum of 5 hidden layers.
-    """
-
-    def __init__(
-        self,
-        V=None,
-        y=None,
-        batch_size=32,
-        missing_mask=None,
-        output_shape=None,
-        n_components=3,
-        weights_initializer="glorot_normal",
-        hidden_layer_sizes="midpoint",
-        num_hidden_layers=1,
-        l1_penalty=0.01,
-        l2_penalty=0.01,
-        dropout_rate=0.2,
-        num_classes=3,
-        phase=1,
-        sample_weight=None,
-    ):
-        super(UBPPhase1, self).__init__()
-
-        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
-        self.binary_accuracy_tracker = tf.keras.metrics.Mean(
-            name="binary_accuracy"
-        )
-
-        nn = NeuralNetworkMethods()
-        self.nn = nn
-
-        if V is None:
-            self._V = nn.init_weights(y.shape[0], n_components)
-        elif isinstance(V, dict):
-            self._V = V[n_components]
-        else:
-            self._V = V
-
-        self._y = y
-
-        hidden_layer_sizes = nn.validate_hidden_layers(
-            hidden_layer_sizes, num_hidden_layers
-        )
-
-        hidden_layer_sizes = nn.get_hidden_layer_sizes(
-            y.shape[1], self._V.shape[1], hidden_layer_sizes
-        )
-
-        nn.validate_model_inputs(y, missing_mask, output_shape)
-
-        self._missing_mask = missing_mask
-        self.weights_initializer = weights_initializer
-        self.phase = phase
-        self.dropout_rate = dropout_rate
-        self._sample_weight = sample_weight
-
-        ### NOTE: I tried using just _V as the input to be refined, but it
-        # wasn't getting updated. So I copy it here and it works.
-        # V_latent is refined during train_step.
-        self.V_latent_ = self._V.copy()
-
-        # Initialize parameters used during train_step() and test_step.
-        # input_with_mask_ is set during the UBPCallbacks() execution.
-        self._batch_idx = 0
-        self._batch_size = batch_size
-        self.n_components = n_components
-
-        if l1_penalty == 0.0 and l2_penalty == 0.0:
-            kernel_regularizer = None
-        else:
-            kernel_regularizer = l1_l2(l1_penalty, l2_penalty)
-
-        self.kernel_regularizer = kernel_regularizer
-        kernel_initializer = weights_initializer
-
-        # Construct single-layer perceptron.
-
-        self.dense1 = Dense(
-            output_shape * num_classes,
-            input_shape=(n_components,),
-            kernel_initializer=kernel_initializer,
-            kernel_regularizer=kernel_regularizer,
-        )
-
-        self.rshp = Reshape((output_shape, num_classes))
-
-    def call(self, inputs):
-        x = self.dense1(inputs)
-        return self.rshp(x)
-
-    def model(self):
-        x = tf.keras.Input(shape=(self.n_components,))
-        return tf.keras.Model(inputs=[x], outputs=self.call(x))
-
-    def set_model_outputs(self):
-        x = tf.keras.Input(shape=(self.n_components,))
-        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,
-        ]
-
-    def train_step(self, data):
-        """Train step function. Parameters are set in the UBPCallbacks callback"""
-        y = self._y
-
-        (
-            v,
-            y_true,
-            sample_weight,
-            missing_mask,
-            batch_start,
-            batch_end,
-        ) = self.nn.prepare_training_batches(
-            self.V_latent_,
-            y,
-            self._batch_size,
-            self._batch_idx,
-            True,
-            self.n_components,
-            self._sample_weight,
-            self._missing_mask,
-        )
-
-        src = [v]
-
-        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
-
-        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=2),
-        )
-
-        # NOTE: Earlier model architectures incorrectly
-        # applied one gradient to all the variables, including
-        # the weights and v. Here we apply them separately, per
-        # the UBP manuscript.
-        with tf.GradientTape(persistent=True) as tape:
-            # Forward pass. Watch input tensor v.
-            tape.watch(v)
-            y_pred = self(v, training=True)
-            y_pred_masked = tf.boolean_mask(
-                y_pred, tf.reduce_any(tf.not_equal(y_true, -1), axis=2)
-            )
-            ### NOTE: If you get the error, "'tuple' object has no attribute
-            ### 'rank'", then convert y_true to a tensor object."
-            loss = self.compiled_loss(
-                y_true_masked,
-                y_pred_masked,
-                sample_weight=sample_weight_masked,
-            )
-
-        # Refine the watched variables with
-        # gradient descent backpropagation
-        gradients = tape.gradient(loss, self.trainable_variables)
-        self.optimizer.apply_gradients(
-            zip(gradients, self.trainable_variables)
-        )
-
-        # Apply separate gradients to v.
-        vgrad = tape.gradient(loss, src)
-        self.optimizer.apply_gradients(zip(vgrad, src))
-
-        del tape
-
-        # NOTE: run_eagerly must be set to True in the compile() method for this
-        # to work. Otherwise it can't convert a Tensor object to a numpy array.
-        # There is really no other way to set v back to V_latent_ in graph
-        # mode as far as I know. eager execution is slower, so it would be nice
-        # to find a way to do this without converting to numpy.
-        self.V_latent_[batch_start:batch_end, :] = v.numpy()
-
-        ### NOTE: If you get the error, "'tuple' object has no attribute
-        ### 'rank', then convert y_true to a tensor object."
-        # history object that gets returned from model.fit().
-        self.total_loss_tracker.update_state(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(),
-        }
-
-    def test_step(self, data):
-        """Train step function. Parameters are set in the UBPCallbacks callback"""
-        y = self._y
-
-        (
-            v,
-            y_true,
-            sample_weight,
-            missing_mask,
-            batch_start,
-            batch_end,
-        ) = self.nn.prepare_training_batches(
-            self.V_latent_,
-            y,
-            self._batch_size,
-            self._batch_idx,
-            True,
-            self.n_components,
-            self._sample_weight,
-            self._missing_mask,
-        )
-
-        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
-
-        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=2),
-        )
-
-        y_pred = self(v, training=False)
-        y_pred_masked = tf.boolean_mask(
-            y_pred, tf.reduce_any(tf.not_equal(y_true, -1), axis=2)
-        )
-        ### NOTE: If you get the error, "'tuple' object has no attribute
-        ### 'rank'", then convert y_true to a tensor object."
-        loss = self.compiled_loss(
-            y_true_masked,
-            y_pred_masked,
-            sample_weight=sample_weight_masked,
-        )
-
-        ### 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(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(),
-        }
-
-    @property
-    def V_latent(self):
-        """Randomly initialized input that gets refined during training.
-        :noindex:
-        """
-        return self.V_latent_
-
-    @property
-    def batch_size(self):
-        """Batch (=step) size per epoch.
-        :noindex:
-        """
-        return self._batch_size
-
-    @property
-    def batch_idx(self):
-        """Current batch (=step) index.
-        :noindex:
-        """
-        return self._batch_idx
-
-    @property
-    def y(self):
-        """Input dataset.
-        :noindex:
-        """
-        return self._y
-
-    @property
-    def missing_mask(self):
-        """Missing mask of shape (y.shape[0], y.shape[1])
-        :noindex:
-        """
-        return self._missing_mask
-
-    @property
-    def sample_weight(self):
-        """Sample weights of shape (y.shape[0], y.shape[1])
-        :noindex:
-        """
-        return self._sample_weight
-
-    @V_latent.setter
-    def V_latent(self, value):
-        """Set randomly initialized input. Gets refined during training.
-        :noindex:
-        """
-        self.V_latent_ = value
-
-    @batch_size.setter
-    def batch_size(self, value):
-        """Set batch_size parameter.
-        :noindex:
-        """
-        self._batch_size = int(value)
-
-    @batch_idx.setter
-    def batch_idx(self, value):
-        """Set current batch (=step) index.
-        :noindex:
-        """
-        self._batch_idx = int(value)
-
-    @y.setter
-    def y(self, value):
-        """Set y after each epoch.
-        :noindex:
-        """
-        self._y = value
-
-    @missing_mask.setter
-    def missing_mask(self, value):
-        """Set missing_mask after each epoch.
-        :noindex:
-        """
-        self._missing_mask = value
-
-    @sample_weight.setter
-    def sample_weight(self, value):
-        """Set sample_weight after each epoch.
-        :noindex:
-        """
-        self._sample_weight = value
-
-
-class UBPPhase2(tf.keras.Model):
-    """UBP Phase 2 model to train and use to predict imputations.
-
-    UBPPhase2 subclasses the tf.keras.Model and overrides the train_step function, which does training for each batch in each epoch.
-
-    Phase 2 does not refine V, it just refines the weights.
-
-    Args:
-        V (numpy.ndarray(float)): V should have been randomly initialized and will be used as the input data that gets refined during training. Defaults to None.
-
-        y (numpy.ndarray): Target values to predict. Actual input data. Defaults to None.
-
-        batch_size (int, optional): Batch size per epoch. Defaults to 32.
-
-        missing_mask (numpy.ndarray): Missing data mask for y. Defaults to None.
-
-        output_shape (int): Output units for n_features dimension. Output will be of shape (batch_size, n_features). Defaults to None.
-
-        n_components (int, optional): Number of features in input V to use. Defaults to 3.
-
-        weights_initializer (str, optional): Kernel initializer to use for initializing model weights. Defaults to "glorot_normal".
-
-        hidden_layer_sizes (NoneType, optional): Output units for each hidden layer. List should be of same length as the number of hidden layers. Not used for UBP Phase 1, but is here for compatibility. Defaults to "midpoint".
-
-        num_hidden_layers (int, optional): Number of hidden layers to use. Not used in UBP Phase 1, but is here for compatibility. Defaults to 1.
-
-        hidden_activation (str, optional): Activation function to use for hidden layers. Defaults to "elu".
-
-        l1_penalty (float, optional): L1 regularization penalty to use to reduce overfitting. Defaults to 0.01.
-
-        l2_penalty (float, optional): L2 regularization penalty to use to reduce overfitting. Defaults to 0.01.
-
-        dropout_rate (float, optional): Dropout rate during training to reduce overfitting. Must be a float between 0 and 1. Defaults to 0.2.
-
-        num_classes (int, optional): Number of classes in output. Corresponds to the 3rd dimension of the output shape (batch_size, n_features, num_classes). Defaults to 3.
-
-        phase (int, optional): Current phase if doing UBP model. Defaults to 1.
-
-        sample_weight (numpy.ndarray, optional): 2D sample weights of shape (n_samples, n_features). Should have values for each class weighted. Defaults to None.
-
-    Example:
-        >>>model = UBPPhase2(V=V, y=y, batch_size=32, missing_mask=missing_mask, output_shape=y_train.shape[1], 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, num_classes=3, phase=3)
-        >>>
-        >>>model.compile(optimizer=optimizer, loss=loss_func, metrics=[my_metrics], run_eagerly=True)
-        >>>
-        >>>history = model.fit(X, y, batch_size=batch_size, epochs=epochs, callbacks=[MyCallback()], validation_split=validation_split, shuffle=False)
-
-    Raises:
-        TypeError: V, y, missing_mask, output_shape must not be NoneType.
-        ValueError: Maximum of 5 hidden layers.
+    The final output from either path is reshaped into a tensor of shape `(batch_size, n_features, n_channels, n_classes)`. The model is optimized using a `MaskedFocalLoss` function, which effectively handles the missing data and class imbalance common in genomic datasets.
     """
 
     def __init__(
         self,
-        V=None,
-        y=None,
-        batch_size=32,
-        missing_mask=None,
-        output_shape=None,
-        n_components=3,
-        weights_initializer="glorot_normal",
-        hidden_layer_sizes="midpoint",
-        num_hidden_layers=1,
-        hidden_activation="elu",
-        l1_penalty=0.01,
-        l2_penalty=0.01,
-        dropout_rate=0.2,
-        num_classes=3,
-        phase=2,
-        sample_weight=None,
+        n_features: int,
+        prefix: str,
+        *,
+        num_classes: int = 3,
+        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,
+        device: Literal["cpu", "gpu", "mps"] = "cpu",
+        verbose: bool = False,
+        debug: bool = False,
     ):
-        super(UBPPhase2, self).__init__()
-
-        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
-        self.binary_accuracy_tracker = tf.keras.metrics.Mean(
-            name="binary_accuracy"
-        )
-
-        nn = NeuralNetworkMethods()
-        self.nn = nn
-
-        if V is None:
-            self._V = nn.init_weights(y.shape[0], n_components)
-        elif isinstance(V, dict):
-            self._V = V[n_components]
-        else:
-            self._V = V
-
-        self._y = y
-
-        hidden_layer_sizes = nn.validate_hidden_layers(
-            hidden_layer_sizes, num_hidden_layers
-        )
-
-        hidden_layer_sizes = nn.get_hidden_layer_sizes(
-            y.shape[1], self._V.shape[1], hidden_layer_sizes
-        )
-
-        nn.validate_model_inputs(y, missing_mask, output_shape)
-
-        self._missing_mask = missing_mask
-        self.weights_initializer = weights_initializer
-        self.phase = phase
-        self.dropout_rate = dropout_rate
-        self._sample_weight = sample_weight
-
-        ### NOTE: I tried using just _V as the input to be refined, but it
-        # wasn't getting updated. So I copy it here and it works.
-        # V_latent is refined during train_step.
-        self.V_latent_ = self._V.copy()
-
-        # Initialize parameters used during train_step.
-        self._batch_idx = 0
-        self._batch_size = batch_size
-        self.n_components = n_components
-
-        if l1_penalty == 0.0 and l2_penalty == 0.0:
-            kernel_regularizer = None
-        else:
-            kernel_regularizer = l1_l2(l1_penalty, l2_penalty)
-
-        self.kernel_regularizer = kernel_regularizer
-        kernel_initializer = weights_initializer
-
-        if hidden_activation.lower() == "leaky_relu":
-            activation = LeakyReLU(alpha=0.01)
-
-        elif hidden_activation.lower() == "prelu":
-            activation = PReLU()
-
-        elif hidden_activation.lower() == "selu":
-            activation = "selu"
-            kernel_initializer = "lecun_normal"
-
-        else:
-            activation = hidden_activation
-
-        if num_hidden_layers > 5:
-            raise ValueError(
-                f"The maximum number of hidden layers is 5, but got "
-                f"{num_hidden_layers}"
-            )
-
-        self.dense2 = None
-        self.dense3 = None
-        self.dense4 = None
-        self.dense5 = None
-
-        # Construct multi-layer perceptron.
-        # Add hidden layers dynamically.
-        self.dense1 = Dense(
-            hidden_layer_sizes[0],
-            input_shape=(n_components,),
-            activation=activation,
-            kernel_initializer=kernel_initializer,
-            kernel_regularizer=kernel_regularizer,
-        )
-
-        if num_hidden_layers >= 2:
-            self.dense2 = Dense(
-                hidden_layer_sizes[1],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-            )
-
-        if num_hidden_layers >= 3:
-            self.dense3 = Dense(
-                hidden_layer_sizes[2],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-            )
-
-        if num_hidden_layers >= 4:
-            self.dense4 = Dense(
-                hidden_layer_sizes[3],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-            )
-
-        if num_hidden_layers == 5:
-            self.dense5 = Dense(
-                hidden_layer_sizes[4],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-                kernel_regularizer=kernel_regularizer,
-            )
-
-        self.output1 = Dense(
-            output_shape * num_classes,
-            kernel_initializer=kernel_initializer,
-            kernel_regularizer=kernel_regularizer,
-        )
-
-        self.rshp = Reshape((output_shape, num_classes))
-
-        self.dropout_layer = Dropout(rate=dropout_rate)
-
-    def call(self, inputs, training=None):
-        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)
+        """Initializes the UBPModel.
 
-        x = self.output1(x)
-        return self.rshp(x)
+        Args:
+            n_features (int): The number of features (SNPs) in the input data.
+            prefix (str): A prefix used for logging.
+            num_classes (int): The number of possible allele classes (e.g., 3 for 0, 1, 2). Defaults to 3.
+            hidden_layer_sizes (list[int] | np.ndarray): A list of integers specifying the size of each hidden layer in the deep (Phase 2/3) decoder. Defaults to [128, 64].
+            latent_dim (int): The dimensionality of the input latent space. Defaults to 2.
+            dropout_rate (float): The dropout rate for regularization in the deep decoder. Defaults to 0.2.
+            activation (str): The non-linear activation function to use in the deep decoder's hidden layers. Defaults to 'relu'.
+            gamma (float): The focusing parameter for the focal loss function. Defaults to 2.0.
+            device (Literal["cpu", "gpu", "mps"]): The PyTorch device to run the model on. Defaults to 'cpu'.
+            verbose (bool): If True, enables detailed logging. Defaults to False.
+            debug (bool): If True, enables debug mode. Defaults to False.
+        """
+        super(UBPModel, self).__init__()
+
+        logman = LoggerManager(
+            name=__name__, prefix=prefix, verbose=verbose, debug=debug
+        )
+        self.logger = logman.get_logger()
+
+        self.n_features = n_features
+        self.num_classes = num_classes
+        self.latent_dim = latent_dim
+        self.gamma = gamma
+        self.device = device
+
+        if isinstance(hidden_layer_sizes, np.ndarray):
+            hidden_layer_sizes = hidden_layer_sizes.tolist()
+
+        # Final layer output size is now n_features * num_classes
+        final_output_size = n_features * num_classes
+
+        # Phase 1 decoder: Simple linear model
+        self.phase1_decoder = nn.Sequential(
+            nn.Linear(latent_dim, final_output_size, device=device),
+        )
+
+        # Phase 2 & 3 uses the Convolutional Decoder
+        act_factory = self._resolve_activation_factory(activation)
+
+        if hidden_layer_sizes[0] > hidden_layer_sizes[-1]:
+            hidden_layer_sizes = list(reversed(hidden_layer_sizes))
+
+        # Phase 2 & 3: Flexible deeper network
+        layers = []
+        input_dim = latent_dim
+        for size in hidden_layer_sizes:
+            layers.append(nn.Linear(input_dim, size))
+            layers.append(nn.BatchNorm1d(size))
+            layers.append(nn.Dropout(dropout_rate))
+            layers.append(act_factory())
+            input_dim = size
+
+        layers.append(nn.Linear(hidden_layer_sizes[-1], final_output_size))
+
+        self.phase23_decoder = nn.Sequential(*layers)
+        self.reshape = (self.n_features, self.num_classes)
+
+    def _resolve_activation_factory(
+        self, activation: Literal["relu", "elu", "selu", "leaky_relu"]
+    ) -> callable:
+        """Resolves an activation function factory from a string name.
+
+        This method acts as a factory, returning a callable (lambda function) that produces the desired PyTorch activation function module when called.
+
+        Args:
+            activation (Literal["relu", "elu", "selu", "leaky_relu"]): The name of the activation function.
+
+        Returns:
+            callable: A factory function that, when called, returns an instance of the specified activation layer.
+
+        Raises:
+            ValueError: If the provided activation name is not supported.
+        """
+        a = activation.lower()
+        if a == "relu":
+            return lambda: nn.ReLU()
+        if a == "elu":
+            return lambda: nn.ELU()
+        if a == "leaky_relu":
+            return lambda: nn.LeakyReLU()
+        if a == "selu":
+            return lambda: nn.SELU()
+
+        msg = f"Activation function {activation} not supported."
+        self.logger.error(msg)
+        raise ValueError(msg)
+
+    def forward(self, x: torch.Tensor, phase: int = 1) -> torch.Tensor:
+        """Performs the forward pass through the UBP model.
+
+        This method routes the input tensor through the appropriate decoder based on the specified training `phase`. The final output is reshaped to match the target data structure of `(batch_size, n_features, n_channels, n_classes)`.
+
+        Args:
+            x (torch.Tensor): The input latent tensor of shape `(batch_size, latent_dim)`.
+            phase (int): The training phase (1, 2, or 3), which determines which decoder path to use.
+
+        Returns:
+            torch.Tensor: The reconstructed output tensor.
 
-    def model(self):
-        x = tf.keras.Input(shape=(self.n_components,))
-        return tf.keras.Model(inputs=[x], outputs=self.call(x))
-
-    def set_model_outputs(self):
-        x = tf.keras.Input(shape=(self.n_components,))
-        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,
-        ]
-
-    def train_step(self, data):
-        """Train step function. Parameters are set in the UBPCallbacks callback"""
-        y = self._y
-
-        (
-            v,
-            y_true,
-            sample_weight,
-            missing_mask,
-            _,
-            __,
-        ) = self.nn.prepare_training_batches(
-            self.V_latent_,
-            y,
-            self._batch_size,
-            self._batch_idx,
-            True,
-            self.n_components,
-            self._sample_weight,
-            self._missing_mask,
-        )
-
-        if sample_weight is not None:
-            sample_weight_masked = tf.convert_to_tensor(
-                sample_weight[~missing_mask], dtype=tf.float32
-            )
+        Raises:
+            ValueError: If an invalid phase is provided.
+        """
+        if phase == 1:
+            # Linear decoder for phase 1
+            x = self.phase1_decoder(x)
+            return x.view(-1, *self.reshape)
+        elif phase in {2, 3}:
+            x = self.phase23_decoder(x)
+            return x.view(-1, *self.reshape)
         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=2),
-        )
-
-        # NOTE: Earlier model architectures incorrectly
-        # applied one gradient to all the variables, including
-        # the weights and v. Here we apply them separately, per
-        # the UBP manuscript.
-        with tf.GradientTape() as tape:
-            # Forward pass
-            y_pred = self(v, training=True)
-            y_pred_masked = tf.boolean_mask(
-                y_pred, tf.reduce_any(tf.not_equal(y_true, -1), axis=2)
-            )
-            ### NOTE: If you get the error, "'tuple' object has no attribute
-            ### 'rank'", then convert y_true to a tensor object."
-            loss = self.compiled_loss(
-                y_true_masked,
-                y_pred_masked,
-                sample_weight=sample_weight_masked,
-            )
-
-        # Refine the watched variables with backpropagation
-        gradients = tape.gradient(loss, self.trainable_variables)
-        self.optimizer.apply_gradients(
-            zip(gradients, self.trainable_variables)
-        )
-
-        ### 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(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(),
-        }
-
-    def test_step(self, data):
-        """Train step function. Parameters are set in the UBPCallbacks callback"""
-        y = self._y
-
-        (
-            v,
-            y_true,
-            sample_weight,
-            missing_mask,
-            _,
-            __,
-        ) = self.nn.prepare_training_batches(
-            self.V_latent_,
-            y,
-            self._batch_size,
-            self._batch_idx,
-            True,
-            self.n_components,
-            self._sample_weight,
-            self._missing_mask,
-        )
-
-        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
-
-        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=2),
-        )
+            msg = f"Invalid phase: {phase}. Expected 1, 2, or 3."
+            self.logger.error(msg)
+            raise ValueError(msg)
 
-        y_pred = self(v, training=False)
-        y_pred_masked = tf.boolean_mask(
-            y_pred, tf.reduce_any(tf.not_equal(y_true, -1), axis=2)
-        )
-        ### NOTE: If you get the error, "'tuple' object has no attribute
-        ### 'rank'", then convert y_true to a tensor object."
-        loss = self.compiled_loss(
-            y_true_masked,
-            y_pred_masked,
-            sample_weight=sample_weight_masked,
-        )
-
-        ### 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(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(),
-        }
-
-    @property
-    def V_latent(self):
-        """Randomly initialized input variable that gets refined during training.
-        :noindex:
-        """
-        return self.V_latent_
-
-    @property
-    def batch_size(self):
-        """Batch (=step) size per epoch.
-        :noindex:
-        """
-        return self._batch_size
-
-    @property
-    def batch_idx(self):
-        """Current batch (=step) index.
-        :noindex:
-        """
-        return self._batch_idx
-
-    @property
-    def y(self):
-        """Full input dataset.
-        :noindex:
-        """
-        return self._y
-
-    @property
-    def missing_mask(self):
-        """Get missing_mask for current epoch.
-        :noindex:
-        """
-        return self._missing_mask
-
-    @property
-    def sample_weight(self):
-        """Get sample_weight for current epoch.
-        :noindex:
-        """
-        return self._sample_weight
-
-    @V_latent.setter
-    def V_latent(self, value):
-        """Set randomly initialized input variable. Gets refined during training.
-        :noindex:
-        """
-        self.V_latent_ = value
-
-    @batch_size.setter
-    def batch_size(self, value):
-        """Set batch_size parameter.
-        :noindex:
-        """
-        self._batch_size = int(value)
-
-    @batch_idx.setter
-    def batch_idx(self, value):
-        """Set current batch (=step) index.
-        :noindex:
-        """
-        self._batch_idx = int(value)
-
-    @y.setter
-    def y(self, value):
-        """Set y after each epoch.
-        :noindex:
-        """
-        self._y = value
-
-    @missing_mask.setter
-    def missing_mask(self, value):
-        """Set missing_mask after each epoch.
-        :noindex:
-        """
-        self._missing_mask = value
-
-    @sample_weight.setter
-    def sample_weight(self, value):
-        """Set sample_weight after each epoch.
-        :noindex:
-        """
-        self._sample_weight = value
-
-
-class UBPPhase3(tf.keras.Model):
-    """UBP Phase 3 model to train and use to predict imputations.
-
-    UBPPhase3 subclasses the tf.keras.Model and overrides the train_step function, which does training and evaluation for each batch in each single epoch.
-
-    Phase 3 Refines both the weights and V.
-
-    Args:
-        V (numpy.ndarray(float)): V should have been randomly initialized and will be used as the input data that gets refined during training. Defaults to None.
-
-        y (numpy.ndarray): Target values to predict. Actual input data. Defaults to None.
-
-        batch_size (int, optional): Batch size per epoch. Defaults to 32.
-
-        missing_mask (numpy.ndarray): Missing data mask for y. Defaults to None.
-
-        output_shape (int): Output units for n_features dimension. Output will be of shape (batch_size, n_features). Defaults to None.
-
-        n_components (int, optional): Number of features in input V to use. Defaults to 3.
-
-        weights_initializer (str, optional): Kernel initializer to use for initializing model weights. Defaults to "glorot_normal".
-
-        hidden_layer_sizes (NoneType, optional): Output units for each hidden layer. List should be of same length as the number of hidden layers. Not used for UBP Phase 1, but is here for compatibility. Defaults to "midpoint".
-
-        num_hidden_layers (int, optional): Number of hidden layers to use. Not used in UBP Phase 1, but is here for compatibility. Defaults to 1.
-
-        hidden_activation (str, optional): Activation function to use for hidden layers. Defaults to "elu".
-
-        l1_penalty (float, optional): L1 regularization penalty to use to reduce overfitting. Defaults to 0.01.
-
-        l2_penalty (float, optional): L2 regularization penalty to use to reduce overfitting. Defaults to 0.01.
-
-        dropout_rate (float, optional): Dropout rate during training to reduce overfitting. Must be a float between 0 and 1. Defaults to 0.2.
-
-        num_classes (int, optional): Number of classes in output. Corresponds to the 3rd dimension of the output shape (batch_size, n_features, num_classes). Defaults to 3.
-
-        phase (int, optional): Current phase if doing UBP model. Defaults to 1.
-
-        sample_weight (numpy.ndarray, optional): 2D sample weights of shape (n_samples, n_features). Should have values for each class weighted. Defaults to None.
-
-    Example:
-        >>>model = UBPPhase3(V=V, y=y, batch_size=32, missing_mask=missing_mask, output_shape=y_train.shape[1], 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, num_classes=3, phase=3)
-        >>>
-        >>>model.compile(optimizer=optimizer, loss=loss_func, metrics=[my_metrics], run_eagerly=True)
-        >>>
-        >>>history = model.fit(X, y, batch_size=batch_size, epochs=epochs, callbacks=[MyCallback()], validation_split=validation_split, shuffle=False)
-
-    Raises:
-        TypeError: V, y, missing_mask, output_shape must not be NoneType.
-        ValueError: Maximum of 5 hidden layers.
-    """
-
-    def __init__(
+    def compute_loss(
         self,
-        V=None,
-        y=None,
-        batch_size=32,
-        missing_mask=None,
-        output_shape=None,
-        n_components=3,
-        weights_initializer="glorot_normal",
-        hidden_layer_sizes="midpoint",
-        num_hidden_layers=1,
-        hidden_activation="elu",
-        dropout_rate=0.2,
-        num_classes=3,
-        phase=3,
-        sample_weight=None,
-    ):
-        super(UBPPhase3, self).__init__()
-
-        self.total_loss_tracker = tf.keras.metrics.Mean(name="total_loss")
-        self.binary_accuracy_tracker = tf.keras.metrics.Mean(
-            name="binary_accuracy"
-        )
-
-        nn = NeuralNetworkMethods()
-        self.nn = nn
-
-        if V is None:
-            self._V = nn.init_weights(y.shape[0], n_components)
-        elif isinstance(V, dict):
-            self._V = V[n_components]
-        else:
-            self._V = V
-
-        self._y = y
-
-        hidden_layer_sizes = nn.validate_hidden_layers(
-            hidden_layer_sizes, num_hidden_layers
+        y: torch.Tensor,
+        outputs: torch.Tensor,
+        mask: torch.Tensor | None = None,
+        class_weights: torch.Tensor | None = None,
+        gamma: float = 2.0,
+    ) -> torch.Tensor:
+        """Computes the masked focal loss between model outputs and ground truth.
+
+        This method calculates the loss value, handling class imbalance with weights and ignoring masked (missing) values in the ground truth tensor.
+
+        Args:
+            y (torch.Tensor): The ground truth tensor of shape `(batch_size, n_features, n_channels)`.
+            outputs (torch.Tensor): The model's raw output (logits) of shape `(batch_size, n_features, n_channels, n_classes)`.
+            mask (torch.Tensor | None): An optional boolean mask indicating which elements should be included in the loss calculation.
+            class_weights (torch.Tensor | None): An optional tensor of weights for each class to address imbalance.
+            gamma (float): The focusing parameter for the focal loss.
+
+        Returns:
+            torch.Tensor: The computed scalar loss value.
+        """
+        if class_weights is None:
+            class_weights = torch.ones(self.num_classes, device=outputs.device)
+
+        if mask is None:
+            mask = torch.ones_like(y, dtype=torch.bool)
+
+        # Explicitly flatten all tensors to the (N, C) and (N,) format.
+        # This creates a clear contract with the new MaskedFocalLoss function.
+        n_classes = outputs.shape[-1]
+        logits_flat = outputs.reshape(-1, n_classes)
+        targets_flat = y.reshape(-1)
+        mask_flat = mask.reshape(-1)
+
+        criterion = MaskedFocalLoss(gamma=gamma, alpha=class_weights)
+
+        return criterion(
+            logits_flat.to(self.device),
+            targets_flat.to(self.device),
+            valid_mask=mask_flat.to(self.device),
         )
-
-        hidden_layer_sizes = nn.get_hidden_layer_sizes(
-            y.shape[1], self._V.shape[1], hidden_layer_sizes
-        )
-
-        nn.validate_model_inputs(y, missing_mask, output_shape)
-
-        self._missing_mask = missing_mask
-        self.weights_initializer = weights_initializer
-        self.phase = phase
-        self.dropout_rate = dropout_rate
-        self._sample_weight = sample_weight
-
-        ### NOTE: I tried using just _V as the input to be refined, but it
-        # wasn't getting updated. So I copy it here and it works.
-        # V_latent is refined during train_step.
-        self.V_latent_ = self._V.copy()
-
-        # Initialize parameters used during train_step.
-        self._batch_idx = 0
-        self._batch_size = batch_size
-        self.n_components = n_components
-
-        # No regularization in phase 3.
-        kernel_regularizer = None
-        self.kernel_regularizer = kernel_regularizer
-        kernel_initializer = None
-
-        if hidden_activation.lower() == "leaky_relu":
-            activation = LeakyReLU(alpha=0.01)
-
-        elif hidden_activation.lower() == "prelu":
-            activation = PReLU()
-
-        elif hidden_activation.lower() == "selu":
-            activation = "selu"
-            kernel_initializer = "lecun_normal"
-
-        else:
-            activation = hidden_activation
-
-        if num_hidden_layers > 5:
-            raise ValueError(
-                f"The maximum number of hidden layers is 5, but got "
-                f"{num_hidden_layers}"
-            )
-
-        self.dense2 = None
-        self.dense3 = None
-        self.dense4 = None
-        self.dense5 = None
-
-        # Construct multi-layer perceptron.
-        # Add hidden layers dynamically.
-        self.dense1 = Dense(
-            hidden_layer_sizes[0],
-            input_shape=(n_components,),
-            activation=activation,
-            kernel_initializer=kernel_initializer,
-        )
-
-        if num_hidden_layers >= 2:
-            self.dense2 = Dense(
-                hidden_layer_sizes[1],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-            )
-
-        if num_hidden_layers >= 3:
-            self.dense3 = Dense(
-                hidden_layer_sizes[2],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-            )
-
-        if num_hidden_layers >= 4:
-            self.dense4 = Dense(
-                hidden_layer_sizes[3],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-            )
-
-        if num_hidden_layers == 5:
-            self.dense5 = Dense(
-                hidden_layer_sizes[4],
-                activation=activation,
-                kernel_initializer=kernel_initializer,
-            )
-
-        self.output1 = Dense(
-            output_shape * num_classes,
-            kernel_initializer=kernel_initializer,
-        )
-
-        self.rshp = Reshape((output_shape, num_classes))
-
-        self.dropout_layer = Dropout(rate=dropout_rate)
-        self.activation = Activation("sigmoid")
-
-    def call(self, inputs, training=None):
-        x = self.dense1(inputs)
-        if self.dense2 is not None:
-            x = self.dense2(x)
-        if self.dense3 is not None:
-            x = self.dense3(x)
-        if self.dense4 is not None:
-            x = self.dense4(x)
-        if self.dense5 is not None:
-            x = self.dense5(x)
-
-        x = self.output1(x)
-        x = self.rshp(x)
-        return self.activation(x)
-
-    def model(self):
-        x = tf.keras.Input(shape=(self.n_components,))
-        return tf.keras.Model(inputs=[x], outputs=self.call(x))
-
-    def set_model_outputs(self):
-        x = tf.keras.Input(shape=(self.n_components,))
-        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,
-        ]
-
-    def train_step(self, data):
-        """Train step function. Parameters are set in the UBPCallbacks callback"""
-        y = self._y
-
-        (
-            v,
-            y_true,
-            sample_weight,
-            missing_mask,
-            batch_start,
-            batch_end,
-        ) = self.nn.prepare_training_batches(
-            self.V_latent_,
-            y,
-            self._batch_size,
-            self._batch_idx,
-            True,
-            self.n_components,
-            self._sample_weight,
-            self._missing_mask,
-        )
-
-        src = [v]
-
-        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
-
-        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=2),
-        )
-
-        # NOTE: Earlier model architectures incorrectly
-        # applied one gradient to all the variables, including
-        # the weights and v. Here we apply them separately, per
-        # the UBP manuscript.
-        with tf.GradientTape(persistent=True) as tape:
-            # Forward pass. Watch input tensor v.
-            tape.watch(v)
-            y_pred = self(v, training=True)
-            y_pred_masked = tf.boolean_mask(
-                y_pred, tf.reduce_any(tf.not_equal(y_true, -1), axis=2)
-            )
-            ### NOTE: If you get the error, "'tuple' object has no attribute
-            ### 'rank'", then convert y_true to a tensor object."
-            loss = self.compiled_loss(
-                y_true_masked,
-                y_pred_masked,
-                sample_weight=sample_weight_masked,
-            )
-
-        # Refine the watched variables with
-        # gradient descent backpropagation
-        gradients = tape.gradient(loss, self.trainable_variables)
-        self.optimizer.apply_gradients(
-            zip(gradients, self.trainable_variables)
-        )
-
-        # Apply separate gradients to v.
-        vgrad = tape.gradient(loss, src)
-        self.optimizer.apply_gradients(zip(vgrad, src))
-
-        del tape
-
-        # NOTE: run_eagerly must be set to True in the compile() method for this
-        # to work. Otherwise it can't convert a Tensor object to a numpy array.
-        # There is really no other way to set v back to V_latent_ in graph
-        # mode as far as I know. eager execution is slower, so it would be nice
-        # to find a way to do this without converting to numpy.
-        self.V_latent_[batch_start:batch_end, :] = v.numpy()
-
-        self.total_loss_tracker.update_state(loss)
-        self.binary_accuracy_tracker.update_state(
-            tf.keras.metrics.binary_accuracy(y_true_masked, y_pred_masked)
-        )
-
-        ### NOTE: If you get the error, "'tuple' object has no attribute
-        ### 'rank', then convert y_true to a tensor object."
-        return {
-            "loss": self.total_loss_tracker.result(),
-            "binary_accuracy": self.binary_accuracy_tracker.result(),
-        }
-
-    def test_step(self, data):
-        """Train step function. Parameters are set in the UBPCallbacks callback"""
-        y = self._y
-
-        (
-            v,
-            y_true,
-            sample_weight,
-            missing_mask,
-            batch_start,
-            batch_end,
-        ) = self.nn.prepare_training_batches(
-            self.V_latent_,
-            y,
-            self._batch_size,
-            self._batch_idx,
-            True,
-            self.n_components,
-            self._sample_weight,
-            self._missing_mask,
-        )
-
-        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
-
-        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=2),
-        )
-
-        y_pred = self(v, training=False)
-        y_pred_masked = tf.boolean_mask(
-            y_pred, tf.reduce_any(tf.not_equal(y_true, -1), axis=2)
-        )
-        ### NOTE: If you get the error, "'tuple' object has no attribute
-        ### 'rank'", then convert y_true to a tensor object."
-        loss = self.compiled_loss(
-            y_true_masked,
-            y_pred_masked,
-            sample_weight=sample_weight_masked,
-        )
-
-        ### 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(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(),
-        }
-
-    @property
-    def V_latent(self):
-        """Randomly initialized input variable that gets refined during training.
-        :noindex:
-        """
-        return self.V_latent_
-
-    @property
-    def batch_size(self):
-        """Batch (=step) size per epoch.
-        :noindex:
-        """
-        return self._batch_size
-
-    @property
-    def batch_idx(self):
-        """Current batch (=step) index.
-        :noindex:
-        """
-        return self._batch_idx
-
-    @property
-    def y(self):
-        """Full input dataset y.
-        :noindex:
-        """
-        return self._y
-
-    @property
-    def missing_mask(self):
-        """Missing mask of shape (y.shape[0], y.shape[1])
-        :noindex:
-        """
-        return self._missing_mask
-
-    @property
-    def sample_weight(self):
-        """Sample weights of shpe (y.shape[0], y.shape[1])
-        :noindex:
-        """
-        return self._sample_weight
-
-    @V_latent.setter
-    def V_latent(self, value):
-        """Set randomly initialized input variable. Refined during training.
-        :noindex:
-        """
-        self.V_latent_ = value
-
-    @batch_size.setter
-    def batch_size(self, value):
-        """Set batch_size parameter.
-        :noindex:
-        """
-        self._batch_size = int(value)
-
-    @batch_idx.setter
-    def batch_idx(self, value):
-        """Set current batch (=step) index.
-        :noindex:
-        """
-        self._batch_idx = int(value)
-
-    @y.setter
-    def y(self, value):
-        """Set y after each epoch.
-        :noindex:
-        """
-        self._y = value
-
-    @missing_mask.setter
-    def missing_mask(self, value):
-        """Set missing_mask after each epoch.
-        :noindex:
-        """
-        self._missing_mask = value
-
-    @sample_weight.setter
-    def sample_weight(self, value):
-        """Set sample_weight after each epoch.
-        :noindex:
-        """
-        self._sample_weight = value