# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + id="F2mVdX5d5xGo" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 153} outputId="a59b5be4-b74d-4b03-f40c-a8389e933eff"
# !unzip DCGRU_github.zip

# + [markdown] id="ctOD3pK-7NvE" colab_type="text"
# # Import libraries

# + id="epKMcPh19JBx" colab_type="code" colab={}
import numpy as np
import networkx as nx
import tensorflow as tf
from tensorflow import keras
from tensorflow.python.keras import backend as K
import matplotlib.pyplot as plt
import pickle

from synth_signal.gen_signal import generate_signal
from dcgru_cell_tf2 import DCGRUCell

# + [markdown] id="lnyTJ1Xq6rDV" colab_type="text"
# # Load signal and adjacency matrix

# + id="n1wHIxlF6t8o" colab_type="code" colab={}
signal = pickle.load(open('signal.pickle','rb'))
G_adj_mx = pickle.load(open('G_adj_mx.pickle','rb'))
num_nodes = G_adj_mx.shape[0]

# + [markdown] id="ulAFKp0-7JR_" colab_type="text"
# # How to build a DCGRU layer

# + id="EBuPXDTf7f5D" colab_type="code" colab={}
# define the dcgru cell
dcgru_cell = DCGRUCell(units=20,adj_mx=G_adj_mx, K_diffusion=2,
                       num_nodes=num_nodes,filter_type="random_walk")

# + [markdown] id="WZlGsjTF_pkT" colab_type="text"
# The state size of the dcgru cell is a vector of dimension (units * num_nodes). In the example we have units = 20 and num_nodes = 24, hence the state size will be 20*24 = 480

# + id="iH76iw6583jz" colab_type="code" colab={}
# wrap the dcgru cell in a keras RNN layer
Dcgru_layer = keras.layers.RNN(dcgru_cell)

# + [markdown] id="pppY7RMl9NbH" colab_type="text"
# # Example of a sequential model

# + id="u5u8heKZ85pU" colab_type="code" colab={}
model_dcgru = keras.models.Sequential([
  keras.Input(shape=(None, num_nodes, 1)),  # Input dimensions: sequence length (None := arbitrary length)
  Dcgru_layer,                              #                   number of nodes in the graph
  keras.layers.Dense(num_nodes)             #                   signal dimensionality (1 in the example)
  ])

# + id="B9wvDmab87qU" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 221} outputId="83808933-825e-4622-fad1-bc21daa7f28b"
model_dcgru.summary()

# + [markdown] id="nH3f-Rkq9-t1" colab_type="text"
# ## Data format

# + id="2dW_CVVu9irW" colab_type="code" colab={}
# create X, y from signal
seq_len = 8
horizon = 4
X = []
y = []

for i in range(100,len(signal)-(seq_len+horizon)):
  X.append(signal[i:i+seq_len])
  y.append(signal[i+seq_len+horizon])
X = np.expand_dims(np.array(X), -1)   # expand dimensionality to match the correct input shape
y = np.array(y)

# + id="00-snk5ZCIya" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 102} outputId="fb57684f-f3b2-46b5-c87b-25082dad4911"
print(X.shape)
print(X.shape[0], ": number of input sequences")
print(X.shape[1], ": sequence length")
print(X.shape[2], ": number of nodhs")
print(X.shape[3], ": signal dimensionality")

# + id="_TQoTD4jDahQ" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="9dae6b3e-21e9-47c3-cb5f-d6a40a09e752"
print(y.shape)

# + [markdown] id="daQrHrfQC2YJ" colab_type="text"
# ## Model training

# + id="gYZIZF2DB4il" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 374} outputId="706ac4fc-117d-4876-be8b-68de08e1706f"
opt = keras.optimizers.Adam(learning_rate=0.001, clipnorm=1.0)
model_dcgru.compile(loss="mse", optimizer=opt)
model_dcgru.fit(X, y, epochs=10, verbose=1)

StarcoderdataJupyterScriptsDedupFiltered