title traffic forecasting using graph neural networks and lstm arash khodadadihttps www linkedin cominarashkhodadadi08a02490 date created 20211228 last modified 20231122 description this example demonstrates how to do timeseries forecasting over graphs accelerator gpu introduction this example shows how to forecast traffic condition using graph neural networks and lstm specifically we are interested in predicting the future values of the traffic speed given a history of the traffic speed for a collection of road segments one popular method to solve this problem is to consider each road segment s traffic speed as a separate timeseries and predict the future values of each timeseries using the past values of the same timeseries this method however ignores the dependency of the traffic speed of one road segment on the neighboring segments to be able to take into account the complex interactions between the traffic speed on a collection of neighboring roads we can define the traffic network as a graph and consider the traffic speed as a signal on this graph in this example we implement a neural network architecture which can process timeseries data over a graph we first show how to process the data and create a tf data datasethttps www tensorflow orgapidocspythontfdatadataset for forecasting over graphs then we implement a model which uses graph convolution and lstm layers to perform forecasting over a graph the data processing and the model architecture are inspired by this paper yu bing haoteng yin and zhanxing zhu spatiotemporal graph convolutional networks a deep learning framework for traffic forecasting proceedings of the 27th international joint conference on artificial intelligence 2018 githubhttps github comveritasyinstgcnijcai18 setup data preparation data description we use a realworld traffic speed dataset named pemsd7 we use the version collected and prepared by yu et al 2018https arxiv orgabs1709 04875 and available herehttps github comveritasyinstgcnijcai18treemasterdataset the data consists of two files pemsd7w228 csv contains the distances between 228 stations across the district 7 of california pemsd7v228 csv contains traffic speed collected for those stations in the weekdays of may and june of 2012 the full description of the dataset can be found in yu et al 2018https arxiv orgabs1709 04875 loading data subsampling roads to reduce the problem size and make the training faster we will only work with a sample of 26 roads out of the 228 roads in the dataset we have chosen the roads by starting from road 0 choosing the 5 closest roads to it and continuing this process until we get 25 roads you can choose any other subset of the roads we chose the roads in this way to increase the likelihood of having roads with correlated speed timeseries sampleroutes contains the ids of the selected roads data visualization here are the timeseries of the traffic speed for two of the routes we can also visualize the correlation between the timeseries in different routes using this correlation heatmap we can see that for example the speed in routes 4 5 6 are highly correlated splitting and normalizing data next we split the speed values array into trainvalidationtest sets and normalize the resulting arrays splits data into trainvaltest sets and normalizes the data args dataarray ndarray of shape numtimesteps numroutes trainsize a float value between 0 0 and 1 0 that represent the proportion of the dataset to include in the train split valsize a float value between 0 0 and 1 0 that represent the proportion of the dataset to include in the validation split returns trainarray valarray testarray creating tensorflow datasets next we create the datasets for our forecasting problem the forecasting problem can be stated as follows given a sequence of the road speed values at times t1 t2 tt we want to predict the future values of the roads speed for times tt1 tth so for each time t the inputs to our model are t vectors each of size n and the targets are h vectors each of size n where n is the number of roads we use the keras builtin function keras utils timeseriesdatasetfromarray the function createtfdataset below takes as input a numpy ndarray and returns a tf data dataset in this function inputsequencelengtht and forecasthorizonh the argument multihorizon needs more explanation assume forecasthorizon3 if multihorizontrue then the model will make a forecast for time steps tt1 tt2 tt3 so the target will have shape t 3 but if multihorizonfalse the model will make a forecast only for time step tt3 and so the target will have shape t 1 you may notice that the input tensor in each batch has shape batchsize inputsequencelength numroutes 1 the last dimension is added to make the model more general at each time step the input features for each raod may contain multiple timeseries for instance one might want to use temperature timeseries in addition to historical values of the speed as input features in this example however the last dimension of the input is always 1 we use the last 12 values of the speed in each road to forecast the speed for 3 time steps ahead creates tensorflow dataset from numpy array this function creates a dataset where each element is a tuple inputs targets inputs is a tensor of shape batchsize inputsequencelength numroutes 1 containing the inputsequencelength past values of the timeseries for each node targets is a tensor of shape batchsize forecasthorizon numroutes containing the forecasthorizon future values of the timeseries for each node args dataarray np ndarray with shape numtimesteps numroutes inputsequencelength length of the input sequence in number of timesteps forecasthorizon if multihorizontrue the target will be the values of the timeseries for 1 to forecasthorizon timesteps ahead if multihorizonfalse the target will be the value of the timeseries forecasthorizon steps ahead only one value batchsize number of timeseries samples in each batch shuffle whether to shuffle output samples or instead draw them in chronological order multihorizon see forecasthorizon returns a tf data dataset instance roads graph as mentioned before we assume that the road segments form a graph the pemsd7 dataset has the road segments distance the next step is to create the graph adjacency matrix from these distances following yu et al 2018https arxiv orgabs1709 04875 equation 10 we assume there is an edge between two nodes in the graph if the distance between the corresponding roads is less than a threshold computes the adjacency matrix from distances matrix it uses the formula in https github comveritasyinstgcnijcai18datapreprocessing to compute an adjacency matrix from the distance matrix the implementation follows that paper args routedistances np ndarray of shape numroutes numroutes entry i j of this array is the distance between roads i j sigma2 determines the width of the gaussian kernel applied to the square distances matrix epsilon a threshold specifying if there is an edge between two nodes specifically ai j1 if np expw2i j sigma2 epsilon and ai j0 otherwise where a is the adjacency matrix and w2routedistances routedistances returns a boolean graph adjacency matrix the function computeadjacencymatrix returns a boolean adjacency matrix where 1 means there is an edge between two nodes we use the following class to store the information about the graph network architecture our model for forecasting over the graph consists of a graph convolution layer and a lstm layer graph convolution layer our implementation of the graph convolution layer resembles the implementation in this keras examplehttps keras ioexamplesgraphgnncitations note that in that example input to the layer is a 2d tensor of shape numnodes infeat but in our example the input to the layer is a 4d tensor of shape numnodes batchsize inputseqlength infeat the graph convolution layer performs the following steps the nodes representations are computed in self computenodesrepresentation by multiplying the input features by self weight the aggregated neighbors messages are computed in self computeaggregatedmessages by first aggregating the neighbors representations and then multiplying the results by self weight the final output of the layer is computed in self update by combining the nodes representations and the neighbors aggregated messages computes each node s representation the nodes representations are obtained by multiplying the features tensor with self weight note that self weight has shape infeat outfeat args features tensor of shape numnodes batchsize inputseqlen infeat returns a tensor of shape numnodes batchsize inputseqlen outfeat forward pass args features tensor of shape numnodes batchsize inputseqlen infeat returns a tensor of shape numnodes batchsize inputseqlen outfeat lstm plus graph convolution by applying the graph convolution layer to the input tensor we get another tensor containing the nodes representations over time another 4d tensor for each time step a node s representation is informed by the information from its neighbors to make good forecasts however we need not only information from the neighbors but also we need to process the information over time to this end we can pass each node s tensor through a recurrent layer the lstmgc layer below first applies a graph convolution layer to the inputs and then passes the results through a lstm layer layer comprising a convolution layer followed by lstm and dense layers def init self infeat outfeat lstmunits int inputseqlen int outputseqlen int graphinfo graphinfo graphconvparams typing optionaldict none kwargs super initkwargs graph conv layer if graphconvparams is none graphconvparams aggregationtype mean combinationtype concat activation none self graphconv graphconv infeat outfeat graphinfo graphconvparams self lstm layers lstmlstmunits activationrelu self dense layers denseoutputseqlen self inputseqlen self outputseqlen inputseqlen outputseqlen def callself inputs convert shape to numnodes batchsize inputseqlen infeat inputs ops transposeinputs 2 0 1 3 gcnout self graphconv inputs gcnout has shape numnodes batchsize inputseqlen outfeat shape ops shapegcnout numnodes batchsize inputseqlen outfeat shape0 shape1 shape2 shape3 lstm takes only 3d tensors as input gcnout ops reshape gcnout batchsize numnodes inputseqlen outfeat lstmout self lstm gcnout lstmout has shape batchsize numnodes lstmunits denseoutput self dense lstmout denseoutput has shape batchsize numnodes outputseqlen output ops reshape denseoutput numnodes batchsize self outputseqlen return ops transpose output 1 2 0 returns tensor of shape batchsize outputseqlen numnodes model training infeat 1 batchsize 64 epochs 20 inputsequencelength 12 forecasthorizon 3 multihorizon false outfeat 10 lstmunits 64 graphconvparams aggregationtype mean combinationtype concat activation none stgcn lstmgc infeat outfeat lstmunits inputsequencelength forecasthorizon graph graphconvparams inputs layers inputinputsequencelength graph numnodes infeat outputs stgcninputs model keras models modelinputs outputs model compile optimizerkeras optimizers rmsproplearningrate0 0002 losskeras losses meansquarederror model fit traindataset validationdatavaldataset epochsepochs callbackskeras callbacks earlystoppingpatience10 making forecasts on test set xtest y nexttestdataset asnumpyiterator ypred model predictxtest plt figurefigsize18 6 plt ploty 0 0 plt plotypred 0 0 plt legendactual forecast naivemse modelmse np squarextest 1 0 y 0 mean np squareypred 0 y 0 mean printfnaive mae naivemse model mae modelmse introduction this example shows how to forecast traffic condition using graph neural networks and lstm specifically we are interested in predicting the future values of the traffic speed given a history of the traffic speed for a collection of road segments one popular method to solve this problem is to consider each road segment s traffic speed as a separate timeseries and predict the future values of each timeseries using the past values of the same timeseries this method however ignores the dependency of the traffic speed of one road segment on the neighboring segments to be able to take into account the complex interactions between the traffic speed on a collection of neighboring roads we can define the traffic network as a graph and consider the traffic speed as a signal on this graph in this example we implement a neural network architecture which can process timeseries data over a graph we first show how to process the data and create a tf data dataset https www tensorflow org api_docs python tf data dataset for forecasting over graphs then we implement a model which uses graph convolution and lstm layers to perform forecasting over a graph the data processing and the model architecture are inspired by this paper yu bing haoteng yin and zhanxing zhu spatio temporal graph convolutional networks a deep learning framework for traffic forecasting proceedings of the 27th international joint conference on artificial intelligence 2018 github https github com veritasyin stgcn_ijcai 18 setup data preparation data description we use a real world traffic speed dataset named pemsd7 we use the version collected and prepared by yu et al 2018 https arxiv org abs 1709 04875 and available here https github com veritasyin stgcn_ijcai 18 tree master dataset the data consists of two files pemsd7_w_228 csv contains the distances between 228 stations across the district 7 of california pemsd7_v_228 csv contains traffic speed collected for those stations in the weekdays of may and june of 2012 the full description of the dataset can be found in yu et al 2018 https arxiv org abs 1709 04875 loading data sub sampling roads to reduce the problem size and make the training faster we will only work with a sample of 26 roads out of the 228 roads in the dataset we have chosen the roads by starting from road 0 choosing the 5 closest roads to it and continuing this process until we get 25 roads you can choose any other subset of the roads we chose the roads in this way to increase the likelihood of having roads with correlated speed timeseries sample_routes contains the ids of the selected roads data visualization here are the timeseries of the traffic speed for two of the routes we can also visualize the correlation between the timeseries in different routes using this correlation heatmap we can see that for example the speed in routes 4 5 6 are highly correlated splitting and normalizing data next we split the speed values array into train validation test sets and normalize the resulting arrays splits data into train val test sets and normalizes the data args data_array ndarray of shape num_time_steps num_routes train_size a float value between 0 0 and 1 0 that represent the proportion of the dataset to include in the train split val_size a float value between 0 0 and 1 0 that represent the proportion of the dataset to include in the validation split returns train_array val_array test_array creating tensorflow datasets next we create the datasets for our forecasting problem the forecasting problem can be stated as follows given a sequence of the road speed values at times t 1 t 2 t t we want to predict the future values of the roads speed for times t t 1 t t h so for each time t the inputs to our model are t vectors each of size n and the targets are h vectors each of size n where n is the number of roads we use the keras built in function keras utils timeseries_dataset_from_array the function create_tf_dataset below takes as input a numpy ndarray and returns a tf data dataset in this function input_sequence_length t and forecast_horizon h the argument multi_horizon needs more explanation assume forecast_horizon 3 if multi_horizon true then the model will make a forecast for time steps t t 1 t t 2 t t 3 so the target will have shape t 3 but if multi_horizon false the model will make a forecast only for time step t t 3 and so the target will have shape t 1 you may notice that the input tensor in each batch has shape batch_size input_sequence_length num_routes 1 the last dimension is added to make the model more general at each time step the input features for each raod may contain multiple timeseries for instance one might want to use temperature timeseries in addition to historical values of the speed as input features in this example however the last dimension of the input is always 1 we use the last 12 values of the speed in each road to forecast the speed for 3 time steps ahead creates tensorflow dataset from numpy array this function creates a dataset where each element is a tuple inputs targets inputs is a tensor of shape batch_size input_sequence_length num_routes 1 containing the input_sequence_length past values of the timeseries for each node targets is a tensor of shape batch_size forecast_horizon num_routes containing the forecast_horizon future values of the timeseries for each node args data_array np ndarray with shape num_time_steps num_routes input_sequence_length length of the input sequence in number of timesteps forecast_horizon if multi_horizon true the target will be the values of the timeseries for 1 to forecast_horizon timesteps ahead if multi_horizon false the target will be the value of the timeseries forecast_horizon steps ahead only one value batch_size number of timeseries samples in each batch shuffle whether to shuffle output samples or instead draw them in chronological order multi_horizon see forecast_horizon returns a tf data dataset instance roads graph as mentioned before we assume that the road segments form a graph the pemsd7 dataset has the road segments distance the next step is to create the graph adjacency matrix from these distances following yu et al 2018 https arxiv org abs 1709 04875 equation 10 we assume there is an edge between two nodes in the graph if the distance between the corresponding roads is less than a threshold computes the adjacency matrix from distances matrix it uses the formula in https github com veritasyin stgcn_ijcai 18 data preprocessing to compute an adjacency matrix from the distance matrix the implementation follows that paper args route_distances np ndarray of shape num_routes num_routes entry i j of this array is the distance between roads i j sigma2 determines the width of the gaussian kernel applied to the square distances matrix epsilon a threshold specifying if there is an edge between two nodes specifically a i j 1 if np exp w2 i j sigma2 epsilon and a i j 0 otherwise where a is the adjacency matrix and w2 route_distances route_distances returns a boolean graph adjacency matrix the function compute_adjacency_matrix returns a boolean adjacency matrix where 1 means there is an edge between two nodes we use the following class to store the information about the graph network architecture our model for forecasting over the graph consists of a graph convolution layer and a lstm layer graph convolution layer our implementation of the graph convolution layer resembles the implementation in this keras example https keras io examples graph gnn_citations note that in that example input to the layer is a 2d tensor of shape num_nodes in_feat but in our example the input to the layer is a 4d tensor of shape num_nodes batch_size input_seq_length in_feat the graph convolution layer performs the following steps the nodes representations are computed in self compute_nodes_representation by multiplying the input features by self weight the aggregated neighbors messages are computed in self compute_aggregated_messages by first aggregating the neighbors representations and then multiplying the results by self weight the final output of the layer is computed in self update by combining the nodes representations and the neighbors aggregated messages 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 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 lstm plus graph convolution by applying the graph convolution layer to the input tensor we get another tensor containing the nodes representations over time another 4d tensor for each time step a node s representation is informed by the information from its neighbors to make good forecasts however we need not only information from the neighbors but also we need to process the information over time to this end we can pass each node s tensor through a recurrent layer the lstmgc layer below first applies a graph convolution layer to the inputs and then passes the results through a lstm layer layer comprising a convolution layer followed by lstm and dense layers graph conv layer forward pass args inputs 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 gcn_out has shape num_nodes batch_size input_seq_len out_feat lstm takes only 3d tensors as input lstm_out has shape batch_size num_nodes lstm_units dense_output has shape batch_size num_nodes output_seq_len returns tensor of shape batch_size output_seq_len num_nodes model training making forecasts on test set now we can use the trained model to make forecasts for the test set below we compute the mae of the model and compare it to the mae of naive forecasts the naive forecasts are the last value of the speed for each node of course the goal here is to demonstrate the method not to achieve the best performance to improve the model s accuracy all model hyperparameters should be tuned carefully in addition several of the lstmgc blocks can be stacked to increase the representation power of the model
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import pandas as pd
import numpy as np
import typing
import matplotlib.pyplot as plt

