models.py

"""This module contains the models used in the experiments.

Classes
------
NeuralNet
    A basic neural network with optional dropout.
NeuralNetEnsemble
    An ensemble of networks used for processing incomplete examples.
AttentionModel
    Attention-based model used for processing incomplete examples.
"""

from typing import List, Optional

import torch
from torch import nn
from torch import Tensor
from torch.nn.functional import one_hot
from torch.autograd import Function
import wandb

from pdi.config import ModelConfig
from pdi.data.constants import N_COLUMNS
from pdi.data.group_id_helpers import group_id_to_binary_array
from pdi.data.types import GroupID

ACTIVATIONS = {"ReLU": nn.ReLU}


def build_model(cfg: ModelConfig, group_ids: list[GroupID]):
    if cfg.architecture == "mlp":
        model = NeuralNet(
            layers=[N_COLUMNS, *cfg.mlp.hidden_layers, 1],
            activation=ACTIVATIONS[cfg.mlp.activation],
            dropout=cfg.mlp.dropout,
        )
    elif cfg.architecture == "ensemble":
        model = NeuralNetEnsemble(
            group_ids=group_ids,
            hidden_layers=cfg.ensemble.hidden_layers,
            activation=ACTIVATIONS[cfg.ensemble.activation],
            dropout=cfg.ensemble.dropout,
        )
    elif cfg.architecture == "attention":
        model = AttentionModel(
            in_dim=N_COLUMNS + 1,  # +1 for value in one hot encoding
            embed_hidden_layers=cfg.attention.embed_hidden_layers,
            embed_dim=cfg.attention.embed_dim,
            ff_hidden_layers=cfg.attention.mlp_hidden_layers,
            encoder_ff_hidden=cfg.attention.encoder_ff_hidden,
            pool_hidden_layers=cfg.attention.pool_hidden_layers,
            num_heads=cfg.attention.num_heads,
            num_blocks=cfg.attention.num_blocks,
            activation=ACTIVATIONS[cfg.attention.activation],
            dropout=cfg.attention.dropout,
        )
    elif cfg.architecture == "attention_dann":
        model = AttentionModelDANN(
            in_dim=N_COLUMNS + 1,  # +1 for value in one hot encoding
            embed_hidden_layers=cfg.attention.embed_hidden_layers,
            embed_dim=cfg.attention.embed_dim,
            ff_hidden_layers=cfg.attention.mlp_hidden_layers,
            encoder_ff_hidden=cfg.attention.encoder_ff_hidden,
            pool_hidden_layers=cfg.attention.pool_hidden_layers,
            num_heads=cfg.attention_dann.attention.num_heads,
            num_blocks=cfg.attention_dann.attention.num_blocks,
            dom_hidden_layers=cfg.attention_dann.dom_hidden_layers,
            activation=ACTIVATIONS[cfg.attention_dann.attention.activation],
            dropout=cfg.attention_dann.attention.dropout,
            alpha=cfg.attention_dann.alpha,
        )
    elif cfg.architecture == "attention_cdan":
        model = AttentionModelCDAN(
            in_dim=N_COLUMNS + 1,  # +1 for value in one hot encoding
            embed_hidden_layers=cfg.attention.embed_hidden_layers,
            embed_dim=cfg.attention.embed_dim,
            ff_hidden_layers=cfg.attention.mlp_hidden_layers,
            encoder_ff_hidden=cfg.attention.encoder_ff_hidden,
            pool_hidden_layers=cfg.attention.pool_hidden_layers,
            num_heads=cfg.attention_cdan.attention.num_heads,
            num_blocks=cfg.attention_cdan.attention.num_blocks,
            dom_hidden_layers=cfg.attention_cdan.dom_hidden_layers,
            activation=ACTIVATIONS[cfg.attention_cdan.attention.activation],
            dropout=cfg.attention_cdan.attention.dropout,
            alpha=cfg.attention_cdan.alpha,
        )
    else:
        raise KeyError(f"Architecture {cfg.architecture} does not exist!")

    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    # Log to wandb and stdout
    if wandb.run:
        wandb.log({"num_parameters": num_params})
    print(f"Model {cfg.architecture} has been initialized:")
    print(f"\tNumber of trainable parameters: {num_params}")

    return model


class NeuralNet(nn.Module):
    """NeuralNet is a basic neural network with variable layer dimensions, activation function and optional dropout."""

    def __init__(
        self,
        layers: list[int],
        activation: type[nn.Module],
        dropout: Optional[float] = None,
    ):
        """__init__

        Args:
            layers (list[int]): list of layer dimensions, including the input layer, hidden layers and output layer.
            activation (nn.Module): activation function
            dropout (Optional[float], optional): dropout rate. Defaults to None.
        """
        super().__init__()
        self.layers = nn.ModuleList()
        for in_f, out_f in zip(layers[:-2], layers[1:-1]):
            self.layers.append(nn.Linear(in_f, out_f))
            self.layers.append(activation())
            if dropout is not None:
                self.layers.append(nn.Dropout(dropout))
        self.layers.append(nn.Linear(layers[-2], layers[-1]))

    def forward(self, x: Tensor) -> Tensor:
        for layer in self.layers:
            x = layer(x)
        return x


