GATGAPModel Class Reference

Inheritance diagram for GATGAPModel:

Public Member Functions
def	__init__ (self, units=128, num_features=8, num_pdg=NUM_PDG, emb_size=8, attention_heads=4, n_layers=5, use_gap=False)
def	forward (self, graph)

Public Attributes
	pdg_embedding
	Embedding layer for PDG IDs.
	gat_layers
	List of GAT modules to update node features.
	fc_output
	Output layer for final prediction.

Detailed Description

Input:
   dgl graph built from decay event

Arguments:
   units(int): Number of units for the output dimension of GAT Convolutional layers
   as well as the dimension of global features
   num_features(int): Number of features attached to each node or particle as NN input
   num_pdg(int): Number of all possible PDG IDs
   emb_size(int): Dimension of embedded PDG space
   attention_heads(int): Number of attention heads for GAT Convolutional layers
   n_layers(int): Number of GAT Convolutional layers
   use_gap(bool): Whether to use Global Attention Pooling (GAP) for the production of global features

Returns:
   logits(float): Indicating the probability of an event being able to pass the
   corresponding skim, need `sigmoid` to be used as a prediction

Definition at line 88 of file gatgap.py.

Constructor & Destructor Documentation

__init__()

def __init__	(		self,
			units = 128,
			num_features = 8,
			num_pdg = NUM_PDG,
			emb_size = 8,
			attention_heads = 4,
			n_layers = 5,
			use_gap = False
	)

Initialise the class.

:param units: Number of units for the output dimension of GAT Convolutional layers
as well as the dimension of global features.
:param num_features: Number of features attached to each node or particle as NN input.
:param num_pdg: Number of all possible PDG IDs.
:param emb_size: Dimension of embedded PDG space.
:param attention_heads: Number of attention heads for GAT Convolutional layers.
:param n_layers: Number of GAT Convolutional layers.
:param use_gap: Whether to use Global Attention Pooling (GAP) for the production of global features.

Definition at line 108 of file gatgap.py.

    ):
        """
        Initialise the class.
 
        :param units: Number of units for the output dimension of GAT Convolutional layers
        as well as the dimension of global features.
        :param num_features: Number of features attached to each node or particle as NN input.
        :param num_pdg: Number of all possible PDG IDs.
        :param emb_size: Dimension of embedded PDG space.
        :param attention_heads: Number of attention heads for GAT Convolutional layers.
        :param n_layers: Number of GAT Convolutional layers.
        :param use_gap: Whether to use Global Attention Pooling (GAP) for the production of global features.
        """
        super().__init__()
        
        self.pdg_embedding = torch.nn.Embedding(num_pdg + 1, emb_size)
        in_feats = num_features + emb_size
        
        self.gat_layers = torch.nn.ModuleList()
        in_feats_glob = 0
        for i in range(n_layers):
            self.gat_layers.append(
                GATModule(
                    in_feats=in_feats,
                    units=units,
                    num_heads=attention_heads,
                    in_feats_glob=in_feats_glob,
                    use_gap=use_gap
                )
            )
            in_feats = units * attention_heads
            in_feats_glob = units
 
        
        self.fc_output = torch.nn.Linear(units, 1)
 

Member Function Documentation

forward()

def forward	(		self,
			graph
	)

Forward pass of the GATGAPModel.

Arguments:
    graph (torch.Tensor): DGLGraph representing the decay tree.

Returns:
    torch.Tensor: the final prediction with size 1.

Definition at line 153 of file gatgap.py.

    def forward(self, graph):
        """
        Forward pass of the GATGAPModel.
 
        Arguments:
            graph (torch.Tensor): DGLGraph representing the decay tree.
 
        Returns:
            torch.Tensor: the final prediction with size 1.
        """
        h_pdg = graph.ndata["x_pdg"]
        h_feat = graph.ndata["x_feature"]
        h_pdg = self.pdg_embedding(h_pdg.long())
        h = torch.cat((h_pdg, h_feat), axis=1)
        hg = None
        for layer in self.gat_layers:
            h, hg = layer(graph, h, hg)
        return self.fc_output(hg)

Member Data Documentation

fc_output

fc_output

Output layer for final prediction.

Definition at line 151 of file gatgap.py.

gat_layers

gat_layers

List of GAT modules to update node features.

Definition at line 135 of file gatgap.py.

pdg_embedding

pdg_embedding

Embedding layer for PDG IDs.

Definition at line 132 of file gatgap.py.

---

The documentation for this class was generated from the following file:

• generators/scripts/smartBKG/models/gatgap.py