# Tutorial: Fraud detection with Graph Attention Networks
By Sing Kwan NG and Anthony TAING as part of the Stanford CS224W course project.
In this tutorial we will show you the fraud detection use case (for a node classification task) on Bitcoin dataset  by using the first version of GAT model and the GATv2 (Graph Attention Network model version2) .We will implement it by using our customs layers with PyG and the buitin PyG layers.

# Setup

## Mount Drive

```python
# project  foldername

FOLDERNAME = "./"
#!mkdir save_results

# load python files from foldername
import sys
sys.path.append(FOLDERNAME)

# for auto-reloading external modules
%load_ext autoreload
%autoreload 2
```

## Load packages

```python
!python -c "import torch; print(torch.__version__)"
```

    1.10.1+cu111

```python
!python -c "import torch; print(torch.version.cuda)"
```

    11.1

```python
try:
    import torch_geometric
except:
    import torch

    CUDA = "cu111"
    TORCH = "1.11.1"

    !pip install torch-scatter torch-sparse torch-cluster torch-geometric -f https://data.pyg.org/whl/torch-1.10.0+cpu.html
    #!pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
    #!pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
    ## !pip install torch-cluster -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
    ## !pip install torch-spline-conv -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
    #!pip install torch-geometric
```

```python
import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt

from torch.nn import Linear, LayerNorm, ReLU, Dropout
import torch.nn.functional as F
from torch_geometric.data import Data, DataLoader
from sklearn.metrics import accuracy_score,f1_score,roc_auc_score, precision_score, recall_score, confusion_matrix

import scipy.sparse as sp

import warnings
warnings.filterwarnings("ignore")
import pandas as pd
from sklearn.model_selection import train_test_split

import torch.nn.functional as F
from sklearn.metrics import roc_auc_score
from torch.nn import Linear
from torch_geometric.nn import GCNConv,GATConv,GATv2Conv

```

# Data ETL

We have hosted the dataset on Gdrive to the public for convenience for one click download. However, the original data is hosted on kaggle as well. If you want to download data from Kaggle using API(if you already have an account) or get it manually from: https://www.kaggle.com/ellipticco/elliptic-data-set.

## Download data

## Load data
We reformat classes into numerical values Unknown:2,Illicit:1 and Licit:0.

```python
DATA_FOLDERNAME ="./data/elliptic_bitcoin_dataset/"
```

```python
# Load data from the folder
df_features = pd.read_csv(DATA_FOLDERNAME+'elliptic_txs_features.csv',header=None)
df_edges = pd.read_csv(DATA_FOLDERNAME+"elliptic_txs_edgelist.csv")
df_classes =  pd.read_csv(DATA_FOLDERNAME+"elliptic_txs_classes.csv")

```

```python
df_features.shape , df_edges.shape , df_classes.shape
```

    ((203769, 167), (234355, 2), (203769, 2))

```python
# Add class names for easy understanding
# reformat classes 0:licit, 1:illicit, 2:unknow
df_classes['class'] = df_classes['class'].map({'unknown': 2, '1':1, '2':0})
```

```python
# See repartition of nodes per class
group_class = df_classes.groupby('class').count()
plt.title("# of nodes per class")
plt.barh([ 'Licit','Illicit', 'Unknown'], group_class['txId'].values, color=['g', 'orange', 'r'] )

```

## Preview of data features, edges and their labels

```python
# View node features
# Data is each node is a transaction ID, and edges are a bit weird, but its like a "previous transaction" edge type
# colume 0 = transaction id
df_features.head()
```

      0
      1
      2
      3
      4
      5
      6
      7
      8
      9
      ...
      157
      158
      159
      160
      161
      162
      163
      164
      165
      166

      0
      230425980
      1
      -0.171469
      -0.184668
      -1.201369
      -0.121970
      -0.043875
      -0.113002
      -0.061584
      -0.162097
      ...
      -0.562153
      -0.600999
      1.461330
      1.461369
      0.018279
      -0.087490
      -0.131155
      -0.097524
      -0.120613
      -0.119792

      1
      5530458
      1
      -0.171484
      -0.184668
      -1.201369
      -0.121970
      -0.043875
      -0.113002
      -0.061584
      -0.162112
      ...
      0.947382
      0.673103
      -0.979074
      -0.978556
      0.018279
      -0.087490
      -0.131155
      -0.097524
      -0.120613
      -0.119792

      2
      232022460
      1
      -0.172107
      -0.184668
      -1.201369
      -0.121970
      -0.043875
      -0.113002
      -0.061584
      -0.162749
      ...
      0.670883
      0.439728
      -0.979074
      -0.978556
      -0.098889
      -0.106715
      -0.131155
      -0.183671
      -0.120613
      -0.119792

      3
      232438397
      1
      0.163054
      1.963790
      -0.646376
      12.409294
      -0.063725
      9.782742
      12.414558
      -0.163645
      ...
      -0.577099
      -0.613614
      0.241128
      0.241406
      1.072793
      0.085530
      -0.131155
      0.677799
      -0.120613
      -0.119792

      4
      230460314
      1
      1.011523
      -0.081127
      -1.201369
      1.153668
      0.333276
      1.312656
      -0.061584
      -0.163523
      ...
      -0.511871
      -0.400422
      0.517257
      0.579382
      0.018279
      0.277775
      0.326394
      1.293750
      0.178136
      0.179117

5 rows × 167 columns

```python
# See edges
df_edges.head()
```

      txId1
      txId2

      0
      230425980
      5530458

      1
      232022460
      232438397

      2
      230460314
      230459870

      3
      230333930
      230595899

      4
      232013274
      232029206

```python
# See classes
df_classes.head()
```

      txId
      class

      0
      230425980
      2

      1
      5530458
      2

      2
      232022460
      2

      3
      232438397
      0

      4
      230460314
      2

## Merge dataset

```python
# Merge features with classes
df_merge = df_features.merge(df_classes, how='left', right_on="txId", left_on=0)
df_merge = df_merge.sort_values(0).reset_index(drop=True)
df_merge.head()
```

      0
      1
      2
      3
      4
      5
      6
      7
      8
      9
      ...
      159
      160
      161
      162
      163
      164
      165
      166
      txId
      class

      0
      1076
      48
      -0.168500
      0.270909
      -0.091383
      -0.046932
      -0.043875
      -0.029140
      -0.061584
      -0.163591
      ...
      1.461330
      1.461369
      0.018279
      0.470019
      1.216796
      1.151607
      1.519700
      1.521399
      1076
      2

      1
      2534
      6
      -0.170834
      -0.131425
      1.018602
      0.028105
      0.055376
      0.054722
      -0.061584
      -0.163572
      ...
      0.955101
      0.459257
      -0.098889
      -0.087490
      -0.099080
      -0.122137
      -0.379970
      -0.379288
      2534
      0

      2
      3181
      34
      1.305212
      -0.210553
      -1.756361
      -0.121970
      97.300650
      -0.113002
      -0.061584
      1.348765
      ...
      0.059948
      0.113967
      -0.098889
      1.969527
      0.037532
      -0.131010
      0.006994
      0.017772
      3181
      0

      3
      3321
      1
      -0.169615
      -0.184668
      -1.201369
      -0.121970
      -0.043875
      -0.113002
      -0.061584
      -0.160199
      ...
      0.241128
      0.241406
      -0.098889
      -0.087490
      -0.084674
      -0.140597
      1.519700
      1.521399
      3321
      2

      4
      3889
      48
      -0.086232
      -0.101835
      -0.646376
      -0.121970
      17.046997
      -0.113002
      -0.061584
      -0.074885
      ...
      0.082065
      0.114773
      -0.098889
      8.948005
      1.024948
      -0.009570
      -0.080708
      -0.123601
      3889
      2

5 rows × 169 columns

## Edge Index: Map trans IDs to Node IDs

We map all nodes to indices, and filter edges data by mapping txid1 and txid2 with the indices. We convert it to tensors, and initialize a weight tensor with the same length of the edge index and with values 1.

```python
# Setup trans ID to node ID mapping
nodes = df_merge[0].values

map_id = {j:i for i,j in enumerate(nodes)} # mapping nodes to indexes

# Create edge df that has transID mapped to nodeIDs
edges = df_edges.copy()
edges.txId1 = edges.txId1.map(map_id) #get nodes idx1 from edges list and filtered data
edges.txId2 = edges.txId2.map(map_id)

edges = edges.astype(int)

edge_index = np.array(edges.values).T #convert into an array
edge_index = torch.tensor(edge_index, dtype=torch.long).contiguous() # create a tensor

print("shape of edge index is {}".format(edge_index.shape))
edge_index

```

    shape of edge index is torch.Size([2, 234355])

    tensor([[138670, 141325, 139232,  ..., 100420,  54833, 101159],
            [  4142, 142201, 139223,  ..., 100419,  81951, 101163]])

### Edge weights
Set all edge weights to be 1

```python
#create weights tensor with same shape of edge_index
weights = torch.tensor([1]* edge_index.shape[1] , dtype=torch.double)
weights
```

    tensor([1., 1., 1.,  ..., 1., 1., 1.], dtype=torch.float64)

## Labels

```python
# Define labels
labels = df_merge['class'].values
print("lables", np.unique(labels))
labels
```

    lables [0 1 2]

    array([2, 0, 0, ..., 2, 2, 2])

## Node Features

```python
# mapping txIds to corresponding indices, to pass node features to the model

node_features = df_merge.drop(['txId'], axis=1).copy()
# node_features[0] = node_features[0].map(map_id) # Convert transaction ID to node ID \
print("unique=",node_features["class"].unique())

# Retain known vs unknown IDs
classified_idx = node_features['class'].loc[node_features['class']!=2].index # filter on known labels
unclassified_idx = node_features['class'].loc[node_features['class']==2].index

classified_illicit_idx = node_features['class'].loc[node_features['class']==1].index # filter on illicit labels
classified_licit_idx = node_features['class'].loc[node_features['class']==0].index # filter on licit labels

# Drop unwanted columns, 0 = transID, 1=time period, class = labels
node_features = node_features.drop(columns=[0, 1, 'class'])

# Convert to tensor
node_features_t = torch.tensor(np.array(node_features.values, dtype=np.double), dtype=torch.double)# drop unused columns
node_features_t
```

    unique= [2 0 1]

    tensor([[-0.1685,  0.2709, -0.0914,  ...,  1.1516,  1.5197,  1.5214],
            [-0.1708, -0.1314,  1.0186,  ..., -0.1221, -0.3800, -0.3793],
            [ 1.3052, -0.2106, -1.7564,  ..., -0.1310,  0.0070,  0.0178],
            ...,
            [-0.1727, -0.1588, -1.2014,  ..., -0.2698, -0.1206, -0.1198],
            [-0.1727, -0.1588, -1.2014,  ..., -0.2698, -0.1206, -0.1198],
            [-0.1433, -0.1588, -1.2014,  ..., -0.0975, -0.1206, -0.1198]],
           dtype=torch.float64)

