The codes for the "Heterogeneous Reinforcement Learning for Defending Power Grids Against Attacks" paper submitted to APML journal. Create an Issue if you have a problem running the code.

TLDR;

1. Install the "environment.yml" environment (Tested in Windows).
2. For training an RL agent in a power grid, run rev_***method***_***condition***_***grid_name***.py files. For example, for training a deep Q agent under attack condition in IEEE-118 run the rev_deepQ_attack_118.py file. Your model and logs will be saved.
3. To inspect your saved logs, run rev_tensorboard.py file.
4. To evaluate your agent, use 'rev_method_eval.py` file.
5. To create .csv timeseries from your evaluations use rev_extract_to_CSV.py file.
6. To train your TGCN model based on the created timeseries use grid_main.py

Long version guide

Welcome to the guide for using the Temporal Graph Convolutional Network (TGCN) Reinforcement Learning project! This guide is designed to provide you with everything you need to get started with using the TGCN analyzer in a power grid reinforcement learning tasks.

The TGCN project is aimed at improving the efficiency of learning unknown power grids under attack by dividing the large action space into smaller action spaces. This project utilizes TGCN, a powerful deep learning architecture capable of analyzing power grid topology and temporal dynamics to create a more compact representation of the grid that is better suited for the reinforcement learning agent.

To facilitate power grid simulation under attack, the project utilizes Grid2Op, an open-source Python package. Grid2Op is specifically designed for building and training reinforcement learning agents in power systems. It provides a comprehensive set of tools for modeling and simulating power grids, including power flow calculations, network topology analysis, and attack scenarios.

The machine learning portion of this project is implemented using Tensorflow, an open-source machine learning framework developed by Google. Tensorflow provides a flexible and scalable platform for building and training deep learning models, including graph convolutional networks like TGCN.

This guide will cover everything from installing the necessary packages and dependencies to running the TGCN model on your power grid simulation data. We will walk you through the entire process of using the TGCN analyzer in your reinforcement learning tasks, from preprocessing the data to training and evaluating the model.

Whether you are an experienced machine learning practitioner or just getting started with reinforcement learning in power systems, this guide will provide you with the knowledge and tools you need to use TGCN to analyze power grid topology and temporal dynamics. By the end of this guide, you will be able to use TGCN to create a more efficient and effective reinforcement learning agent for your power grid simulation under attack.

So, let's get started!

Overview

The project consists of several code files, each serving a specific purpose. The primary file, "grid_main.py", is responsible for training the Temporal Graph Convolutional Network (TGCN) on generated attack data.

• "grid_gen.py" generates data for normal conditions where there is no attack, while "grid_gen_adversarial.py" generates data for attack conditions. "grid_data_extractor.py" extracts data from the stored episodes generated by the previous two files.
• "grid_adj.py" sets the adjacency matrix of the specific power grid and stores it in a .csv file.
• "grid_distributed_RL.py" creates a gym loop and serves as an example of how the trained TGCN individual agents will be utilized in a gym loop.
• Finally, "custom_functions.py" contains functions defined for pre-processing data and other utilities that facilitate the code.

By utilizing these files in conjunction with the necessary packages and dependencies, users can efficiently train and evaluate the TGCN analyzer for power grid reinforcement learning tasks.

Dependencies

In order to successfully implement the TGCN analyzer in a power grid reinforcement learning tasks, you will need to install several packages and dependencies. These packages and dependencies are essential for working with the Grid2Op Python package and Tensorflow machine learning framework, as well as for manipulating and visualizing data. The following packages are required for this project:

• matplotlib-inline=0.1.6: This package is used for inline plotting in Jupyter notebooks and provides a convenient way to visualize data.
• python=3.8.16: This is the programming language used for implementing the TGCN analyzer.
• spyder-kernels=2.4.2: This package is required for running the Python IDE, Spyder.
• zeromq=4.3.4: This package provides a messaging library used for high-performance asynchronous communication between processes.
• grid2op==1.8.1: This is the open-source Python package used for simulating power grids under attack.
• networkx==3.0: This package is used for graph analysis and manipulation, which is an essential component of the TGCN model.
• numpy==1.24.2: This package is used for scientific computing and is a fundamental dependency for many Python packages.
• pandapower==2.11.1: This package is used for power system analysis and is required for working with power grid simulation data.
• pandas==1.5.3: This package is used for data manipulation and analysis, which is essential for preprocessing data for use with the TGCN model.
• scipy==1.10.1: This package is used for scientific computing and is required for working with power grid simulation data.
• tqdm==4.65.0: This package provides a progress bar for iterative processes and is useful for monitoring the progress of model training and evaluation.
• numba==0.56.4: This package is a just-in-time compiler for Python that translates Python functions into optimized machine code at runtime, resulting in significant speedup for numerical computations and other performance-critical tasks.
• stable-baselines3==2.2.1: This package is a set of high-quality implementations of reinforcement learning algorithms in Python, built on top of the PyTorch deep learning framework. It provides a user-friendly interface for training and evaluating reinforcement learning agents, making it easier for researchers and practitioners to experiment with and deploy RL algorithms in various applications.