class NeuralNetEnsemble(nn.Module):
    """NeuralNetEnsemble is an ensemble of networks used for processing incomplete examples.
    Each group has a separate neural network in the ensemble, responsible for processing examples belonging to that group.
    """

    def __init__(
        self,
        group_ids: list[GroupID],
        hidden_layers: list[int],
        activation: type[nn.Module] = nn.ReLU,
        dropout: float = 0.4,
    ):
        """__init__

        Args:
            group_ids (list[GroupID]): list of ids of groups in the dataset.
                Group ids should be binary numbers based on the missing values, e.g.
                an example with 5 values, missing the 1st value, should be given id 0b11110.

            hidden_layers (list[int]): list of hidden + output layer dimensions.
                All neural networks use the same dimensions.

            activation (nn.Module, optional): Activation function. Defaults to nn.ReLU.
            dropout (float, optional): dropout rate. Defaults to 0.4.
        """
        super().__init__()
        self.models = nn.ModuleDict({
            str(g_id): NeuralNet(
                [bin(g_id).count("1")] + hidden_layers + [1], activation, dropout
            )
            for g_id in group_ids
        })

        # caching tuple -> group_id to save few transformations during inference
        self.group_id_map = {
            tuple(bool(b) for b in group_id_to_binary_array(g_id)): str(g_id)
            for g_id in group_ids
        }

    def forward(self, x: Tensor):
        nan_mask = torch.isnan(x[0])
        group_id = self.group_id_map[tuple((~nan_mask).tolist())]
        x = x[:, ~nan_mask]
        return self.models[str(group_id)](x)


class AttentionModel(nn.Module):
    """AttentionModel is an attention-based model used for processing incomplete examples."""

    class _AttentionPooling(nn.Module):
        def __init__(
            self, in_dim: int, pool_dim: List[int], activation: type[nn.Module]
        ):
            super().__init__()
            self.net = NeuralNet([in_dim, *pool_dim, in_dim], activation)
            self.softmax = nn.Softmax(dim=1)

        def forward(self, x):
            alpha = self.net(x)

            attention_weights = self.softmax(alpha)
            out = torch.mul(attention_weights, x)
            out = torch.sum(out, dim=-2)
            return out

    class _VecToMat(nn.Module):
        def __init__(self):
            super().__init__()

        def forward(self, x: Tensor) -> Tensor:
            not_missing = ~torch.isnan(x)
            not_missing_idx = not_missing.nonzero()[:, -1]
            values = x[not_missing]

            if x.dim() == 1:
                feat_count = not_missing.nonzero().shape[0]
                rows = 1
            else:
                feat_count = not_missing[0].nonzero().shape[0]
                rows, _ = x.shape

            oh = one_hot(not_missing_idx, N_COLUMNS)  # pylint: disable=not-callable
            one_hot_indices = oh.reshape(rows, feat_count, N_COLUMNS)

            values = values.reshape(rows, feat_count, 1)
            result = torch.cat((one_hot_indices, values), -1)
            return result

    def __init__(
        self,
        in_dim: int,
        embed_hidden_layers: List[int],
        embed_dim: int,
        encoder_ff_hidden: int,
        ff_hidden_layers: List[int],
        pool_hidden_layers: List[int],
        num_heads: int,
        num_blocks: int,
        activation: type[nn.Module] = nn.ReLU,
        dropout: float = 0.4,
    ):
        """__init__

        Args:
            in_dim (int): input dimension.
            embed_hidden (int): hidden dimension in the embedding layer.
            embed_dim (int): output dimension of the embedding layer.
            ff_hidden (int): hidden dimension in the feed forward layer in the Transformer.
            pool_hidden (int): pooling hidden dimension.
            num_heads (int): number of heads in the Transformer.
            num_blocks (int): number of blocks in the Transformer.
            activation (nn.Module, optional): activation function. Defaults to nn.ReLU.
            dropout (float, optional): dropout rate. Defaults to 0.4.
        """
        super().__init__()
        self.to_feature_set = AttentionModel._VecToMat()
        self.emb = NeuralNet([in_dim, *embed_hidden_layers, embed_dim], activation)
        self.drop = nn.Dropout(dropout)
        encoder_layer = nn.TransformerEncoderLayer(
            embed_dim,
            num_heads,
            encoder_ff_hidden,
            dropout,
            activation(),
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer, num_blocks, enable_nested_tensor=False
        )
        self.pool = AttentionModel._AttentionPooling(
            embed_dim, pool_hidden_layers, activation
        )
        self.net = NeuralNet([embed_dim, *ff_hidden_layers, 1], activation)

    def forward(self, x: Tensor) -> Tensor:
        x = self.to_feature_set(x)
        x = self.emb(x)
        x = self.drop(x)
        x = self.encoder(x)
        x = self.pool(x)
        x = self.net(x)
        return x

    def predict(self, incomplete_tensor: Tensor) -> Tensor:
        return self.forward(incomplete_tensor)


