classify_characters_neighbors.py

import os
import string
import random

import numpy as np
from scipy.io import loadmat
from scipy.ndimage import gaussian_filter1d
from sklearn.decomposition import PCA
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import confusion_matrix
from matplotlib import pyplot as plt


MATPLOTLIB_COLORS = plt.rcParams["axes.prop_cycle"].by_key()["color"]

ALL_CHARS = [
    *string.ascii_lowercase,
    "greaterThan",
    "tilde",
    "questionMark",
    "apostrophe",
    "comma",
]
CHAR_REPLACEMENTS = {
    "greaterThan": ">",
    "tilde": "~",
    "questionMark": "?",
    "apostrophe": "'",
    "comma": ",",
}
CHAR_TO_CLASS_MAP = {char: idx for idx, char in enumerate(ALL_CHARS)}
CLASS_TO_CHAR_MAP = {idx: char for idx, char in enumerate(ALL_CHARS)}

REACTION_TIME_BINS = 10
TRAINING_WINDOW_BINS = 150

NUM_ELECTRODES = 192
NUM_PCS = 15

WARPING_ALPHAS = np.linspace(1.0, 2.0, 11)

OUTPUTS_DIR = os.path.abspath("./outputs")


########################################################################################
# Main function.
########################################################################################


def main_KNN():
    ## Load the data.

    data_dicts = load_data()

    ## Run the whole process multiple times to get a series of results.

    accuracy_KNN = run_multiple_KNN(data_dicts)
    print(f"accuracy_KNN: {accuracy_KNN}")

    ## Vary number of electrodes.

    NUM_ELECTRODES_TO_TRY = [24, 48, 72, 96, 120, 144, 168, 192]

    accuracies_by_electrodes_KNN = [
        run_multiple_KNN(data_dicts, num_electrodes=num_electrodes)
        for num_electrodes in NUM_ELECTRODES_TO_TRY
    ]
    print(f"accuracies_by_electrodes_KNN: {accuracies_by_electrodes_KNN}")

    ## Vary number of training trials.

    NUM_TRAIN_TRIALS_TO_TRY = [300, 600, 900, 1200, 1500, 1800, 2100, 2400, 2700]

    accuracies_by_train_trials_KNN = [
        run_multiple_KNN(data_dicts, num_train_trials=num_train_trials)
        for num_train_trials in NUM_TRAIN_TRIALS_TO_TRY
    ]
    print(f"accuracies_by_train_trials_KNN: {accuracies_by_train_trials_KNN}")


########################################################################################
# Helper functions.
########################################################################################


def load_data():
    """"""
    DATA_DIR = os.path.abspath("./handwritingBCIData/Datasets/")
    letters_filepaths = []
    for root, _, filenames in os.walk(DATA_DIR):
        for filename in filenames:
            filepath = os.path.join(root, filename)
            if filename == "singleLetters.mat":
                letters_filepaths.append(filepath)

    letters_filepaths = sorted(letters_filepaths)

    data_dicts = []
    for filepath in letters_filepaths:
        print(f"Loading {filepath} ...")
        data_dict = loadmat(filepath)
        data_dicts.append(data_dict)

    return data_dicts