By installing these packages and dependencies, you will have the necessary tools and frameworks to work with power grid simulation data and to train and evaluate the TGCN model.

- matplotlib-inline=0.1.6=py38haa95532_0
- python=3.8.16=h6244533_3
- spyder-kernels=2.4.2=py38haa95532_0
- zeromq=4.3.4=hd77b12b_0
- pip:
   - grid2op==1.8.1
   - networkx==3.0
   - numpy==1.24.2
   - pandapower==2.11.1
   - pandas==1.5.3
   - scipy==1.10.1
   - tqdm==4.65.0

Generating data

This Python code is used to generate attack or normal datasets using the Grid2Op library, which is a Reinforcement Learning (RL) framework for power systems. In summary, this code generates an attack or normal dataset for power systems using the Grid2Op RL framework. It defines an environment, specifies an agent class, and runs the environment to generate the dataset.

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

import os
import sys
import numpy as np
import grid2op
from grid2op.PlotGrid import PlotMatplot
from grid2op.Runner import Runner
from grid2op.Agent import DoNothingAgent, PowerLineSwitch, TopologyGreedy
from grid2op.Action import PowerlineSetAction
from grid2op.Opponent import RandomLineOpponent, BaseActionBudget, GeometricOpponent
from grid2op.Episode import EpisodeReplay
from tqdm import tqdm

First, some warnings are ignored, and the necessary libraries are imported. The grid2op module is imported to use its functionality to create an environment. The PlotMatplot module is imported to plot the grid, and the Runner module is used to define an agent and run an instance of the environment. The DoNothingAgent, PowerLineSwitch, and TopologyGreedy classes are used as agent classes in the code which can be the custom RL agent defined by the user.

# env_name = 'rte_case5_example'
# env_name = 'rte_case14_realistic'
env_name = 'l2rpn_wcci_2022'
# env_name = 'educ_case14_storage'

env = grid2op.make(env_name, test = False)
line_list = ['91_101_55',
 '31_113_76',
 '87_88_29',
 '11_13_57',
 '7_29_100',
 '48_68_170',
 '62_63_160',
 '11_15_83',
 '55_57_148',
 '2_4_96',
 '63_64_163',
 '8_9_140',
 '83_84_25',
 '104_107_66',
 '27_28_99',
 '36_38_115',
 '15_16_86',
 '13_14_79',
 '102_104_61',
 '26_27_98'] 

env_with_opponent = grid2op.make(env_name,
                                  opponent_attack_cooldown= 12*3, # 12 * Hours
                                  opponent_attack_duration= 12*0.5, # 12 * Hours 
                                  opponent_budget_per_ts= 0.1,
                                  opponent_init_budget= 0,
                                  opponent_action_class=PowerlineSetAction,
                                  opponent_class=RandomLineOpponent,
                                  opponent_budget_class=BaseActionBudget,
                                  # kwargs_opponent={"lines_attacked":
                                  #    ["1_3_3", "1_4_4", "3_6_15", "9_10_12", "11_12_13", "12_13_14"]}
                                  kwargs_opponent={"lines_attacked": line_list}
                                  )


agentClass = PowerLineSwitch  
path_name = f"Dataset/Attack_dataset/{env_name}/PLS/" 
# path_name = f"Dataset/Normal_dataset/{env_name}/PLS/" 

# agentClass = TopologyGreedy 
# path_name = f"Dataset/Attack_dataset/{env_name}/TG/" 
# path_name = f"Dataset/Normal_dataset/{env_name}/TG/" 

In the code, an environment name is specified, which can be set to one of the available environment names in Grid2Op. The line_list variable specifies the lines to be attacked during training, and RandomLineOpponent is used as an opponent class.

