Classification with Neural Networks

import copy

import torch
import torch.nn as nn
import torch.optim as optim

from sklearn.datasets import make_moons
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
from sklearn.datasets import make_blobs

from ipywidgets import interact, IntSlider
import owncloud

from sklearn.svm import SVC

def train(model: nn.Module, features: torch.Tensor, labels: torch.Tensor, nepochs: int = 100) -> None:
    """Train model on features from data"""

    loss_function = nn.CrossEntropyLoss()
    optimizer=torch.optim.RMSprop(model.parameters(), lr=0.1)
    for epoch in range(nepochs):
        optimizer.zero_grad()
        output = model.forward(features)
        loss = loss_function(output, labels)
        loss.backward()
        optimizer.step()

def calculate_accuracy(model: nn.Module, features: torch.Tensor, labels: torch.Tensor, verbose = True) -> torch.Tensor:
    """Calculate accuracy of model's classification"""

    output = model(features)
    accuracy = (labels == torch.argmax(output, dim=1)).sum()/len(labels)
    if verbose:
        print(f'Accuracy: {(accuracy*100).item()}')
    return (accuracy*100).item()

def do_binary_classification(model, X, labels, nepochs=40, step_size = 1, optimizer_class = optim.RMSprop):
    """Binary classification"""

    loss_function=nn.BCEWithLogitsLoss()

    optimizer = optimizer_class(model.parameters(), lr=0.1)

    fig, axes = plt.subplots(nrows=2, ncols=4, figsize = (12,6))
    max_subplots = 8
    isubplot = 0
    epochs_to_plot = set(np.linspace(0, nepochs - 1, max_subplots, dtype=int))

    # train
    for epoch in range(nepochs):
        optimizer.zero_grad()
        y_pred = model(X)
        loss = loss_function(y_pred.squeeze(), labels.float())
        loss.backward()
        optimizer.step()

        if epoch in epochs_to_plot and isubplot < max_subplots:

            # plot the result at this epoch
            with torch.no_grad():
                y_pred = model(X)
                probs = torch.sigmoid(y_pred)

                # create meshgrid for decision boundary
                x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5
                y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5
                xx, yy = np.meshgrid(np.linspace(x_min, x_max, 200),
                                        np.linspace(y_min, y_max, 200))
                grid_points = torch.tensor(np.c_[xx.ravel(), yy.ravel()], dtype=torch.float32)
                grid_preds = torch.sigmoid(model(grid_points)).reshape(xx.shape)

            irow = isubplot // 4
            icol = isubplot % 4
            axes[irow, icol].set_title(f"epoch={epoch}, loss={loss.item():.4f}")
            # plot decision boundary as a contour line at probability = 0.5
            axes[irow, icol].contour(xx, yy, grid_preds.numpy(), 'k--', levels=[0.5], linewidths=2, linestyles='dashed')
            axes[irow, icol].scatter(*X.T, s=20, c=(probs > 0.5).squeeze(), cmap='Set1')
            fig.tight_layout()

            isubplot += 1

def generate_clusters_nonlinear_classification(n_samples=256, noise=0.1, random_state=42):
    """Generate and plot two interleaving moon-shaped clusters with decision boundary."""
    features_orig, labels_orig = make_moons(n_samples, noise=noise, random_state=random_state)
    features = torch.tensor(features_orig, dtype=torch.float32)
    labels = torch.tensor(labels_orig, dtype=torch.long)

    # Fit SVM to approximate the true decision boundary
    svm = SVC(kernel='rbf', C=10, gamma='scale')
    svm.fit(features_orig, labels_orig)

    # Create meshgrid for decision boundary
    x_min, x_max = features[:, 0].min() - 0.5, features[:, 0].max() + 0.5
    y_min, y_max = features[:, 1].min() - 0.5, features[:, 1].max() + 0.5
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 200),
                        np.linspace(y_min, y_max, 200))
    Z = svm.decision_function(np.c_[xx.ravel(), yy.ravel()])
    Z = Z.reshape(xx.shape)

    plt.scatter(*features.T, s=20, c=labels, cmap='Dark2')
    plt.contour(xx, yy, Z, levels=[0], colors='black', linestyles='dashed', linewidths=2)
    plt.title('Clusters with a Non-Linear Decision Boundary')
    plt.xlim(-1.5, 2.5)
    plt.ylim(-1, 1.5)
    
    return features, labels


