Week 47: Recurrent neural networks and Autoencoders

"""
Key points:
1. The input images (28x28 pixels) are treated as sequences of 28 timesteps with 28 features each
2. The LSTM layer processes this sequential data
3. A final dense layer with softmax activation handles the classification
4. Typical accuracy ranges between 95-98% (lower than CNNs but reasonable for demonstration)

Note: LSTMs are not typically used for image classification (CNNs are more efficient), but this demonstrates how to adapt them for such tasks. Training might take longer compared to CNN architectures.

To improve performance, you could:
1. Add more LSTM layers
2. Use Bidirectional LSTMs
3. Increase the number of units
4. Add dropout for regularization
5. Use learning rate scheduling
"""

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from tensorflow.keras.utils import to_categorical

# Load and preprocess data
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# Normalize pixel values to [0, 1]
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

# Reshape data for LSTM (samples, timesteps, features)
# MNIST images are 28x28, so we treat each image as 28 timesteps of 28 features
x_train = x_train.reshape((-1, 28, 28))
x_test = x_test.reshape((-1, 28, 28))

# Convert labels to one-hot encoding
y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)

# Build LSTM model
model = Sequential()
model.add(LSTM(128, input_shape=(28, 28)))  # 128 LSTM units
model.add(Dense(10, activation='softmax'))

# Compile the model
model.compile(loss='categorical_crossentropy',
             optimizer='adam',
             metrics=['accuracy'])

# Display model summary
model.summary()

# Train the model
history = model.fit(x_train, y_train,
                   batch_size=64,
                   epochs=10,
                   validation_split=0.2)

# Evaluate on test data
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2)
print(f'\nTest accuracy: {test_acc:.4f}')