Two agent classes, PowerLineSwitch and TopologyGreedy, are commented out in the code. The PowerLineSwitch connect or reconnects the powerline, while the TopologyGreedy class changes the topology of a substation.

if not os.path.exists(path_name):
   os.makedirs(path_name)
path_to_save = os.path.abspath(path_name)

nb_episode = 1000
max_iter = 12*24 # 12 * Hours 

# pbar = tqdm(total=nb_episode, disable=False)

runner = Runner(**env_with_opponent.get_params_for_runner(), agentClass=agentClass) 
# runner = Runner(**env.get_params_for_runner(), agentClass=agentClass) 

res=runner.run(nb_episode = nb_episode, max_iter = max_iter, path_save = path_to_save, pbar = True) 

# for i in range(nb_episode):
#     plot_epi = EpisodeReplay(path_to_save)
#     plot_epi.replay_episode(res[i][1], gif_name='this_episode')

The path_name variable defines the path to store the dataset generated. The nb_episode variable defines the number of episodes to generate in the dataset. The max_iter variable is used to specify the maximum number of iterations per episode, with each iteration being one time step in the simulation. The Runner class is then used to generate the dataset.

env_with_opponent.get_params_for_runner() is used to get the parameters of the environment with the specified opponent, while agentClass is set to one of the agent classes defined earlier. The generated dataset is saved in the path specified by path_to_save. The pbar argument is set to True to display a progress bar during the generation of the dataset.

Extracting the generated data

This code snippet is used to extract data from the saved episodes generated by previous files. Specifically, it analyzes the correlation between the load flow parameter rho for each powerline in the simulation.

from grid2op.Episode import EpisodeData
import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

env_name = 'rte_case14_realistic'
path_ = f'Dataset/Normal_dataset/{env_name}/TG/'
path_agent = os.path.abspath(path_)
episode_studied = EpisodeData.list_episode(path_agent)
n_episodes = len(episode_studied)

this_episode = EpisodeData.from_disk(*episode_studied[0])

First, the required packages are imported: EpisodeData from grid2op, os, matplotlib.pyplot, numpy, and pandas.

Next, the variable env_name is set to the name of the power grid environment being analyzed, and the path to the directory containing the saved episodes is defined. The list of saved episodes is then obtained using EpisodeData.list_episode, and the number of episodes is determined using len.

for i in range(n_episodes):
    this_episode = EpisodeData.from_disk(*episode_studied[i])
    nstep = len(this_episode.observations)
    for t, obs in enumerate(this_episode.observations):
        nlines = obs.n_line
        break

    rho = np.zeros((nstep, nlines))

    for t, obs in enumerate(this_episode.observations):
        rho[t, :] = obs.rho

    rho = rho[:-1,:]
    if i==0:
        rho_Total=rho 
    else:
        rho_Total=np.concatenate((rho_Total,rho),axis=0)

A for loop is then used to iterate through each episode. Within this loop, the variable this_episode is set to the data for the current episode using the variable EpisodeData.from_disk. The number of time steps and the number of powerlines in the observation for the current episode are then determined using len and the n_line attribute of the observation, respectively.

The rho variable is then initialized as a 2D numpy array of zeros with dimensions (nstep, nlines), where nstep is the number of time steps in the episode and nlines is the number of powerlines. Another for loop is used to iterate through each time step, and for each time step, the load flow parameter rho for each powerline is extracted from the observation and stored in rho.

fig = plt.figure(figsize=(8, 8))
plt.matshow(np.corrcoef(rho_Total.T), 0)
plt.clim(-1,1)
plt.colorbar()
plt.xlabel("powerline number")
plt.ylabel("powerline number")
fig.savefig(f'rho_normal_{env_name}_TG_correlation',dpi=600)

df = pd.DataFrame(rho_Total)
df.to_csv(f'rho_normal_{env_name}_TG.csv')

After all episodes have been analyzed, the correlation between the load flow parameter rho for each powerline is computed using numpy.corrcoef and plotted using matplotlib.pyplot.matshow. The resulting figure is saved using matplotlib.pyplot.savefig.

In addition, the rho data for all episodes is saved to a CSV file using pandas.DataFrame.to_csv.

To use this code for a different environment, simply set the env_name variable to the name of the desired environment and ensure that the directory containing the saved episodes for that environment exists.

Adjacency Matrix