def organize_data(data_dicts, limit_electrodes=None, limit_train_trials=None):
    """"""

    print("Preparing data ...")

    X_train = []
    X_test = []
    y_train = []
    y_test = []

    # Random electrode order to let us limit electrodes.
    rand_electrode_order = list(range(NUM_ELECTRODES))
    random.shuffle(rand_electrode_order)

    # Iterate through the sessions.
    # NUM_SESSIONS = 1
    NUM_SESSIONS = None
    for data_dict in data_dicts[:NUM_SESSIONS]:
        neural = data_dict["neuralActivityTimeSeries"]
        go_cue_bins = data_dict["goPeriodOnsetTimeBin"].ravel().astype(int)
        delay_cue_bins = data_dict["delayPeriodOnsetTimeBin"].ravel().astype(int)
        prompts = np.array([a[0] for a in data_dict["characterCues"].ravel()])
        block_by_bin = data_dict["blockNumsTimeSeries"].ravel()
        block_nums = data_dict["blockList"].ravel()

        # If specified, only use a random subset of electrodes.
        if limit_electrodes is not None:
            neural = neural[:, rand_electrode_order[:limit_electrodes]]

        # Iterate through each block in this session.
        for block_num in block_nums:
            # Get means and stddevs from a random set of train trials in the block, and
            # the rest of the trials can be used for test.
            block_trial_mask = [block_by_bin[b] == block_num for b in go_cue_bins]
            num_trials_in_block = sum(block_trial_mask)
            random_trial_idxs = list(range(num_trials_in_block))
            random.shuffle(random_trial_idxs)
            train_end_idx = int(num_trials_in_block * 0.8)
            train_trial_idxs = random_trial_idxs[:train_end_idx]
            block_go_cue_bins = go_cue_bins[block_trial_mask]
            block_delay_cue_bins = delay_cue_bins[block_trial_mask]
            block_prompts = prompts[block_trial_mask]
            # Loop through the train trials and add the neural data to our list.
            neural_to_zscore_based_on = []
            for trial_idx in train_trial_idxs:
                # For convenience, ignore the last trial in the block.
                if trial_idx + 1 >= len(block_delay_cue_bins):
                    continue
                start_bin = block_delay_cue_bins[trial_idx]
                end_bin = block_delay_cue_bins[trial_idx + 1]
                neural_to_zscore_based_on.extend(neural[start_bin:end_bin])
            neural_to_zscore_based_on = np.array(neural_to_zscore_based_on)
            block_means = np.mean(neural_to_zscore_based_on, axis=0)
            block_stddevs = np.std(neural_to_zscore_based_on, axis=0)

            for trial_idx in range(num_trials_in_block):
                # Get the training window for this trial.
                go_cue_bin = block_go_cue_bins[trial_idx]
                window_start_bin = int(go_cue_bin) + REACTION_TIME_BINS
                window_end_bin = window_start_bin + TRAINING_WINDOW_BINS
                # Get the neural data in this window.
                trial_neural = neural[window_start_bin:window_end_bin]
                # Z-score the neural data using the block-specific means and stddevs.
                with np.errstate(divide="ignore", invalid="ignore"):
                    trial_zscored_neural = (trial_neural - block_means) / block_stddevs
                trial_zscored_neural = np.nan_to_num(
                    trial_zscored_neural, nan=0, posinf=0, neginf=0
                )

                # Get the character for this trial.
                trial_label = block_prompts[trial_idx]

                # Skip rest trials.
                if trial_label == "doNothing":
                    continue

                # Add the trial to the appropriate set of data (train or test).
                if trial_idx in train_trial_idxs:
                    X_train.append(trial_zscored_neural)
                    y_train.append(trial_label)
                else:
                    X_test.append(trial_zscored_neural)
                    y_test.append(trial_label)

    # Smooth the neural data over time.
    SMOOTHING_STDDEV = 3.0
    X_train = np.array(
        [gaussian_filter1d(w, sigma=SMOOTHING_STDDEV, axis=0) for w in X_train]
    )
    X_test = np.array(
        [gaussian_filter1d(w, sigma=SMOOTHING_STDDEV, axis=0) for w in X_test]
    )

    # Fit a PCA model (only on the training data).

    print("Fitting PCA model ...")

    pca_model = PCA(n_components=NUM_PCS).fit(np.concatenate(X_train))

    # Get PCA-transformed data (all of training and test).
    X_train = np.array([pca_model.transform(w) for w in X_train])
    X_test = np.array([pca_model.transform(w) for w in X_test])

    # Flatten the each trial's PCs so KNN can operate on 1D vectors.
    X_train = np.reshape(X_train, (X_train.shape[0], -1))
    X_test = np.reshape(X_test, (X_test.shape[0], -1))

    # Convert the characters to ints, for compatibility with pytorch.
    y_train = np.array([CHAR_TO_CLASS_MAP[ch] for ch in y_train])
    y_test = np.array([CHAR_TO_CLASS_MAP[ch] for ch in y_test])

    # If specified, only use a random subset of train trials.
    if limit_train_trials:
        X_train = X_train[:limit_train_trials]
        y_train = y_train[:limit_train_trials]

    print(f"X_train.shape: {X_train.shape}")
    print(f"X_test.shape: {X_test.shape}")
    print(f"y_train.shape: {y_train.shape}")
    print(f"y_test.shape: {y_test.shape}")

    return X_train, X_test, y_train, y_test


def time_lengthen_and_slice_vector(y, alpha):
    """"""
    # Sample the series at a smaller interval than originally (alpha is expected to be
    # >= 1 because we're lengthening).
    t = np.arange(len(y))
    new_y = np.interp(t / alpha, t, y)
    # This new series has the same number of elements as the original but covers a
    # smaller portion in time. If this new series is returned and treated as if the
    # values correspond to the original time points, the original series has effectively
    # been stretched and then sliced at its original time length.
    return new_y


def time_lengthen_and_slice_matrix(m, alpha):
    """"""
    # If each row is a time bin, each electrode's time series is a column. Iterate
    # through the columns, lengthen each one.
    cols = m.T
    new_cols = np.array(
        [time_lengthen_and_slice_vector(series, alpha) for series in cols]
    )
    new_m = new_cols.T
    return new_m