```python
# See node features again
node_features
```

      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      ...
      157
      158
      159
      160
      161
      162
      163
      164
      165
      166

      0
      -0.168500
      0.270909
      -0.091383
      -0.046932
      -0.043875
      -0.029140
      -0.061584
      -0.163591
      -0.164980
      -0.009283
      ...
      0.073047
      -0.039637
      1.461330
      1.461369
      0.018279
      0.470019
      1.216796
      1.151607
      1.519700
      1.521399

      1
      -0.170834
      -0.131425
      1.018602
      0.028105
      0.055376
      0.054722
      -0.061584
      -0.163572
      -0.167757
      -0.038545
      ...
      1.228858
      0.379357
      0.955101
      0.459257
      -0.098889
      -0.087490
      -0.099080
      -0.122137
      -0.379970
      -0.379288

      2
      1.305212
      -0.210553
      -1.756361
      -0.121970
      97.300650
      -0.113002
      -0.061584
      1.348765
      1.321754
      -0.049707
      ...
      1.348450
      1.590664
      0.059948
      0.113967
      -0.098889
      1.969527
      0.037532
      -0.131010
      0.006994
      0.017772

      3
      -0.169615
      -0.184668
      -1.201369
      -0.121970
      -0.043875
      -0.113002
      -0.061584
      -0.160199
      -0.166062
      -0.049707
      ...
      -0.577099
      -0.500080
      0.241128
      0.241406
      -0.098889
      -0.087490
      -0.084674
      -0.140597
      1.519700
      1.521399

      4
      -0.086232
      -0.101835
      -0.646376
      -0.121970
      17.046997
      -0.113002
      -0.061584
      -0.074885
      -0.081943
      -0.049707
      ...
      0.501062
      0.362510
      0.082065
      0.114773
      -0.098889
      8.948005
      1.024948
      -0.009570
      -0.080708
      -0.123601

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      203764
      -0.172978
      -0.172527
      0.463609
      -0.121970
      -0.043875
      -0.113002
      -0.061584
      -0.163640
      -0.169455
      -0.049707
      ...
      -0.577099
      -0.600999
      0.241128
      0.241406
      0.018279
      -0.087490
      -0.131155
      -0.097524
      -0.120613
      -0.119792

      203765
      -0.172669
      -0.158783
      -1.201369
      -0.121970
      -0.063725
      -0.113002
      -0.061584
      -0.163323
      -0.169142
      -0.049707
      ...
      -0.577099
      -0.626229
      0.241128
      0.241406
      -0.216057
      -0.125939
      -0.131155
      -0.269818
      -0.120613
      -0.119792

      203766
      -0.172669
      -0.158783
      -1.201369
      -0.121970
      -0.063725
      -0.113002
      -0.061584
      -0.163323
      -0.169142
      -0.049707
      ...
      -0.577099
      -0.626229
      0.241128
      0.241406
      -0.216057
      -0.125939
      -0.131155
      -0.269818
      -0.120613
      -0.119792

      203767
      -0.172669
      -0.158783
      -1.201369
      -0.121970
      -0.063725
      -0.113002
      -0.061584
      -0.163323
      -0.169142
      -0.049707
      ...
      -0.577099
      -0.626229
      0.241128
      0.241406
      -0.216057
      -0.125939
      -0.131155
      -0.269818
      -0.120613
      -0.119792

      203768
      -0.143292
      -0.158783
      -1.201369
      -0.121970
      -0.043875
      -0.113002
      -0.061584
      -0.133266
      -0.139507
      -0.049707
      ...
      -0.577099
      -0.613614
      0.241128
      0.241406
      0.018279
      -0.087490
      -0.131155
      -0.097524
      -0.120613
      -0.119792

203769 rows × 165 columns

## Train Test splits

```python
# Create a known vs unknown mask
train_idx, valid_idx = train_test_split(classified_idx.values, test_size=0.15)
print("train_idx size {}".format(len(train_idx)))
print("tets_idx size {}".format(len(valid_idx)))
```

    train_idx size 39579
    tets_idx size 6985

# Create a PyG dataset

Each PyG dataset stores a list of `torch_geometric.data.Data` objects, where each `torch_geometric.data.Data` object represents a graph. We can easily get the `Data` object by indexing into the dataset.