class utils:
    train = train
    calculate_accuracy = calculate_accuracy
    do_binary_classification = do_binary_classification
    generate_clusters_nonlinear_classification = generate_clusters_nonlinear_classification

Path('data').mkdir(exist_ok=True, parents=True)

owncloud.Client.from_public_link('https://uni-bonn.sciebo.de/s/9bZdbJc2JX22KTQ', folder_password="ibots").get_file('/', f'data/spike_count_and_stimulus_type_per_trial.pt')

# generate two linearly separable clusters
features_linear, labels_linear = make_blobs(
    n_samples=256, 
    centers=[[-1, 0.7], [1, -0.7]],
    cluster_std=0.4,
    random_state=2026
)
features_linear = torch.tensor(features_linear, dtype=torch.float32)
labels_linear = torch.tensor(labels_linear, dtype=torch.long)

plt.scatter(*features_linear.T, s=20, c=labels_linear, cmap='Dark2');
plt.title('Two Clusters Separated by a Linear Decision Boundary');

# Add illustration of decision boundary line (perpendicular bisector of the line connecting centers)
x_line = np.linspace(-2.5, 2.5, 100)
y_line = x_line / 0.7 - 0.3  # passes through origin with slope 1/0.7
plt.plot(x_line, y_line, 'k--', linewidth=2, label='Linear Decision Boundary')

plt.xlim(-2.5, 2.5)
plt.ylim(-2.5, 2.5)
plt.legend();

# create model
torch.manual_seed(2026)
model = nn.Linear(2,1)

utils.do_binary_classification(model, features_linear, labels_linear)

features, labels = utils.generate_clusters_nonlinear_classification()

utils.do_binary_classification(model, features, labels)

utils.do_binary_classification(model, features, labels, nepochs = 100)

nn.BCEWithLogitsLoss()

# simulated model outputs (logits) for 4 samples
# Positive values = model thinks "active", negative = model thinks "inactive"
model_output = torch.tensor([2.0, -1.0, 0.5, -2.0])

# True labels: 1 = active, 0 = inactive
labels_true = torch.tensor([1.0, 0.0, 1.0, 0.0])

print(f"Model predictions (logits): {model_output}")
print(f"True labels: {labels_true}")

# create the binary cross-entropy loss function
loss_function = nn.BCEWithLogitsLoss()
# calculate binary cross-entropy loss
loss = loss_function(model_output, labels_true)

print(f"Binary cross-entropy loss: {loss.item():.4f}")

# bad predictions - signs often don't match true labels
model_output = torch.tensor([-1.0, 1.0, -0.5, 1.0])

# True labels (same as before)
labels_true = torch.tensor([1.0, 0.0, 1.0, 0.0])

# model outputs for 3 samples, with 3 classes each (stimulus left, stimulus middle, stimulus right)
# each row is one sample, each column is the score for one class
model_output = torch.tensor([
    [2.0, 0.5, -1.0],   # Sample 1: highest score for stimulus left (index 0)
    [-1.0, 2.5, 0.0],   # Sample 2: highest score for stimulus middle (index 1)
    [0.0, -0.5, 3.0]    # Sample 3: highest score for stimulus right (index 2)
])

# true labels (class indices): stimulus left=0, stimulus middle=1, stimulus right=2
labels_true = torch.tensor([0, 1, 2])

print(f"Model predictions (logits):\n{model_output}")
print(f"True labels: {labels_true}")

# calculate cross-entropy loss
loss_function = nn.CrossEntropyLoss()
loss = loss_function(model_output, labels_true)

print(f"Cross-entropy loss: {loss.item():.4f}")