import tensorflow as tf
import keras
from keras import layers
from keras import ops







url = "https://github.com/VeritasYin/STGCN_IJCAI-18/raw/master/dataset/PeMSD7_Full.zip"
data_dir = keras.utils.get_file(origin=url, extract=True, archive_format="zip")
data_dir = data_dir.rstrip("PeMSD7_Full.zip")

route_distances = pd.read_csv(
    os.path.join(data_dir, "PeMSD7_W_228.csv"), header=None
).to_numpy()
speeds_array = pd.read_csv(
    os.path.join(data_dir, "PeMSD7_V_228.csv"), header=None
).to_numpy()

print(f"route_distances shape={route_distances.shape}")
print(f"speeds_array shape={speeds_array.shape}")



sample_routes = [
    0,
    1,
    4,
    7,
    8,
    11,
    15,
    108,
    109,
    114,
    115,
    118,
    120,
    123,
    124,
    126,
    127,
    129,
    130,
    132,
    133,
    136,
    139,
    144,
    147,
    216,
]
route_distances = route_distances[np.ix_(sample_routes, sample_routes)]
speeds_array = speeds_array[:, sample_routes]

print(f"route_distances shape={route_distances.shape}")
print(f"speeds_array shape={speeds_array.shape}")



plt.figure(figsize=(18, 6))
plt.plot(speeds_array[:, [0, -1]])
plt.legend(["route_0", "route_25"])



