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

# + [markdown] colab_type="text" id="P7sX-B4aUbRN" pycharm={"name": "#%% md\n"}
# # Introduction to RNNs
# This notebook is part of the [SachsLab Workshop for Intracranial Neurophysiology and Deep Learning](https://github.com/SachsLab/IntracranialNeurophysDL).
#
# ### Normalize Environments
# Run the first two cells to normalize Local / Colab environments, then proceed below for the lesson.
#
# <table class="tfo-notebook-buttons" align="left">
#   <td>
#     <a target="_blank" href="https://colab.research.google.com/github/SachsLab/IntracranialNeurophysDL/blob/master/notebooks/06_01_intro_to_RNNs.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
#   </td>
#   <td>
#     <a target="_blank" href="https://github.com/SachsLab/IntracranialNeurophysDL/blob/master/notebooks/06_01_intro_to_RNNs.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
#   </td>
# </table>

# + colab_type="code" id="aL4QcKd5UbRT" outputId="a4accba4-507a-4eaf-cb68-c5a2263fd8aa" pycharm={"is_executing": false, "name": "#%%\n"} colab={"base_uri": "https://localhost:8080/", "height": 67}
from pathlib import Path
import os
try:
    # See if we are running on google.colab
    import google.colab
    IN_COLAB = True
    # Setup tensorflow 2.0
    # !pip install -q tensorflow-gpu==2.0.0-rc0
except ModuleNotFoundError:
    IN_COLAB = False
    import sys
    if Path.cwd().stem == 'notebooks':
        os.chdir(Path.cwd().parent)

# + colab_type="code" id="rXM9gHqLUbRP" pycharm={"is_executing": false, "name": "#%%\n"} colab={}
# Common imports
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

plt.rcParams.update({'font.size': 22})
# %matplotlib inline
# %load_ext autoreload
# %autoreload 2

# + [markdown] colab_type="text" id="k9Og3lOfUbRV" pycharm={"name": "#%% md\n"}
# ## RNN Step-by-Step
# A recurrent neural network (RNN) is an artificial neural network that maintains an internal state (memory)
# as it processes items in a sequence. RNNs are most often applied to sequence data such as text
# (sequences of characters; sequences of words) or time series (e.g., stock prices, weather data, neurophysiology).
# RNNs can provide an output at each step of the sequence to predict different sequences (e.g., text translation) or
# the next item in the same sequence (e.g., assistive typing, stock prediction, prosthetic control).
# The loss does not back-propagate to update weights until the sequence is complete. RNNs can also be configured
# to produce a single output at the end of a sequence (e.g., classify a tweet sentiment, decode categorical intention).
#
# We will create our own RNN step-by-step. We will train it using toy data that we generate.

# + [markdown] colab_type="text" id="_DAvkmycUbRV" pycharm={"name": "#%% md\n"}
# ### Generate data
# X is a multi-channel time series and Y is a different multi-channel time series
# constructed from a linear combination of a delayed version of X.

# + colab_type="code" id="kA9Fqu_qUbRW" pycharm={"is_executing": false, "name": "#%%\n"} colab={}
PEAK_FREQS = [10, 22, 75]  # Hz
FS = 1000.0  # Hz
DURATION = 5.0  # seconds
DELAY = 0.01  # seconds
N_OUT = 2  # Y dimensions

n_samples = int(DURATION * FS)
delay_samples = int(DELAY * FS)
t_vec = np.arange(n_samples + delay_samples) / FS

X = np.zeros((n_samples + delay_samples, len(PEAK_FREQS)), dtype=np.float32)
for chan_ix, pf in enumerate(PEAK_FREQS):
    X[:, chan_ix] = (1 / (chan_ix + 1) ) * np.sin(t_vec * 2 * np.pi * pf)

# Create mixing matrix that mixes inputs into outputs
W = 2 * np.random.rand(N_OUT, len(PEAK_FREQS)) - 1
# W = W / np.sum(W, axis=1, keepdims=True)

b = np.random.rand(N_OUT) - 0.5

Y = W @ X[:-delay_samples, :].T + b[:, None]
Y = Y.T
X = X[delay_samples:, :]
t_vec = t_vec[delay_samples:]

X += 0.05*np.random.rand(*X.shape)
Y += 0.05*np.random.rand(*Y.shape)

# + colab_type="code" id="A4PXkiSEUbRY" outputId="9fcdf1da-e60b-4493-f8d8-41cc5007c14a" pycharm={"is_executing": false, "name": "#%%\n"} colab={"base_uri": "https://localhost:8080/", "height": 378}
plt.figure(figsize=(8, 6), facecolor='white')
plt.plot(t_vec[:100], X[:100, :])
plt.plot(t_vec[:100], Y[:100, :], 'k')
plt.show()