The code provided sets the adjacency matrix of a specific power grid and stores it in a CSV file. Let's go through the code step by step:

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

import numpy as np
import pandas as pd
import grid2op
from grid2op.PlotGrid import PlotMatplot
from grid2op.Runner import Runner
import networkx as nx
import matplotlib.pyplot as plt 

The code first imports some necessary libraries. The "warnings" library is used to ignore future warnings, while "numpy" and "pandas" are used for data manipulation. "grid2op" library provides the necessary tools to interact with the power grid simulation environment. "networkx" is used for creating, manipulating, and studying complex networks. Finally, "matplotlib" is a visualization library used for plotting the network.

env_name = 'rte_case14_realistic'

env = grid2op.make(env_name, test=True) # default = "rte_case14_realistic" environment 

obs = env.reset() 

graph = obs.as_networkx()

linegraph=nx.line_graph(graph)
oldAdj = nx.adjacency_matrix(linegraph)
oldAdj = oldAdj.todense()

The variable "env_name" is set to the name of the environment we want to use. In this case, it is "rte_case14_realistic". The next line creates an instance of the power grid simulation environment using the "make()" method of the "grid2op" library. The "test" argument is set to "True" to run the environment in test mode. The "reset()" method of the environment object is called to reset the environment and return an initial observation. The "as_networkx()" method of the observation object is called to obtain the network graph of the power grid. The "line_graph()" method of "networkx" is used to convert the network graph into a line graph. The adjacency matrix of the line graph is obtained using the "adjacency_matrix()" method of "networkx". The resulting matrix is in sparse format. The sparse matrix is converted to a dense matrix using the "todense()" method.

dict_={'0':18, '1':3, '2':16, '3':2, '4':15, '5':7, 
       '6':8, '7':1, '8':6, '9':9, '10':4, '11':19, 
       '12':13, '13':5, '14':14, '15':11, '16':12, '17':17, 
       '18':0, '19':10}

Adj=np.zeros((20,20))
for i in range(20):
    for j in range(20):
        if oldAdj[i,j]==1:
            m = dict_[str(i)]
            n = dict_[str(j)]
            Adj[m,n] = 1

df = pd.DataFrame(Adj)
df.to_csv('Adj.csv')

A dictionary is created to map the original node indices to a new index range from 0 to 19. The adjacency matrix is updated by mapping the node indices using the dictionary created in the previous step. The updated adjacency matrix is stored in a pandas DataFrame object. The pandas DataFrame object is saved to a CSV file named "Adj.csv" using the "to_csv()" method.

Distributed RL

The code is a Python script that creates a gym loop for the grid2op power network simulation environment. The gym loop is designed to simulate the behavior of individual agents in the environment using a trained model.

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

import grid2op
from grid2op.Agent import RecoPowerlineAgent, PowerLineSwitch, TopologyGreedy
from grid2op.Runner import Runner
import tensorflow as tf
import random
import numpy as np 
from ttictoc import tic, toc
import os

test_number = 1
attack_time = 24
decision_interval = 10
nb_episode = 50
max_iter = 50
costly  = 0.7

env = grid2op.make('rte_case14_realistic')
attack_list = ["1_3_3", "1_4_4", "3_6_15", "9_10_12", "11_12_13", "12_13_14"]

TG_model = tf.keras.models.load_model(f'%% YOUR SAVED MODEL PATH %%')
PLS_model = tf.keras.models.load_model(f'%% YOUR SAVED MODEL PATH %%')

The first section of the code imports the required libraries and sets some parameters for the simulation. These parameters include the test number, attack time, decision interval, number of episodes, maximum iterations, and the costly factor. The warnings.simplefilter() function is used to ignore any FutureWarnings that might be generated during the execution of the code.

def extreme_finder(pred, th):
    pred_extreme = np.where(pred>th,1,0)    
    extreme_degree = sum(sum(pred_extreme))
    return extreme_degree   

def agent_selection(PLS_pred, TG_pred, th):
    PLS_ex_deg = extreme_finder(PLS_pred, th)
    TG_ex_deg = extreme_finder(TG_pred, th)
    if costly * PLS_ex_deg > TG_ex_deg:
        selected_agent = 'TopologyGreedy'
    else:
        selected_agent = 'PowerLineSwitch'
    if PLS_ex_deg == 0 and PLS_ex_deg == 0:
        selected_agent = 'RecoPowerlineAgent'
    return selected_agent 

