particle_transfomer.py

import torch
import torch.nn as nn
import torch.nn.functional as F


class Embed(nn.Module):
    def __init__(self, input_dim, output_dim, normalize_input=False, event_level=False, activation='gelu'):
        super().__init__()

        self.input_bn = nn.BatchNorm1d(input_dim) if normalize_input else None
        self.fc1 = nn.Linear(input_dim, output_dim)
        self.fc2 = nn.Linear(output_dim, output_dim)
        self.fc3 = nn.Linear(output_dim, output_dim)
        self.event_level = event_level
        

    def forward(self, x):
        if self.input_bn is not None:
            # x: (batch, embed_dim, seq_len)
            x = self.input_bn(x)
            if not self.event_level: 
                x = x.permute(2, 0, 1).contiguous()

        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))

        return x
    

class AttBlock(nn.Module):
    def __init__(self, embed_dims, linear_dims1, linear_dims2, num_heads=8, activation='relu'):
        super(AttBlock, self).__init__()

        self.layer_norm1 = nn.LayerNorm(embed_dims)
        self.multihead_attention = nn.MultiheadAttention(embed_dims, num_heads)
        self.layer_norm2 = nn.LayerNorm(embed_dims)
        self.linear1 = nn.Linear(embed_dims, linear_dims1)
        self.activation = nn.ReLU() if activation == 'relu' else nn.GELU()
        self.layer_norm3 = nn.LayerNorm(linear_dims1)
        self.linear2 = nn.Linear(linear_dims1, linear_dims2)

    def forward(self, x, padding_mask=None):
        # Layer normalization 1
        x = self.layer_norm1(x)

        if padding_mask is not None:
        # Assuming mask is 0 for non-padded and 1 for padded elements,
        # convert it to a boolean tensor with `True` for padded locations.
            padding_mask = padding_mask.bool()

        
        x_att, attention = self.multihead_attention(x, x, x, key_padding_mask=padding_mask, need_weights=True, average_attn_weights=True)
        
        # Skip connection
        x = x + x_att # Skip connection
        # Layer normalization 2
        x = self.layer_norm2(x)
        # Linear layer and activation
        x_linear1 = self.activation(self.linear1(x))
        # Skip connection for the first linear layer
        x = x + x_linear1
        # Layer normalization 3
        x = self.layer_norm3(x_linear1)
        # Linear layer with specified output dimensions
        x_linear2 = self.linear2(x)
        # Skip connection for the second linear layer
        x = x + x_linear2
        return x, attention

class ClassBlock(nn.Module):
    def __init__(self, embed_dims, linear_dims1, linear_dims2, num_heads=8, activation='relu'):
        super(ClassBlock, self).__init__()

        self.layer_norm1 = nn.LayerNorm(embed_dims)
        self.multihead_attention = nn.MultiheadAttention(embed_dims, num_heads)
        self.layer_norm2 = nn.LayerNorm(embed_dims)
        self.linear1 = nn.Linear(embed_dims, linear_dims1)
        self.activation = nn.ReLU() if activation == 'relu' else nn.GELU()
        self.layer_norm3 = nn.LayerNorm(linear_dims1)
        self.linear2 = nn.Linear(linear_dims1, linear_dims2)

    def forward(self, x, class_token, padding_mask=None):
        # Concatenate the class token to the input sequence along the sequence length dimension  
        x = torch.cat((class_token, x), dim=0)  # (seq_len+1, batch, embed_dim)
        # Layer normalization 1
        x = self.layer_norm1(x)

        # Multihead Attention
        if padding_mask is not None:
            # Ensure mask has the correct shape for attention
            padding_mask = torch.cat((torch.zeros_like(padding_mask[:, :1]), padding_mask), dim=1)
            padding_mask = padding_mask.bool()
        
        
        x_att, attention = self.multihead_attention(class_token, x, x, key_padding_mask=padding_mask, need_weights=True, average_attn_weights=False)
        # Layer normalization 2
        x = self.layer_norm2(x_att)
        x = class_token + x  # Skip connection
        # Linear layer and activation
        x_linear1 = self.activation(self.linear1(x))
        # Layer normalization 3
        x_linear1 = self.layer_norm3(x_linear1)
        # Linear layer with specified output dimensions
        x_linear2 = self.linear2(x_linear1 )
        # Skip connection for the second linear layer
        x = x + x_linear2
        return x, attention

class MLPHead(nn.Module):
    def __init__(self, input_dim, hidden_dim1, hidden_dim2, output_dim):
        super(MLPHead, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim1)
        self.fc2 = nn.Linear(hidden_dim1, hidden_dim2)
        self.fc3 = nn.Linear(hidden_dim2, output_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
    
    
class AnalysisObjectTransformer(nn.Module):
    def __init__(self, input_dim_obj, input_dim_event, embed_dims, linear_dims1, linear_dims2, mlp_hidden_1, mlp_hidden_2, num_heads=8):
        super(AnalysisObjectTransformer, self).__init__()

        self.embed_dims = embed_dims
       
        # Embedding layer (assumed to be external)
        self.embedding_layer = Embed(input_dim_obj, embed_dims)
        self.embedding_layer_event_level = Embed(input_dim_event, embed_dims, event_level=True)

        # Three blocks of self-attention
        self.block1 = AttBlock(embed_dims, linear_dims1, linear_dims1, num_heads)
        self.block2 = AttBlock(linear_dims1, linear_dims1, linear_dims1, num_heads)
        self.block3 = AttBlock(linear_dims1, linear_dims2, linear_dims2, num_heads)
        self.block5 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads)
        self.block6 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads)
        self.block7 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads)

        # Output linear layer and sigmoid activation

        self.mlp = MLPHead(embed_dims + input_dim_event, mlp_hidden_1, mlp_hidden_2, output_dim=1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, event_level, mask=None):

        cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims), requires_grad=True)
        cls_token = nn.init.trunc_normal_(cls_token, std=.02)
        # Embedding layer

        x = self.embedding_layer(x)
        x = x.permute(1, 0, 2)
        
        attention_weights = []

        # Three blocks of self-attention
        x, attention = self.block1(x, padding_mask=mask)

        attention_weights.append(attention)
        x, attention  = self.block2(x, padding_mask=mask)
        attention_weights.append(attention)
        x, attention = self.block3(x, padding_mask=mask)
        attention_weights.append(attention)

        cls_tokens  = cls_token.expand(1, x.size(1), -1)  # (1, N, C)
        cls_tokens, attention  = self.block5(x, cls_tokens, padding_mask=mask)
        cls_tokens, attention  = self.block6(x, cls_tokens, padding_mask=mask)
        cls_tokens, attention  = self.block7(x, cls_tokens, padding_mask=mask)

        x = torch.cat((cls_tokens.squeeze(0), event_level), dim=-1)
        x = self.mlp(x)
        output_probabilities = self.sigmoid(x)
        return output_probabilities, attention_weights