For more information such as what is stored in the `Data` object, please refer to the [documentation](https://pytorch-geometric.readthedocs.io/en/latest/modules/data.html#torch_geometric.data.Data).

We convert data to PyGeometric graph data format by using node_features, edge_index, weights and labels initialized before.

```python
weights
```

```python
data_train = Data(x=node_features_t, edge_index=edge_index, edge_attr=weights,
                               y=torch.tensor(labels, dtype=torch.double))
# Add in the train and valid idx
data_train.train_idx = train_idx
data_train.valid_idx = valid_idx
data_train.test_idx = unclassified_idx
```

# Model setup

```python
import torch
import torch_scatter
import torch.nn as nn
import torch.nn.functional as F

import torch_geometric.nn as pyg_nn
import torch_geometric.utils as pyg_utils

from torch import Tensor
from typing import Union, Tuple, Optional
from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
                                    OptTensor)

from torch.nn import Parameter, Linear
from torch_sparse import SparseTensor, set_diag
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.utils import remove_self_loops, add_self_loops, softmax, degree

from torch_geometric.nn import GCNConv,GATConv,GATv2Conv
import pickle
```

## Prebuilt with PyG's built-in  layers
In this section, we are using PyG's built-in layers with 2 layers and an additional MLP layers for the GATs models.

### GCNConv

```python
class GCN(torch.nn.Module):
    def __init__(self):
        super(GCN, self).__init__()
        torch.manual_seed(12345)
        self.conv1 = GCNConv(165, 128)
        self.conv2 = GCNConv(128, 2)
        self.classifier = Linear(2, 1)

    def forward(self, data, adj=None):
        x, edge_index = data.x, data.edge_index
        h = self.conv1(x, edge_index)
        h = h.tanh()
        h = self.conv2(h, edge_index)
        embeddings = h.tanh()  # Final GNN embedding space.

        # Apply a final (linear) classifier.
        out = self.classifier(embeddings)

        # return out, embeddings
        return F.sigmoid(out)
```

### GATConv

```python

class GAT(torch.nn.Module):
    def __init__(self,input_dim, hidden_dim, output_dim,args):
        super(GAT, self).__init__()
        #use our gat message passing
        self.conv1 = GATConv(input_dim, hidden_dim, heads=args['heads'])
        self.conv2 = GATConv(args['heads'] * hidden_dim, hidden_dim, heads=args['heads'])

        self.post_mp = nn.Sequential(
            nn.Linear(args['heads'] * hidden_dim, hidden_dim), nn.Dropout(args['dropout'] ),
            nn.Linear(hidden_dim, output_dim))

    def forward(self, data, adj=None):
        x, edge_index = data.x, data.edge_index
        # Layer 1
        x = self.conv1(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)
        # Layer 2
        x = self.conv2(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)
        # MLP output
        x = self.post_mp(x)
        return F.sigmoid(x)

```

### GATv2Conv

```python

class GATv2(torch.nn.Module):
    def __init__(self,input_dim, hidden_dim, output_dim,args):
        super(GATv2, self).__init__()
        #use our gat message passing
        self.conv1 = GATv2Conv(input_dim, hidden_dim, heads=args['heads'])
        self.conv2 = GATv2Conv(args['heads'] * hidden_dim, hidden_dim, heads=args['heads'])

        self.post_mp = nn.Sequential(
            nn.Linear(args['heads'] * hidden_dim, hidden_dim), nn.Dropout(args['dropout'] ),
            nn.Linear(hidden_dim, output_dim))

    def forward(self, data, adj=None):
        x, edge_index = data.x, data.edge_index
        # Layer 1
        x = self.conv1(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)
        # Layer 2
        x = self.conv2(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)
        # MLP output
        x = self.post_mp(x)
        return F.sigmoid(x)

```

## Custom GAT Implementation

Attention mechanisms have become the state-of-the-art in many sequence-based tasks such as machine translation and learning sentence representations. One of the major benefits of attention-based mechanisms is their ability to focus on the most relevant parts of the input to make decisions. In this problem, we will see how attention mechanisms can be used to perform node classification over graph-structured data through the usage of Graph Attention Networks (GATs) ([Veličković et al. (2018)](https://arxiv.org/abs/1710.10903)).

The building block of the Graph Attention Network is the graph attention layer, which is a variant of the aggregation function. Let $N$ be the number of nodes and $F$ be the dimension of the feature vector for each node. The input to each graph attentional layer is a set of node features: $\mathbf{h} = \{\overrightarrow{h_1}, \overrightarrow{h_2}, \dots, \overrightarrow{h_N}$\}, $\overrightarrow{h_i} \in R^F$. The output of each graph attentional layer is a new set of node features, which may have a new dimension $F'$: $\mathbf{h'} = \{\overrightarrow{h_1'}, \overrightarrow{h_2'}, \dots, \overrightarrow{h_N'}\}$, with $\overrightarrow{h_i'} \in \mathbb{R}^{F'}$.

We will now describe how this transformation is performed for each graph attention layer. First, a shared linear transformation parametrized by the weight matrix $\mathbf{W} \in \mathbb{R}^{F' \times F}$ is applied to every node.

Next, we perform self-attention on the nodes. We use a shared attention function $a$:
\begin{equation}
a : \mathbb{R}^{F'} \times \mathbb{R}^{F'} \rightarrow \mathbb{R}.
\end{equation}

that computes the attention coefficients capturing the importance of node $j$'s features to node $i$:
\begin{equation}
e_{ij} = a(\mathbf{W_l}\overrightarrow{h_i}, \mathbf{W_r} \overrightarrow{h_j})
\end{equation}

The most general formulation of self-attention allows every node to attend to all other nodes which drops all structural information. However, to utilize graph structure in the attention mechanisms, we use **masked attention**. In masked attention, we only compute attention coefficients $e_{ij}$ for nodes $j \in \mathcal{N}_i$ where $\mathcal{N}_i$ is some neighborhood of node $i$ in the graph.

To easily compare coefficients across different nodes, we normalize the coefficients across $j$ using a softmax function:
\begin{equation}
\alpha_{ij} = \text{softmax}_j(e_{ij}) = \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}_i} \exp(e_{ik})}
\end{equation}

For this problem, our attention mechanism $a$ will be a single-layer feedforward neural network parametrized by a weight vectors $\overrightarrow{a_l} \in \mathbb{R}^{F'}$ and $\overrightarrow{a_r} \in \mathbb{R}^{F'}$, followed by a LeakyReLU nonlinearity (with negative input slope 0.2). Let $\cdot^T$ represent transposition and $||$ represent concatenation. The coefficients computed by our attention mechanism may be expressed as:

\begin{equation}
\alpha_{ij} = \frac{\exp\Big(\text{LeakyReLU}\Big(\overrightarrow{a_l}^T \mathbf{W_l} \overrightarrow{h_i} + \overrightarrow{a_r}^T\mathbf{W_r}\overrightarrow{h_j}\Big)\Big)}{\sum_{k\in \mathcal{N}_i} \exp\Big(\text{LeakyReLU}\Big(\overrightarrow{a_l}^T \mathbf{W_l} \overrightarrow{h_i} + \overrightarrow{a_r}^T\mathbf{W_r}\overrightarrow{h_k}\Big)\Big)}
\end{equation}

For the following questions, we denote `alpha_l` = $\alpha_l = [...,\overrightarrow{a_l}^T \mathbf{W_l} \overrightarrow{h_i},...] \in \mathcal{R}^n$ and `alpha_r` = $\alpha_r = [..., \overrightarrow{a_r}^T \mathbf{W_r} \overrightarrow{h_j}, ...] \in \mathcal{R}^n$.

At every layer of GAT, after the attention coefficients are computed for that layer, the aggregation function can be computed by a weighted sum of neighborhood messages, where weights are specified by $\alpha_{ij}$.

Now, we use the normalized attention coefficients to compute a linear combination of the features corresponding to them. These aggregated features will serve as the final output features for every node.

\begin{equation}
h_i' = \sum_{j \in \mathcal{N}_i} \alpha_{ij} \mathbf{W_r} \overrightarrow{h_j}.
\end{equation}

### Multi-Head Attention
To stabilize the learning process of self-attention, we use multi-head attention. To do this we use $K$ independent attention mechanisms, or ``heads'' compute output features as in the above equations. Then, we concatenate these output feature representations:

\begin{equation}
    \overrightarrow{h_i}' = ||_{k=1}^K \Big(\sum_{j \in \mathcal{N}_i} \alpha_{ij}^{(k)} \mathbf{W_r}^{(k)} \overrightarrow{h_j}\Big)
\end{equation}

where $||$ is concentation, $\alpha_{ij}^{(k)}$ are the normalized attention coefficients computed by the $k$-th attention mechanism $(a^k)$, and $\mathbf{W}^{(k)}$ is the corresponding input linear transformation's weight matrix. Note that for this setting, $\mathbf{h'} \in \mathbb{R}^{KF'}$.

### Custom GAT model
PyG provides the MessagePassing base class, which helps in creating such kinds of message passing graph neural networks by automatically taking care of message propagation. \\
-**MessagePassing()**: Defines the aggregation scheme to use ("add", "mean" or "max") and the flow direction of message passing (either "source_to_target" or "target_to_source").

-**MessagePassing.propagate(edge_index,size=None,kwargs)**: The initial call to start propagating messages. Takes in the edge indices and all additional data which is needed to construct messages and to update node embeddings.

-**MessagePassing.message(...)**: Constructs messages to node  for each edge. It can take any argument which was initially passed to propagate().  In addition, tensors passed to propagate() can be mapped to the respective nodes and by appending _i or _j to the variable name.
We generally refer to i as the central nodes that aggregates information, and refer to j as the neighboring nodes, since this is the most common notation.

 Further details can be found at the official PyG documentation.

```python
class myGAT(MessagePassing):
    def __init__(self, in_channels, out_channels, heads = 1,
                 negative_slope = 0.2, dropout = 0., **kwargs):
        super(myGAT, self).__init__(node_dim=0, **kwargs)

        self.in_channels = in_channels # node features input dimension
        self.out_channels = out_channels # node level output dimension
        self.heads = heads # No. of attention heads
        self.negative_slope = negative_slope
        self.dropout = dropout

        self.lin_l = None
        self.lin_r = None
        self.att_l = None
        self.att_r = None

        # Initialization
        self.lin_l = Linear(in_channels, heads*out_channels)
        self.lin_r = self.lin_l
        self.att_l = Parameter(torch.Tensor(1, heads, out_channels).float())
        self.att_r = Parameter(torch.Tensor(1, heads, out_channels).float())
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.lin_l.weight)
        nn.init.xavier_uniform_(self.lin_r.weight)
        nn.init.xavier_uniform_(self.att_l)
        nn.init.xavier_uniform_(self.att_r)

    def forward(self, x, edge_index, size = None):

        H, C = self.heads, self.out_channels # DIM：H, outC

        #Linearly transform node feature matrix.
        x_source = self.lin_l(x).view(-1,H,C) # DIM: [Nodex x In] [in x H * outC] => [nodes x H * outC] => [nodes, H, outC]
        x_target = self.lin_r(x).view(-1,H,C) # DIM: [Nodex x In] [in x H * outC] => [nodes x H * outC] => [nodes, H, outC]

        # Alphas will be used to calculate attention later
        alpha_l = (x_source * self.att_l).sum(dim=-1) # DIM: [nodes, H, outC] x [H, outC] => [nodes, H]
        alpha_r = (x_target * self.att_r).sum(dim=-1) # DIM: [nodes, H, outC] x [H, outC] => [nodes, H]

        #  Start propagating messages (runs message and aggregate)
        out = self.propagate(edge_index, x=(x_source, x_target), alpha=(alpha_l, alpha_r),size=size) # DIM: [nodes, H, outC]
        out = out.view(-1, self.heads * self.out_channels) # DIM: [nodes, H * outC]

        return out

    def message(self, x_j, alpha_j, alpha_i, index, ptr, size_i):
        # Calculate attention for edge pairs
        attention = F.leaky_relu((alpha_j + alpha_i), self.negative_slope) # EQ(1) DIM: [Edges, H]
        attention = softmax(attention, index, ptr, size_i) # EQ(2) DIM: [Edges, H] | This softmax only calculates it over all neighbourhood nodes
        attention = F.dropout(attention, p=self.dropout, training=self.training) # DIM: [Edges, H]

        # Multiple attention with node features for all edges
        out = x_j * attention.unsqueeze(-1)  # EQ(3.1) [Edges, H, outC] x [Edges, H] = [Edges, H, outC];

        return out

    def aggregate(self, inputs, index, dim_size = None):
        # EQ(3.2) For each node, aggregate messages for all neighbourhood nodes
        out = torch_scatter.scatter(inputs, index, dim=self.node_dim,
                                    dim_size=dim_size, reduce='sum') # inputs (from message) DIM: [Edges, H, outC] => DIM: [Nodes, H, outC]
        return out

```

### Custom GATv2 model
This new model GATv2 introduces a simple modification by modifying the order of operations. \\
They simply apply "a" layer after the non-linearity and the "W" layer after the concatenation, effectively applying an MLP to compute the score for each query-key pair. \\
GAT( Velickovic et al.2018):
\begin{equation}
    e(h_i,h_j)=LeakyReLU(a^T \dot [Wh_i||Wh_j])
\end{equation}
GATv2( fixed version):
\begin{equation}
    e(h_i,h_j)=a^T LeakyReLU(W\dot [h_i||h_j])
\end{equation}

```python
class myGATv2(MessagePassing):
    def __init__(self, in_channels, out_channels, heads = 1,
                 negative_slope = 0.2, dropout = 0., **kwargs):
        super(myGATv2, self).__init__(node_dim=0, **kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.heads = heads
        self.negative_slope = negative_slope
        self.dropout = dropout

        self.lin_l = None
        self.lin_r = None
        self.att_l = None
        self.att_r = None
        self._alpha = None
        # self.lin_l is the linear transformation that you apply to embeddings
        # BEFORE message passing.
        self.lin_l =  Linear(in_channels, heads*out_channels)
        self.lin_r = self.lin_l

        self.att = Parameter(torch.Tensor(1, heads, out_channels))
        self.reset_parameters()

    #initialize parameters with xavier uniform
    def reset_parameters(self):
        nn.init.xavier_uniform_(self.lin_l.weight)
        nn.init.xavier_uniform_(self.lin_r.weight)
        nn.init.xavier_uniform_(self.att)

    def forward(self, x, edge_index, size = None):

        H, C = self.heads, self.out_channels # DIM：H, outC
        #Linearly transform node feature matrix.
        x_source = self.lin_l(x).view(-1,H,C) # DIM: [Nodex x In] [in x H * outC] => [nodes x H * outC] => [nodes, H, outC]
        x_target = self.lin_r(x).view(-1,H,C) # DIM: [Nodex x In] [in x H * outC] => [nodes x H * outC] => [nodes, H, outC]

        #  Start propagating messages (runs message and aggregate)
        out= self.propagate(edge_index, x=(x_source,x_target),size=size) # DIM: [nodes, H, outC]
        out= out.view(-1, self.heads * self.out_channels)       # DIM: [nodes, H * outC]
        alpha = self._alpha
        self._alpha = None
        return out

    #Process a message passing
    def message(self, x_j,x_i,  index, ptr, size_i):
        #computation using previous equationss
        x = x_i + x_j
        x  = F.leaky_relu(x, self.negative_slope)   # See Equation above: Apply the non-linearty function
        alpha = (x * self.att).sum(dim=-1)          # Apply attnention "a" layer after the non-linearity
        alpha = softmax(alpha, index, ptr, size_i)  # This softmax only calculates it over all neighbourhood nodes
        self._alpha = alpha
        alpha= F.dropout(alpha,p=self.dropout,training=self.training)
        # Multiple attention with node features for all edges
        out= x_j*alpha.unsqueeze(-1)

        return out
    #Aggregation of messages
    def aggregate(self, inputs, index, dim_size = None):
        out = torch_scatter.scatter(inputs, index, dim=self.node_dim,
                                    dim_size=dim_size, reduce='sum')
        return out
```

### GATCustom
We use PyG and define two models by using our previous GAT and GATv2. \\
Each layer consists of running the GAT layer, followed by a relu nonlinearity function and dropout. They can be stacked multiple times with some output layers.We had a linear MLP output layer as a post-processing.

```python
class GATmodif(torch.nn.Module):
    def __init__(self,input_dim, hidden_dim, output_dim,args):
        super(GATmodif, self).__init__()
        #use our gat message passing
        self.conv1 = myGAT(input_dim, hidden_dim,heads=args['heads'])
        self.conv2 = myGAT(args['heads']  *hidden_dim, hidden_dim,heads=args['heads'])

        self.post_mp = nn.Sequential(
            nn.Linear(args['heads'] * hidden_dim, hidden_dim), nn.Dropout(args['dropout'] ),
            nn.Linear(hidden_dim, output_dim))

    def forward(self, data, adj=None):
        x, edge_index = data.x, data.edge_index
        # Layer 1
        x = self.conv1(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)

        # Layer 2
        x = self.conv2(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)

        # MLP output
        x = self.post_mp(x)
        return F.sigmoid(x)

class GATv2modif(torch.nn.Module):
    def __init__(self,input_dim, hidden_dim, output_dim,args):
        super(GATv2modif, self).__init__()
        #use our gat message passing
        self.conv1 = myGATv2(input_dim, hidden_dim,heads=args['heads'])
        self.conv2 = myGATv2(args['heads'] *hidden_dim, hidden_dim,heads=args['heads'])

        self.post_mp = nn.Sequential(
            nn.Linear(args['heads']  * hidden_dim, hidden_dim), nn.Dropout(args['dropout'] ),
            nn.Linear(hidden_dim, output_dim))

    def forward(self, data, adj=None):
        x, edge_index = data.x, data.edge_index
        # Layer 1
        x = self.conv1(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)

        # Layer 2
        x = self.conv2(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)

        # MLP output
        x = self.post_mp(x)
        return F.sigmoid(x)

```

Here, we are implementing a custom GAT with 3 layers

```python
class GATmodif_3layer(torch.nn.Module):
    def __init__(self,input_dim, hidden_dim, output_dim, args):
        super(GATmodif_3layer, self).__init__()
        #use our gat message passing
        self.conv1 = myGAT(input_dim, hidden_dim,heads=args['heads'])
        self.conv2 = myGAT(args['heads'] * hidden_dim, hidden_dim,heads=args['heads'])
        self.conv3 = myGAT(args['heads'] * hidden_dim, hidden_dim,heads=args['heads'])

        self.post_mp = nn.Sequential(
            nn.Linear(args['heads'] * hidden_dim, hidden_dim), nn.Dropout(args['dropout'] ),
            nn.Linear(hidden_dim, output_dim))

    def forward(self, data, adj=None):
        x, edge_index = data.x, data.edge_index
        # Layer 1
        x = self.conv1(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)

        # Layer 2
        x = self.conv2(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)

        # Layer 3
        x = self.conv3(x, edge_index)
        x = F.dropout(F.relu(x), p=args['dropout'], training=self.training)

        # MLP output
        x = self.post_mp(x)
        return F.sigmoid(x)
```

# Model training

## GNNTrainer object

```python
class GnnTrainer(object):

  def __init__(self, model):
    self.model = model
    self.metric_manager = MetricManager(modes=["train", "val"])

  def train(self, data_train, optimizer, criterion, scheduler, args):

    self.data_train = data_train
    for epoch in range(args['epochs']):
        self.model.train()
        optimizer.zero_grad()
        out = self.model(data_train)

        out = out.reshape((data_train.x.shape[0]))
        loss = criterion(out[data_train.train_idx], data_train.y[data_train.train_idx])
        ## Metric calculations
        # train data
        target_labels = data_train.y.detach().cpu().numpy()[data_train.train_idx]
        pred_scores = out.detach().cpu().numpy()[data_train.train_idx]
        train_acc, train_f1,train_f1macro, train_aucroc, train_recall, train_precision, train_cm = self.metric_manager.store_metrics("train", pred_scores, target_labels)

        ## Training Step
        loss.backward()
        optimizer.step()

        # validation data
        self.model.eval()
        target_labels = data_train.y.detach().cpu().numpy()[data_train.valid_idx]
        pred_scores = out.detach().cpu().numpy()[data_train.valid_idx]
        val_acc, val_f1,val_f1macro, val_aucroc, val_recall, val_precision, val_cm = self.metric_manager.store_metrics("val", pred_scores, target_labels)

        if epoch%5 == 0:
          print("epoch: {} - loss: {:.4f} - accuracy train: {:.4f} -accuracy valid: {:.4f}  - val roc: {:.4f}  - val f1micro: {:.4f}".format(epoch, loss.item(), train_acc, val_acc, val_aucroc,val_f1))

  # To predict labels
  def predict(self, data=None, unclassified_only=True, threshold=0.5):
    # evaluate model:
    self.model.eval()
    if data is not None:
      self.data_train = data

    out = self.model(self.data_train)
    out = out.reshape((self.data_train.x.shape[0]))

    if unclassified_only:
      pred_scores = out.detach().cpu().numpy()[self.data_train.test_idx]
    else:
      pred_scores = out.detach().cpu().numpy()

    pred_labels = pred_scores > threshold

    return {"pred_scores":pred_scores, "pred_labels":pred_labels}

  # To save metrics
  def save_metrics(self, save_name, path="./save/"):
    file_to_store = open(path + save_name, "wb")
    pickle.dump(self.metric_manager, file_to_store)
    file_to_store.close()

  # To save model
  def save_model(self, save_name, path="./save/"):
    torch.save(self.model.state_dict(), path + save_name)
```

## Metric manager
We can use mutiple metrics, but we are interested on auroc, f1micro score. \\
Actually 2 types of f1 score can be used:  \\
-'micro': Calculate metrics globally by counting the total true positives, false negatives and false positives.

-'macro':  Calculate metrics for each label, and find their unweighted mean. This does not take label imbalance into account.

```python
class MetricManager(object):
  def __init__(self, modes=["train", "val"]):

    self.output = {}

    for mode in modes:
      self.output[mode] = {}
      self.output[mode]["accuracy"] = []
      self.output[mode]["f1micro"] = []
      self.output[mode]["f1macro"] = []
      self.output[mode]["aucroc"] = []
      #new
      self.output[mode]["precision"] = []
      self.output[mode]["recall"] = []
      self.output[mode]["cm"] = []

  def store_metrics(self, mode, pred_scores, target_labels, threshold=0.5):

    # calculate metrics
    pred_labels = pred_scores > threshold
    accuracy = accuracy_score(target_labels, pred_labels)
    f1micro = f1_score(target_labels, pred_labels,average='micro')
    f1macro = f1_score(target_labels, pred_labels,average='macro')
    aucroc = roc_auc_score(target_labels, pred_scores)
    #new
    recall = recall_score(target_labels, pred_labels)
    precision = precision_score(target_labels, pred_labels)
    cm = confusion_matrix(target_labels, pred_labels)

    # Collect results
    self.output[mode]["accuracy"].append(accuracy)
    self.output[mode]["f1micro"].append(f1micro)
    self.output[mode]["f1macro"].append(f1macro)
    self.output[mode]["aucroc"].append(aucroc)
    #new
    self.output[mode]["recall"].append(recall)
    self.output[mode]["precision"].append(precision)
    self.output[mode]["cm"].append(cm)

    return accuracy, f1micro,f1macro, aucroc,recall,precision,cm

  # Get best results
  def get_best(self, metric, mode="val"):

    # Get best results index
    best_results = {}
    i = np.array(self.output[mode][metric]).argmax()

    # Output
    for m in self.output[mode].keys():
      best_results[m] = self.output[mode][m][i]

    return best_results
```

## Training and validation
We set parameters and choose the model. Set prebuild parameters to True, if you want to use builtin PyG layers, otherwise False to use customs layers.

```python
# Set training arguments, set prebuild=True to use builtin PyG models otherwise False
args={"epochs":100,
      'lr':0.01,
      'weight_decay':1e-5,
      'prebuild':True,
      'heads':2,
      'hidden_dim': 128,
      'dropout': 0.5
      }

device = "cpu"
```

```python
# Model selector GAT or GATv2
net = "GAT"

if net == "GAT":
    if args['prebuild']==True:
      model = GAT(data_train.num_node_features, args['hidden_dim'], 1, args)
      print("Prebuilt GAT from PyG ")
    else:
      model = GATmodif(data_train.num_node_features, args['hidden_dim'], 1,args)
      print("Custom GAT implemented")
elif net == "GATv2":
    # args['heads'] = 1
    if args['prebuild']==True:
      model = GATv2(data_train.num_node_features, args['hidden_dim'], 1, args)
      print("Prebuilt GATv2 from PyG ")
    else:
      model = GATv2modif(data_train.num_node_features,  args['hidden_dim'], 1,args)
      print("Custom GATv2 implemented")

model.double().to(device)
```

    Prebuilt GAT from PyG

    GAT(
      (conv1): GATConv(165, 128, heads=2)
      (conv2): GATConv(256, 128, heads=2)
      (post_mp): Sequential(
        (0): Linear(in_features=256, out_features=128, bias=True)
        (1): Dropout(p=0.5, inplace=False)
        (2): Linear(in_features=128, out_features=1, bias=True)
      )
    )

```python
# Push data to GPU
data_train = data_train.to(device)

# Setup training settings
optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')
criterion = torch.nn.BCELoss()

# Train
gnn_trainer_gat = GnnTrainer(model)
gnn_trainer_gat.train(data_train, optimizer, criterion, scheduler, args)

gnn_trainer_gat.save_metrics("GATprebuilt.results", path=FOLDERNAME + "/save_results/")
gnn_trainer_gat.save_model("GATprebuilt.pth", path=FOLDERNAME + "/save_results/")
```

```python
#here we run GATv2 model
model = GATv2(data_train.num_node_features, args['hidden_dim'], 1, args)
model.double()
# Push data to GPU
data_train = data_train.to(device)

# Setup training settings
optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')
criterion = torch.nn.BCELoss()
# Train
gnn_trainer_gatv2 = GnnTrainer(model)
gnn_trainer_gatv2.train(data_train, optimizer, criterion, scheduler, args)

gnn_trainer_gatv2.save_metrics("GATv2prebuilt.results", path=FOLDERNAME + "/save_results/")
gnn_trainer_gatv2.save_model("GATv2prebuilt.pth", path=FOLDERNAME + "/save_results/")
```

```python
#here we run GATmodif model
model = GATmodif(data_train.num_node_features, args['hidden_dim'], 1, args)
model.double()
# Push data to GPU
data_train = data_train.to(device)

# Setup training settings
optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')
criterion = torch.nn.BCELoss()
# Train
gnn_trainer_gatmodif = GnnTrainer(model)
gnn_trainer_gatmodif.train(data_train, optimizer, criterion, scheduler, args)

gnn_trainer_gatmodif.save_metrics("GATmodif.results", path=FOLDERNAME + "/save_results/")
gnn_trainer_gatmodif.save_model("GATmodif.pth", path=FOLDERNAME + "/save_results/")
```

    epoch: 0 - loss: 0.7125 - accuracy train: 0.3300 -accuracy valid: 0.3359  - val roc: 0.5196  - val f1micro: 0.3359
    epoch: 5 - loss: 0.4592 - accuracy train: 0.9020 -accuracy valid: 0.9047  - val roc: 0.5681  - val f1micro: 0.9047
    epoch: 10 - loss: 0.3773 - accuracy train: 0.9020 -accuracy valid: 0.9047  - val roc: 0.6035  - val f1micro: 0.9047
    epoch: 15 - loss: 0.3451 - accuracy train: 0.9017 -accuracy valid: 0.9044  - val roc: 0.6135  - val f1micro: 0.9044
    epoch: 20 - loss: 0.3200 - accuracy train: 0.9018 -accuracy valid: 0.9044  - val roc: 0.6576  - val f1micro: 0.9044
    epoch: 25 - loss: 0.3052 - accuracy train: 0.9020 -accuracy valid: 0.9047  - val roc: 0.6619  - val f1micro: 0.9047
    epoch: 30 - loss: 0.2976 - accuracy train: 0.9021 -accuracy valid: 0.9049  - val roc: 0.6814  - val f1micro: 0.9049
    epoch: 35 - loss: 0.2898 - accuracy train: 0.9032 -accuracy valid: 0.9081  - val roc: 0.6865  - val f1micro: 0.9081
    epoch: 40 - loss: 0.2825 - accuracy train: 0.9051 -accuracy valid: 0.9099  - val roc: 0.7193  - val f1micro: 0.9099
    epoch: 45 - loss: 0.2747 - accuracy train: 0.9067 -accuracy valid: 0.9098  - val roc: 0.7331  - val f1micro: 0.9098
    epoch: 50 - loss: 0.2694 - accuracy train: 0.9094 -accuracy valid: 0.9110  - val roc: 0.7328  - val f1micro: 0.9110
    epoch: 55 - loss: 0.2623 - accuracy train: 0.9114 -accuracy valid: 0.9157  - val roc: 0.7661  - val f1micro: 0.9157
    epoch: 60 - loss: 0.2607 - accuracy train: 0.9134 -accuracy valid: 0.9167  - val roc: 0.7603  - val f1micro: 0.9167
    epoch: 65 - loss: 0.2538 - accuracy train: 0.9155 -accuracy valid: 0.9165  - val roc: 0.7530  - val f1micro: 0.9165
    epoch: 70 - loss: 0.2512 - accuracy train: 0.9166 -accuracy valid: 0.9168  - val roc: 0.7647  - val f1micro: 0.9168
    epoch: 75 - loss: 0.2488 - accuracy train: 0.9173 -accuracy valid: 0.9181  - val roc: 0.7651  - val f1micro: 0.9181
    epoch: 80 - loss: 0.2452 - accuracy train: 0.9196 -accuracy valid: 0.9178  - val roc: 0.7735  - val f1micro: 0.9178
    epoch: 85 - loss: 0.2421 - accuracy train: 0.9209 -accuracy valid: 0.9174  - val roc: 0.7730  - val f1micro: 0.9174
    epoch: 90 - loss: 0.2405 - accuracy train: 0.9212 -accuracy valid: 0.9195  - val roc: 0.7781  - val f1micro: 0.9195
    epoch: 95 - loss: 0.2387 - accuracy train: 0.9224 -accuracy valid: 0.9200  - val roc: 0.7959  - val f1micro: 0.9200

You can also run our custom models GATmodif and GATv2modif by setting in args, prebuild=False or run with GATmodif_3layer: \\
```model = GATmodif_3layer(data_train.num_node_features, args['hidden_dim'], 1, args) ``` \\
After training few models, you can load your saved outputs. Here we load the outputs of GAT and GATv2 from builtin PyG.

```python
# Fetch results from saved
fn="GATprebuilt.results"
mmGATprebuilt = pickle.load(open(FOLDERNAME + "/save_results/" + fn, "rb"))
mmGATprebuilt.get_best("aucroc")

fn="GATv2prebuilt.results"
mmGATv2prebuilt = pickle.load(open(FOLDERNAME + "/save_results/" + fn, "rb"))
mmGATv2prebuilt.get_best("aucroc")

```

    {'accuracy': 0.9733715103793844,
     'aucroc': 0.9837897717309015,
     'cm': array([[6237,   82],
            [ 104,  562]]),
     'f1macro': 0.9216616620238052,
     'f1micro': 0.9733715103793844,
     'precision': 0.8726708074534162,
     'recall': 0.8438438438438438}

# Performance visualization
Then we can compare the performance of our models by plotting the results with a selected metric.

```python
import plotly.express as px
import plotly.graph_objects as go

# Validation accuracy comparisons
def plot_training_comp(metric_manager_list, names, metric="aucroc", version="val",title="Val set accuracy comparison"):

  fig = go.Figure()
  fig = fig.update_layout(title=title,
                          xaxis_title="Epoch",
                          yaxis_title=metric)

  for i, metric_manager in enumerate(metric_manager_list):
    epochs = np.arange(len(metric_manager.output[version][metric]))
    fig.add_trace(go.Scatter(x=epochs,
                             y=metric_manager.output[version][metric],
                             name=names[i]))
  fig.show()

# Load saved outputs
def load_results(fn):
  mm1 = pickle.load(open(FOLDERNAME + "/save_results/" + fn, "rb"))
  return mm1

```

```python
# Compare GAT vs GATv2 with 2 heads, prebuilt PyGs models
plot_training_comp([mmGATprebuilt,mmGATv2prebuilt], ["GAT" ,"GATv2" ],"f1micro",title="Val set f1micro, heads=2")
plot_training_comp([mmGATprebuilt,mmGATv2prebuilt], ["GAT" ,"GATv2" ],"accuracy",title="Val set accuracy ,heads=2")
plot_training_comp([mmGATprebuilt,mmGATv2prebuilt], ["GAT" ,"GATv2" ],"f1macro",title="Val set f1macro ,heads=2")
plot_training_comp([mmGATprebuilt,mmGATv2prebuilt], ["GAT" ,"GATv2" ],"aucroc",title="Val set aucroc ,heads=2")
plot_training_comp([mmGATprebuilt,mmGATv2prebuilt], ["GAT" ,"GATv2" ],"recall",title="Val set recall ,heads=2")
plot_training_comp([mmGATprebuilt,mmGATv2prebuilt], ["GAT" ,"GATv2" ],"precision",title="Val set precision ,heads=2")

```

<meta charset="utf-8" />

<meta charset="utf-8" />

<meta charset="utf-8" />

<meta charset="utf-8" />

<meta charset="utf-8" />

<meta charset="utf-8" />

## Graph visualization
Now after training previous models, we can also visualize our predicted labels nodes as licit or illicit.We are going to show the graph for one time period.

```python
import networkx as nx
import matplotlib.pyplot as plt

# Load one model
m1 = GATv2(data_train.num_node_features, args['hidden_dim'], 1, args).to(device).double()
m1.load_state_dict(torch.load(FOLDERNAME + "/save_results/" + "GATv2prebuilt.pth"))
gnn_t2 = GnnTrainer(m1)
output = gnn_t2.predict(data=data_train, unclassified_only=False)
output

```

    {'pred_labels': array([False, False, False, ..., False, False, False]),
     'pred_scores': array([6.70029005e-04, 6.34844810e-06, 0.00000000e+00, ...,
            1.76779501e-30, 2.24113257e-01, 2.16711998e-04])}

```python
# Get index for one time period
time_period = 28
sub_node_list = df_merge.index[df_merge.loc[:, 1] == time_period].tolist()

# Fetch list of edges for that time period
edge_tuples = []
for row in data_train.edge_index.view(-1, 2).numpy():
  if (row[0] in sub_node_list) | (row[1] in sub_node_list):
    edge_tuples.append(tuple(row))
len(edge_tuples)

# Fetch predicted results for that time period
node_color = []
for node_id in sub_node_list:
  if node_id in classified_illicit_idx: #
     label = "red" # fraud
  elif node_id in classified_licit_idx:
     label = "green" # not fraud
  else:
    if output['pred_labels'][node_id]:
      label = "orange" # Predicted fraud
    else:
      label = "blue" # Not fraud predicted

  node_color.append(label)

# Setup networkx graph
G = nx.Graph()
G.add_edges_from(edge_tuples)

# Plot the graph
plt.figure(3,figsize=(16,16))
plt.title("Time period:"+str(time_period))
nx.draw_networkx(G, nodelist=sub_node_list, node_color=node_color, node_size=6, with_labels=False)
```

The diagram above has the following legend: Green = Not illicit (not fraud), Red = illicit (Fraud), Blue = Predicted not illicit, Orange = Predicted illicit.

## Training Cross comparisons
Here,we show you how to load few results from previous trained models , we should assign names then plots all results. The current visualization has additional models that we ran as well.

```python
# Take names of the saved files
model_name_list = ["GATv2prebuilt.results","GATprebuilt.results"]
#Assign names to plots
names = ["GATv2_prebuilt", "GAT_prebuilt"]

#Iterate to load saved outputs and plots
mm_list0 = [load_results(fn) for fn in model_name_list]
plot_training_comp(mm_list0, names, "aucroc",title="Val set aucroc comparison")

```

<meta charset="utf-8" />

## Best metrics

Next, we show you the best results of each model into a table.

```python
# Plot best results of each into a table
aucroc_l = []
accuracy_l = []
f1micro_l = []
f1macro_l = []
model_versions = []
for c1, mm in enumerate(mm_list0):
  model_versions.append(names[c1])
  best = mm.get_best("aucroc")

  aucroc_l.append(best["aucroc"])
  accuracy_l.append(best["accuracy"])
  f1micro_l.append(best["f1micro"])
  f1macro_l.append(best["f1macro"])

d = {"model_version": model_versions, "aucroc":aucroc_l, "accuracy": accuracy_l, "f1_macro": f1macro_l, "f1_micro": f1micro_l}
df = pd.DataFrame(data=d)
```

```python
df.round(3)
```

      model_version
      aucroc
      accuracy
      f1_macro
      f1_micro

      0
      GATv2_prebuilt
      0.984
      0.973
      0.922
      0.973

      1
      GAT_prebuilt
      0.980
      0.974
      0.924
      0.974

      2
      GATcustom
      0.796
      0.920
      0.658
      0.920

      3
      GATv2custom
      0.786
      0.919
      0.655
      0.919

      4
      GAT3layers
      0.745
      0.915
      0.610
      0.915


# <center> AcM HW04 </center>
<center> Yuanxing Cheng, Xavier, 程远星, 810925466 </center>

## Question 1
$\newcommand{\ffrac}{\displaystyle \frac}
\newcommand{\Tran}[1]{{#1}^{\mathrm{T}}}
\newcommand{\d}[1]{\displaystyle{#1}}
\newcommand{\Var}[1]{\mathrm{Var}\left[#1\right]}
\newcommand{\using}[1]{\stackrel{\mathrm{#1}}{=}}
\newcommand{\EE}[1]{\mathbb{E}\left[ #1 \right]}
\newcommand{\I}[1]{\mathrm{I}\left( #1 \right)}$

($\mathrm{a}$)

$\begin{align}
\mathrm{PV} &= \sum\limits_{t=0}^{47-45} \upsilon^{t}\cdot S_{t+1} \cdot {_{t}}p_{45} \\
&= S_1 + \upsilon \cdot S_2 \cdot \left( 1 - q_{45}\right) + \upsilon^2 \cdot S_3 \cdot p_{45}\cdot p_{46} \\
&= \$ 1000 + \ffrac{1} {1.07} \times \$ 1200 \cdot \left( 1 - 0.015\right) + \left(\ffrac{1} {1.07}\right)^2 \times \$ 2000 \cdot \left( 1 - 0.015\right)\left( 1 - 0.018\right) \\
&= \$ 3794.37505458992
\end{align}$

$$\\[1em]$$

($\mathrm{b}$)

$\begin{align}
\mathrm{PV} &= \sum\limits_{t=0}^{47-45} \upsilon^{t+1}\cdot S_{t+1} \cdot {_{t+1}}p_{45} \\
&= \upsilon \cdot S_1 \cdot \left( 1 - q_{45}\right) + \upsilon^2 \cdot S_2 \cdot \cdot p_{45} \cdot p_{46} + \upsilon^3 \cdot S_3 \cdot p_{45} \cdot p_{46} \cdot p_{47}\\
&= \ffrac{1} {1.07} \times \$ 1000 \cdot \left( 1 - 0.015\right) + \left( \ffrac{1} {1.07}\right)^2 \cdot \$ 1200
\cdot 0.985 \cdot \left( 1 - 0.018\right) \\
&\;\;\;\;+ \left(\ffrac{1} {1.07}\right)^3 \times \$ 2000 \cdot 0.985 \cdot 0.982 \cdot \left( 1 - 0.021\right) \\
&= \$ 3480.3805580702056
\end{align}$

$$\\[22em]$$

## Question 2

$\begin{align}
\newcommand{\d}[1]{\displaystyle{#1}}
\bar{a}_{x:\overline{1}\d{|}} &= \int_{0}^{1}e^{-\delta t} {_{t}}p_{x} \; \mathrm{d}t \\
&= \int_{0}^{1}e^{-\delta t} \left(\exp\left\{-\int_{0}^{t}\mu_{x+s}\;\mathrm{d}s\right\}\right) \; \mathrm{d}t \\
&= \int_{0}^{1}e^{-\delta t} e^{-\mu t} \; \mathrm{d}t \\
&= -\ffrac{1} {\delta + \mu} e^{-\left( \delta + \mu \right)t} \;\Big.\Big|_{0}^{1} \\
&= -\ffrac{1} {\delta + \mu} e^{-\left( \delta + \mu \right)} + \ffrac{1} {\delta + \mu}
\end{align}$

$$\\[46em]$$

## Question 3

$\begin{align}
\ddot{a}_{x:\overline{4}\d{|}} &= \sum_{k=0}^{4-1} \ddot{a}_{\overline{k+1}\d{|}} \cdot {_{k|}}q_{x} + {_{4}}p_{x} \cdot \ddot{a}_{\overline{4}\d{|}} \\
&= \sum_{k=0}^{4-1} \ddot{a}_{\overline{k+1}\d{|}} \cdot {_{k|}}q_{x} + \left[1-\big( \left(1 - {_{1}}p_{x}\right) + \left({_{1}}p_{x} - {_{2}}p_{x}\right) + \left({_{2}}p_{x} - {_{3}}p_{x}\right) + \left({_{3}}p_{x} - {_{4}}p_{x}\right) \big) \right] \cdot \ddot{a}_{\overline{4}\d{|}} \\
&= \ddot{a}_{\overline{1}\d{|}} \cdot {_{0|}}q_{x} + \ddot{a}_{\overline{2}\d{|}} \cdot {_{1|}}q_{x} + \ddot{a}_{\overline{3}\d{|}} \cdot {_{2|}}q_{x} + \ddot{a}_{\overline{4}\d{|}} \cdot {_{3|}}q_{x} + \left(1- {_{0|}}p_{x} - {_{1|}}p_{x} - {_{2|}}p_{x} - {_{3|}}p_{x} \right) \cdot \ddot{a}_{\overline{4}\d{|}} \\[0.5em]
&= 1.00 \times 0.33 + 1.93 \times 0.24 + 2.80 \times 0.16 + 3.62 \times 0.11 + \left( 1 - 0.33 - 0.24 - 0.16 - 0.11\right) \times 3.62\\[0.5em]
&= 2.2186
\end{align}$

$$\\[48em]$$

## Question 4

Since $
\Var{\bar{a}_{\overline{T_x}\d{|}}} = \Var{\ffrac{1 - \upsilon^{T_x}} {\delta}} = \ffrac{{^{2}}\bar{A}_x - \bar{A}_x^2} {\delta^2} = 4$, we have

$\bar{A}_x = \sqrt{{^{2}}\bar{A}_x - \delta^2 \cdot \Var{\bar{a}_{\overline{T_x}\d{|}}}} = \sqrt{0.08 - 0.1^2 \cdot 4} = 0.2$, so that finally

$\bar{a}_x = \ffrac{1 - \bar{A}_x} {\delta} = \ffrac{1 - 0.2} {0.1} = 8$

$$\\[50em]$$

## Question 5

$\begin{align}
\bar{a}_{\overline{x:\overline{10}\d{|}}} &= \bar{a}_{\overline{10}\d{|}} + {_{10}}E_{x} \cdot \bar{a}_{x+10} \\
&= \frac{1-\upsilon^{10}} {\delta} + {_{10}}p_{x} \cdot \upsilon^{10} \cdot \frac{1-\bar{A}_{x+10}} {\delta} \\
&= \frac{1-\left(\ffrac{1} {1+i}\right)^{10}} {\log(1+i)} + {_{10}}p_{x} \cdot \left(\frac{1} {1+i}\right)^{10} \cdot \frac{1-\left(\ffrac{\bar{A}_{x} - \bar{A}_{x:\overline{10}\d{|}}^{1}} {{_{10}}p_{x} \cdot \left(\frac{1} {1+i}\right)^{10}}\right)} {\log(1+i)} \\
&= 17.207417
\end{align}$


```python
try:
    import openmdao.api as om
    import dymos as dm
except ImportError:
    !python -m pip install openmdao[notebooks]
    !python -m pip install dymos[docs]
    import openmdao.api as om
    import dymos as dm
```

# Commercial Aircraft Range Maximization by Differential Inclusion

In this example we seek to maximize the range of a commercial aircraft
flying at a constant Mach number with a fixed initial fuel load. The
initial and final altitudes are constrained to 10000 ft (that is, in
this simplified example we're ignoring the aircraft behavior near the
ground).

This example is simplified in several ways. The Mach number is fixed at
0.8. When only considering fuel burn, the aircraft will tend to fly
slower to conserve fuel, ignoring the running costs of being airborne.

For a more detailed optimization of commercial aircraft performance, see
Betts {cite}`betts1995application`.

## Differential Inclusion

In this case we also use the steady flight assumption to allow the
problem to be parameterized with a different control set. Rather than
providing throttle/thrust and angle of attack as control variables, we
use a _differential inclusion_ approach. Whereas altitude and velocity
are often treated as state variables to be integrated, here we take a
slightly different approach.

Rather than using the differential equations to compute rates for the
flight path angle and true airspeed, we use a nonlinear solver to
determine the alpha and thrust time-history which is needed to make the
given altitude and airspeed time-history possible. This technique is
known as differential inclusion. It was demonstrated by Seywald {cite}`Seywald1994`.
Since the collocation
techniques in Dymos are based on polynomial approximations, rates of the
control variables can be easily calculated, which can then be applied to
the reordered dynamics equations to solve for angle of attack and
thrust.

The use of collocation/psuedospectral techniques in solving problems via
differential inclusion has been examined by Fahroo and Ross {cite}`fahroo2001second`
and Kumar and Seywald {cite}`Kumar1996`. Since the differential inclusion
approach demonstrated here relies on nonlinear solvers running within
the ODE model, obtaining accurate derivatives is paramount. Computing
derivatives with finite differences in the presence of nonlinear solvers
can lead to inaccurate derivatives based on the properties of the
solver. With OpenMDAO's unified derivatives framework, the sensitivy to
solver tolerances is greatly reduced.

### Pitfalls of Differential Inclusion

Differential inclusion works well in this case but it is highly
dependent on the parameterization of the problem. For instance, it is
possible for the optimizer to suggest a flight profile that requires
thrust well beyond the capabilities of the vehicle. In our case, the
throttle parameter $\tau$ is path constrained within the optimizer, so
that the optimizer can make the solution physically attainable by our
model.

We have two options for how to tackle this problem. We could specify
altitude as a control variable, and use it's approximate derivative
($alt\_rate$) to provide the altitude rate for the flight
equilibrium conditions. In practice, however, this noisy derivative of
altitude can cause numerical difficulties. Other times, the optimizer
will shape the control to satisfy our equilibrium conditions at the
collocation nodes but be wildly wrong in between these nodes. One simple
way to get around this is to still integrate altitude, but to provide
$climb\_rate$ as the control variable. This makes it easy to
constraint $climb\_rate$ with bounds constraints, and the
time-history of altitude will be smoother since it is the time integral
of the rate of climb.

#### Solver Options

In general, _openmdao.api.AnalysisError_ should be raised
when the design variables provided by the optimizer result in
nonconvergence of the solvers. This informs pyoptsparse to backtrack
along its line search for a solution. This error can be raised by the
_compute_ method of a component. The nonlinear solver can
also raise this error if it fails to converge by setting the following
options:

```python
self.nonlinear_solver.options['maxiter'] = 150
self.nonlinear_solver.options['err_on_maxiter'] = True
```

In the course of solving, an outer NewtonSolver won't necessarily force
subsystems to be solved, instead relying on the derivative at the at the
previous solution. In the context of OpenMDAO, even explicit components
constitute subsystems. This can lead to poorly converged solutions.

Setting the _solve\_subystems_ option to True will force
subsystems to be updated, making sure that all derivative information is
up-to-date.

```python
self.nonlinear_solver.options['solve_subsystems'] = True
self.nonlinear_solver.options['max_sub_solves'] = 10
```

#### The Impact of Choice of Transcription

The high-order Gauss-Lobatto method used by Dymos evaluates the ordinary
differential equations in a two-step process. First, the ODEs are
evaluated at the *state discretization* nodes. Using the state values
and derivatives, values and rates at the *collocation* nodes are then
interpolated. If a large rate is found at one of the state
discretization nodes, the interpolation at the collocation node can be
outside of the expected range, leading to nonconvergence of solvers
within the ODEs.

There is also an overhead cost to invoking solvers. While differential
inclusion reduces the number of design variables for the optimizer, it
also slows down each evaluation of the ODE somewhat. Since evaluating
the collocation constraints in the Gauss-Lobatto method requires two
evaluations of the ODE, it typically suffers greater slowdown that the
Radau Pseudospectral method, which evaluates the defects of the ODE with
a single evaluation that encompasses all the nodes.

For these reasons the Radau Pseudospectral transcription is typically
preferable when using a differential inclusion approach to problems.

## Problem Formulation

### State Variables

|Name        |Description                   |Fixed Initial Value  |Fixed Final Value|
|------------|------------------------------|---------------------|-----------------|
|range       |distance flown along ground   |True (0 NM)          |False            |
|mass\_fuel  |mass of fuel aboard aircraft  |True (30000 lbm)     |True (0 lbm)     |
|alt         |aircraft altitude             |True (10000 ft)      |True (10000 ft)  |

### Dynamic Controls

|Name         |Description             |Optimized  |Fixed Initial Value  |Fixed Final Value|
|-------------|------------------------|-----------|---------------------|-----------------|
|climb\_rate  |aircraft rate of climb  |True       |False                |False            |

### Parameters

|Name           |Description                 |Optimized             |
|---------------|----------------------------|----------------------|
|mach           |mach number                 |False (0.8)           |
|S              |aerodynamic reference area  |False (427.8 m\*\*2)  |
|mass\_empty    |aircraft empty mass         |False (330693.393 lbm)|
|mass\_payload  |aircraft payload mass       |False (74100 lbm)     |

### Objective

|Name   | Description | Location | Minimized or Maximized |
|-------|-------------|----------|------------------------|
|r      | range       | final    | Maximized              |

### Nonlinear Path Constraints

|Name  | Description                 |Lower  |Upper|
|------|-----------------------------|-------|-----|
|tau   | engine throttle parameter   |0.01   |1.0  |

### Nonlinear Boundary Constraints

None

## Models

### Atmosphere

This problem uses an analytic fit to the 1976 standard atmosphere.

|Name  |Description           |Input or Output|
|------|----------------------|---------------|
|alt   |altitude (m)          |input          |
|pres  |static pressure (Pa)  |output         |
|temp  |temperature (K)       |output         |
|sos   |speed of sound (m/s)  |output         |
|rho   |density (kg/m\*\*3)   |output         |

### True Airspeed

_TrueAirspeedComp_ uses the Mach number, provided as a
control, and the speed of sound from the atmosphere model to compute the
true airspeed of the aircraft.

\begin{align}
  TAS &= mach \cdot sos
\end{align}

|Name  |Description           |Input or Output|
|------|----------------------|---------------|
|mach  |Mach number           |input          |
|sos   |speed of sound (m/s)  |input          |
|TAS   |true airspeed (m/s)   |output         |

### Flight Path Angle

_SteadyFlightPathAngleComp_ uses the true airspeed and the
climb rate, obtained by differentiating the altitude time history at the
nodes, to compute the flight path angle.

\begin{align}
  \gamma &= \arctan \frac{TAS}{\dot h}
\end{align}

|Name       |Description              |Input or Output|
|-----------|-------------------------|---------------|
|TAS        |true airspeed (m/s)      |input          |
|alt\_rate  |climb rate (m/s)         |input          |
|gamma      |flight path angle (rad)  |output         |

### Range Rate

_RangeRateComp_ uses the true airspeed and the flight path
angle to determine the velocity projected along the ground. This is the
derivative of the state variable _range_.

\begin{align}
  \dot{range} &= TAS \cos \gamma
\end{align}

|Name        |Description              |Input or Output|
|------------|-------------------------|---------------|
|TAS         |true airspeed (m/s)      |input          |
|gamma       |flight path angle (rad)  |input          |
|dXdt:range  |range rate (m/s)         |output         |

### Mass

The component _MassComp_ defined in
`mass_comp.py` computes the aircraft total mass based on
its empty mass, payload mass, and current fuel mass. It also computes
total weight which simplifies some equations later on.

\begin{align}
  mass_{total} &= mass_{empty} + mass_{payload} + mass_{fuel} W_{total} = 9.80665 \, mass_{total}
\end{align}

|Name           |Description                |Input or Output|
|---------------|---------------------------|---------------|
|mass\_empty    |aircraft empty mass (kg)   |input          |
|mass\_payload  |payload mass (kg)          |input          |
|mass\_fuel     |fuel mass (kg)             |input          |
|mass\_total    |total aircraft mass (kg)   |output         |
|W\_total       |total aircraft weight (N)  |output         |

### Dynamic Pressure

The _DynamicPressureComp_ computes the dynamic pressure from
true airspeed and atmospheric density.

\begin{align}
  q &= \frac{1}{2} \rho TAS^2
\end{align}

|Name  |Description                      |Input or Output|
|------|---------------------------------|---------------|
|TAS   |true airspeed (m/s)              |input          |
|rho   |atmospheric density (kg/m\*\*3)  |input          |
|q     |dynamic pressure (Pa)            |output         |

### Aerodynamics

The aerodynamics group computes the aerodynamic coefficients and forces
on the vehicle. It consists of an interpolation component which outputs
lift, drag, and moment coefficients as a function of Mach number, angle
of attack, altitude, and tail rotation angle. A second component then
uses these coefficients, along with dynamic pressure and aerodynamic
reference area, to compute the lift and drag forces on the vehicle.

The aerodynamics group resides within the flight equilibrium group. As
that group iterates to find the combination of thrust coefficient, angle
of attack, and tail rotation angle, aerodynamics needs to update the
values of the interpolated coefficients and resulting forces.

Organizationally speaking, we logically could have put the dynamic
pressure component within the aerodynamics group. However, since that
group doesn't need to be updated with changes in alpha and tail angle,
it's more efficient to leave it outside of flight equilibrium group.

  |Name   |Description                |Input or Output|
  |-------|---------------------------|---------------|
  |mach   |mach number                |input          |
  |alt    |altitude (m)               |input          |
  |alpha  |angle of attack (deg)      |input          |
  |eta    |tail rotation angle (deg)  |input          |
  |CL     |lift coefficient           |output         |
  |CD     |drag coefficient           |output         |
  |CM     |moment coefficient         |output         |
  |L      |lift force (N)             |output         |
  |D      |drag force (N)             |output         |

```{Note}
This example uses _openmdao.api.MetaModelStructuredComp_ to
interpolate aerodynamic properties of the vehicle. This component is
somewhat easier to use since it is distributed as part of OpenMDAO, but
it can be significantly slower than alternatives such as
[MBI](https://github.com/OpenMDAO/MBI) and
[SMT](https://github.com/SMTorg/smt). The Aerodynamics group includes an
implementation of the aerodynamics coefficients interpolant which uses
MBI that is commented out. By switching to MBI for this component, this
problem will solve roughly 20 times faster.
```

### Flight Equilibrium

The steady flight equilibrium group uses balances to solve for the angle
of attack and tail plane rotation angle such that the aircraft is in
steady flight (the rates of change in flight path angle and true
airspeed are zero) and the aerodynamic moment in the pitch axis is zero.

Of course, in reality the vehicle will accelerate, but the flight
profile being modeled is so benign that assuming steady flight at
discrete points (nodes) in the trajectory is not terribly inaccurate.

The thrust necessary for steady flight is computed by balancing the drag
equation

\begin{align}
  C_T &= W_{total} \frac{sin \gamma}{q \cdot S \cdot \cos \alpha} + \frac{C_D}{\cos \alpha}
\end{align}

The lift coefficient required for steady flight is found by balancing
lift and weight:

\begin{align}
  \tilde{C_L} &= W_{total} \frac{cos \gamma}{q \cdot S} - C_T \sin \alpha
\end{align}

Using coefficients in the balance equations is better scaled from a numerical standpoint.

### Propulsion

Having determined the thrust, the propulsion group then computes the
rate of fuel burn. In addition, by normalizing thrust at any point by
the maximum possible thrust, we obtain the throttle parameter $\tau$.
The propulsion group uses a number of components to perform these
calculations.

\begin{align}
  T &= C_T \cdot q \cdot S
\end{align}

Maximum thrust is computed by multiplying sea-level thrust by the ratio
of pressure to sea-level atmospheric pressure.

\begin{align}
  T_{max} &= T_{max,sl} \frac{P}{P_{sl}}
\end{align}

The throttle parameter is then the ratio current thrust to maximum
possible thrust.

\begin{align}
  \tau &= \frac{T}{T_{max}}
\end{align}

The thrust specific fuel consumption is computed as follows:

\begin{align}
  TSFC &= TSFC_{sl} - 1.5 E - 10 \cdot 9.80665 \cdot alt
\end{align}

Finally, fuel burn rate is:

\begin{align}
  \dot{mass_{fuel}} &= -TSFC \frac{T}{9.80665}
\end{align}

## The ODE System: aircraft_ode.py

```python
class AircraftODE(om.Group):

    def initialize(self):
        self.options.declare('num_nodes', types=int,
                             desc='Number of nodes to be evaluated in the RHS')

    def setup(self):
        nn = self.options['num_nodes']

        self.add_subsystem(name='mass_comp',
                           subsys=MassComp(num_nodes=nn),
                           promotes_inputs=['mass_fuel'])

        self.connect('mass_comp.W_total', 'flight_equilibrium.W_total')

        self.add_subsystem(name='atmos',
                           subsys=USatm1976Comp(num_nodes=nn),
                           promotes_inputs=[('h', 'alt')])

        self.connect('atmos.pres', 'propulsion.pres')
        self.connect('atmos.sos', 'tas_comp.sos')
        self.connect('atmos.rho', 'q_comp.rho')

        self.add_subsystem(name='tas_comp',
                           subsys=TrueAirspeedComp(num_nodes=nn))

        self.connect('tas_comp.TAS',
                     ('gam_comp.TAS', 'q_comp.TAS', 'range_rate_comp.TAS'))

        self.add_subsystem(name='gam_comp',
                           subsys=SteadyFlightPathAngleComp(num_nodes=nn))

        self.connect('gam_comp.gam', ('flight_equilibrium.gam', 'range_rate_comp.gam'))

        self.add_subsystem(name='q_comp',
                           subsys=DynamicPressureComp(num_nodes=nn))

        self.connect('q_comp.q', ('aero.q', 'flight_equilibrium.q', 'propulsion.q'))

        self.add_subsystem(name='flight_equilibrium',
                           subsys=SteadyFlightEquilibriumGroup(num_nodes=nn),
                           promotes_inputs=['aero.*', 'alt'],
                           promotes_outputs=['aero.*'])

        self.connect('flight_equilibrium.CT', 'propulsion.CT')

        self.add_subsystem(name='propulsion', subsys=PropulsionGroup(num_nodes=nn),
                           promotes_inputs=['alt'])

        self.add_subsystem(name='range_rate_comp', subsys=RangeRateComp(num_nodes=nn))

        # We promoted multiple inputs to the group named 'alt'.
        # In order to automatically create a source variable for 'alt', we must specify their units
        # since they have different units in different components.
        self.set_input_defaults('alt', val=np.ones(nn,), units='m')
```

In this case the system has only two integrated states:
_range_ and _mass\_fuel_. There are six parameters.
Two of them (_alt_ and _climb\_rate_) will be varied dynamically in the phase
The other four (_mach_, _S_, _mass\_empty_, and _mass\_payload_), will be set to fixed values as non-optimized parameters.
More details on the various models involved can be found in the examples code.

## Building and running the problem

In the following code we define and solve the optimal control problem.
Note that we demonstrate the use of externally-connected design parameters in this case.
The four parameters are connected to a source provided by the _assumptions_ IndepVarComp.

```python
om.display_source("dymos.examples.aircraft_steady_flight.aircraft_ode")
```

```python
import matplotlib.pyplot as plt

import openmdao.api as om
import dymos as dm

from dymos.examples.aircraft_steady_flight.aircraft_ode import AircraftODE
from dymos.examples.plotting import plot_results
from dymos.utils.lgl import lgl

p = om.Problem(model=om.Group())
p.driver = om.pyOptSparseDriver()
p.driver.options['optimizer'] = 'SLSQP'
p.driver.declare_coloring()

num_seg = 15
seg_ends, _ = lgl(num_seg + 1)

traj = p.model.add_subsystem('traj', dm.Trajectory())

phase = traj.add_phase('phase0',
                       dm.Phase(ode_class=AircraftODE,
                                transcription=dm.Radau(num_segments=num_seg,
                                                       segment_ends=seg_ends,
                                                       order=3, compressed=False)))

# Pass Reference Area from an external source
assumptions = p.model.add_subsystem('assumptions', om.IndepVarComp())
assumptions.add_output('S', val=427.8, units='m**2')
assumptions.add_output('mass_empty', val=1.0, units='kg')
assumptions.add_output('mass_payload', val=1.0, units='kg')

phase.set_time_options(fix_initial=True,
                       duration_bounds=(300, 10000),
                       duration_ref=5600)

phase.add_state('range', units='NM',
                rate_source='range_rate_comp.dXdt:range',
                fix_initial=True, fix_final=False, ref=1e-3,
                defect_ref=1e-3, lower=0, upper=2000)

phase.add_state('mass_fuel', units='lbm',
                rate_source='propulsion.dXdt:mass_fuel',
                fix_initial=True, fix_final=True,
                upper=1.5E5, lower=0.0, ref=1e2, defect_ref=1e2)

phase.add_state('alt', units='kft',
                rate_source='climb_rate',
                fix_initial=True, fix_final=True,
                lower=0.0, upper=60, ref=1e-3, defect_ref=1e-3)

phase.add_control('climb_rate', units='ft/min', opt=True, lower=-3000, upper=3000,
                  targets=['gam_comp.climb_rate'],
                  rate_continuity=True, rate2_continuity=False)

phase.add_control('mach', targets=['tas_comp.mach', 'aero.mach'], units=None, opt=False)

phase.add_parameter('S',
                    targets=['aero.S', 'flight_equilibrium.S', 'propulsion.S'],
                    units='m**2')

phase.add_parameter('mass_empty', targets=['mass_comp.mass_empty'], units='kg')
phase.add_parameter('mass_payload', targets=['mass_comp.mass_payload'], units='kg')

phase.add_path_constraint('propulsion.tau', lower=0.01, upper=2.0)

p.model.connect('assumptions.S', 'traj.phase0.parameters:S')
p.model.connect('assumptions.mass_empty', 'traj.phase0.parameters:mass_empty')
p.model.connect('assumptions.mass_payload', 'traj.phase0.parameters:mass_payload')

phase.add_objective('range', loc='final', ref=-1.0e-4)

phase.add_timeseries_output('aero.CL')
phase.add_timeseries_output('aero.CD')

p.setup()

p['traj.phase0.t_initial'] = 0.0
p['traj.phase0.t_duration'] = 3600.0
p['traj.phase0.states:range'] = phase.interp('range', ys=(0, 724.0))
p['traj.phase0.states:mass_fuel'] = phase.interp('mass_fuel', ys=(30000, 1e-3))
p['traj.phase0.states:alt'][:] = 10.0

p['traj.phase0.controls:mach'][:] = 0.8

p['assumptions.S'] = 427.8
p['assumptions.mass_empty'] = 0.15E6
p['assumptions.mass_payload'] = 84.02869 * 400

dm.run_problem(p)

exp_out = traj.simulate()

plot_results([('traj.phase0.timeseries.states:range', 'traj.phase0.timeseries.states:alt',
               'range (NM)', 'altitude (kft)'),
              ('traj.phase0.timeseries.time', 'traj.phase0.timeseries.states:mass_fuel',
               'time (s)', 'fuel mass (lbm)'),
              ('traj.phase0.timeseries.time', 'traj.phase0.timeseries.CL',
               'time (s)', 'lift coefficient'),
              ('traj.phase0.timeseries.time', 'traj.phase0.timeseries.CD',
               'time (s)', 'drag coefficient')],
             title='Commercial Aircraft Optimization',
             p_sol=p, p_sim=exp_out)

plt.show()
```

```python
from openmdao.utils.assert_utils import assert_near_equal

assert_near_equal(p.get_val('traj.phase0.timeseries.states:range', units='NM')[-1],
                  726.85, tolerance=1.0E-2)

assert_near_equal(p.get_val('traj.phase0.timeseries.CL')[0],
                  exp_out.get_val('traj.phase0.timeseries.CL')[0],
                  tolerance=1.0E-9)

assert_near_equal(p.get_val('traj.phase0.timeseries.CD')[0],
                  exp_out.get_val('traj.phase0.timeseries.CD')[0],
                  tolerance=1.0E-9)
```

## References

```{bibliography}
:filter: docname in docnames
```