# bad predictions - highest scores don't match true labels
model_output = torch.tensor([
    [0.5, 2.0, -1.0],   # Sample 1: highest score for stimulus middle, but true label is stimulus left
    [-1.0, 2.5, 0.0],   # Sample 2: highest score for stimulus middle (correct)
    [2.0, -0.5, 1.0]    # Sample 3: highest score for stimulus left, but true label is stimulus right
])

# True labels (same as before)
labels_true = torch.tensor([0, 1, 2])

features, labels = utils.generate_clusters_nonlinear_classification()

torch.manual_seed(2026)
# create model
model = nn.Sequential(
    nn.Linear(2,16),
    nn.ReLU(),
    nn.Linear(16,2)
)

# Get model predictions
with torch.no_grad():
    output = model(features)

print(f'Output shape: {output.shape}')
print(f'First 5 rows of output:\n{output[:5]}')

predicted_labels = torch.argmax(output, dim=1)
predicted_labels[:10]

utils.train(model, features, labels, nepochs=50)

# Get model predictions
with torch.no_grad():
    output = model(features)

print(f'Output shape: {output.shape}')
print(f'First 5 rows of output:\n{output[:5]}')

data = torch.load('data/spike_count_and_stimulus_type_per_trial.pt')

features = data['features']
labels = data['labels']

print(f'Number of trials: {features.shape[0]}')
print(f'Number of cells: {features.shape[1]}')
print(f'len(labels): {len(labels)}')
print(f'labels: {labels}')

n_features = 1268

# create model
torch.manual_seed(2026)
model = nn.Sequential(
    nn.Linear(n_features, 64),
    nn.ReLU(),
    nn.Linear(64, 3)
)

# create optimizer for model 
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# create loss function
loss_fn = nn.CrossEntropyLoss()

# pass the input to the model and get a prediction
model_pred = model(features)

# calculate the loss
loss = loss_fn(model_pred, labels)

accuracy = utils.calculate_accuracy(model, features, labels, verbose=False)
print(f"Epoch {0}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.4f}")

for epoch in range(1):
    # reset the gradients
    optimizer.zero_grad()

    # pass the input to the model
    model_pred = model(features)

    # calculate the loss
    loss = loss_fn(model_pred, labels)

    # perform backward propagation to calculate the gradients
    loss.backward()

    # update the parameters
    optimizer.step()

for epoch in range(10):
    # reset the gradients
    optimizer.zero_grad()

    # pass the input to the model
    model_pred = model(features)

    # calculate the loss
    loss = loss_fn(model_pred, labels)

    # perform backward propagation to calculate the gradients
    loss.backward()

    # update the parameters
    optimizer.step()

# create model
torch.manual_seed(2026)
model = nn.Sequential(
    nn.Linear(n_features, 64),
    nn.ReLU(),
    nn.Linear(64, 3)
)

# create optimizer for model 
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# initialize list to store loss
losses = []
for epoch in range(20):
    # reset the gradients
    optimizer.zero_grad()

    # pass the input to the model
    model_pred = model(features)

    # calculate the loss
    loss = loss_fn(model_pred, labels)

    # perform backward propagation to calculate the gradients
    loss.backward()

    # update the parameters
    optimizer.step()

    # store loss in this epoch
    losses.append(loss.item())

plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Cross-entropy Loss')
plt.title('Loss Throughout Training');

# create model
torch.manual_seed(2026)
model = nn.Sequential(
    nn.Linear(n_features, 64),
    nn.ReLU(),
    nn.Linear(64, 3)
)

# create optimizer for model 
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)


___ # UPDATE ME: initialize list to store loss
for epoch in range(100):
    # reset the gradients
    optimizer.zero_grad()

    # pass the input to the model
    model_pred = model(features)

    # calculate the loss
    loss = loss_fn(model_pred, labels)

    # perform backward propagation to calculate the gradients
    loss.backward()

    # update the parameters
    optimizer.step()

    
    ___ # UPDATE ME: store loss in this epoch

plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Cross-entropy Loss')
plt.title('Loss Throughout Training');