"""
Key components:
1. **Data Handling**: Uses PyTorch DataLoader with MNIST dataset
2. **LSTM Architecture**:
  - Input sequence of 28 timesteps (image rows)
  - 128 hidden units in LSTM layer
  - Fully connected layer for classification
3. **Training**:
  - Cross-entropy loss
  - Adam optimizer
  - Automatic GPU utilization if available

This implementation typically achieves **97-98% accuracy** after 10 epochs. The main differences from the TensorFlow/Keras version:
- Explicit device management (CPU/GPU)
- Manual training loop
- Different data loading pipeline
- More explicit tensor reshaping

To improve performance, you could:
1. Add dropout regularization
2. Use bidirectional LSTM
3. Implement learning rate scheduling
4. Add batch normalization
5. Increase model capacity (more layers/units)
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Hyperparameters
input_size = 28     # Number of features (pixels per row)
hidden_size = 128   # LSTM hidden state size
num_classes = 10    # Digits 0-9
num_epochs = 10     # Training iterations
batch_size = 64     # Batch size
learning_rate = 0.001

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# MNIST dataset
transform = transforms.Compose([
   transforms.ToTensor(),
   transforms.Normalize((0.1307,), (0.3081,))  # MNIST mean and std
])

train_dataset = datasets.MNIST(root='./data',
                              train=True,
                              transform=transform,
                              download=True)

test_dataset = datasets.MNIST(root='./data',
                             train=False,
                             transform=transform)

train_loader = DataLoader(dataset=train_dataset,
                         batch_size=batch_size,
                         shuffle=True)

test_loader = DataLoader(dataset=test_dataset,
                        batch_size=batch_size,
                        shuffle=False)

# LSTM model
class LSTMModel(nn.Module):
   def __init__(self, input_size, hidden_size, num_classes):
       super(LSTMModel, self).__init__()
       self.hidden_size = hidden_size
       self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
       self.fc = nn.Linear(hidden_size, num_classes)

   def forward(self, x):
       # Reshape input to (batch_size, sequence_length, input_size)
       x = x.reshape(-1, 28, 28)

       # Forward propagate LSTM
       out, _ = self.lstm(x)  # out: (batch_size, seq_length, hidden_size)

       # Decode the hidden state of the last time step
       out = out[:, -1, :]
       out = self.fc(out)
       return out

# Initialize model
model = LSTMModel(input_size, hidden_size, num_classes).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Training loop
total_step = len(train_loader)
for epoch in range(num_epochs):
   model.train()
   for i, (images, labels) in enumerate(train_loader):
       images = images.to(device)
       labels = labels.to(device)

       # Forward pass
       outputs = model(images)
       loss = criterion(outputs, labels)

       # Backward and optimize
       optimizer.zero_grad()
       loss.backward()
       optimizer.step()

       if (i+1) % 100 == 0:
           print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_step}], Loss: {loss.item():.4f}')

   # Test the model
   model.eval()
   with torch.no_grad():
       correct = 0
       total = 0
       for images, labels in test_loader:
           images = images.to(device)
           labels = labels.to(device)
           outputs = model(images)
           _, predicted = torch.max(outputs.data, 1)
           total += labels.size(0)
           correct += (predicted == labels).sum().item()

       print(f'Test Accuracy: {100 * correct / total:.2f}%')

print('Training finished.')

import numpy as np
import pandas as pd
from math import *
import matplotlib.pyplot as plt
import os

# Where to save the figures and data files
PROJECT_ROOT_DIR = "Results"
FIGURE_ID = "Results/FigureFiles"
DATA_ID = "DataFiles/"

if not os.path.exists(PROJECT_ROOT_DIR):
    os.mkdir(PROJECT_ROOT_DIR)

if not os.path.exists(FIGURE_ID):
    os.makedirs(FIGURE_ID)

if not os.path.exists(DATA_ID):
    os.makedirs(DATA_ID)

def image_path(fig_id):
    return os.path.join(FIGURE_ID, fig_id)

def data_path(dat_id):
    return os.path.join(DATA_ID, dat_id)

def save_fig(fig_id):
    plt.savefig(image_path(fig_id) + ".png", format='png')


def SpringForce(v,x,t):
#   note here that we have divided by mass and we return the acceleration
    return  -2*gamma*v-x+Ftilde*cos(t*Omegatilde)


def RK4(v,x,t,n,Force):
    for i in range(n-1):
# Setting up k1
        k1x = DeltaT*v[i]
        k1v = DeltaT*Force(v[i],x[i],t[i])
# Setting up k2
        vv = v[i]+k1v*0.5
        xx = x[i]+k1x*0.5
        k2x = DeltaT*vv
        k2v = DeltaT*Force(vv,xx,t[i]+DeltaT*0.5)
# Setting up k3
        vv = v[i]+k2v*0.5
        xx = x[i]+k2x*0.5
        k3x = DeltaT*vv
        k3v = DeltaT*Force(vv,xx,t[i]+DeltaT*0.5)
# Setting up k4
        vv = v[i]+k3v
        xx = x[i]+k3x
        k4x = DeltaT*vv
        k4v = DeltaT*Force(vv,xx,t[i]+DeltaT)
# Final result
        x[i+1] = x[i]+(k1x+2*k2x+2*k3x+k4x)/6.
        v[i+1] = v[i]+(k1v+2*k2v+2*k3v+k4v)/6.
        t[i+1] = t[i] + DeltaT


# Main part begins here

DeltaT = 0.001
#set up arrays 
tfinal = 20 # in dimensionless time
n = ceil(tfinal/DeltaT)
# set up arrays for t, v, and x
t = np.zeros(n)
v = np.zeros(n)
x = np.zeros(n)
# Initial conditions (can change to more than one dim)
x0 =  1.0 
v0 = 0.0
x[0] = x0
v[0] = v0
gamma = 0.2
Omegatilde = 0.5
Ftilde = 1.0
# Start integrating using Euler's method
# Note that we define the force function as a SpringForce
RK4(v,x,t,n,SpringForce)

# Plot position as function of time    
fig, ax = plt.subplots()
ax.set_ylabel('x[m]')
ax.set_xlabel('t[s]')
ax.plot(t, x)
fig.tight_layout()
save_fig("ForcedBlockRK4")
plt.show()

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt


# Newton's equation for harmonic oscillations with external force

# Global parameters
gamma = 0.2        # damping
Omegatilde = 0.5   # driving frequency
Ftilde = 1.0       # driving amplitude

def spring_force(v, x, t):
    """
    SpringForce:
    note: divided by mass => returns acceleration
    a = -2*gamma*v - x + Ftilde*cos(Omegatilde * t)
    """
    return -2.0 * gamma * v - x + Ftilde * np.cos(Omegatilde * t)


def rk4_trajectory(DeltaT=0.001, tfinal=20.0, x0=1.0, v0=0.0):
    """
    Returns t, x, v arrays.
    """
    n = int(np.ceil(tfinal / DeltaT))

    t = np.zeros(n, dtype=np.float32)
    x = np.zeros(n, dtype=np.float32)
    v = np.zeros(n, dtype=np.float32)

    x[0] = x0
    v[0] = v0

    for i in range(n - 1):
        # k1
        k1x = DeltaT * v[i]
        k1v = DeltaT * spring_force(v[i], x[i], t[i])

        # k2
        vv = v[i] + 0.5 * k1v
        xx = x[i] + 0.5 * k1x
        k2x = DeltaT * vv
        k2v = DeltaT * spring_force(vv, xx, t[i] + 0.5 * DeltaT)

        # k3
        vv = v[i] + 0.5 * k2v
        xx = x[i] + 0.5 * k2x
        k3x = DeltaT * vv
        k3v = DeltaT * spring_force(vv, xx, t[i] + 0.5 * DeltaT)

        # k4
        vv = v[i] + k3v
        xx = x[i] + k3x
        k4x = DeltaT * vv
        k4v = DeltaT * spring_force(vv, xx, t[i] + DeltaT)

        # Update
        x[i + 1] = x[i] + (k1x + 2.0 * k2x + 2.0 * k3x + k4x) / 6.0
        v[i + 1] = v[i] + (k1v + 2.0 * k2v + 2.0 * k3v + k4v) / 6.0
        t[i + 1] = t[i] + DeltaT

    return t, x, v


# Sequence generation for RNN training

def create_sequences(x, seq_len):
    """
    Given a 1D array x (e.g., position as a function of time),
    create input/target sequences for next-step prediction.

    Inputs:  [x_i, x_{i+1}, ..., x_{i+seq_len-1}]
    Targets: [x_{i+1}, ..., x_{i+seq_len}]
    """
    xs = []
    ys = []
    for i in range(len(x) - seq_len):
        seq_x = x[i : i + seq_len]
        seq_y = x[i + 1 : i + seq_len + 1]  # shifted by one step
        xs.append(seq_x)
        ys.append(seq_y)

    xs = np.array(xs, dtype=np.float32)      # shape: (num_samples, seq_len)
    ys = np.array(ys, dtype=np.float32)      # shape: (num_samples, seq_len)
    # Add feature dimension (1 feature: the position x)
    xs = np.expand_dims(xs, axis=-1)         # (num_samples, seq_len, 1)
    ys = np.expand_dims(ys, axis=-1)         # (num_samples, seq_len, 1)
    return xs, ys


class OscillatorDataset(Dataset):
    def __init__(self, seq_len=50, DeltaT=0.001, tfinal=20.0, x0=1.0, v0=0.0):
        t, x, v = rk4_trajectory(DeltaT=DeltaT, tfinal=tfinal, x0=x0, v0=v0)
        self.t = t
        self.x = x
        self.v = v
        xs, ys = create_sequences(x, seq_len=seq_len)
        self.inputs = torch.from_numpy(xs)  # (N, seq_len, 1)
        self.targets = torch.from_numpy(ys) # (N, seq_len, 1)

    def __len__(self):
        return self.inputs.shape[0]

    def __getitem__(self, idx):
        return self.inputs[idx], self.targets[idx]


# RNN model (LSTM-based in this example)

class RNNPredictor(nn.Module):
    def __init__(self, input_size=1, hidden_size=32, num_layers=1, output_size=1):
        super().__init__()
        self.lstm = nn.LSTM(input_size=input_size,
                            hidden_size=hidden_size,
                            num_layers=num_layers,
                            batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # x: (batch, seq_len, input_size)
        out, _ = self.lstm(x)   # out: (batch, seq_len, hidden_size)
        out = self.fc(out)      # (batch, seq_len, output_size)
        return out


# Training part

def train_model(
    seq_len=50,
    DeltaT=0.001,
    tfinal=20.0,
    batch_size=64,
    num_epochs=10,
    hidden_size=64,
    lr=1e-3,
    device=None,
):
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Dataset & DataLoader
    dataset = OscillatorDataset(seq_len=seq_len, DeltaT=DeltaT, tfinal=tfinal)
    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    # Model, loss, optimizer
    model = RNNPredictor(input_size=1, hidden_size=hidden_size, output_size=1)
    model.to(device)

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    # Training loop
    model.train()
    for epoch in range(num_epochs):
        epoch_loss = 0.0
        for batch_x, batch_y in train_loader:
            batch_x = batch_x.to(device)
            batch_y = batch_y.to(device)

            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item() * batch_x.size(0)

        epoch_loss /= len(train_loader.dataset)
        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.6f}")

    return model, dataset


# Evaluation / visualization

def evaluate_and_plot(model, dataset, seq_len=50, device=None):
    if device is None:
        device = next(model.parameters()).device

    model.eval()
    with torch.no_grad():
        # Take a single sequence from the dataset
        x_seq, y_seq = dataset[0]  # shapes: (seq_len, 1), (seq_len, 1)
        x_input = x_seq.unsqueeze(0).to(device)  # (1, seq_len, 1)
        # Model prediction (next-step for whole sequence)
        y_pred = model(x_input).cpu().numpy().squeeze(-1).squeeze(0)  # (seq_len,)
        # True target
        y_true = y_seq.numpy().squeeze(-1)  # (seq_len,)
        # Plot comparison
        plt.figure()
        plt.plot(y_true, label="True x(t+Δt)", linestyle="-")
        plt.plot(y_pred, label="Predicted x(t+Δt)", linestyle="--")
        plt.xlabel("Time step in sequence")
        plt.ylabel("Position")
        plt.legend()
        plt.title("RNN next-step prediction on oscillator trajectory")
        plt.tight_layout()
        plt.show()

# This is the main part of the code where we define the network

if __name__ == "__main__":
    # Hyperparameters can be tweaked as you like
    seq_len = 50
    DeltaT = 0.001
    tfinal = 20.0
    num_epochs = 10
    batch_size = 64
    hidden_size = 64
    lr = 1e-3

    model, dataset = train_model(
        seq_len=seq_len,
        DeltaT=DeltaT,
        tfinal=tfinal,
        batch_size=batch_size,
        num_epochs=num_epochs,
        hidden_size=hidden_size,
        lr=lr,
    )

    evaluate_and_plot(model, dataset, seq_len=seq_len)