def dist_with_time_warp(m0, m1, cache={}):
    """"""
    # The args have to be vectors, so reshape them to the matrices they represent.
    m0 = m0.reshape(-1, NUM_PCS)
    m1 = m1.reshape(-1, NUM_PCS)

    # Lookup in our cache the warped versions of the inputs. If not there, calculate
    # them and add them to our cache first.
    m0.flags.writeable = False
    m0_id = hash(m0.tobytes())
    if m0_id not in cache:
        warped_m0s = {
            alpha: time_lengthen_and_slice_matrix(m0, alpha) for alpha in WARPING_ALPHAS
        }
        cache[m0_id] = warped_m0s
    else:
        warped_m0s = cache[m0_id]

    m1.flags.writeable = False
    m1_id = hash(m1.tobytes())
    if m1_id not in cache:
        warped_m1s = {
            alpha: time_lengthen_and_slice_matrix(m1, alpha) for alpha in WARPING_ALPHAS
        }
        cache[m1_id] = warped_m1s
    else:
        warped_m1s = cache[m1_id]

    # Iterate through a list of scaling factors alpha.

    min_dist_so_far = np.inf

    for alpha in WARPING_ALPHAS:
        # Keep the first matrix the same but stretch the other matrix in time by a
        # factor of alpha (using linear interpolation), and take the distance between
        # the first matrix and the warped other matrix.

        dist_0 = np.linalg.norm(m0 - warped_m1s[alpha])
        # If this is the lowest distance we've seen, keep it.
        if dist_0 < min_dist_so_far:
            min_dist_so_far = dist_0

        # Do the same thing, but lengthen the first matrix instead.

        dist_1 = np.linalg.norm(m1 - warped_m0s[alpha])
        if dist_1 < min_dist_so_far:
            min_dist_so_far = dist_1

    return min_dist_so_far


def plot_confusion_matrix(y_test, y_pred_test, accuracy_str):
    """"""

    fig, confusion_ax = plt.subplots()

    confusion_results = confusion_matrix(y_test, y_pred_test, normalize="true")

    heatmap = confusion_ax.imshow(confusion_results, origin="lower")

    fig.colorbar(heatmap, ax=confusion_ax)

    confusion_ax.set_xticks(np.arange(len(ALL_CHARS)))
    confusion_ax.set_xticklabels(ALL_CHARS, rotation=45, ha="right")
    confusion_ax.set_xlabel("predicted character")

    confusion_ax.set_yticks(np.arange(len(ALL_CHARS)))
    confusion_ax.set_yticklabels(ALL_CHARS)
    confusion_ax.set_ylabel("true character")

    model_str = "KNearestNeighbors"
    confusion_ax.set_title(
        f"{model_str} on single-letter instructed-delay task (accuracy: {accuracy_str})"
    )

    plt.tight_layout()

    plt.show()


def run_multiple_KNN(
    data_dicts, num_electrodes=None, num_train_trials=None, num_runs=5
):
    """"""

    accuracy_results = []

    for run_idx in range(num_runs):
        print(f"KNearestNeighbors run {run_idx + 1} / {num_runs}")

        ## Preprocess and label the data.

        X_train, X_test, y_train, y_test = organize_data(
            data_dicts,
            limit_electrodes=num_electrodes,
            limit_train_trials=num_train_trials,
        )

        ## Train a k-nearest neighbors model on the preprocessed training data.

        print("Training k-nearest neighbors model ...")

        NUM_NEIGHBORS = 10
        knn_model = KNeighborsClassifier(
            n_neighbors=NUM_NEIGHBORS, metric=dist_with_time_warp
        )
        knn_model.fit(X_train, y_train)

        ## Evaluate the k-nearest neighbors model by calculating accuracy on the test
        ## set.

        print("Evaluating k-nearest neighbors model ...")

        y_pred_test = knn_model.predict(X_test)

        test_accuracy = np.sum(y_pred_test == y_test) / len(y_test)

        # Store this run's result.
        accuracy_results.append(test_accuracy)

        accuracy_str = f"{round(test_accuracy, 3):.3f}"
        print(f"accuracy: {accuracy_str}")

        ## Optionally plot the confusion matrix.

        show_confusion_matrix = False
        if show_confusion_matrix:
            plot_confusion_matrix(y_test, y_pred_test, accuracy_str)

    mean_accuracy = np.mean(accuracy_results)

    print(f"accuracies: {accuracy_results}")
    print(f"mean accuracy: {np.mean(accuracy_results)}")

    return mean_accuracy


if __name__ == "__main__":
    main_KNN()