# + [markdown] colab_type="text" id="8ymuKRXPUbRc" pycharm={"name": "#%% md\n"}
# ### Define forward pass and loss
# We aren't actually going to run this next cell.
# This is just to give you an indication of what the forward pass looks like.

# + colab_type="code" id="spnNLauUUbRd" outputId="595f796f-e37a-4a3e-8eca-7da000931293" pycharm={"is_executing": false, "name": "#%%\n"} colab={"base_uri": "https://localhost:8080/", "height": 35}
activation_fn = lambda x: x
# activation_fn = np.tanh

# Initialize parameters
state_t = np.zeros((N_OUT,))
_W = np.random.random((N_OUT, X.shape[1]))  # Mixes input with output
_U = np.random.random((N_OUT, N_OUT))       # Mixes old state with output
_b = np.random.random((N_OUT,))             # Bias term

# Create variable to store output
successive_outputs = []

# Iterate over each timestep in X
for x_t in X:
    y_t = activation_fn(np.dot(_W, x_t) + np.dot(_U, state_t) + _b)
    successive_outputs.append(y_t)
    state_t = y_t
    
final_output_sequence = np.stack(successive_outputs, axis=0)
loss = np.mean( (Y - final_output_sequence)**2 )
print(loss)

# + [markdown] colab_type="text" id="ukwchqVWUbRf"
# ## RNN in Tensorflow
# [Tutorial](https://www.tensorflow.org/tutorials/sequences/text_generation) (text generation w/ eager)

# + [markdown] colab_type="text" id="cOHeBypcUbRg" pycharm={"name": "#%% md\n"}
# ### Prepare data for Tensorflow
# In the tutorial linked above, the `batch` transformation is used to convert a continuous sequence into
# many sequences, then the batch transform is used AGAIN to get batches of sequences.

# + colab_type="code" id="kW2toraOUbRg" pycharm={"is_executing": false, "name": "#%%\n"} colab={}
SEQ_LENGTH = 200
BATCH_SIZE = 2
BUFFER_SIZE = 10000

dataset = tf.data.Dataset.from_tensor_slices((X, Y))
dataset = dataset.batch(SEQ_LENGTH, drop_remainder=True)  # Continuous to segmented sequences
dataset = dataset.shuffle(BUFFER_SIZE)
dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)  # Segmented sequences to batches of seg. seq.

# + [markdown] colab_type="text" id="iVYi9e43UbRj" pycharm={"name": "#%% md\n"}
# ### Define RNN model in Tensorflow

# + colab_type="code" id="kxhHmWZXUbRk" outputId="d6324dfd-db59-44fa-f60c-79bf25979cd6" pycharm={"is_executing": false, "name": "#%%\n"} colab={"base_uri": "https://localhost:8080/", "height": 218}
inputs = tf.keras.layers.Input(shape=(SEQ_LENGTH, X.shape[-1]))
outputs = tf.keras.layers.SimpleRNN(N_OUT, return_sequences=True,
                                    activation='linear')(inputs)
model = tf.keras.Model(inputs, outputs)
model.compile(optimizer='rmsprop', loss='mean_squared_error')
model.summary()

# + colab_type="code" id="L9nwtlv8UbRn" pycharm={"is_executing": false, "name": "#%%\n"} colab={}
# Just to save us some training time, let's cheat and init the model with good weights
old_init_weights = model.layers[-1].get_weights()
model.layers[-1].set_weights([W.T, old_init_weights[1], b])

# + colab_type="code" id="Xk1qktbEUbRo" outputId="a506a658-ffab-4bc5-de35-39c513a757df" pycharm={"is_executing": false, "name": "#%%\n"} colab={"base_uri": "https://localhost:8080/", "height": 1000}
EPOCHS = 50
history = model.fit(x=dataset, epochs=EPOCHS, verbose=1)

# + colab_type="code" id="buU5s_XlUbRr" outputId="c6ca8690-e302-485a-8f82-ae88217eb2d0" pycharm={"is_executing": false} colab={"base_uri": "https://localhost:8080/", "height": 378}
_Y1 = model.predict(X[:SEQ_LENGTH, :][None, :, :])

plt.figure(figsize=(8, 6), facecolor='white')
plt.plot(Y[:SEQ_LENGTH, :], label='real')
plt.plot(_Y1[0], label='pred')
plt.legend()
plt.show()

# + colab_type="code" id="VQeoT-SlUbRu" outputId="43725caf-8679-42f5-d27c-f479d28f0fd1" pycharm={"is_executing": false, "name": "#%%\n"} colab={"base_uri": "https://localhost:8080/", "height": 305}
_W, _U, _b = model.layers[-1].get_weights()

print(W)
print(_W.T)
print(b, _b)

plt.figure(figsize=(10, 6), facecolor='white')
plt.subplot(1, 2, 1)
plt.imshow(W)
plt.subplot(1, 2, 2)
plt.imshow(_W.T)
plt.show()

StarcoderdataJupyterScriptsDedupFiltered