plt.figure(figsize=(8, 8))
plt.matshow(np.corrcoef(speeds_array.T), 0)
plt.xlabel("road number")
plt.ylabel("road number")





train_size, val_size = 0.5, 0.2


def preprocess(data_array: np.ndarray, train_size: float, val_size: float):
    

    num_time_steps = data_array.shape[0]
    num_train, num_val = (
        int(num_time_steps * train_size),
        int(num_time_steps * val_size),
    )
    train_array = data_array[:num_train]
    mean, std = train_array.mean(axis=0), train_array.std(axis=0)

    train_array = (train_array - mean) / std
    val_array = (data_array[num_train : (num_train + num_val)] - mean) / std
    test_array = (data_array[(num_train + num_val) :] - mean) / std

    return train_array, val_array, test_array


train_array, val_array, test_array = preprocess(
    speeds_array, train_size, val_size
)

print(f"train set size: {train_array.shape}")
print(f"validation set size: {val_array.shape}")
print(f"test set size: {test_array.shape}")





batch_size = 64
input_sequence_length = 12
forecast_horizon = 3
multi_horizon = False


def create_tf_dataset(
    data_array: np.ndarray,
    input_sequence_length: int,
    forecast_horizon: int,
    batch_size: int = 128,
    shuffle=True,
    multi_horizon=True,
):
    

    inputs = keras.utils.timeseries_dataset_from_array(
        np.expand_dims(data_array[:-forecast_horizon], axis=-1),
        None,
        sequence_length=input_sequence_length,
        shuffle=False,
        batch_size=batch_size,
    )

    target_offset = (
        input_sequence_length
        if multi_horizon
        else input_sequence_length + forecast_horizon - 1
    )
    target_seq_length = forecast_horizon if multi_horizon else 1
    targets = keras.utils.timeseries_dataset_from_array(
        data_array[target_offset:],
        None,
        sequence_length=target_seq_length,
        shuffle=False,
        batch_size=batch_size,
    )

    dataset = tf.data.Dataset.zip((inputs, targets))
    if shuffle:
        dataset = dataset.shuffle(100)

    return dataset.prefetch(16).cache()


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

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





