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

pgsui/impute/unsupervised/models/ubp_model.py

@@ -1,1126 +1,200 @@
-import logging
-import os
-import sys
-import warnings
+from typing import Callable, 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
+from pgsui.utils.logging_utils import configure_logger
 
-# 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) decoder for genotype logits.
 
-from tensorflow.python.util import deprecation
+    The model reconstructs locus-level genotype probabilities (two states for haploid data or three for diploid data) from a latent vector. It exposes two decoding branches so the training schedule can follow the UBP recipe:
 
-# 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 shallow linear layer that co-trains with latent codes.
+    2. **Phase 2/3 decoder** - a deeper MLP with batch normalization and dropout that is first trained in isolation and later fine-tuned jointly with the latents.
 
-
-# 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,
-    LeakyReLU,
-    PReLU,
-)
-
-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.
+    Both paths ultimately reshape their logits to ``(batch_size, n_features, num_classes)`` and training uses ``MaskedFocalLoss`` to focus on hard examples while masking missing entries.
     """
 
     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,
+        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(UBPPhase1, self).__init__()
-
-        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
-
-    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,
-                regularization_losses=self.losses,
-            )
-
-        # 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: If you get the error, "'tuple' object has no attribute
-        ### 'rank', then convert y_true to a tensor object."
-        self.compiled_metrics.update_state(
-            y_true_masked,
-            y_pred_masked,
-            sample_weight=sample_weight_masked,
-        )
-
-        # 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()
-
-        # history object that gets returned from model.fit().
-        return {m.name: m.result() for m in self.metrics}
-
-    @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
+        """Initializes the UBPModel.
 
-    @V_latent.setter
-    def V_latent(self, value):
-        """Set randomly initialized input. Gets refined during training.
-        :noindex:
+        Args:
+            n_features (int): The number of features (SNPs) in the input data.
+            prefix (str): A prefix used for logging.
+            num_classes (int): Number of genotype states per locus (typically 2 or 3). 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.
         """
-        self.V_latent_ = value
+        super(UBPModel, self).__init__()
+
+        logman = LoggerManager(
+            name=__name__, prefix=prefix, verbose=verbose, debug=debug
+        )
+        self.logger = configure_logger(
+            logman.get_logger(), verbose=verbose, debug=debug
+        )
+
+        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[[], nn.Module]:
+        """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[[], nn.Module]: 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`` and reshapes the logits to the
+        `(batch_size, n_features, num_classes)` grid expected by the loss.
+
+        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: Logits shaped as `(batch_size, n_features, num_classes)`.
 
-    @batch_size.setter
-    def batch_size(self, value):
-        """Set batch_size parameter.
-        :noindex:
+        Raises:
+            ValueError: If an invalid phase is provided.
         """
-        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.
-    """
-
-    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,
-    ):
-        super(UBPPhase2, self).__init__()
-
-        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"
-
+        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:
-            activation = hidden_activation
+            msg = f"Invalid phase: {phase}. Expected 1, 2, or 3."
+            self.logger.error(msg)
+            raise ValueError(msg)
 
-        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)
-
-        x = self.output1(x)
-        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
-
-    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
-            )
-        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,
-                regularization_losses=self.losses,
-            )
-
-        # 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.compiled_metrics.update_state(
-            y_true_masked,
-            y_pred_masked,
-            sample_weight=sample_weight_masked,
-        )
-
-        # history object that gets returned from fit().
-        return {m.name: m.result() for m in self.metrics}
-
-    @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__()
-
-        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
-
-        # 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,
+        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): Integer ground-truth genotypes of shape `(batch_size, n_features)`.
+            outputs (torch.Tensor): Logits of shape `(batch_size, n_features, num_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),
         )
-
-        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)
-
-    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)
-        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
-
-    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,
-                regularization_losses=self.losses,
-            )
-
-        # 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: If you get the error, "'tuple' object has no attribute
-        ### 'rank', then convert y_true to a tensor object."
-        self.compiled_metrics.update_state(
-            y_true_masked,
-            y_pred_masked,
-            sample_weight=sample_weight_masked,
-        )
-
-        # 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()
-
-        # history object that gets returned from fit().
-        return {m.name: m.result() for m in self.metrics}
-
-    @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