This section defines two utility functions that will be used later in the code. The first function extreme_finder() is used to count the number of extreme values (values above a certain threshold) in the input array. The second function agent_selection() selects the appropriate agent based on the predictions made by the TG and PLS models.

selected_agent = 'RecoPowerlineAgent'
rho = np.zeros((nb_episode, max_iter + 1, env.n_line))
agent_list = []

for th in np.arange(0.7,1,0.1):
    th = round(th,2)
    file_name = f'Results/CompareTH/reults_{th}.txt'
    os.makedirs(os.path.dirname(file_name), exist_ok=True)

    for i in range(nb_episode):
        message = f'episde {i} started ...\n'
        print(message)
        with open(file_name,'a+') as f:
            f.write(message)
            f.close()
        t = 0 
        obs = env.reset()
        rho[i,t,:] = obs.rho 
        done = False
        reward = env.reward_range[0]
        while not done:
            if selected_agent == 'RecoPowerlineAgent':
                agent = RecoPowerlineAgent(env.action_space)
            elif selected_agent == 'TopologyGreedy': 
                agent = TopologyGreedy(env.action_space)
            elif selected_agent == 'PowerLineSwitch': 
                agent = PowerLineSwitch(env.action_space)

            act = agent.act(obs, reward, done)
            obs, reward, done, info = env.step(act)
            t = t + 1
            rho[i,t,:] = obs.rho
            if t == attack_time :
                attacked_line = random.choice(attack_list)
                act = env.action_space.disconnect_powerline(line_name=attacked_line)
            if t == attack_time + decision_interval :
                window = rho[i,t-23:t+1,:]
                window = np.expand_dims(window, axis = 0)
                window = np.expand_dims(window, axis = 3)
                TG_pred = TG_model.predict(window)
                TG_pred = TG_pred[0,:,:]
                PLS_pred = PLS_model.predict(window)
                PLS_pred = PLS_pred[0,:,:]
                selected_agent  = agent_selection(PLS_pred, TG_pred, th)
                agent_change = True
                tic()
            if t == max_iter:
                done = True
        if done:
            agent_list.append(selected_agent)
            message = f'episde {i} ended. Selected agent was {selected_agent}.\n' 
            print(message)
            with open(file_name,'a+') as f:
                f.write(message)
                f.close()
            if agent_change:
                each_step_comp = toc()/(t-(attack_time + decision_interval))
                message = f'it costed {each_step_comp} seconds per step\n'
                print(message)
                with open(file_name,'a+') as f:
                    f.write(message)
                    f.close()
                agent_change = False

This section is the main gym loop. The first part of the loop initializes some variables and starts a loop to iterate over different threshold values. The threshold values are used in the agent_selection() function to determine the appropriate agent to use. The loop also initializes a list to keep track of the selected agents for each episode.

TGCN

The primary file, ```grid_main.py`'', is responsible for training the Temporal Graph Convolutional Network (TGCN) on generated attack data.

import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

import pandas as pd
import numpy as np
import os
import typing
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.preprocessing import timeseries_dataset_from_array
from custom_functions import func_create_tf_dataset, func_preprocess, func_compute_adjacency_matrix

This code imports the necessary libraries and sets an environment variable to avoid issues with the KMP_DUPLICATE_LIB_OK flag. It also imports custom functions that will be used later in the code.

adjacency_matrix = pd.read_csv('Adj.csv')
# adjacency_matrix = pd.read_csv('connections.csv')
adjacency_matrix = np.array(adjacency_matrix)[:,1:]


rho = pd.read_csv('%% YOUR DATASET PATH %%')
rho = np.array(rho)[:,1:]


fig = plt.figure(figsize=(8, 8))
plt.matshow(np.corrcoef(rho.T), 0)
plt.colorbar()
plt.clim(-1,1)
plt.xlabel("line number")
plt.ylabel("line number")
fig.savefig('rho_correlation',dpi=600)

This code section reads the adjacency matrix and attack data (in the form of the matrix of PLS values) from csv files, converts them into numpy arrays, and visualizes the correlation between the columns of the attack data matrix using a heatmap. The adjacency matrix is used to represent the network topology, and the attack data matrix represents the line-wise values of PLS during a power system attack. The correlation matrix provides an idea of how different lines in the power system are affected during the attack. The resulting heatmap is saved as an image file.

adjacency_matrix = np.corrcoef(rho.T)