def compute_adjacency_matrix(
    route_distances: np.ndarray, sigma2: float, epsilon: float
):
    
    num_routes = route_distances.shape[0]
    route_distances = route_distances / 10000.0
    w2, w_mask = (
        route_distances * route_distances,
        np.ones([num_routes, num_routes]) - np.identity(num_routes),
    )
    return (np.exp(-w2 / sigma2) >= epsilon) * w_mask





class GraphInfo:
    def __init__(self, edges: typing.Tuple[list, list], num_nodes: int):
        self.edges = edges
        self.num_nodes = num_nodes


sigma2 = 0.1
epsilon = 0.5
adjacency_matrix = compute_adjacency_matrix(route_distances, sigma2, epsilon)
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])}"
)






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 = self.add_weight(
            initializer=keras.initializers.GlorotUniform(),
            shape=(in_feat, out_feat),
            dtype="float32",
            trainable=True,
        )
        self.activation = layers.Activation(activation)

    def aggregate(self, neighbour_representations):
        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):
        
        return ops.matmul(features, self.weight)

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

    def update(
        self, nodes_representation, aggregated_messages
    ):
        if self.combination_type == "concat":
            h = ops.concatenate([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):
        
        nodes_representation = self.compute_nodes_representation(features)
        aggregated_messages = self.compute_aggregated_messages(features)
        return self.update(nodes_representation, aggregated_messages)





class LSTMGC(layers.Layer):
    

    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)

        
        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):
        

        
        inputs = ops.transpose(inputs, [2, 0, 1, 3])

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

        
        gcn_out = ops.reshape(
            gcn_out, (batch_size * num_nodes, input_seq_len, out_feat)
        )
        lstm_out = self.lstm(
            gcn_out
        )  

        dense_output = self.dense(
            lstm_out
        )  
        output = ops.reshape(
            dense_output, (num_nodes, batch_size, self.output_seq_len)
        )
        return ops.transpose(
            output, [1, 2, 0]
        )  




in_feat = 1
batch_size = 64
epochs = 20
input_sequence_length = 12
forecast_horizon = 3
multi_horizon = False
out_feat = 10
lstm_units = 64
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)

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



x_test, y = next(test_dataset.as_numpy_iterator())
y_pred = model.predict(x_test)
plt.figure(figsize=(18, 6))
plt.plot(y[:, 0, 0])
plt.plot(y_pred[:, 0, 0])
plt.legend(["actual", "forecast"])

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}")