# DANN
# pylint: disable=abstract-method,arguments-differ
class ReverseLayerF(Function):
    @staticmethod
    def forward(ctx, x, alpha):
        ctx.alpha = alpha
        return x.view_as(x)

    @staticmethod
    def backward(ctx, grad_output):
        output = grad_output.neg() * ctx.alpha
        return output, None
# pylint: enable=abstract-method,arguments-differ

class AttentionModelDANN(AttentionModel):
    """AttentionModelDANN is an attention-based model with Domain Adversarial Neural Network (DANN) capabilities."""

    def __init__(
        self,
        in_dim: int,
        embed_hidden_layers: List[int],
        embed_dim: int,
        encoder_ff_hidden: int,
        ff_hidden_layers: List[int],
        pool_hidden_layers: List[int],
        num_heads: int,
        num_blocks: int,
        dom_hidden_layers: List[int],
        activation: type[nn.Module] = nn.ReLU,
        dropout: float = 0.4,
        alpha: float = 1.0,
    ):
        super().__init__(
            in_dim,
            embed_hidden_layers,
            embed_dim,
            encoder_ff_hidden,
            ff_hidden_layers,
            pool_hidden_layers,
            num_heads,
            num_blocks,
            activation,
            dropout,
        )
        self.domain_classifier = NeuralNet(
            [embed_dim, *dom_hidden_layers, 1], activation
        )
        self.alpha = alpha

    def feature_extraction(self, x: Tensor) -> Tensor:
        feature = self.to_feature_set(x)
        feature = self.emb(feature)
        feature = self.drop(feature)
        feature = self.encoder(feature)
        feature = self.pool(feature)

        return feature

    def forward(self, x: Tensor) -> Tensor:
        # Feature extraction part
        feature = self.feature_extraction(x)

        # Domain adversarial part
        reverse_x = ReverseLayerF.apply(feature, self.alpha)
        domain_posterior = self.domain_classifier(reverse_x)

        # Classifier part
        class_posterior = self.net(feature)

        return torch.cat((class_posterior, domain_posterior), dim=1)

class AttentionModelCDAN(AttentionModel):
    """AttentionModelCDAN is an attention-based model with Conditional Domain Adversarial Neural Network (CDAN) capabilities."""

    def __init__(
        self,
        in_dim: int,
        embed_hidden_layers: List[int],
        embed_dim: int,
        encoder_ff_hidden: int,
        ff_hidden_layers: List[int],
        pool_hidden_layers: List[int],
        num_heads: int,
        num_blocks: int,
        dom_hidden_layers: List[int],
        activation: type[nn.Module] = nn.ReLU,
        dropout: float = 0.4,
        alpha: float = 1.0,
    ):
        super().__init__(
            in_dim,
            embed_hidden_layers,
            embed_dim,
            encoder_ff_hidden,
            ff_hidden_layers,
            pool_hidden_layers,
            num_heads,
            num_blocks,
            activation,
            dropout,
        )
        self.domain_classifier = NeuralNet(
            [embed_dim * embed_dim, *dom_hidden_layers, 1], activation
        )
        self.alpha = alpha

    def feature_extraction(self, x: Tensor) -> Tensor:
        feature = self.to_feature_set(x)
        feature = self.emb(feature)
        feature = self.drop(feature)
        feature = self.encoder(feature)
        feature = self.pool(feature)

        return feature

    def forward(self, x: Tensor) -> Tensor:
        # Feature extraction part
        feature = self.feature_extraction(x)

        # Classifier part
        class_posterior = self.net(feature)

        # Conditional domain adversarial part
        # Compute the outer product of feature and class_posterior
        outer_product = torch.bmm(
            class_posterior.unsqueeze(2), feature.unsqueeze(1)
        ).view(feature.size(0), -1)

        # Reverse gradient and pass through domain classifier
        reverse_outer_product = ReverseLayerF.apply(outer_product, self.alpha)
        domain_posterior = self.domain_classifier(reverse_outer_product)

        return torch.cat((class_posterior, domain_posterior), dim=1)