adjacency_matrix = np.random.rand(20,20)
adjacency_matrix  = adjacency_matrix * np.transpose(adjacency_matrix)
for i in range(len(adjacency_matrix)):
    adjacency_matrix[i][i]=1

adjacency_matrix = np.identity(20)

The code initializes the adjacency matrix with different values in each block. You should choose the one that works for you the best.

train_size, val_size = 0.7, 0.2

train_array, val_array, test_array = func_preprocess(rho, train_size, val_size)

in_feat = 1
epochs = 20
out_feat = 10
lstm_units = 64
batch_size = 64
input_sequence_length = 12
forecast_horizon = 1
multi_horizon = True
learning_rate = 0.01
patience = 10


train_dataset, val_dataset = (func_create_tf_dataset(
    data_array, 
    input_sequence_length, 
    forecast_horizon, 
    batch_size)
    for data_array in [train_array, val_array]
)

test_dataset = func_create_tf_dataset(
    test_array,
    input_sequence_length,
    forecast_horizon,
    batch_size=test_array.shape[0],
    shuffle=False,
    multi_horizon=multi_horizon,
)

This code block sets up the parameters and creates training, validation, and test datasets for the TGCN model. train_size and val_size are defined as the proportions of the dataset used for training and validation, respectively. func_preprocess() function is used to split the dataset into training, validation, and test sets. It takes rho (the input dataset), train_size, and val_size as input parameters and returns train_array, val_array, and test_array.

Next, several parameters are defined for the TGCN model:

• in_feat: The number of input features (in this case, it is set to 1)
• epochs: The number of epochs to train the model
• out_feat: The number of output features (in this case, it is set to 10)
• lstm_units: The number of LSTM units in the TGCN model
• batch_size: The batch size used for training and validation datasets
• input_sequence_length: The length of the input sequence for each batch
• forecast_horizon: The number of time steps to forecast
• multi_horizon: A Boolean variable indicating whether the model is predicting a single or multiple horizons
• learning_rate: The learning rate used in the optimizer
• patience: The number of epochs to wait before early stopping if there is no improvement in validation loss.

Finally, the func_create_tf_dataset() function is used to create train_dataset, val_dataset, and test_dataset. It takes the data arrays, input_sequence_length, forecast_horizon, and batch_size as input parameters and returns a TensorFlow dataset object. multi_horizon is set to True for training and validation datasets but set to False for the test dataset.

class GraphInfo:
    def __init__(self, edges: typing.Tuple[list, list], num_nodes: int):
        self.edges = edges
        self.num_nodes = num_nodes
        
node_indices, neighbor_indices = np.where(adjacency_matrix == 1)
graph = GraphInfo(
    edges=(node_indices.tolist(), neighbor_indices.tolist()),
    num_nodes=adjacency_matrix.shape[0],
)
print(f"number of nodes: {graph.num_nodes}, number of edges: {len(graph.edges[0])}")

This code defines a class called GraphInfo with two attributes: edges and num_nodes.

• edges is a tuple of two lists, containing the node indices and their corresponding neighbors' indices.
• num_nodes is an integer representing the total number of nodes in the graph.

The where function of NumPy is used to find the indices of non-zero elements in the adjacency_matrix, which correspond to the edges of the graph. These indices are then converted into two lists and passed as arguments to initialize an instance of the GraphInfo class. Finally, the number of nodes and edges in the graph are printed.

class GraphConv(layers.Layer):
    def __init__(
        self,
        in_feat,
        out_feat,
        graph_info: GraphInfo,
        aggregation_type="mean",
        combination_type="concat",
        activation: typing.Optional[str] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.in_feat = in_feat
        self.out_feat = out_feat
        self.graph_info = graph_info
        self.aggregation_type = aggregation_type
        self.combination_type = combination_type
        self.weight = tf.Variable(
            initial_value=keras.initializers.glorot_uniform()(
                shape=(in_feat, out_feat), dtype="float32"
            ),
            trainable=True,
        )
        self.activation = layers.Activation(activation)

    def aggregate(self, neighbour_representations: tf.Tensor):
        aggregation_func = {
            "sum": tf.math.unsorted_segment_sum,
            "mean": tf.math.unsorted_segment_mean,
            "max": tf.math.unsorted_segment_max,
        }.get(self.aggregation_type)

        if aggregation_func:
            return aggregation_func(
                neighbour_representations,
                self.graph_info.edges[0],
                num_segments=self.graph_info.num_nodes,
            )

        raise ValueError(f"Invalid aggregation type: {self.aggregation_type}")

    def compute_nodes_representation(self, features: tf.Tensor):
        """Computes each node's representation.

        The nodes' representations are obtained by multiplying the features tensor with
        `self.weight`. Note that
        `self.weight` has shape `(in_feat, out_feat)`.

        Args:
            features: Tensor of shape `(num_nodes, batch_size, input_seq_len, in_feat)`

        Returns:
            A tensor of shape `(num_nodes, batch_size, input_seq_len, out_feat)`
        """
        return tf.matmul(features, self.weight)

    def compute_aggregated_messages(self, features: tf.Tensor):
        neighbour_representations = tf.gather(features, self.graph_info.edges[1])
        aggregated_messages = self.aggregate(neighbour_representations)
        return tf.matmul(aggregated_messages, self.weight)

    def update(self, nodes_representation: tf.Tensor, aggregated_messages: tf.Tensor):
        if self.combination_type == "concat":
            h = tf.concat([nodes_representation, aggregated_messages], axis=-1)
        elif self.combination_type == "add":
            h = nodes_representation + aggregated_messages
        else:
            raise ValueError(f"Invalid combination type: {self.combination_type}.")

        return self.activation(h)

    def call(self, features: tf.Tensor):
        """Forward pass.

        Args:
            features: tensor of shape `(num_nodes, batch_size, input_seq_len, in_feat)`

        Returns:
            A tensor of shape `(num_nodes, batch_size, input_seq_len, out_feat)`
        """
        nodes_representation = self.compute_nodes_representation(features)
        aggregated_messages = self.compute_aggregated_messages(features)
        return self.update(nodes_representation, aggregated_messages)

This is a custom Keras layer for performing graph convolution. The GraphConv layer takes in some parameters during initialization such as in_feat, out_feat, graph_info, aggregation_type, and combination_type. During a forward pass, this layer computes each node's representation by multiplying the features tensor with a weight tensor of shape (in_feat, out_feat). The node representations are then used to compute aggregated messages, which is a function of the neighbor representations. The neighbor representations are obtained by gathering the features tensor along the second axis, using the neighbor indices from the graph information. The messages are then aggregated using one of three functions, specified by the aggregation_type parameter, which can be "sum", "mean", or "max". The aggregated messages are combined with the nodes' representations using one of two methods, specified by the combination_type parameter, which can be "concat" or "add". Finally, the combined messages are passed through an activation function and returned as output.

class LSTMGC(layers.Layer):
    """Layer comprising a convolution layer followed by LSTM and dense layers."""

    def __init__(
        self,
        in_feat,
        out_feat,
        lstm_units: int,
        input_seq_len: int,
        output_seq_len: int,
        graph_info: GraphInfo,
        graph_conv_params: typing.Optional[dict] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        # graph conv layer
        if graph_conv_params is None:
            graph_conv_params = {
                "aggregation_type": "mean",
                "combination_type": "concat",
                "activation": None,
            }
        self.graph_conv = GraphConv(in_feat, out_feat, graph_info, **graph_conv_params)

        self.lstm = layers.LSTM(lstm_units, activation="relu")
        self.dense = layers.Dense(output_seq_len)

        self.input_seq_len, self.output_seq_len = input_seq_len, output_seq_len

    def call(self, inputs):
        """Forward pass.

        Args:
            inputs: tf.Tensor of shape `(batch_size, input_seq_len, num_nodes, in_feat)`

        Returns:
            A tensor of shape `(batch_size, output_seq_len, num_nodes)`.
        """

        # convert shape to  (num_nodes, batch_size, input_seq_len, in_feat)
        inputs = tf.transpose(inputs, [2, 0, 1, 3])

        gcn_out = self.graph_conv(
            inputs
        )  # gcn_out has shape: (num_nodes, batch_size, input_seq_len, out_feat)
        shape = tf.shape(gcn_out)
        num_nodes, batch_size, input_seq_len, out_feat = (
            shape[0],
            shape[1],
            shape[2],
            shape[3],
        )

        # LSTM takes only 3D tensors as input
        gcn_out = tf.reshape(gcn_out, (batch_size * num_nodes, input_seq_len, out_feat))
        lstm_out = self.lstm(
            gcn_out
        )  # lstm_out has shape: (batch_size * num_nodes, lstm_units)

        dense_output = self.dense(
            lstm_out
        )  # dense_output has shape: (batch_size * num_nodes, output_seq_len)
        output = tf.reshape(dense_output, (num_nodes, batch_size, self.output_seq_len))
        return tf.transpose(
            output, [1, 2, 0]
        )  # returns Tensor of shape (batch_size, output_seq_len, num_nodes)

This is a custom layer called LSTMGC which is composed of a GraphConv layer followed by a LSTM layer and then a Dense layer. The GraphConv layer takes in in_feat input features and applies graph convolution to compute out_feat output features. It takes in a GraphInfo object, which contains information about the graph structure such as the edges and the number of nodes. The LSTM layer takes in the output of the GraphConv layer and applies LSTM computation to the input. The Dense layer then takes in the LSTM output and computes the final output. The call method of this layer takes in inputs of shape (batch_size, input_seq_len, num_nodes, in_feat) and returns a tensor of shape (batch_size, output_seq_len, num_nodes). The graph_conv_params argument is optional and can be used to set the parameters for the GraphConv layer such as aggregation_type, combination_type, and activation. If not provided, default values are used.

graph_conv_params = {
    "aggregation_type": "mean",
    "combination_type": "concat",
    "activation": None,
}

st_gcn = LSTMGC(
    in_feat,
    out_feat,
    lstm_units,
    input_sequence_length,
    forecast_horizon,
    graph,
    graph_conv_params,
)
inputs = layers.Input((input_sequence_length, graph.num_nodes, in_feat))
outputs = st_gcn(inputs)

This code defines an instance of the LSTMGC class and builds a Keras model using it. The LSTMGC class is a custom Keras layer that represents a neural network layer with three sub-layers: a graph convolution layer, an LSTM layer, and a dense layer. The graph convolution layer applies convolution operations to the input graph data, the LSTM layer processes the output of the graph convolution layer, and the dense layer produces the final output of the layer. The graph_conv_params parameter specifies the configuration of the graph convolution layer, including the type of aggregation operation to use (mean), the type of combination operation to use (concatenate), and the activation function to use (none). The Keras model is constructed by defining the input tensor shape and then applying the LSTMGC layer to the input tensor. The resulting output tensor represents the output of the layer, which is fed to subsequent layers in the model.

model = keras.models.Model(inputs, outputs)
model.compile(
    optimizer=keras.optimizers.RMSprop(learning_rate=learning_rate),
    loss=keras.losses.MeanSquaredError(),
)
model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=epochs,
    callbacks=[keras.callbacks.EarlyStopping(patience=patience)],
)

x_test, y = next(test_dataset.as_numpy_iterator())
y_pred = model.predict(x_test)

In the first block of code, a custom LSTMGC layer is defined that consists of a graph convolutional layer followed by an LSTM layer and a dense layer. Then, an instance of the LSTMGC layer is created with the provided parameters and an input placeholder layer is defined using layers.Input(). Finally, a model is created using keras.models.Model() by specifying the input and output layers and is compiled with an optimizer and loss function.

In the second block of code, the model is trained on the training dataset and evaluated on the validation dataset using model.fit(). Finally, the trained model is used to predict on the test dataset using model.predict(). Note that train_dataset, val_dataset, and test_dataset are instances of tf.data.Dataset class which is used to represent a potentially large set of elements (e.g. image data, text data, etc.) and provide an efficient way to stream these elements in batches for training and evaluation.

for i in range(20):
    fig = plt.figure(figsize=(18, 6))
    plt.plot(y[:, 0, i])
    plt.plot(y_pred[:, 0, i])
    plt.legend(["actual", "forecast"])
    # fig.savefig('pred_results_{}'.format(i))
    
naive_mse, model_mse = (
    np.square(x_test[:, -1, :, 0] - y[:, 0, :]).mean(),
    np.square(y_pred[:, 0, :] - y[:, 0, :]).mean(),
)
print(f"naive MAE: {naive_mse}, model MAE: {model_mse}")

The model is compiled using the mean squared error as the loss function and RMSprop as the optimizer. The fit function is then used to train the model on train_dataset with validation on val_dataset. Early stopping is implemented to prevent overfitting. After training, the model is used to predict on the test dataset test_dataset, and the resulting forecasts are compared with the actual values using mean squared error. Finally, a loop is run to plot the forecasts and actual values for the first 20 nodes in the test set, and the mean squared errors of the model and a naive baseline are printed.